zipkin2.collector.activemq.ActiveMQSpanConsumer Maven / Gradle / Ivy

Go to download
/*
 * Copyright The OpenZipkin Authors
 * SPDX-License-Identifier: Apache-2.0
 */
package zipkin2.collector.activemq;

import java.io.Closeable;
import java.io.IOException;
import java.util.LinkedHashMap;
import java.util.Map;
import javax.jms.BytesMessage;
import javax.jms.JMSException;
import javax.jms.Message;
import javax.jms.MessageListener;
import javax.jms.Queue;
import javax.jms.QueueReceiver;
import javax.jms.QueueSession;
import javax.jms.Session;
import javax.jms.TextMessage;
import org.apache.activemq.ActiveMQConnection;
import org.apache.activemq.transport.TransportListener;
import zipkin2.Callback;
import zipkin2.CheckResult;
import zipkin2.collector.Collector;
import zipkin2.collector.CollectorMetrics;

import static java.nio.charset.StandardCharsets.UTF_8;

/**
 * Consumes spans from messages on a ActiveMQ queue. Malformed messages will be discarded. Errors in
 * the storage component will similarly be ignored, with no retry of the message.
 */
final class ActiveMQSpanConsumer implements TransportListener, MessageListener, Closeable {
  static final Callback NOOP = new Callback() {
    @Override public void onSuccess(Void value) {
    }

    @Override public void onError(Throwable t) {
    }
  };

  static final CheckResult
    CLOSED = CheckResult.failed(new IllegalStateException("Collector intentionally closed")),
    INTERRUPTION = CheckResult.failed(new IOException("Recoverable error on ActiveMQ connection"));

  final Collector collector;
  final CollectorMetrics metrics;

  final ActiveMQConnection connection;
  final Map sessionToReceiver = new LinkedHashMap<>();

  volatile CheckResult checkResult = CheckResult.OK;

  ActiveMQSpanConsumer(Collector collector, CollectorMetrics metrics, ActiveMQConnection conn) {
    this.collector = collector;
    this.metrics = metrics;
    this.connection = conn;
    connection.addTransportListener(this);
  }

  /** JMS contract is one session per thread: we need a new session up to our concurrency level. */
  void registerInNewSession(ActiveMQConnection connection, String queue) throws JMSException {
    // Pass redundant info as we can't use default method in activeMQ
    QueueSession session = connection.createQueueSession(false, Session.AUTO_ACKNOWLEDGE);
    // No need to do anything on ActiveMQ side as physical queues are created on demand
    Queue destination = session.createQueue(queue);
    QueueReceiver receiver = session.createReceiver(destination);
    receiver.setMessageListener(this);
    sessionToReceiver.put(session, receiver);
  }

  @Override public void onCommand(Object o) {
  }

  @Override public void onException(IOException error) {
    checkResult = CheckResult.failed(error);
  }

  @Override public void transportInterupted() {
    checkResult = INTERRUPTION;
  }

  @Override public void transportResumed() {
    checkResult = CheckResult.OK;
  }

  @Override public void onMessage(Message message) {
    metrics.incrementMessages();
    byte[] serialized; // TODO: consider how to reuse buffers here
    try {
      if (message instanceof BytesMessage bytesMessage) {
        serialized = new byte[(int) bytesMessage.getBodyLength()];
        bytesMessage.readBytes(serialized);
      } else if (message instanceof TextMessage textMessage) {
        String text = textMessage.getText();
        serialized = text.getBytes(UTF_8);
      } else {
        metrics.incrementMessagesDropped();
        return;
      }
    } catch (Exception e) {
      metrics.incrementMessagesDropped();
      return;
    }

    metrics.incrementBytes(serialized.length);
    if (serialized.length == 0) return; // lenient on empty messages
    collector.acceptSpans(serialized, NOOP);
  }

  @Override public void close() {
    if (checkResult == CLOSED) return;
    checkResult = CLOSED;
    connection.removeTransportListener(this);
    try {
      for (Map.Entry sessionReceiver : sessionToReceiver.entrySet()) {
        sessionReceiver.getValue().setMessageListener(null); // deregister this
        sessionReceiver.getKey().close();
      }
      connection.close();
    } catch (JMSException ignored) {
      // EmptyCatch ignored
    }
  }
}