org.atmosphere.plugin.rabbitmq.RabbitMQBroadcaster.java Source code

Introduction

Here is the source code for org.atmosphere.plugin.rabbitmq.RabbitMQBroadcaster.java

Source

/*
 * Copyright 2013 Jean-Francois Arcand
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy of
 * the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations under
 * the License.
 */
package org.atmosphere.plugin.rabbitmq;

import com.rabbitmq.client.AMQP;
import com.rabbitmq.client.Channel;
import com.rabbitmq.client.DefaultConsumer;
import com.rabbitmq.client.Envelope;
import com.rabbitmq.client.MessageProperties;
import org.atmosphere.cpr.AtmosphereConfig;
import org.atmosphere.cpr.BroadcasterFuture;
import org.atmosphere.cpr.Entry;
import org.atmosphere.util.SimpleBroadcaster;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;

/**
 * Simple {@link org.atmosphere.cpr.Broadcaster} implementation based on RabbitMQ
 *
 * @author Thibault Normand
 * @author Jean-Francois Arcand
 */
public class RabbitMQBroadcaster extends SimpleBroadcaster {

    private static final Logger logger = LoggerFactory.getLogger(RabbitMQBroadcaster.class);

    private String queueName;
    private String consumerTag;
    private final RabbitMQConnectionFactory factory;
    private final Channel channel;
    private final String exchangeName;

    public RabbitMQBroadcaster(String id, AtmosphereConfig config) {
        super(id, config);
        factory = RabbitMQConnectionFactory.getFactory(config);
        channel = factory.channel();
        exchangeName = factory.exchangeName();
    }

    @Override
    public void setID(String id) {
        super.setID(id);
        restartConsumer();
    }

    @Override
    public String getID() {
        String id = super.getID();
        if (id.startsWith("/*")) {
            id = "atmosphere";
        }
        return id;
    }

    @Override
    protected void push(Entry entry) {
        if (destroyed.get()) {
            return;
        }

        outgoingBroadcast(entry.message);
    }

    public void outgoingBroadcast(Object message) {
        try {
            String id = getID();
            if (message instanceof String) {
                logger.trace("Outgoing broadcast : {}", message);

                channel.basicPublish(exchangeName, id, MessageProperties.PERSISTENT_TEXT_PLAIN,
                        message.toString().getBytes());
            } else {
                throw new IOException("Message is not a string, so could not be handled !");
            }

        } catch (IOException e) {
            logger.warn("Failed to send message over RabbitMQ", e);
        }
    }

    void restartConsumer() {
        try {
            final String id = getID();

            if (consumerTag != null) {
                logger.debug("Delete consumer {}", consumerTag);
                channel.basicCancel(consumerTag);
                consumerTag = null;
            }

            if (queueName != null) {
                logger.debug("Delete queue {}", queueName);
                channel.queueUnbind(queueName, exchangeName, id);
                channel.queueDelete(queueName);
                queueName = null;
            }

            queueName = channel.queueDeclare().getQueue();
            channel.queueBind(queueName, exchangeName, id);

            logger.info("Create AMQP consumer on queue {}, for routing key {}", queueName, id);

            DefaultConsumer queueConsumer = new DefaultConsumer(channel) {
                @Override
                public void handleDelivery(String consumerTag, Envelope envelope, AMQP.BasicProperties properties,
                        byte[] body) throws IOException {

                    // Not for us.
                    if (!envelope.getRoutingKey().equalsIgnoreCase(id))
                        return;

                    String message = new String(body);
                    try {
                        Object newMsg = filter(message);
                        // if newSgw == null, that means the message has been filtered.
                        if (newMsg != null) {
                            deliverPush(new Entry(newMsg,
                                    new BroadcasterFuture<Object>(newMsg, RabbitMQBroadcaster.this), message),
                                    true);
                        }
                    } catch (Throwable t) {
                        logger.error("failed to push message: " + message, t);
                    }
                }
            };

            consumerTag = channel.basicConsume(queueName, true, queueConsumer);
            logger.info("Consumer " + consumerTag + " for queue {}, on routing key {}", queueName, id);

        } catch (Throwable ex) {
            String msg = "Unable to initialize RabbitMQBroadcaster";
            logger.error(msg, ex);
            throw new IllegalStateException(msg, ex);
        }
    }

    @Override
    public synchronized void releaseExternalResources() {
        try {
            if (channel != null && channel.isOpen()) {
                if (consumerTag != null) {
                    channel.basicCancel(consumerTag);
                }
            }
        } catch (Exception ex) {
            logger.trace("", ex);
        }
    }
}