Suppress capturing responses

TrafficCapture/captureKafkaOffloader/build.gradle

@@ -9,20 +9,24 @@ repositories {
 }
 
 dependencies {
-    api 'io.netty:netty-buffer:4.1.100.Final'
+    implementation platform("io.opentelemetry:opentelemetry-bom:1.32.0")
+
     implementation project(':captureOffloader')
     implementation project(':coreUtilities')
-    implementation 'org.projectlombok:lombok:1.18.26'
-    implementation 'com.google.protobuf:protobuf-java:3.22.2'
-    implementation 'org.apache.kafka:kafka-clients:3.6.0'
-    implementation 'software.amazon.msk:aws-msk-iam-auth:1.1.9'
-    implementation 'org.slf4j:slf4j-api:2.0.7'
+    implementation group: 'com.google.protobuf', name:'protobuf-java', version:'3.22.2'
+    api            group:'io.netty', name:'netty-buffer', version: '4.1.100.Final'
+    implementation group: 'io.opentelemetry', name:'opentelemetry-api'
+    implementation group: 'org.projectlombok', name:'lombok', version:'1.18.26'
+    implementation group: 'org.apache.kafka', name:'kafka-clients', version:'3.6.0'
+    implementation group: 'org.slf4j', name:'slf4j-api', version:'2.0.7'
+    implementation group: 'software.amazon.msk', name:'aws-msk-iam-auth', version:'1.1.9'
+
     testImplementation project(':captureProtobufs')
-    testImplementation 'org.mockito:mockito-core:4.6.1'
-    testImplementation 'org.mockito:mockito-junit-jupiter:4.6.1'
     testImplementation group: 'org.apache.logging.log4j', name: 'log4j-api', version: '2.20.0'
     testImplementation group: 'org.apache.logging.log4j', name: 'log4j-core', version: '2.20.0'
     testImplementation group: 'org.apache.logging.log4j', name: 'log4j-slf4j2-impl', version: '2.20.0'
+    testImplementation group: 'org.mockito', name: 'mockito-core', version: '4.6.1'
+    testImplementation group: 'org.mockito', name: 'mockito-junit-jupiter', version: '4.6.1'
     testImplementation group: 'org.slf4j', name: 'slf4j-api', version: '2.0.7'
 }
 

TrafficCapture/captureKafkaOffloader/src/main/java/org/opensearch/migrations/trafficcapture/kafkaoffloader/KafkaCaptureFactory.java

@@ -10,21 +10,31 @@
 import org.opensearch.migrations.coreutils.MetricsAttributeKey;
 import org.opensearch.migrations.coreutils.MetricsEvent;
 import org.apache.kafka.clients.producer.RecordMetadata;
+import org.opensearch.migrations.tracing.SimpleMeteringClosure;
+import org.opensearch.migrations.tracing.commoncontexts.IConnectionContext;
 import org.opensearch.migrations.trafficcapture.CodedOutputStreamHolder;
 import org.opensearch.migrations.trafficcapture.IChannelConnectionCaptureSerializer;
 import org.opensearch.migrations.trafficcapture.IConnectionCaptureFactory;
 import org.opensearch.migrations.trafficcapture.OrderedStreamLifecyleManager;
 import org.opensearch.migrations.trafficcapture.StreamChannelConnectionCaptureSerializer;
 import org.opensearch.migrations.coreutils.MetricsLogger;
+import org.opensearch.migrations.trafficcapture.kafkaoffloader.tracing.KafkaRecordContext;
+import org.opensearch.migrations.trafficcapture.tracing.ConnectionContext;
 
+import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.time.Duration;
+import java.time.Instant;
 import java.util.Arrays;
 import java.util.concurrent.CompletableFuture;
 
 
 @Slf4j
 public class KafkaCaptureFactory implements IConnectionCaptureFactory<RecordMetadata> {
 
+    private static final SimpleMeteringClosure METERING_CLOSURE = new SimpleMeteringClosure("KafkaCapture");
+
+
     private static final MetricsLogger metricsLogger = new MetricsLogger("BacksideHandler");
 
     private static final String DEFAULT_TOPIC_NAME_FOR_TRAFFIC = "logging-traffic-topic";
@@ -51,8 +61,10 @@
     }
 
     @Override
