Add the metrics collector abstraction

Author: acogoluegnes

pom.xml

@@ -130,6 +130,13 @@
             <scope>test</scope>
         </dependency>
 
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-junit-jupiter</artifactId>
+            <version>${mockito.version}</version>
+            <scope>test</scope>
+        </dependency>
+
         <dependency>
             <groupId>io.dropwizard.metrics</groupId>
             <artifactId>metrics-core</artifactId>
@@ -209,7 +216,7 @@
                     <excludes>
                         <exclude>**/*TestSuite.java</exclude>
                     </excludes>
-                    <argLine>${test-arguments}</argLine>
+                    <argLine>-Xshare:off ${test-arguments}</argLine>
                     <systemPropertyVariables>
                         <net.bytebuddy.experimental>true</net.bytebuddy.experimental>
                         <rabbitmqctl.bin>DOCKER:rabbitmq</rabbitmqctl.bin>

src/main/java/com/rabbitmq/model/amqp/AmqpConnection.java

@@ -23,6 +23,7 @@
 import static java.util.Collections.singletonMap;
 
 import com.rabbitmq.model.*;
+import com.rabbitmq.model.metrics.MetricsCollector;
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.List;
@@ -111,19 +112,16 @@ class AmqpConnection extends ResourceBase implements Connection {
     } else {
       disconnectHandler =
           (c, e) -> {
-            if (this.closed.compareAndSet(false, true)) {
-              ModelException failureCause = convert(e.failureCause(), "Connection disconnected");
-              this.state(CLOSING, failureCause);
-              this.releaseManagementResources();
-              this.state(CLOSED, failureCause);
-            }
+            ModelException failureCause = convert(e.failureCause(), "Connection disconnected");
+            this.close(failureCause);
           };
       this.recoveryRequestQueue = null;
       this.recoveryLoop = null;
     }
     this.nativeConnection = connect(this.connectionSettings, builder.name(), disconnectHandler);
     this.management = createManagement();
     this.state(OPEN);
+    this.environment.metricsCollector().openConnection();
   }
 
   @Override
@@ -166,38 +164,7 @@ public RpcServerBuilder rpcServerBuilder() {
 
   @Override
   public void close() {
-    if (this.closed.compareAndSet(false, true)) {
-      this.state(CLOSING);
-      if (this.recoveryLoop != null) {
-        this.recoveryLoop.cancel(true);
-      }
-      if (this.topologyListener instanceof AutoCloseable) {
-        try {
-          ((AutoCloseable) this.topologyListener).close();
-        } catch (Exception e) {
-          LOGGER.info("Error while closing topology listener", e);
-        }
-      }
-      this.closeManagement();
-      for (RpcClient rpcClient : this.rpcClients) {
-        rpcClient.close();
-      }
-      for (RpcServer rpcServer : this.rpcServers) {
-        rpcServer.close();
-      }
-      for (AmqpPublisher publisher : this.publishers) {
-        publisher.close();
-      }
-      for (AmqpConsumer consumer : this.consumers) {
-        consumer.close();
-      }
-      try {
-        this.nativeConnection.close();
-      } catch (Exception e) {
-        LOGGER.warn("Error while closing native connection", e);
-      }
-      this.state(CLOSED);
-    }
+    this.close(null);
   }
 
   // internal API
@@ -508,6 +475,10 @@ Clock clock() {
     return this.environment.clock();
   }
 
+  MetricsCollector metricsCollector() {
+    return this.environment.metricsCollector();
+  }
+
   Publisher createPublisher(AmqpPublisherBuilder builder) {
     // TODO copy the builder properties to create the publisher
     AmqpPublisher publisher = new AmqpPublisher(builder);
@@ -573,6 +544,42 @@ private String currentConnectionLabel() {
     }
   }
 
+  private void close(Throwable cause) {
+    if (this.closed.compareAndSet(false, true)) {
+      this.state(CLOSING, cause);
+      if (this.recoveryLoop != null) {
+        this.recoveryLoop.cancel(true);
+      }
+      if (this.topologyListener instanceof AutoCloseable) {
+        try {
+          ((AutoCloseable) this.topologyListener).close();
+        } catch (Exception e) {
+          LOGGER.info("Error while closing topology listener", e);
+        }
+      }
+      this.closeManagement();
+      for (RpcClient rpcClient : this.rpcClients) {
+        rpcClient.close();
+      }
+      for (RpcServer rpcServer : this.rpcServers) {
+        rpcServer.close();
+      }
+      for (AmqpPublisher publisher : this.publishers) {
+        publisher.close();
+      }
+      for (AmqpConsumer consumer : this.consumers) {
+        consumer.close();
+      }
+      try {
+        this.nativeConnection.close();
+      } catch (Exception e) {
+        LOGGER.warn("Error while closing native connection", e);
+      }
+      this.state(CLOSED, cause);
+      this.environment.metricsCollector().closeConnection();
+    }
+  }
+
   @Override
   public String toString() {
     return this.environment.toString() + "-" + this.id;

src/main/java/com/rabbitmq/model/amqp/AmqpConsumer.java

@@ -21,6 +21,7 @@
 
 import com.rabbitmq.model.Consumer;
 import com.rabbitmq.model.ModelException;
+import com.rabbitmq.model.metrics.MetricsCollector;
 import java.lang.reflect.Field;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
@@ -59,6 +60,7 @@ class AmqpConsumer extends ResourceBase implements Consumer {
   private final AmqpConnection connection;
   private final AtomicBoolean paused = new AtomicBoolean(false);
   private final AtomicReference<CountDownLatch> echoedFlowAfterPauseLatch = new AtomicReference<>();
+  private final MetricsCollector metricsCollector;
   // native receiver internal state, accessed only in the native executor/scheduler
   private ProtonReceiver protonReceiver;
   private Scheduler protonExecutor;
@@ -76,7 +78,9 @@ class AmqpConsumer extends ResourceBase implements Consumer {
     this.initStateFromNativeReceiver(this.nativeReceiver);
     this.connection = builder.connection();
     this.startReceivingLoop();
+    this.metricsCollector = this.connection.metricsCollector();
     this.state(OPEN);
+    this.metricsCollector.openConsumer();
   }
 
   @Override
@@ -107,6 +111,7 @@ private Runnable createReceiveTask(Receiver receiver, MessageHandler messageHand
         while (!Thread.currentThread().isInterrupted()) {
           Delivery delivery = receiver.receive(100, TimeUnit.MILLISECONDS);
           if (delivery != null) {
+            this.metricsCollector.consume();
             AmqpMessage message = new AmqpMessage(delivery.message());
             // TODO make disposition idempotent
             Consumer.Context context =
@@ -117,6 +122,8 @@ public void accept() {
                     try {
                       protonExecutor.execute(() -> replenishCreditIfNeeded());
                       delivery.disposition(DeliveryState.accepted(), true);
+                      metricsCollector.consumeDisposition(
+                          MetricsCollector.ConsumeDisposition.ACCEPTED);
                     } catch (ClientIllegalStateException | ClientIOException e) {
                       LOGGER.debug("message accept failed: {}", e.getMessage());
                     } catch (ClientException e) {
@@ -129,6 +136,8 @@ public void discard() {
                     try {
                       protonExecutor.execute(() -> replenishCreditIfNeeded());
                       delivery.disposition(DeliveryState.rejected("", ""), true);
+                      metricsCollector.consumeDisposition(
+                          MetricsCollector.ConsumeDisposition.DISCARDED);
                     } catch (ClientIllegalStateException | ClientIOException e) {
                       LOGGER.debug("message discard failed: {}", e.getMessage());
                     } catch (ClientException e) {
@@ -141,6 +150,8 @@ public void requeue() {
                     try {
                       protonExecutor.execute(() -> replenishCreditIfNeeded());
                       delivery.disposition(DeliveryState.released(), true);
+                      metricsCollector.consumeDisposition(
+                          MetricsCollector.ConsumeDisposition.REQUEUED);
                     } catch (ClientIllegalStateException | ClientIOException e) {
                       LOGGER.debug("message requeue failed: {}", e.getMessage());
                     } catch (ClientException e) {
@@ -196,6 +207,7 @@ private void close(Throwable cause) {
         LOGGER.warn("Error while closing receiver", e);
       }
       this.state(CLOSED, cause);
+      this.metricsCollector.closeConsumer();
     }
   }
 

src/main/java/com/rabbitmq/model/amqp/AmqpEnvironment.java

@@ -18,6 +18,8 @@
 package com.rabbitmq.model.amqp;
 
 import com.rabbitmq.model.*;
+import com.rabbitmq.model.metrics.MetricsCollector;
+import com.rabbitmq.model.metrics.NoOpMetricsCollector;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
@@ -46,9 +48,12 @@ class AmqpEnvironment implements Environment {
   private final ScheduledExecutorService scheduledExecutorService;
   private volatile ScheduledFuture<?> clockRefreshFuture;
   private final AtomicBoolean clockRefreshSet = new AtomicBoolean(false);
+  private final MetricsCollector metricsCollector;
 
   AmqpEnvironment(
-      ExecutorService executorService, DefaultConnectionSettings<?> connectionSettings) {
+      ExecutorService executorService,
+      DefaultConnectionSettings<?> connectionSettings,
+      MetricsCollector metricsCollector) {
     this.id = ID_SEQUENCE.getAndIncrement();
     connectionSettings.copyTo(this.connectionSettings);
     this.connectionSettings.consolidate();
@@ -64,6 +69,8 @@ class AmqpEnvironment implements Environment {
     }
     this.scheduledExecutorService =
         Executors.newScheduledThreadPool(0, Utils.defaultThreadFactory());
+    this.metricsCollector =
+        metricsCollector == null ? NoOpMetricsCollector.INSTANCE : metricsCollector;
   }
 
   DefaultConnectionSettings<?> connectionSettings() {
@@ -117,6 +124,10 @@ ScheduledExecutorService scheduledExecutorService() {
     return this.scheduledExecutorService;
   }
 
+  MetricsCollector metricsCollector() {
+    return this.metricsCollector;
+  }
+
   void addConnection(AmqpConnection connection) {
     this.connections.add(connection);
   }

src/main/java/com/rabbitmq/model/amqp/AmqpEnvironmentBuilder.java

@@ -20,13 +20,16 @@
 import com.rabbitmq.model.ConnectionSettings;
 import com.rabbitmq.model.Environment;
 import com.rabbitmq.model.EnvironmentBuilder;
+import com.rabbitmq.model.metrics.MetricsCollector;
+import com.rabbitmq.model.metrics.NoOpMetricsCollector;
 import java.util.concurrent.ExecutorService;
 
 public class AmqpEnvironmentBuilder implements EnvironmentBuilder {
 
   private final DefaultEnvironmentConnectionSettings connectionSettings =
       new DefaultEnvironmentConnectionSettings(this);
   private ExecutorService executorService;
+  private MetricsCollector metricsCollector = NoOpMetricsCollector.INSTANCE;
 
   public AmqpEnvironmentBuilder() {}
 
@@ -35,13 +38,18 @@ public AmqpEnvironmentBuilder executorService(ExecutorService executorService) {
     return this;
   }
 
+  public AmqpEnvironmentBuilder metricsCollector(MetricsCollector metricsCollector) {
+    this.metricsCollector = metricsCollector;
+    return this;
+  }
+
   public EnvironmentConnectionSettings connectionSettings() {
     return this.connectionSettings;
   }
 
   @Override
   public Environment build() {
-    return new AmqpEnvironment(executorService, connectionSettings);
+    return new AmqpEnvironment(executorService, connectionSettings, metricsCollector);
   }
 
   public interface EnvironmentConnectionSettings

src/main/java/com/rabbitmq/model/amqp/AmqpManagement.java

@@ -156,9 +156,21 @@ public UnbindSpecification unbind() {
   public void close() {
     if (this.closed.compareAndSet(false, true) && this.initialized.get()) {
       this.releaseResources();
-      this.receiver.close();
-      this.sender.close();
-      this.session.close();
+      try {
+        this.receiver.close();
+      } catch (Exception e) {
+        LOGGER.debug("Error while closing management receiver: {}", e.getMessage());
+      }
+      try {
+        this.sender.close();
+      } catch (Exception e) {
+        LOGGER.debug("Error while closing management sender: {}", e.getMessage());
+      }
+      try {
+        this.session.close();
+      } catch (Exception e) {
+        LOGGER.debug("Error while closing management session: {}", e.getMessage());
+      }
     }
   }
 

src/main/java/com/rabbitmq/model/amqp/AmqpPublisher.java

@@ -21,6 +21,7 @@
 
 import com.rabbitmq.model.Message;
 import com.rabbitmq.model.Publisher;
+import com.rabbitmq.model.metrics.MetricsCollector;
 import java.util.concurrent.*;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
@@ -42,6 +43,7 @@ class AmqpPublisher extends ResourceBase implements Publisher {
   private final String address;
   private final AmqpConnection connection;
   private final AtomicBoolean closed = new AtomicBoolean(false);
+  private final MetricsCollector metricsCollector;
 
   AmqpPublisher(AmqpPublisherBuilder builder) {
     super(builder.listeners());
@@ -50,7 +52,9 @@ class AmqpPublisher extends ResourceBase implements Publisher {
     this.address = builder.address();
     this.connection = builder.connection();
     this.sender = this.createSender(builder.connection().nativeSession(), this.address);
+    this.metricsCollector = this.connection.metricsCollector();
     this.state(OPEN);
+    this.metricsCollector.openPublisher();
   }
 
   @Override
@@ -82,10 +86,14 @@ public void publish(Message message, Callback callback) {
             } catch (InterruptedException | ExecutionException e) {
               status = Status.FAILED;
             }
-            @SuppressWarnings("rawtypes")
             DefaultContext defaultContext = new DefaultContext(message, status);
+            this.metricsCollector.publishDisposition(
+                status == Status.ACCEPTED
+                    ? MetricsCollector.PublishDisposition.ACCEPTED
+                    : MetricsCollector.PublishDisposition.FAILED);
             callback.handle(defaultContext);
           });
+      this.metricsCollector.publish();
     } catch (ClientIllegalStateException e) {
       // the link is closed
       this.close(ExceptionUtils.convert(e));
@@ -129,6 +137,7 @@ private void close(Throwable cause) {
         LOGGER.warn("Error while closing sender", e);
       }
       this.state(State.CLOSED, cause);
+      this.metricsCollector.closePublisher();
     }
   }