uk.ac.sanger.cgp.wwdocker.messages.Messaging.java Source code

Java tutorial

Introduction

Here is the source code for uk.ac.sanger.cgp.wwdocker.messages.Messaging.java

Source

/**
 * Copyright (c) 2015 Genome Research Ltd.
 *
 * Author: Cancer Genome Project cgpit@sanger.ac.uk
 *
 * This file is part of WwDocker.
 *
 * WwDocker is free software: you can redistribute it and/or modify it under the
 * terms of the GNU Affero General Public License as published by the Free
 * Software Foundation; either version 3 of the License, or (at your option) any
 * later version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
 * details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program. If not, see <http://www.gnu.org/licenses/>.
 *
 * 1. The usage of a range of years within a copyright statement contained
 * within this distribution should be interpreted as being equivalent to a list
 * of years including the first and last year specified and all consecutive
 * years between them. For example, a copyright statement that reads 'Copyright
 * (c) 2005, 2007- 2009, 2011-2012' should be interpreted as being identical to
 * a statement that reads 'Copyright (c) 2005, 2007, 2008, 2009, 2011, 2012' and
 * a copyright statement that reads "Copyright (c) 2005-2012' should be
 * interpreted as being identical to a statement that reads 'Copyright (c) 2005,
 * 2006, 2007, 2008, 2009, 2010, 2011, 2012'."
 */
package uk.ac.sanger.cgp.wwdocker.messages;

import com.rabbitmq.client.AMQP.BasicProperties;
import com.rabbitmq.client.AMQP.BasicProperties.Builder;
import com.rabbitmq.client.Channel;
import com.rabbitmq.client.Connection;
import com.rabbitmq.client.ConnectionFactory;
import com.rabbitmq.client.MessageProperties;
import com.rabbitmq.client.QueueingConsumer;
import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.TimeoutException;
import org.apache.commons.configuration.BaseConfiguration;
import org.apache.commons.io.FileUtils;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import uk.ac.sanger.cgp.wwdocker.actions.Utils;
import uk.ac.sanger.cgp.wwdocker.beans.WorkerState;

/**
 *
 * @author kr2
 */
public class Messaging {

    private static final Logger logger = LogManager.getLogger();

    BaseConfiguration config;

    public Messaging(BaseConfiguration config) {
        this.config = config;
    }

    public Connection getRmqConn() throws IOException {
        ConnectionFactory factory = new ConnectionFactory();
        factory.setHost(config.getString("rabbit_host"));
        factory.setPort(config.getInt("rabbit_port", 5672));
        factory.setUsername(config.getString("rabbit_user"));
        factory.setPassword(config.getString("rabbit_pw"));
        factory.setNetworkRecoveryInterval(60000); // retry every 60 seconds
        factory.setAutomaticRecoveryEnabled(true);
        return factory.newConnection();
    }

    public void closeRmqConn(Connection conn) {
        try {
            conn.close(-1); // 0.5 sec
        } catch (IOException e) {
            throw new RuntimeException(e.toString(), e);
        }
    }

    public void sendMessage(String queue, Object in) throws IOException, InterruptedException, TimeoutException {
        sendMessage(queue, in, null);
    }

    /**
     * Sends a message to the specified queue
     * @param queue
     * @param in
     * @param host
     * @throws IOException
     * @throws InterruptedException 
     * @throws TimeoutException
     */
    public void sendMessage(String queue, Object in, String host)
            throws IOException, InterruptedException, TimeoutException {
        String message;
        Builder propBuilder = MessageProperties.MINIMAL_PERSISTENT_BASIC.builder();
        if (in.getClass().equals(String.class)) {
            message = (String) in;
        } else {
            message = Utils.objectToJson(in);
            if (in.getClass().equals(WorkerState.class)) {
                host = ((WorkerState) in).getResource().getHostName();
            }
        }
        if (host != null) {
            Map<String, Object> headers = new HashMap();
            headers.put("host", host);
            propBuilder.headers(headers);
        }
        BasicProperties mProp = propBuilder.build();

        Connection connectionSend = getRmqConn();
        Channel channel = connectionSend.createChannel();
        channel.confirmSelect();
        channel.queueDeclare(queue, true, false, false, null);
        channel.basicPublish("", queue, mProp, message.getBytes());
        channel.waitForConfirmsOrDie(5000);
        logger.info(queue + " sent: " + message);
        channel.close();
        closeRmqConn(connectionSend);
    }