-    public IChannelConnectionCaptureSerializer<RecordMetadata> createOffloader(String connectionId) {
-        return new StreamChannelConnectionCaptureSerializer<>(nodeId, connectionId, new StreamManager(connectionId));
+    public IChannelConnectionCaptureSerializer<RecordMetadata> createOffloader(IConnectionContext ctx,
+                                                                               String connectionId) {
+        return new StreamChannelConnectionCaptureSerializer<>(nodeId, connectionId,
+                new StreamManager(ctx, connectionId));
     }
 
     @AllArgsConstructor
@@ -65,12 +77,33 @@
         }
     }
 
-    @AllArgsConstructor
     class StreamManager extends OrderedStreamLifecyleManager<RecordMetadata> {
+        IConnectionContext telemetryContext;
         String connectionId;
+        Instant startTime;
+
+        public StreamManager(IConnectionContext ctx, String connectionId) {
+            this.telemetryContext = ctx;
+            METERING_CLOSURE.meterIncrementEvent(telemetryContext, "offloader_created");
+            METERING_CLOSURE.meterDeltaEvent(telemetryContext, "offloaders_active", 1);
+
+            this.connectionId = connectionId;
+            this.startTime = Instant.now();
+        }
+
+        @Override
+        public void close() throws IOException {
+            log.atInfo().setMessage(() -> "factory.close()").log();
+            METERING_CLOSURE.meterHistogramMillis(telemetryContext, "offloader_stream_lifetime",
+                    Duration.between(startTime, Instant.now()));
+            METERING_CLOSURE.meterDeltaEvent(telemetryContext, "offloaders_active", -1);
+            METERING_CLOSURE.meterIncrementEvent(telemetryContext, "offloader_closed");
+        }
 
         @Override
         public CodedOutputStreamWrapper createStream() {
+            METERING_CLOSURE.meterIncrementEvent(telemetryContext, "stream_created");
+
             ByteBuffer bb = ByteBuffer.allocate(bufferSize);
             return new CodedOutputStreamWrapper(CodedOutputStream.newInstance(bb), bb);
         }
@@ -93,8 +126,14 @@
                 // Used to essentially wrap Future returned by Producer to CompletableFuture
                 var cf = new CompletableFuture<RecordMetadata>();
                 log.debug("Sending Kafka producer record: {} for topic: {}", recordId, topicNameForTraffic);
+
+                var flushContext = new KafkaRecordContext(telemetryContext,
+                        METERING_CLOSURE.makeSpanContinuation("flushRecord"),
+                        topicNameForTraffic, recordId, kafkaRecord.value().length);
+                METERING_CLOSURE.meterIncrementEvent(telemetryContext, "stream_flush_called");
+
                 // Async request to Kafka cluster
-                producer.send(kafkaRecord, handleProducerRecordSent(cf, recordId));
+                producer.send(kafkaRecord, handleProducerRecordSent(cf, recordId, flushContext));
                 metricsLogger.atSuccess(MetricsEvent.RECORD_SENT_TO_KAFKA)
                         .setAttribute(MetricsAttributeKey.CHANNEL_ID, connectionId)
                         .setAttribute(MetricsAttributeKey.TOPIC_NAME, topicNameForTraffic)
@@ -108,29 +147,42 @@
                 throw e;
             }
         }
