The following diagram illustrates the interactions between the client and server WebSocket peers and the RabbitMQ exchanges and queues:

The client and server WebSockets connecting

In the previous diagram, the publication of messages towards the direct exchange for user-to-user messages, and the topic exchange for group messages, is not different than before. What is different is that instead of getting the messages from RabbitMQ queues, you will consume them. What's the difference? When you consume messages, you register a listener that new messages arriving in the queue will be automatically delivered to. So, unlike the synchronous basicGet operations you were performing before, you will now be using an asynchronous consumer to receive the queued messages.

Since you're still working on the CCM's application server that is in Java, you will look at RabbitMQ's Java API. Registering a queue consumer is as simple as follows:

channel.basicConsume(queue, autoAck, consumer)

Here, the consumer is an implementation of the com.rabbitmq.client.Consumer interface. The interface defines the contract between a queue consumer and RabbitMQ. It sports several methods, some related to receiving error notifications, but the main method you will focus on is the following:

void handleDelivery(String consumerTag,
                    Envelope envelope,
                    AMQP.BasicProperties properties,
                    byte[] body)
    throws IOException;

The handleDelivery method is called whenever a message is received from the queue. Hence, when this method is called, you will want to push the message back to the frontend via the WebSocket server.

The consumer is bound to the channel that was used to consume a particular queue. If this channel is closed, the consumer will stop receiving messages. Since a channel cannot be reopened and has to be recreated from scratch, the implication is that both the channel and its consumer must be re-established in case of problems. CCM decides to tackle this problem by wrapping the consumer in a class that supports the reconnection mechanism.

public class Subscription
{
    private final static Logger LOGGER = Logger.getLogger(Subscription.class.getName());

    private final String queue;
    private final SubscriptionDeliverHandler handler;

    private volatile DefaultConsumer consumer;

    public Subscription(final String queue, final SubscriptionDeliverHandler handler)
    {
        this.queue = queue;
        this.handler = handler;
    }

public interface SubscriptionDeliverHandler
{
    void handleDelivery(Channel channel,
                        Envelope envelope,
                        AMQP.BasicProperties properties,
                        byte[] body);
}

public String start(final Channel channel) throws IOException
{
    consumer = null;

    if (channel != null)
    {
        try
        {
            consumer = new DefaultConsumer(channel)
            {
                @Override
                public void handleDelivery(final String consumerTag,
                    final Envelope envelope,
                    final BasicProperties properties,
                    final byte[] body) throws IOException
                {
                  handler.handleDelivery(channel, envelope, properties, body);
                }

            };

            final boolean autoAck = false;
            final String consumerTag = channel.basicConsume(queue, autoAck, consumer);
            LOGGER.info("Consuming queue: " + queue + ": with tag: " + consumerTag + " on channel: "
                        + channel);

            return consumerTag;
        }
        catch (final Exception e)
        {
            LOGGER.log(Level.SEVERE, "Failed to start consuming queue: " + queue, e);
            consumer = null;
        }
    }

    return null;
}

public void stop()
{
    final Channel channel = getChannel();
    if (channel == null)
    {
        return;
    }

    LOGGER.log(Level.INFO, "Stopping subscription: " + this);

    try
    {
        channel.basicCancel(consumer.getConsumerTag());
    }
    catch (final Exception e)
    {
        LOGGER.log(Level.SEVERE, "Failed to cancel subscription: " + this, e);
    }

    try
    {
        channel.close();
    }
    catch (final Exception e)
    {
        LOGGER.log(Level.SEVERE, "Failed to close channel: " + channel, e);
    }
    finally
    {
        consumer = null;
    }
}

public Channel getChannel()
{
    return consumer == null ? null : consumer.getChannel();
}

@Override
protected void finalize() throws Throwable
{
    stop();
}

@Override
public String toString()
{
    final ToStringHelper tsh = Objects.toStringHelper(this).addValue(hashCode()).add("queue", queue);
    if (consumer != null)
    {
        tsh.add("channel", getChannel());
        tsh.add("consumerTag", consumer.getConsumerTag());
    }
    return tsh.toString();
}

private final Set<Subscription> subscriptions;
 
public RabbitMqManager(final ConnectionFactory factory)
{
    // ... existing code omitted

    subscriptions = synchronizedSet(new HashSet<Subscription>());
}

public Subscription createSubscription(final String queue, final SubscriptionDeliverHandler handler)
{
    final Subscription subscription = new Subscription(queue, handler);
    subscriptions.add(subscription);
    startSubscription(subscription);
    return subscription;
}

private void startSubscription(final Subscription subscription)
{
    final Channel channel = createChannel();

    if (channel != null)
    {
        try
        {
            subscription.start(channel);
        }
        catch (final Exception e)
        {
            LOGGER.log(Level.SEVERE, "Failed to start subscription: " + subscription + " on channel: "
                                     + channel, e);
        }
    }
}

private void restartSubscriptions()
{
    LOGGER.info("Restarting " + subscriptions.size() + " subscriptions");

    for (final Subscription subscription : subscriptions)
    {
        startSubscription(subscription);
    }
}

You first need to refactor the UserMessageManager class to expose variants of sendUserMessage and sendTopicMessage that take a channel argument. Indeed, since you will have an active channel associated with a subscription, you will use it not only to consume messages, but also to produce them.

On top of these basic refactorings, you also need to add the following to the UserMessageManager method to allow creating a subscription for a particular user inbox:

public Subscription subscribeToUserInbox(final long userId, final SubscriptionDeliverHandler handler)
{
    final String queue = getUserInboxQueue(userId);
    return rabbitMqManager.createSubscription(queue, handler);
}

CCM uses a JSR-356-compliant implementation of server-side WebSocket in its Java application backend. In this model, an application has to expose WebSocket endpoints, so CCM will create one endpoint dedicated to user messaging.

We will look at the main methods of UserMessageServerEndpoint, keeping in mind that WebSocket authentication will not be discussed. Let's first look at what happens when a user browser connects to the WebSocket server:

@OnOpen
public void startSubscription(@PathParam("user-id") final long userId, final Session session)
{
    session.setMaxIdleTimeout(0);

    final Subscription subscription = userMessageManager.subscribeToUserInbox(userId,
        new SubscriptionDeliverHandler()
        {
            @Override
            public void handleDelivery(final Channel channel,
              final Envelope envelope,
              final BasicProperties properties,
              final byte[] body)
            {
              try
              {
                final String contentEncoding = properties.getContentEncoding();
                session.getBasicRemote().sendText(new String(body, contentEncoding));
                channel.basicAck(envelope.getDeliveryTag(), false);
                }
                catch (final Exception e)
                {
                  LOGGER.log(Level.SEVERE,
                    "Failed to push over websocket message ID: " + properties.getMessageId(), e);

                  try
                  {
                      final boolean requeue = true;
                      channel.basicReject(envelope.getDeliveryTag(), requeue);
                  }
                  catch (final Exception e2)
                  {
                      LOGGER.log(Level.SEVERE,
                        "Failed to reject and requeue message ID: " + properties.getMessageId(), e);
                  }
                }
            }
        });

    session.getUserProperties().put(RABBITMQ_SUBSCRIPTION, subscription);
}

The important aspect of this method is that it uses userMessageManager received by the dependency injection, to subscribe SubscriptionDeliverHandler that is in charge of sending the messages consumed from the user-specific queue on WebSocket. Do you see how the channel instance passed in CCM's custom handleDelivery method comes handy? It is required needed to perform manual message acknowledgement with the basicAck channel action. You perform this single message acknowledgment if and only if the sendText WebSocket operation has succeeded (that is, it didn't throw an exception). Otherwise, you use basicReject to actively reject and requeue the delivered message. If you don't acknowledge and fail to reject a message, the RabbitMQ broker will eventually redeliver the message once the subscription is re-established with a new channel.

