package io.lumify.model.rabbitmq;
import com.google.inject.Inject;
import com.rabbitmq.client.Channel;
import com.rabbitmq.client.Connection;
import com.rabbitmq.client.QueueingConsumer;
import io.lumify.core.config.Configuration;
import io.lumify.core.exception.LumifyException;
import io.lumify.core.ingest.WorkerSpout;
import io.lumify.core.ingest.WorkerTuple;
import io.lumify.core.ingest.graphProperty.GraphPropertyWorkerTuple;
import io.lumify.core.util.LumifyLogger;
import io.lumify.core.util.LumifyLoggerFactory;
import org.json.JSONObject;
import java.io.IOException;
public class RabbitMQWorkQueueSpout extends WorkerSpout {
private static final LumifyLogger LOGGER = LumifyLoggerFactory.getLogger(RabbitMQWorkQueueSpout.class);
public static final int DEFAULT_RABBITMQ_PREFETCH_COUNT = 10;
private final String queueName;
private Channel channel;
private QueueingConsumer consumer;
private Connection connection;
private Configuration configuration;
public RabbitMQWorkQueueSpout(String queueName) {
this.queueName = queueName;
}
@Override
public void open() {
try {
this.connection = RabbitMQUtils.openConnection(configuration);
this.channel = RabbitMQUtils.openChannel(this.connection);
this.channel.queueDeclare(queueName, true, false, false, null);
this.consumer = new QueueingConsumer(channel);
Integer prefetchCount = configuration.getInt(Configuration.RABBITMQ_PREFETCH_COUNT, DEFAULT_RABBITMQ_PREFETCH_COUNT);
this.channel.basicQos(prefetchCount, false);
this.channel.basicConsume(this.queueName, false, consumer);
} catch (IOException ex) {
throw new LumifyException("Could not startup RabbitMQ", ex);
}
}
@Override
public void close() {
super.close();
try {
LOGGER.debug("Closing RabbitMQ channel");
this.channel.close();
LOGGER.debug("Closing RabbitMQ connection");
this.connection.close();
} catch (IOException ex) {
LOGGER.error("Could not close RabbitMQ connection and channel", ex);
}
}
@Override
public WorkerTuple nextTuple() throws InterruptedException {
QueueingConsumer.Delivery delivery = this.consumer.nextDelivery(100);
if (delivery == null) {
return null;
}
JSONObject json = new JSONObject(new String(delivery.getBody()));
LOGGER.debug("emit (%s): %s", this.queueName, json.toString());
return new GraphPropertyWorkerTuple(delivery.getEnvelope().getDeliveryTag(), json);
}
@Override
public void ack(Object msgId) {
super.ack(msgId);
long deliveryTag = (Long) msgId;
try {
this.channel.basicAck(deliveryTag, false);
} catch (IOException ex) {
LOGGER.error("Could not ack: %d", deliveryTag, ex);
}
}
@Override
public void fail(Object msgId) {
super.fail(msgId);
long deliveryTag = (Long) msgId;
try {
this.channel.basicNack(deliveryTag, false, false);
} catch (IOException ex) {
LOGGER.error("Could not ack: %d", deliveryTag, ex);
}
}
protected QueueingConsumer getConsumer() {
return this.consumer;
}
@Inject
public void setConfiguration(Configuration configuration) {
this.configuration = configuration;
}
}