-
-        /**
-         * The default KafkaProducer comes with built-in retry and error-handling logic that suits many cases. From the
-         * documentation here for retry: https://kafka.apache.org/35/javadoc/org/apache/kafka/clients/producer/KafkaProducer.html
-         * "If the request fails, the producer can automatically retry. The retries setting defaults to Integer.MAX_VALUE,
-         * and it's recommended to use delivery.timeout.ms to control retry behavior, instead of retries."
-         * <p>
-         * Apart from this the KafkaProducer has logic for deciding whether an error is transient and should be
-         * retried or not retried at all: https://kafka.apache.org/35/javadoc/org/apache/kafka/common/errors/RetriableException.html
-         * as well as basic retry backoff
-         */
-        private Callback handleProducerRecordSent(CompletableFuture<RecordMetadata> cf, String recordId) {
-            return (metadata, exception) -> {
-                if (exception != null) {
-                    log.error("Error sending producer record: {}", recordId, exception);
-                    cf.completeExceptionally(exception);
-                } else {
-                    log.debug("Kafka producer record: {} has finished sending for topic: {} and partition {}",
-                            recordId, metadata.topic(), metadata.partition());
-                    cf.complete(metadata);
-                }
-            };
-        }
     }
 
+    /**
+     * The default KafkaProducer comes with built-in retry and error-handling logic that suits many cases. From the
+     * documentation here for retry: https://kafka.apache.org/35/javadoc/org/apache/kafka/clients/producer/KafkaProducer.html
+     * "If the request fails, the producer can automatically retry. The retries setting defaults to Integer.MAX_VALUE,
+     * and it's recommended to use delivery.timeout.ms to control retry behavior, instead of retries."
+     * <p>
+     * Apart from this the KafkaProducer has logic for deciding whether an error is transient and should be
+     * retried or not retried at all: https://kafka.apache.org/35/javadoc/org/apache/kafka/common/errors/RetriableException.html
+     * as well as basic retry backoff
+     */
+    private Callback handleProducerRecordSent(CompletableFuture<RecordMetadata> cf, String recordId,
+                                              KafkaRecordContext flushContext) {
+        // Keep this out of the inner class because it is more unsafe to include it within
+        // the inner class since the inner class has context that shouldn't be used.  This keeps
+        // that field out of scope.
+        return (metadata, exception) -> {
+            log.atInfo().setMessage(()->"kafka completed sending a record").log();
+            METERING_CLOSURE.meterHistogramMicros(flushContext,
+                    exception==null ? "stream_flush_success_ms" : "stream_flush_failure_ms");
+            METERING_CLOSURE.meterIncrementEvent(flushContext,
+                    exception==null ? "stream_flush_success" : "stream_flush_failure");
+            METERING_CLOSURE.meterIncrementEvent(flushContext,
+                    exception==null ? "stream_flush_success_bytes" : "stream_flush_failure_bytes",
+                    flushContext.getRecordSize());
+            flushContext.currentSpan.end();
+
+            if (exception != null) {
+                log.error("Error sending producer record: {}", recordId, exception);
+                cf.completeExceptionally(exception);
+            } else {
+                log.debug("Kafka producer record: {} has finished sending for topic: {} and partition {}",
+                        recordId, metadata.topic(), metadata.partition());
+                cf.complete(metadata);
+            }
+        };
+    }
 }

TrafficCapture/captureKafkaOffloader/src/main/java/org/opensearch/migrations/trafficcapture/kafkaoffloader/tracing/KafkaRecordContext.java

