Generify ObservationCollector

Type parameter is the context, e.g. Observation for Micrometer.

References #379

Author: acogoluegnes

src/main/java/com/rabbitmq/stream/EnvironmentBuilder.java

@@ -222,7 +222,7 @@ public interface EnvironmentBuilder {
    */
   EnvironmentBuilder metricsCollector(MetricsCollector metricsCollector);
 
-  EnvironmentBuilder observationCollector(ObservationCollector observationCollector);
+  EnvironmentBuilder observationCollector(ObservationCollector<?> observationCollector);
 
   /**
    * The maximum number of producers allocated to a single connection.

src/main/java/com/rabbitmq/stream/ObservationCollector.java

@@ -13,27 +13,27 @@
 // [email protected].
 package com.rabbitmq.stream;
 
-public interface ObservationCollector {
+public interface ObservationCollector<T> {
 
-  ObservationCollector NO_OP =
-      new ObservationCollector() {
+  ObservationCollector<Void> NO_OP =
+      new ObservationCollector<Void>() {
         @Override
-        public Object prePublish(String stream, Message message) {
+        public Void prePublish(String stream, Message message) {
           return null;
         }
 
         @Override
-        public void published(Object context, Message message) {}
+        public void published(Void context, Message message) {}
 
         @Override
         public MessageHandler subscribe(MessageHandler handler) {
           return handler;
         }
       };
 
-  void published(Object context, Message message);
+  void published(T context, Message message);
 
-  Object prePublish(String stream, Message message);
+  T prePublish(String stream, Message message);
 
   MessageHandler subscribe(MessageHandler handler);
 }

src/main/java/com/rabbitmq/stream/impl/SimpleMessageAccumulator.java

@@ -32,8 +32,9 @@ class SimpleMessageAccumulator implements MessageAccumulator {
   private final ToLongFunction<Message> publishSequenceFunction;
   private final Function<Message, String> filterValueExtractor;
   final String stream;
-  final ObservationCollector observationCollector;
+  final ObservationCollector<Object> observationCollector;
 
+  @SuppressWarnings("unchecked")
   SimpleMessageAccumulator(
       int capacity,
       Codec codec,
@@ -42,7 +43,7 @@ class SimpleMessageAccumulator implements MessageAccumulator {
       Function<Message, String> filterValueExtractor,
       Clock clock,
       String stream,
-      ObservationCollector observationCollector) {
+      ObservationCollector<?> observationCollector) {
     this.capacity = capacity;
     this.messages = new LinkedBlockingQueue<>(capacity);
     this.codec = codec;
@@ -52,7 +53,7 @@ class SimpleMessageAccumulator implements MessageAccumulator {
         filterValueExtractor == null ? NULL_FILTER_VALUE_EXTRACTOR : filterValueExtractor;
     this.clock = clock;
     this.stream = stream;
-    this.observationCollector = observationCollector;
+    this.observationCollector = (ObservationCollector<Object>) observationCollector;
   }
 
   public boolean add(Message message, ConfirmationHandler confirmationHandler) {

src/main/java/com/rabbitmq/stream/impl/StreamEnvironment.java

@@ -85,7 +85,7 @@ class StreamEnvironment implements Environment {
   private final Runnable locatorInitializationSequence;
   private final List<Locator> locators = new CopyOnWriteArrayList<>();
   private final ExecutorServiceFactory executorServiceFactory;
-  private final ObservationCollector observationCollector;
+  private final ObservationCollector<?> observationCollector;
 
   StreamEnvironment(
       ScheduledExecutorService scheduledExecutorService,
@@ -102,7 +102,7 @@ class StreamEnvironment implements Environment {
       boolean lazyInit,
       Function<ClientConnectionType, String> connectionNamingStrategy,
       Function<Client.ClientParameters, Client> clientFactory,
-      ObservationCollector observationCollector) {
+      ObservationCollector<?> observationCollector) {
     this.recoveryBackOffDelayPolicy = recoveryBackOffDelayPolicy;
     this.topologyUpdateBackOffDelayPolicy = topologyBackOffDelayPolicy;
     this.byteBufAllocator = byteBufAllocator;
@@ -646,7 +646,7 @@ CompressionCodecFactory compressionCodecFactory() {
     return this.clientParametersPrototype.compressionCodecFactory;
   }
 
-  ObservationCollector observationCollector() {
+  ObservationCollector<?> observationCollector() {
     return this.observationCollector;
   }
 

src/main/java/com/rabbitmq/stream/impl/StreamEnvironmentBuilder.java

@@ -63,7 +63,7 @@ public class StreamEnvironmentBuilder implements EnvironmentBuilder {
   private CompressionCodecFactory compressionCodecFactory;
   private boolean lazyInit = false;
   private Function<Client.ClientParameters, Client> clientFactory = Client::new;
-  private ObservationCollector observationCollector = ObservationCollector.NO_OP;
+  private ObservationCollector<?> observationCollector = ObservationCollector.NO_OP;
 
   public StreamEnvironmentBuilder() {}
 
@@ -282,7 +282,7 @@ StreamEnvironmentBuilder clientFactory(Function<Client.ClientParameters, Client>
   }
 
   @Override
-  public EnvironmentBuilder observationCollector(ObservationCollector observationCollector) {
+  public EnvironmentBuilder observationCollector(ObservationCollector<?> observationCollector) {
     this.observationCollector = observationCollector;
     return this;
   }

src/main/java/com/rabbitmq/stream/observation/micrometer/MicrometerObservationCollector.java

@@ -17,7 +17,7 @@
 import io.micrometer.observation.Observation;
 import io.micrometer.observation.ObservationRegistry;
 
-class MicrometerObservationCollector implements ObservationCollector {
+class MicrometerObservationCollector implements ObservationCollector<Observation> {
 
   private final ObservationRegistry registry;
   private final PublishObservationConvention customPublishConvention, defaultPublishConvention;
@@ -37,19 +37,16 @@ public MicrometerObservationCollector(
   }
 
   @Override
-  public void published(Object context, Message message) {
-    if (context instanceof Observation) {
-      Observation observation = (Observation) context;
-      try {
-        observation.stop();
-      } catch (Exception e) {
-        // TODO log error
-      }
+  public void published(Observation observation, Message message) {
+    try {
+      observation.stop();
+    } catch (Exception e) {
+      // TODO log error
     }
   }
 
   @Override
-  public Object prePublish(String stream, Message message) {
+  public Observation prePublish(String stream, Message message) {
     PublishContext context = new PublishContext(stream, message);
     Observation observation =
         StreamObservationDocumentation.PUBLISH_OBSERVATION.observation(

src/test/java/com/rabbitmq/stream/impl/StreamProducerUnitTest.java

@@ -113,7 +113,7 @@ void init() {
     when(env.locatorOperation(any())).thenCallRealMethod();
     when(env.clock()).thenReturn(clock);
     when(env.codec()).thenReturn(new SimpleCodec());
-    when(env.observationCollector()).thenReturn(ObservationCollector.NO_OP);
+    when(env.observationCollector()).thenAnswer(invocation -> ObservationCollector.NO_OP);
     doAnswer(
             (Answer<Runnable>)
                 invocationOnMock -> {