Add TrafficListener interface

[#161109942]

Fixes #411

(cherry picked from commit 0223b61)

Author: acogoluegnes

src/main/java/com/rabbitmq/client/ConnectionFactory.java

@@ -200,6 +200,15 @@ public class ConnectionFactory implements Cloneable {
      */
     private RetryHandler topologyRecoveryRetryHandler;
 
+    /**
+     * Traffic listener notified of inbound and outbound {@link Command}s.
+     * <p>
+     * Useful for debugging purposes. Default is no-op.
+     *
+     * @since 5.5.0
+     */
+    private TrafficListener trafficListener = TrafficListener.NO_OP;
+
     /** @return the default host to use for connections */
     public String getHost() {
         return host;
@@ -1149,6 +1158,7 @@ public ConnectionParams params(ExecutorService consumerWorkServiceExecutor) {
         result.setTopologyRecoveryFilter(topologyRecoveryFilter);
         result.setConnectionRecoveryTriggeringCondition(connectionRecoveryTriggeringCondition);
         result.setTopologyRecoveryRetryHandler(topologyRecoveryRetryHandler);
+        result.setTrafficListener(trafficListener);
         return result;
     }
 
@@ -1501,6 +1511,7 @@ public void setErrorOnWriteListener(ErrorOnWriteListener errorOnWriteListener) {
 
     /**
      * Set filter to include/exclude entities from topology recovery.
+     *
      * @since 4.8.0
      */
     public void setTopologyRecoveryFilter(TopologyRecoveryFilter topologyRecoveryFilter) {
@@ -1510,6 +1521,7 @@ public void setTopologyRecoveryFilter(TopologyRecoveryFilter topologyRecoveryFil
     /**
      * Allows to decide on automatic connection recovery is triggered.
      * Default is for shutdown not initiated by application or missed heartbeat errors.
+     *
      * @param connectionRecoveryTriggeringCondition
      */
     public void setConnectionRecoveryTriggeringCondition(Predicate<ShutdownSignalException> connectionRecoveryTriggeringCondition) {
@@ -1519,10 +1531,26 @@ public void setConnectionRecoveryTriggeringCondition(Predicate<ShutdownSignalExc
     /**
      * Set retry handler for topology recovery.
      * Default is no retry.
+     *
      * @param topologyRecoveryRetryHandler
      * @since 5.4.0
      */
     public void setTopologyRecoveryRetryHandler(RetryHandler topologyRecoveryRetryHandler) {
         this.topologyRecoveryRetryHandler = topologyRecoveryRetryHandler;
     }
+
+    /**
+     * Traffic listener notified of inbound and outbound {@link Command}s.
+     * <p>
+     * Useful for debugging purposes, e.g. logging all sent and received messages.
+     * Default is no-op.
+     *
+     * @param trafficListener
+     * @see TrafficListener
+     * @see com.rabbitmq.client.impl.LogTrafficListener
+     * @since 5.5.0
+     */
+    public void setTrafficListener(TrafficListener trafficListener) {
+        this.trafficListener = trafficListener;
+    }
 }

src/main/java/com/rabbitmq/client/TrafficListener.java

@@ -0,0 +1,40 @@
+package com.rabbitmq.client;
+
+/**
+ * Contract to log outbound and inbound {@link Command}s.
+ *
+ * @see ConnectionFactory#setTrafficListener(TrafficListener)
+ * @since 5.5.0
+ */
+public interface TrafficListener {
+
+    /**
+     * No-op {@link TrafficListener}.
+     */
+    TrafficListener NO_OP = new TrafficListener() {
+
+        @Override
+        public void write(Command outboundCommand) {
+
+        }
+
+        @Override
+        public void read(Command inboundCommand) {
+
+        }
+    };
+
+    /**
+     * Notified for each outbound {@link Command}.
+     *
+     * @param outboundCommand
+     */
+    void write(Command outboundCommand);
+
+    /**
+     * Notified for each inbound {@link Command}.
+     *
+     * @param inboundCommand
+     */
+    void read(Command inboundCommand);
+}

src/main/java/com/rabbitmq/client/impl/AMQChannel.java

@@ -75,6 +75,8 @@ public abstract class AMQChannel extends ShutdownNotifierComponent {
 
     private final boolean _checkRpcResponseType;
 
+    private final TrafficListener _trafficListener;
+
     /**
      * Construct a channel on the given connection, with the given channel number.
      * @param connection the underlying connection for this channel
@@ -88,6 +90,7 @@ public AMQChannel(AMQConnection connection, int channelNumber) {
         }
         this._rpcTimeout = connection.getChannelRpcTimeout();
         this._checkRpcResponseType = connection.willCheckRpcResponseType();
+        this._trafficListener = connection.getTrafficListener();
     }
 
     /**
@@ -175,6 +178,7 @@ public void handleCompleteInboundCommand(AMQCommand command) throws IOException
         // asynchronous commands (deliveries/returns/other events),
         // and false for commands that should be passed on to some
         // waiting RPC continuation.
+        this._trafficListener.read(command);
         if (!processAsync(command)) {
             // The filter decided not to handle/consume the command,
             // so it must be a response to an earlier RPC.
@@ -444,6 +448,7 @@ public void quiescingTransmit(AMQCommand c) throws IOException {
                     ensureIsOpen();
                 }
             }
+            this._trafficListener.write(c);
             c.transmit(this);
         }
     }

src/main/java/com/rabbitmq/client/impl/AMQConnection.java

@@ -141,6 +141,7 @@ public static Map<String, Object> defaultClientProperties() {
     protected final MetricsCollector metricsCollector;
     private final int channelRpcTimeout;
     private final boolean channelShouldCheckRpcResponseType;
+    private final TrafficListener trafficListener;
 
     /* State modified after start - all volatile */
 
@@ -237,6 +238,7 @@ public AMQConnection(ConnectionParams params, FrameHandler frameHandler, Metrics
         this.channelRpcTimeout = params.getChannelRpcTimeout();
         this.channelShouldCheckRpcResponseType = params.channelShouldCheckRpcResponseType();
 
+        this.trafficListener = params.getTrafficListener() == null ? TrafficListener.NO_OP : params.getTrafficListener();
         this._channel0 = new AMQChannel(this, 0) {
             @Override public boolean processAsync(Command c) throws IOException {
                 return getConnection().processControlCommand(c);
@@ -1125,4 +1127,8 @@ public int getChannelRpcTimeout() {
     public boolean willCheckRpcResponseType() {
         return channelShouldCheckRpcResponseType;
     }
+
+    TrafficListener getTrafficListener() {
+        return trafficListener;
+    }
 }

src/main/java/com/rabbitmq/client/impl/ConnectionParams.java

@@ -20,6 +20,7 @@
 import com.rabbitmq.client.RecoveryDelayHandler.DefaultRecoveryDelayHandler;
 import com.rabbitmq.client.SaslConfig;
 import com.rabbitmq.client.ShutdownSignalException;
+import com.rabbitmq.client.TrafficListener;
 import com.rabbitmq.client.impl.recovery.RetryHandler;
 import com.rabbitmq.client.impl.recovery.TopologyRecoveryFilter;
 
@@ -57,6 +58,8 @@ public class ConnectionParams {
     private ExceptionHandler exceptionHandler;
     private ThreadFactory threadFactory;
 
+    private TrafficListener trafficListener;
+
     public ConnectionParams() {}
 
     public CredentialsProvider getCredentialsProvider() {
@@ -266,4 +269,12 @@ public void setTopologyRecoveryRetryHandler(RetryHandler topologyRecoveryRetryHa
     public RetryHandler getTopologyRecoveryRetryHandler() {
         return topologyRecoveryRetryHandler;
     }
+
+    public void setTrafficListener(TrafficListener trafficListener) {
+        this.trafficListener = trafficListener;
+    }
+
+    public TrafficListener getTrafficListener() {
+        return trafficListener;
+    }
 }

src/main/java/com/rabbitmq/client/impl/LogTrafficListener.java

@@ -0,0 +1,41 @@
+package com.rabbitmq.client.impl;
+
+import com.rabbitmq.client.Command;
+import com.rabbitmq.client.TrafficListener;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * {@link TrafficListener} that logs {@link Command} at <code>TRACE</code> level.
+ * <p>
+ * This implementation checks whether the <code>TRACE</code> log level
+ * is enabled before logging anything. This {@link TrafficListener}
+ * should only be activated for debugging purposes, not in a production
+ * environment.
+ *
+ * @see TrafficListener
+ * @see com.rabbitmq.client.ConnectionFactory#setTrafficListener(TrafficListener)
+ * @since 5.5.0
+ */
+public class LogTrafficListener implements TrafficListener {
+
+    private static final Logger LOGGER = LoggerFactory.getLogger(LogTrafficListener.class);
+
+    @Override
+    public void write(Command outboundCommand) {
+        if (shouldLog(outboundCommand)) {
+            LOGGER.trace("Outbound command: {}", outboundCommand);
+        }
+    }
+
+    @Override
+    public void read(Command inboundCommand) {
+        if (shouldLog(inboundCommand)) {
+            LOGGER.trace("Inbound command: {}", inboundCommand);
+        }
+    }
+
+    protected boolean shouldLog(Command command) {
+        return LOGGER.isTraceEnabled();
+    }
+}

src/test/java/com/rabbitmq/client/test/TrafficListenerTest.java

@@ -0,0 +1,102 @@
+// Copyright (c) 2018 Pivotal Software, Inc.  All rights reserved.
+//
+// This software, the RabbitMQ Java client library, is triple-licensed under the
+// Mozilla Public License 1.1 ("MPL"), the GNU General Public License version 2
+// ("GPL") and the Apache License version 2 ("ASL"). For the MPL, please see
+// LICENSE-MPL-RabbitMQ. For the GPL, please see LICENSE-GPL2.  For the ASL,
+// please see LICENSE-APACHE2.
+//
+// This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY KIND,
+// either express or implied. See the LICENSE file for specific language governing
+// rights and limitations of this software.
+//
+// If you have any questions regarding licensing, please contact us at
+// [email protected].
+
+package com.rabbitmq.client.test;
+
+import com.rabbitmq.client.AMQP;
+import com.rabbitmq.client.Channel;
+import com.rabbitmq.client.Command;
+import com.rabbitmq.client.Connection;
+import com.rabbitmq.client.ConnectionFactory;
+import com.rabbitmq.client.TrafficListener;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+
+import java.util.List;
+import java.util.UUID;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+/**
+ *
+ */
+@RunWith(Parameterized.class)
+public class TrafficListenerTest {
+
+    @Parameterized.Parameter
+    public Consumer<ConnectionFactory> configurator;
+
+    @Parameterized.Parameters
+    public static Object[] data() {
+        return new Object[] { automaticRecoveryEnabled(), automaticRecoveryDisabled() };
+    }
+
+    static Consumer<ConnectionFactory> automaticRecoveryEnabled() {
+        return cf -> cf.setAutomaticRecoveryEnabled(true);
+    }
+
+    static Consumer<ConnectionFactory> automaticRecoveryDisabled() {
+        return cf -> cf.setAutomaticRecoveryEnabled(false);
+    }
+
+    @Test
+    public void trafficListenerIsCalled() throws Exception {
+        ConnectionFactory cf = TestUtils.connectionFactory();
+        TestTrafficListener testTrafficListener = new TestTrafficListener();
+        cf.setTrafficListener(testTrafficListener);
+        configurator.accept(cf);
+        try (Connection c = cf.newConnection()) {
+            Channel ch = c.createChannel();
+            String queue = ch.queueDeclare().getQueue();
+            CountDownLatch latch = new CountDownLatch(1);
+            ch.basicConsume(queue, true,
+                (consumerTag, message) -> latch.countDown(), consumerTag -> {
+                });
+            String messageContent = UUID.randomUUID().toString();
+            ch.basicPublish("", queue, null, messageContent.getBytes());
+            assertTrue(latch.await(5, TimeUnit.SECONDS));
+            assertEquals(1, testTrafficListener.outboundContent.size());
+            assertEquals(messageContent, testTrafficListener.outboundContent.get(0));
+            assertEquals(1, testTrafficListener.inboundContent.size());
+            assertEquals(messageContent, testTrafficListener.inboundContent.get(0));
+        }
+    }
+
+    private static class TestTrafficListener implements TrafficListener {
+
+        final List<String> outboundContent = new CopyOnWriteArrayList<>();
+        final List<String> inboundContent = new CopyOnWriteArrayList<>();
+
+        @Override
+        public void write(Command outboundCommand) {
+            if (outboundCommand.getMethod() instanceof AMQP.Basic.Publish) {
+                outboundContent.add(new String(outboundCommand.getContentBody()));
+            }
+        }
+
+        @Override
+        public void read(Command inboundCommand) {
+            if (inboundCommand.getMethod() instanceof AMQP.Basic.Deliver) {
+                inboundContent.add(new String(inboundCommand.getContentBody()));
+            }
+        }
+    }
+}