@@ -0,0 +1,50 @@
+package org.opensearch.migrations.trafficcapture.kafkaoffloader.tracing;
+
+import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.api.common.AttributesBuilder;
+import io.opentelemetry.api.trace.Span;
+import lombok.AllArgsConstructor;
+import lombok.Getter;
+import org.opensearch.migrations.tracing.ISpanWithParentGenerator;
+import org.opensearch.migrations.tracing.commoncontexts.IConnectionContext;
+import org.opensearch.migrations.tracing.IWithAttributes;
+import org.opensearch.migrations.tracing.IWithStartTime;
+
+import java.time.Instant;
+
+@AllArgsConstructor
+public class KafkaRecordContext implements IWithAttributes<IConnectionContext>, IWithStartTime {
+    static final AttributeKey<String> TOPIC_ATTR = AttributeKey.stringKey("topic");
+    static final AttributeKey<String> RECORD_ID_ATTR = AttributeKey.stringKey("recordId");
+    static final AttributeKey<Long> RECORD_SIZE_ATTR = AttributeKey.longKey("recordSize");
+
+    @Getter
+    public final IConnectionContext enclosingScope;
+    @Getter
+    public final Span currentSpan;
+    @Getter
+    public final Instant startTime;
+    @Getter
+    public final String topic;
+    @Getter
+    public final String recordId;
+    @Getter
+    public final int recordSize;
+
+    public KafkaRecordContext(IConnectionContext enclosingScope, ISpanWithParentGenerator incomingSpan,
+                              String topic, String recordId, int recordSize) {
+        this.enclosingScope = enclosingScope;
+        this.topic = topic;
+        this.recordId = recordId;
+        this.recordSize = recordSize;
+        this.startTime = Instant.now();
+        currentSpan = incomingSpan.apply(this.getPopulatedAttributes(), enclosingScope.getCurrentSpan());
+    }
+
+    @Override
+    public AttributesBuilder fillAttributes(AttributesBuilder builder) {
+        return builder.put(TOPIC_ATTR, getTopic())
+                .put(RECORD_ID_ATTR, getRecordId())
+                .put(RECORD_SIZE_ATTR, getRecordSize());
+    }
+}

TrafficCapture/captureKafkaOffloader/src/test/java/org/opensearch/migrations/trafficcapture/kafkaoffloader/KafkaCaptureFactoryTest.java

@@ -1,6 +1,7 @@
 package org.opensearch.migrations.trafficcapture.kafkaoffloader;
 
 import io.netty.buffer.Unpooled;
+import io.opentelemetry.api.GlobalOpenTelemetry;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.kafka.clients.ApiVersions;
 import org.apache.kafka.clients.producer.Callback;
@@ -18,7 +19,9 @@
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
+import org.opensearch.migrations.tracing.EmptyContext;
 import org.opensearch.migrations.trafficcapture.IChannelConnectionCaptureSerializer;
+import org.opensearch.migrations.trafficcapture.tracing.ConnectionContext;
 
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
@@ -51,7 +54,7 @@ public void testLargeRequestIsWithinKafkaMessageSizeLimit() throws IOException,
         MockProducer<String, byte[]> producer = new MockProducer<>(true, new StringSerializer(), new ByteArraySerializer());
         KafkaCaptureFactory kafkaCaptureFactory =
             new KafkaCaptureFactory(TEST_NODE_ID_STRING, producer, maxAllowableMessageSize);
-        IChannelConnectionCaptureSerializer serializer = kafkaCaptureFactory.createOffloader(connectionId);
+        IChannelConnectionCaptureSerializer serializer = kafkaCaptureFactory.createOffloader(createCtx(), connectionId);
 
         StringBuilder sb = new StringBuilder();
         for (int i = 0; i < 15000; i++) {
@@ -73,6 +76,11 @@ public void testLargeRequestIsWithinKafkaMessageSizeLimit() throws IOException,
         producer.close();
     }
 
