net.leanix.dropkit.amqp.ConsumerRegistry Maven / Gradle / Ivy

Go to download
package net.leanix.dropkit.amqp;

import java.io.IOException;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.inject.Inject;
import com.google.inject.Singleton;
import com.rabbitmq.client.Channel;
import com.rabbitmq.client.ShutdownSignalException;

@Singleton
public class ConsumerRegistry {

    private final static Logger log = LoggerFactory.getLogger(ConsumerRegistry.class);

    static class RemoveOldConsumerRunnable implements Runnable {

        private final ConsumerRegistry consumerRegistry;

        RemoveOldConsumerRunnable(ConsumerRegistry consumerRegistry) {
            this.consumerRegistry = consumerRegistry;
        }

        @Override
        public void run() {
            while (true) {
                consumerRegistry.removeOldConsumers(10L * 60 * 1000);
                try {
                    Thread.sleep(60 * 1000);
                } catch (InterruptedException e) {
                    log.info("Stopping RemoveOldConsumerRunnable loop used to clean up unused consumers.");
                    break;
                }
            }
        }
    };

    // interval that the queue will survive having no consumers
    public static final long QUEUE_X_EXPIRY_MILLIS = 2 * 60 * 1000;// 24L * 60 * 60 * 1000;

    private final ConnectionHolder connectionHolder;
    private final Map consumerMap = new HashMap<>();

    private final QueueConsumerFactory consumerFactory;

    @Inject
    public ConsumerRegistry(ConnectionHolder connectionHolder, QueueConsumerFactory queueingService) {
        this.connectionHolder = connectionHolder;
        this.consumerFactory = queueingService;

        // add simple thread that checks for consumers that are no more used
        ExecutorService executorService = Executors.newFixedThreadPool(1);
        executorService.submit(new RemoveOldConsumerRunnable(this));
    }

    /**
     * Unregisters the consumer and close its channel.
     *
     * @param consumer
     *            the consumer to unregister.
     * @throws IllegalStateException
     *             if another consumer is stored in the registry under the passed in consumer's queue name
     */
    public void unregister(QueueConsumer consumer) {
        log.info("unregistering consumer for " + consumer.getQueueName());
        synchronized (consumerMap) {
            QueueConsumer registered = consumerMap.remove(consumer.getQueueName());
            if (registered != consumer) {
                throw new IllegalStateException("serious kuddelmuddel if this really happens");
            }

            internalPostRemove(consumer);
        }
    }

    // must be called after removing the consumer from consumerMap
    // and with a lock held on the consumerMap monitor (synchronized (this.consumerMap))
    private void internalPostRemove(QueueConsumer consumer) {
        consumer.setUnregistering();

        try {
            // we use a channel only for one consumer, so close the channel.
            if (consumer.getChannel().isOpen()) {
                log.info("closing channel for consumer '{}'", consumer.getQueueName());
                consumer.getChannel().close();
            }
        } catch (IOException | ShutdownSignalException e) {
            log.error("could not close channel - ignoring", e);
        }
    }

    /**
     * Gets the registered consumer with the given queue name; creates a new one if none is registered so far, together with a new channel,
     * starts consuming from the queue name, and puts it into this registry.
     * 
     * 
     * The consumer will consume the queue in exclusive mode.
     *
     * @param queueName
     *            name of the queue the consumer listens on
     * @return the registered consumer
     * @throws IOException
     */
    public QueueConsumer consumerPresto(String queueName) throws IOException {
        QueueConsumer consumer;
        synchronized (consumerMap) {
            consumer = consumerMap.get(queueName);
            if (consumer == null) {
                log.info("for queue name '{}' creating a channel to AMQP server and a consumer using it", queueName);
                Channel channel = connectionHolder.createNewChannel();
                consumer = consumerFactory.createConsumer(queueName, channel, this);
                consumerMap.put(queueName, consumer);

                // declare (create if missing) the queue as durable,
                // which is not removed when the connection closes (exclusive flag)
                // nor when there is temporarily no consumer (autodelete flag)
                // but after a expiry timeout without consumers (x-expires)
                Map args = Collections. singletonMap("x-expires", QUEUE_X_EXPIRY_MILLIS);
                channel.queueDeclare(queueName, true, false, false, args);

                // make consumer exclusive to detect some programming fault when an old consumer is not correctly
                // removed/shut down/cancelled/waddayacallit
                // use empty string as consumerTag to make the server generate one
                try {
                    String consumerTag = consumer.getChannel().basicConsume(queueName, false, "", false, true, null, consumer);
                    consumer.setRegisteredConsumerTag(consumerTag);
                } catch (Exception e) {
                    log.warn("Internal error subscribing new consumer to queue '{}'", queueName);
                    throw e;
                }
            }
        }
        return consumer;
    }

    public void removeOldConsumers(long inactivityMillis) {
        long now = System.currentTimeMillis();

        synchronized (consumerMap) {
            log.debug("consumer remover: removing old consumers from {} total consumers", consumerMap.size());

            int count = 0;
            Iterator iter = consumerMap.values().iterator();
            while (iter.hasNext()) {
                QueueConsumer consumer = iter.next();
                if (now - consumer.getLastUsed() > inactivityMillis) {
                    log.info("removing old consumer with queue name " + consumer.getQueueName());
                    iter.remove();
                    internalPostRemove(consumer);
                    ++count;
                }
            }
            log.debug("{} old consumers removed", count);
        }
    }
}