    public void sendFile(String queue, String host, File f)
            throws IOException, InterruptedException, TimeoutException {
        Map<String, Object> headers = new HashMap();
        headers.put("host", host);
        BasicProperties mProp = MessageProperties.MINIMAL_PERSISTENT_BASIC.builder().headers(headers).build();
        Connection connectionSend = getRmqConn();
        Channel channel = connectionSend.createChannel();
        channel.confirmSelect();
        channel.queueDeclare(queue, true, false, false, null);
        channel.basicPublish("", queue, mProp, Files.readAllBytes(f.toPath()));
        channel.waitForConfirmsOrDie(5000);
        logger.info(queue + " file: " + f.getAbsolutePath());
        channel.close();
        closeRmqConn(connectionSend);
    }

    public void cleanQueue(String queue, String match) throws IOException, InterruptedException {
        logger.trace(queue.concat(": ").concat(match));
        Connection connectionRcv = getRmqConn();
        Channel channel = connectionRcv.createChannel();
        channel.queueDeclare(queue, true, false, false, null);

        QueueingConsumer consumer = new QueueingConsumer(channel);
        channel.basicConsume(queue, false, consumer);
        QueueingConsumer.Delivery delivery = consumer.nextDelivery(1000);
        Set seen = new HashSet();
        while (delivery != null) {
            String body = new String(delivery.getBody());
            if (seen.contains(body)) {
                break;
            }
            if (body.contains(match)) {
                channel.basicAck(delivery.getEnvelope().getDeliveryTag(), false);
            } else {
                channel.basicNack(delivery.getEnvelope().getDeliveryTag(), false, true);
                seen.add(body);
            }
            delivery = consumer.nextDelivery(1000);
        }
        channel.close();
        closeRmqConn(connectionRcv);
    }

    public List<File> getFiles(String queue, Path outFolder, boolean ack) throws IOException, InterruptedException {
        List files = new ArrayList();
        Connection connectionRcv = getRmqConn();
        Channel channel = connectionRcv.createChannel();
        channel.queueDeclare(queue, true, false, false, null);

        QueueingConsumer consumer = new QueueingConsumer(channel);
        channel.basicConsume(queue, false, consumer);
        QueueingConsumer.Delivery delivery = consumer.nextDelivery(1000);

        Set seen = new HashSet();

        while (delivery != null) {
            String host = delivery.getProperties().getHeaders().get("host").toString();
            File outTo = Paths.get(outFolder.toString(), host + ".tar.gz").toFile();
            FileUtils.writeByteArrayToFile(outTo, delivery.getBody());
            if (ack) {
                channel.basicAck(delivery.getEnvelope().getDeliveryTag(), false);
            } else {
                channel.basicNack(delivery.getEnvelope().getDeliveryTag(), false, true);
            }
            if (seen.contains(delivery.getProperties().getHeaders().get("host"))) {
                break;
            }
            seen.add(delivery.getProperties().getHeaders().get("host"));
            files.add(outTo);
            logger.info(queue + " retrieved: " + outTo.getAbsolutePath());
            delivery = consumer.nextDelivery(1000);
        }
        logger.warn("getFiles done");
        channel.close();
        closeRmqConn(connectionRcv);
        return files;
    }

    public WorkerState getWorkerState(String queue, long wait) throws IOException, InterruptedException {
        WorkerState ws = null;
        Connection connectionRcv = getRmqConn();
        Channel channel = connectionRcv.createChannel();
        channel.queueDeclare(queue, true, false, false, null);

        QueueingConsumer consumer = new QueueingConsumer(channel);
        channel.basicConsume(queue, false, consumer);
        QueueingConsumer.Delivery delivery;
        if (wait == -1) {
            delivery = consumer.nextDelivery(); // will block until response
        } else {
            delivery = consumer.nextDelivery(wait);
        }
        if (delivery != null) {
            String message = new String(delivery.getBody());
            ws = (WorkerState) Utils.jsonToObject(message, WorkerState.class);
            channel.basicAck(delivery.getEnvelope().getDeliveryTag(), false);
        }
        channel.close();
        closeRmqConn(connectionRcv);
        return ws;
    }

    public Object getMessageObject(String queue, Class objClass, long wait)
            throws IOException, InterruptedException {
        String message = getMessageString(queue, wait);
        Object result = null;
        if (message != null) {
            result = Utils.jsonToObject(message, objClass);
        }
        return result;
    }