+    private static ConnectionContext createCtx() {
+        return new ConnectionContext("test", "test",
+                x->GlobalOpenTelemetry.getTracer("test").spanBuilder("test").startSpan());
+    }
+
     /**
      * This size calculation is based off the KafkaProducer client request size validation check done when Producer
      * records are sent. This validation appears to be consistent for several versions now, here is a reference to
@@ -97,7 +105,7 @@ private int calculateRecordSize(ProducerRecord<String, byte[]> record, String re
     public void testLinearOffloadingIsSuccessful() throws IOException {
         KafkaCaptureFactory kafkaCaptureFactory =
                 new KafkaCaptureFactory(TEST_NODE_ID_STRING, mockProducer, 1024*1024);
-        IChannelConnectionCaptureSerializer offloader = kafkaCaptureFactory.createOffloader(connectionId);
+        IChannelConnectionCaptureSerializer offloader = kafkaCaptureFactory.createOffloader(createCtx(), connectionId);
 
         List<Callback> recordSentCallbacks = new ArrayList<>(3);
         when(mockProducer.send(any(), any())).thenAnswer(invocation -> {

TrafficCapture/captureOffloader/build.gradle

@@ -20,19 +20,22 @@ sourceSets {
     }
 }
 dependencies {
+    implementation platform("io.opentelemetry:opentelemetry-bom:1.32.0")
     api group: 'io.netty', name: 'netty-buffer', version: '4.1.100.Final'
 
     implementation project(':captureProtobufs')
-    implementation "com.google.protobuf:protobuf-java:3.22.2"
-    implementation 'org.projectlombok:lombok:1.18.26'
+    implementation project(':coreUtilities')
+    implementation group: 'io.opentelemetry', name:'opentelemetry-api'
+    implementation group: 'com.google.protobuf', name: 'protobuf-java', version: '3.22.2'
+    implementation group: 'org.projectlombok', name: 'lombok', version: '1.18.26'
     implementation group: 'org.slf4j', name: 'slf4j-api', version: '2.0.7'
 
     testImplementation group: 'org.apache.logging.log4j', name: 'log4j-api', version: '2.20.0'
     testImplementation group: 'org.apache.logging.log4j', name: 'log4j-core', version: '2.20.0'
     testImplementation group: 'org.apache.logging.log4j', name: 'log4j-slf4j2-impl', version: '2.20.0'
     testImplementation group: 'org.slf4j', name: 'slf4j-api', version: '2.0.7'
 
-    testFixturesImplementation "com.google.protobuf:protobuf-java:3.22.2"
     testFixturesImplementation project(':captureProtobufs')
-
+    testFixturesImplementation project(':coreUtilities')
+    testFixturesImplementation group: 'com.google.protobuf', name: 'protobuf-java', version: '3.22.2'
 }

TrafficCapture/captureOffloader/src/main/java/org/opensearch/migrations/trafficcapture/FileConnectionCaptureFactory.java

@@ -3,6 +3,8 @@
 import lombok.AllArgsConstructor;
 import lombok.Lombok;
 import lombok.extern.slf4j.Slf4j;
+import org.opensearch.migrations.tracing.commoncontexts.IConnectionContext;
+import org.opensearch.migrations.trafficcapture.tracing.ConnectionContext;
 
 import java.io.FileNotFoundException;
 import java.io.FileOutputStream;
@@ -51,6 +53,9 @@
     @AllArgsConstructor
     class StreamManager extends OrderedStreamLifecyleManager<Void> {
         String connectionId;
+        @Override
+        public void close() {}
+
         @Override
         public CodedOutputStreamAndByteBufferWrapper createStream() {
             return new CodedOutputStreamAndByteBufferWrapper(bufferSize);
@@ -80,7 +85,7 @@
     }
 
     @Override
-    public IChannelConnectionCaptureSerializer createOffloader(String connectionId) {
-        return new StreamChannelConnectionCaptureSerializer<Void>(nodeId, connectionId, new StreamManager(connectionId));
+    public IChannelConnectionCaptureSerializer createOffloader(IConnectionContext ctx, String connectionId) {
+        return new StreamChannelConnectionCaptureSerializer<>(nodeId, connectionId, new StreamManager(connectionId));
     }
 }

TrafficCapture/captureOffloader/src/main/java/org/opensearch/migrations/trafficcapture/IConnectionCaptureFactory.java

@@ -1,7 +1,10 @@
 package org.opensearch.migrations.trafficcapture;
 
+import org.opensearch.migrations.tracing.commoncontexts.IConnectionContext;
+import org.opensearch.migrations.trafficcapture.tracing.ConnectionContext;
+
 import java.io.IOException;
 
 public interface IConnectionCaptureFactory<T> {
-    IChannelConnectionCaptureSerializer<T> createOffloader(String connectionId) throws IOException;
+    IChannelConnectionCaptureSerializer<T> createOffloader(IConnectionContext ctx, String connectionId) throws IOException;
 }