com.vmware.vhadoop.vhm.rabbit.RabbitConnection.java Source code

Introduction

Here is the source code for com.vmware.vhadoop.vhm.rabbit.RabbitConnection.java

Source

/***************************************************************************
* Copyright (c) 2013 VMware, Inc. All Rights Reserved.
* 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.
***************************************************************************/

/***************************************************************************
 * Copyright (c) 2012 VMware, Inc. All Rights Reserved.
 * 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 com.vmware.vhadoop.vhm.rabbit;

import java.io.IOException;
import java.util.logging.Level;
import java.util.logging.Logger;

import com.rabbitmq.client.Channel;
import com.rabbitmq.client.Connection;
import com.rabbitmq.client.ConnectionFactory;
import com.rabbitmq.client.QueueingConsumer;
import com.rabbitmq.client.ShutdownListener;
import com.rabbitmq.client.ShutdownSignalException;
import com.vmware.vhadoop.api.vhm.QueueClient.CannotConnectException;

/**
 * Encapsulates the details of connecting to a RabbitMQ queue
 *
 */
public class RabbitConnection {
    private static final Logger _log = Logger.getLogger(RabbitConnection.class.getName());

    interface RabbitCredentials {
        public String getHostName();

        public String getExchangeName();

        public String getRouteKeyCommand();

        public String getRouteKeyStatus();
    }

    private ConnectionFactory _connectionFactory;
    private RabbitCredentials _credentials;

    private Channel _channel = null;
    private Object _channelLock = new Object();
    private Connection _connection = null;
    private Object _connectionLock = new Object();
    private String _queueName = null;
    private QueueingConsumer _consumer = null;
    private Object _consumerLock = new Object();
    private volatile boolean _started = false;

    /* For testing only */
    protected RabbitConnection() {
    }

    public RabbitConnection(RabbitCredentials credentials) {
        _credentials = credentials;
        _connectionFactory = new ConnectionFactory();
        _connectionFactory.setHost(credentials.getHostName());
    }

    private Connection getConnection() throws IOException {
        synchronized (_connectionLock) {
            if (_connection == null || !_connection.isOpen()) {
                _connection = null;
                _connection = _connectionFactory.newConnection();
                _connection.addShutdownListener(new ShutdownListener() {
                    @Override
                    public void shutdownCompleted(ShutdownSignalException cause) {
                        _log.info("Connection shut down");
                        _log.log(Level.FINE, "{0}", cause.getReason());
                        synchronized (_connectionLock) {
                            _connection = null;
                            _channel = null;
                            _queueName = null;
                        }
                        _started = false;
                    }
                });
                _started = true;
                _log.fine("Created new connection");
            }

            return _connection;
        }
    }

    private Channel getChannel() throws IOException {
        synchronized (_channelLock) {
            if (_channel == null || !_channel.isOpen()) {
                _log.fine("Creating new channel");
                _channel = null;
                Connection connection = getConnection();
                _channel = connection.createChannel();
                _channel.addShutdownListener(new ShutdownListener() {
                    @Override
                    public void shutdownCompleted(ShutdownSignalException cause) {
                        _log.info("Channel shut down");
                        _log.log(Level.FINE, "{0}", cause.getReason());
                        synchronized (_channelLock) {
                            _channel = null;
                            _queueName = null;
                        }
                        _started = false;
                    }
                });
            }

            return _channel;
        }
    }

    private synchronized String getQueueName() throws IOException {
        if (_queueName == null) {
            String exchangeName = _credentials.getExchangeName();
            Channel channel = getChannel();
            channel.exchangeDeclare(exchangeName, "direct", true, false, null); /* TODO: Externalize? */
            _queueName = channel.queueDeclare().getQueue();
            channel.queueBind(_queueName, exchangeName, _credentials.getRouteKeyCommand());
            _log.fine("Created transient queue: " + _queueName);
        }

        return _queueName;
    }

    /* TODO: Can we cache the consumer object? */
    QueueingConsumer getConsumer() throws CannotConnectException {
        synchronized (_consumerLock) {
            if (_consumer == null) {
                _log.fine("Creating new consumer");
                try {
                    Channel channel = getChannel();
                    _consumer = new QueueingConsumer(channel) {
                        @Override
                        public void handleShutdownSignal(java.lang.String consumerTag,
                                ShutdownSignalException sig) {
                            super.handleShutdownSignal(consumerTag, sig);
                            _log.info("Consumer received shutdown notification");
                            _log.log(Level.FINE, "{0}", sig.getReason());

                            synchronized (_consumerLock) {
                                _consumer = null;
                            }
                        }
                    };
                    channel.basicConsume(getQueueName(), true, _consumer);
                } catch (Exception e) {
                    throw new CannotConnectException("Unable to get message consumer", e);
                }
            }

            return _consumer;
        }
    }

    boolean connect() {
        Connection connection = null;
        try {
            connection = getConnection();
        } catch (IOException e) {
            /* squash */
        }

        return connection != null;
    }

    protected void sendMessage(String routeKey, byte[] data) throws CannotConnectException {
        boolean retry = false;
        do {
            try {
                Channel channel = getChannel();
                if (channel != null) {
                    channel.basicPublish(_credentials.getExchangeName(), routeKey, null, data);
                }
                retry = false;
            } catch (Exception e) {
                if (!retry) {
                    connect();
                    retry = true;
                } else {
                    throw new CannotConnectException("Unable to send message", e);
                }
            }
        } while (retry);
    }

    protected void sendMessage(byte[] data) throws CannotConnectException {
        sendMessage(_credentials.getRouteKeyStatus(), data);
    }

    public boolean isShutdown() {
        return (_started == false);
    }
}