Also, note how the session's user properties are used to store the subscription so that it can be used in other methods. Indeed, you are required to gracefully terminate the subscription in case of disconnection of the WebSocket, as shown in the following code:

@OnClose
public void stopSubscription(final Session session)
{
    final Subscription subscription = (Subscription) session.getUserProperties().get(
        RABBITMQ_SUBSCRIPTION);

    if (subscription != null)
    {
        subscription.stop();
    }
}

Of course, the Subscription instance is also required to publish messages because, as you may remember, it acts as a channel provider, as shown in the following code:

@OnMessage
public void publishMessage(final String jsonMessage, final Session session)
    throws IOException, EncodeException
{
    final Subscription subscription = (Subscription) session.getUserProperties().get(
        RABBITMQ_SUBSCRIPTION);

    final Channel channel = subscription == null ? null : subscription.getChannel();
    if (channel == null)
    {
        LOGGER.log(Level.SEVERE, "No active channel to dispatch message: " + jsonMessage);
        return;
    }

    // inspect the message to find out where to route it
    final Message message = OBJECT_MAPPER.readValue(jsonMessage, Message.class);
    if (message.getAddresseeId() != null)
    {
        userMessageManager.sendUserMessage(message.getAddresseeId(), jsonMessage, channel);
    }
    else if (!Strings.isNullOrEmpty(message.getTopic()))
    {
        userMessageManager.sendTopicMessage(message.getTopic(), jsonMessage, channel);
    }
    else
    {
        LOGGER.log(Level.SEVERE, "Received unroutable message: " + jsonMessage);
    }
}

Did you see how you now use the sendUserMessage and sendTopicMessage method variants that take a channel as a third argument? There is now no reason to keep using the somewhat wasteful method of creating and closing the channel each time since you now have access to a channel, which itself benefits from the reconnection mechanism you've created.

CCM is now ready to activate server-push for its user messaging feature!

Running the application

CCM has tested whether the server-push mechanism successfully withstands connection issues with the RabbitMQ broker. Its application can start even if it can't connect to the broker, and it can recover if the broker is restarted.

Tip

The robustness of a distributed system isn't the sole responsibility of one actor but a combination of the effort of all its members. No matter how highly available one of its members is, accounting for its potential failure in other members that depend on it will ensure a smooth ride and avoid the proverbial wake-up call.

Let's take a glance at the management console to see how switching to server-push has affected things. The following figure shows the channel's view:

Active consumers keeping channels open

Remember how with the polling approach, where channels were open and closed very quickly, no channel was visible on the view. Now, because each consumer keeps its channel open, you can see active channels in the management console. You can also see the associated usage rates.

Note

There is no logical limit to the number of channels a RabbitMQ broker can handle; the limiting factors are the available memory on the broker as each channel mobilizes memory, and the actual network bandwidth available for pushing messages on all these channels.

Now, let's take a look at the queue view of the management console:

Manual acknowledgements showing up in rates

Did you see how the ack column shows a non-zero rate? This is because you are now using manual acknowledgment. Thus, the RabbitMQ client now sends ack messages over the wire to the broker. This definitely has a cost in terms of bandwidth usage and general performance; however, if in your case, you value the guarantee of successful message processing over speed, it is perfectly acceptable.

After hearing about the resounding success of server-pushed user messages, the customer support team at Clever Coney Media came up with a new requirement: being able to message all users in the system. Let's see how this new feature can be implemented!