    /**
     * Gets a single message from a queue, ideal for getting an item of work.
     * @param queue
     * @param wait
     * @return A JSON string representing an object, you need to know what type of object the queue will return and handle this outside of here
     * @throws IOException
     * @throws InterruptedException 
     */
    public String getMessageString(String queue, long wait) throws IOException, InterruptedException {
        String message = null;
        Connection connectionRcv = getRmqConn();
        Channel channel = connectionRcv.createChannel();
        channel.queueDeclare(queue, true, false, false, null);

        QueueingConsumer consumer = new QueueingConsumer(channel);
        channel.basicConsume(queue, false, consumer);
        QueueingConsumer.Delivery delivery;
        if (wait == -1) {
            delivery = consumer.nextDelivery(); // will block until response
        } else {
            delivery = consumer.nextDelivery(wait);
        }
        if (delivery != null) {
            message = new String(delivery.getBody());
            channel.basicAck(delivery.getEnvelope().getDeliveryTag(), false);
            logger.info(queue + " recieved: " + message);
        }
        channel.close();
        closeRmqConn(connectionRcv);
        return message;
    }

    public boolean queryGaveResponse(String queryQueue, String responseQueue, String query, long wait)
            throws IOException, InterruptedException, TimeoutException {
        boolean response = false;
        // clean up queue we send to first
        getMessageStrings(queryQueue, 100);
        this.sendMessage(queryQueue, query);
        if (getMessageString(responseQueue, wait) != null) {
            response = true;
        } else {
            // clean up queue we sent the query
            getMessageStrings(queryQueue, 100);
        }
        return response;
    }

    public void removeFromStateQueue(String queue, String hostToRemove) throws IOException, InterruptedException {
        logger.trace(queue.concat(": ").concat(hostToRemove));
        Connection connectionRcv = getRmqConn();
        Channel channel = connectionRcv.createChannel();
        int maxTries = channel.queueDeclare(queue, true, false, false, null).getMessageCount();
        logger.trace("Queue : messages\t" + queue + " : " + maxTries);

        QueueingConsumer consumer = new QueueingConsumer(channel);
        channel.basicConsume(queue, false, consumer);
        QueueingConsumer.Delivery delivery = consumer.nextDelivery(200);
        logger.trace(maxTries);
        while (delivery != null && maxTries > 0) {
            // the toString in the middle of this is needed as it is wrapped with another type that can hold 4GB
            Map<String, Object> headers = delivery.getProperties().getHeaders();
            if (headers != null && headers.get("host").toString().equals(hostToRemove)) {
                logger.trace(headers.get("host").toString().concat(" remove"));
                channel.basicAck(delivery.getEnvelope().getDeliveryTag(), false);
            } else {
                channel.basicNack(delivery.getEnvelope().getDeliveryTag(), false, true);
            }
            delivery = consumer.nextDelivery(200);
            maxTries--;
        }
        channel.close();
        closeRmqConn(connectionRcv);
    }

    /**
     * Gets all the messages in a queue, best for queues which receive status updates.
     * @param queue
     * @param wait
     * @return List of JSON strings representing objects, you need to know what type of object the queue will return and handle this outside of here
     * @throws IOException
     * @throws InterruptedException 
     */
    public List<String> getMessageStrings(String queue, long wait) throws IOException, InterruptedException {
        List<String> responses = new ArrayList();
        Connection connectionRcv = getRmqConn();
        Channel channel = connectionRcv.createChannel();
        channel.queueDeclare(queue, true, false, false, null);
        QueueingConsumer consumer = new QueueingConsumer(channel);
        channel.basicConsume(queue, true, consumer);
        while (true) {
            QueueingConsumer.Delivery delivery = consumer.nextDelivery(wait);
            if (delivery == null) {
                break;
            }
            String message = new String(delivery.getBody());
            logger.info(queue + " recieved: " + message);
            responses.add(message);
        }
        channel.close();
        closeRmqConn(connectionRcv);
        return responses;
    }

    @Override
    public String toString() {
        StringBuilder result = new StringBuilder();
        String NEW_LINE = System.getProperty("line.separator");
        result.append(this.getClass().getName()).append(" Object {").append(NEW_LINE);
        result.append(" config: ").append(config).append(NEW_LINE);
        result.append("}");
        return result.toString();
    }

}