Refactor traffic replayer into three files

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

@@ -85,7 +85,7 @@
 
         @Override
         public CodedOutputStreamWrapper createStream() {
-            telemetryContext.getCurrentSpan().addEvent("streamCreated");
+            telemetryContext.addEvent("streamCreated");
 
             ByteBuffer bb = ByteBuffer.allocate(bufferSize);
             return new CodedOutputStreamWrapper(CodedOutputStream.newInstance(bb), bb);
@@ -123,7 +123,7 @@
             return sendFullyAsync(producer, kafkaRecord)
                 .whenComplete(((recordMetadata, throwable) -> {
                     if (throwable != null) {
-                        flushContext.addException(throwable, true);
+                        flushContext.addTraceException(throwable, true);
                         log.error("Error sending producer record: {}", recordId, throwable);
                     } else {
                         log.debug("Kafka producer record: {} has finished sending for topic: {} and partition {}",

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

@@ -30,7 +30,7 @@ public KafkaRecordContext(IRootKafkaOffloaderContext rootScope, IConnectionConte
         this.topic = topic;
         this.recordId = recordId;
         initializeSpan();
-        getCurrentSpan().setAttribute(RECORD_SIZE_ATTR, recordSize);
+        this.setTraceAttribute(RECORD_SIZE_ATTR, recordSize);
     }
 
     public static class MetricInstruments extends CommonScopedMetricInstruments {

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

@@ -29,7 +29,7 @@ public TestRootKafkaOffloaderContext() {
     }
 
     public TestRootKafkaOffloaderContext(InMemoryInstrumentationBundle inMemoryInstrumentationBundle) {
-        super("tests", inMemoryInstrumentationBundle == null ? null :
+        super("tests", DO_NOTHING_TRACKER, inMemoryInstrumentationBundle == null ? null :
                 inMemoryInstrumentationBundle.openTelemetrySdk);
         this.inMemoryInstrumentationBundle = inMemoryInstrumentationBundle;
         final var meter = getMeterProvider().get("test");

TrafficCapture/captureOffloader/src/main/java/org/opensearch/migrations/trafficcapture/tracing/ConnectionContext.java

@@ -27,7 +27,7 @@ public ConnectionContext(IRootOffloaderContext rootInstrumentationScope, String
         super(rootInstrumentationScope);
         this.connectionId = connectionId;
         this.nodeId = nodeId;
-        initializeSpan();
+        initializeSpan(rootInstrumentationScope);
         meterDeltaEvent(getMetrics().activeConnectionsCounter, 1);
     }
 

TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/ActiveContextTracker.java

@@ -0,0 +1,34 @@
+package org.opensearch.migrations.tracing;
+
+import java.util.Comparator;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentSkipListSet;
+import java.util.stream.Stream;
+
+public class ActiveContextTracker implements IContextTracker {
+    final ConcurrentSkipListSet<IScopedInstrumentationAttributes> orderedScopes;
+
+    public ActiveContextTracker() {
+        orderedScopes = makeScopeSkipList();
+    }
+
+    static ConcurrentSkipListSet<IScopedInstrumentationAttributes> makeScopeSkipList() {
+        return new ConcurrentSkipListSet<>(Comparator
+                .comparingLong(IWithStartTimeAndAttributes::getStartNanoTime)
+                .thenComparingInt(System::identityHashCode));
+    }
+
+    @Override
+    public void onContextCreated(IScopedInstrumentationAttributes scopedContext) {
+        orderedScopes.add(scopedContext);
+    }
+
+    @Override
+    public void onContextClosed(IScopedInstrumentationAttributes scopedContext) {
+        orderedScopes.remove(scopedContext);
+    }
+
+    public Stream<IScopedInstrumentationAttributes> getOldestActiveScopes() {
+        return orderedScopes.stream();
+    }
+}

TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/ActiveContextTrackerByActivityType.java

@@ -0,0 +1,48 @@
+package org.opensearch.migrations.tracing;
+
+import java.util.Collection;
+import java.util.Comparator;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentSkipListSet;
+import java.util.stream.Stream;
+
+public class ActiveContextTrackerByActivityType implements IContextTracker {
+    final ConcurrentHashMap<Class<IScopedInstrumentationAttributes>,
+            ConcurrentSkipListSet<IScopedInstrumentationAttributes>> orderedScopesByScopeType;
+
+    public ActiveContextTrackerByActivityType() {
+        orderedScopesByScopeType = new ConcurrentHashMap<>();
+    }
+
+    @Override
+    @SuppressWarnings("unchecked")
+    public void onContextCreated(IScopedInstrumentationAttributes scopedContext) {
+        orderedScopesByScopeType
+                .computeIfAbsent((Class<IScopedInstrumentationAttributes>)scopedContext.getClass(),
+                        c->ActiveContextTracker.makeScopeSkipList())
+                .add(scopedContext);
+    }
+
+    @Override
+    public void onContextClosed(IScopedInstrumentationAttributes scopedContext) {
+        final var skipListByType = orderedScopesByScopeType.get(scopedContext.getClass());
+        assert skipListByType != null : "expected to have already added the scope to the collection, " +
+                "so the top-level class mapping should be present";
+        skipListByType.remove(scopedContext);
+    }
+
+    public Stream<IScopedInstrumentationAttributes>
+    getOldestActiveScopes(Class<IScopedInstrumentationAttributes> activityType) {
+        return Optional.ofNullable(orderedScopesByScopeType.getOrDefault(activityType, null))
+                .stream()
+                .flatMap(Collection::stream);
+    }
+
+    public Stream<Class<IScopedInstrumentationAttributes>> getActiveScopeTypes() {
+        return orderedScopesByScopeType.entrySet().stream()
+                .filter(kvp->!kvp.getValue().isEmpty())
+                .map(Map.Entry::getKey);
+    }
+}

TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/BaseNestedSpanContext.java

@@ -1,15 +1,5 @@
 package org.opensearch.migrations.tracing;
 
-import io.opentelemetry.api.common.AttributeKey;
-import io.opentelemetry.api.common.Attributes;
-import io.opentelemetry.api.common.AttributesBuilder;
-import io.opentelemetry.api.trace.Span;
-import lombok.Getter;
-import lombok.NonNull;
-
-import java.util.Optional;
-import java.util.stream.Stream;
-
 public abstract class BaseNestedSpanContext
         <S extends IInstrumentConstructor, T extends IScopedInstrumentationAttributes>
         extends BaseSpanContext<S>
@@ -21,6 +11,10 @@ protected BaseNestedSpanContext(S rootScope, T enclosingScope) {
         this.enclosingScope = enclosingScope;
     }
 
+    protected void initializeSpan() {
+        initializeSpan(rootInstrumentationScope);
+    }
+
     @Override
     public IScopedInstrumentationAttributes getEnclosingScope() {
         return enclosingScope;

TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/BaseSpanContext.java

@@ -10,7 +10,7 @@
 import java.util.stream.Stream;
 
 public abstract class BaseSpanContext<S extends IInstrumentConstructor>
-        implements IScopedInstrumentationAttributes, IWithStartTimeAndAttributes, IHasRootInstrumentationScope<S>, AutoCloseable {
+        implements IScopedInstrumentationAttributes, IHasRootInstrumentationScope<S>, AutoCloseable {
     @Getter
     protected final S rootInstrumentationScope;
     @Getter
@@ -20,10 +20,9 @@ public abstract class BaseSpanContext<S extends IInstrumentConstructor>
     @Getter
     private Span currentSpan;
 
-    public BaseSpanContext(S rootScope) {
+    protected BaseSpanContext(S rootScope) {
         this.startNanoTime = System.nanoTime();
         this.rootInstrumentationScope = rootScope;
-        rootScope.onContextCreated(this);
     }
 
     protected static <T> AttributesBuilder addAttributeIfPresent(AttributesBuilder attributesBuilder,
@@ -32,27 +31,28 @@ protected static <T> AttributesBuilder addAttributeIfPresent(AttributesBuilder a
     }
 
     @Override
-    public void endSpan() {
-        IScopedInstrumentationAttributes.super.endSpan();
-        rootInstrumentationScope.onContextClosed(this);
+    public @NonNull IContextTracker getContextTracker() {
+        return rootInstrumentationScope;
     }
 
-    protected void initializeSpan() {
-        initializeSpanWithLinkedSpans(null);
+    protected void initializeSpan(@NonNull IInstrumentConstructor constructor) {
+        initializeSpanWithLinkedSpans(constructor, null);
     }
 
-    protected void initializeSpanWithLinkedSpans(Stream<Span> linkedSpans) {
-        initializeSpan(rootInstrumentationScope.buildSpan(this, getActivityName(), linkedSpans));
+    protected void initializeSpanWithLinkedSpans(@NonNull IInstrumentConstructor constructor,
+                                                 Stream<Span> linkedSpans) {
+        initializeSpan(constructor, rootInstrumentationScope.buildSpan(this, getActivityName(), linkedSpans));
     }
 
-    public void initializeSpan(@NonNull Span s) {
+    public void initializeSpan(@NonNull IInstrumentConstructor constructor, @NonNull Span s) {
         assert currentSpan == null : "only expect to set the current span once";
         currentSpan = s;
+        constructor.onContextCreated(this);
     }
 
     @Override
-    public void addException(Throwable e, boolean isPropagating) {
-        IScopedInstrumentationAttributes.super.addException(e, isPropagating);
+    public void addTraceException(Throwable e, boolean isPropagating) {
+        IScopedInstrumentationAttributes.super.addTraceException(e, isPropagating);
         observedExceptionToIncludeInMetrics = e;
     }
 

TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/CompositeContextTracker.java

@@ -0,0 +1,32 @@
+package org.opensearch.migrations.tracing;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+public class CompositeContextTracker implements IContextTracker {
+    private final List<IContextTracker> trackers;
+
+    public CompositeContextTracker(IContextTracker...trackers) {
+        this.trackers = Arrays.stream(trackers).collect(Collectors.toUnmodifiableList());
+    }
+    public CompositeContextTracker(List<IContextTracker> trackers) {
+        this.trackers = new ArrayList<>(trackers);
+    }
+
+    @Override
+    public void onContextCreated(IScopedInstrumentationAttributes scopedContext) {
+        trackers.forEach(ct->ct.onContextCreated(scopedContext));
+    }
+
+    @Override
+    public void onContextClosed(IScopedInstrumentationAttributes scopedContext) {
+        trackers.forEach(ct->ct.onContextClosed(scopedContext));
+    }
+
+    public Stream<IContextTracker> getTrackers() {
+        return trackers.stream();
+    }
+}

TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/IContextTracker.java

@@ -0,0 +1,16 @@
+package org.opensearch.migrations.tracing;
+
+/**
+ * For debugging or observability purposes, this interface allows for tracking the
+ * creation and termination of activities (such as those with spans).
+ */
+public interface IContextTracker {
+    default void onContextCreated(IScopedInstrumentationAttributes newScopedContext) {}
+
+    /**
+     * This can be overridden to track creation and termination of spans
+     */
+    default void onContextClosed(IScopedInstrumentationAttributes newScopedContext) {}
+
+    final static IContextTracker DO_NOTHING_TRACKER = new IContextTracker() {};
+}

TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/IInstrumentConstructor.java

@@ -6,17 +6,6 @@
 
 import java.util.stream.Stream;
 
-public interface IInstrumentConstructor {
+public interface IInstrumentConstructor extends IContextTracker {
     @NonNull Span buildSpan(IScopedInstrumentationAttributes forScope, String spanName, Stream<Span> linkedSpans);
-
-    /**
-     * For debugging, this will be overridden to track creation and termination of spans
-     */
-    default void onContextCreated(IScopedInstrumentationAttributes newScopedContext) {}
-
-    /**
-     * For debugging, this will be overridden to track creation and termination of spans
-     */
-    default void onContextClosed(IScopedInstrumentationAttributes newScopedContext) {}
-
 }

TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/IInstrumentationAttributes.java

@@ -7,12 +7,9 @@
 import io.opentelemetry.api.metrics.LongCounter;
 import io.opentelemetry.api.metrics.LongHistogram;
 import io.opentelemetry.api.metrics.LongUpDownCounter;
-import io.opentelemetry.api.trace.Span;
 import lombok.NonNull;
-import org.opensearch.migrations.Utils;
 
 import java.time.Duration;
-import java.util.ArrayDeque;
 
 public interface IInstrumentationAttributes {
     AttributeKey<Boolean> HAD_EXCEPTION_KEY = AttributeKey.booleanKey("hadException");
@@ -29,10 +26,10 @@ public interface IInstrumentationAttributes {
     }
 
     default void addCaughtException(Throwable e) {
-        addException(e, false);
+        addTraceException(e, false);
     }
 
-    default void addException(Throwable e, boolean exceptionIsPropagating) {
+    default void addTraceException(Throwable e, boolean exceptionIsPropagating) {
         meterIncrementEvent(getMetrics().exceptionCounter);
     }
 

TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/IScopedInstrumentationAttributes.java

@@ -1,5 +1,6 @@
 package org.opensearch.migrations.tracing;
 
+import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.common.AttributesBuilder;
 import io.opentelemetry.api.metrics.DoubleHistogram;
@@ -13,8 +14,7 @@
 
 import java.util.ArrayDeque;
 
-public interface IScopedInstrumentationAttributes
-        extends IWithStartTimeAndAttributes, AutoCloseable {
+public interface IScopedInstrumentationAttributes extends IWithStartTimeAndAttributes, AutoCloseable {
 
     String getActivityName();
 
@@ -26,6 +26,8 @@ public interface IScopedInstrumentationAttributes
 
     @NonNull Span getCurrentSpan();
 
+    @NonNull IContextTracker getContextTracker();
+
     default Attributes getPopulatedSpanAttributes() {
         return getPopulatedSpanAttributesBuilder().build();
     }
@@ -35,9 +37,7 @@ default AttributesBuilder getPopulatedSpanAttributesBuilder() {
         // reverse the order so that the lowest attribute scopes will overwrite the upper ones if there were conflicts
         var stack = new ArrayDeque<IScopedInstrumentationAttributes>();
         while (currentObj != null) {
-            if (currentObj instanceof IScopedInstrumentationAttributes) {
-                stack.addFirst((IScopedInstrumentationAttributes) currentObj);
-            }
+            stack.addFirst((IScopedInstrumentationAttributes) currentObj);
             currentObj = currentObj.getEnclosingScope();
         }
         var builder = stack.stream()
@@ -61,10 +61,11 @@ default DoubleHistogram getEndOfScopeDurationMetric() {
         return getMetrics().contextDuration;
     }
 
-    default void endSpan() {
+    default void endSpan(IContextTracker contextTracker) {
         var span = getCurrentSpan();
         span.setAllAttributes(getPopulatedSpanAttributes());
         span.end();
+        contextTracker.onContextClosed(this);
     }
 
     default void sendMeterEventsForEnd() {
@@ -73,13 +74,13 @@ default void sendMeterEventsForEnd() {
     }
 
     default void close() {
-        endSpan();
+        endSpan(getContextTracker());
         sendMeterEventsForEnd();
     }
 
     @Override
-    default void addException(Throwable e, boolean isPropagating) {
-        IWithStartTimeAndAttributes.super.addException(e, isPropagating);
+    default void addTraceException(Throwable e, boolean isPropagating) {
+        IWithStartTimeAndAttributes.super.addTraceException(e, isPropagating);
         final var span = getCurrentSpan();
         if (isPropagating) {
             span.recordException(e, Attributes.of(SemanticAttributes.EXCEPTION_ESCAPED, true));
@@ -115,4 +116,20 @@ default void meterHistogram(LongHistogram histogram, long value, AttributesBuild
             IWithStartTimeAndAttributes.super.meterHistogram(histogram, value, attributesBuilder);
         }
     }
+
+    default void addEvent(String eventName) {
+        getCurrentSpan().addEvent(eventName);
+    }
+
+    default void setTraceAttribute(AttributeKey<Long> attributeKey, long attributeValue) {
+        getCurrentSpan().setAttribute(attributeKey, attributeValue);
+    }
+
+    default void setAttribute(AttributeKey<String> attributeKey, String attributeValue) {
+        getCurrentSpan().setAttribute(attributeKey, attributeValue);
+    }
+
+    default void setAllAttributes(Attributes allAttributes) {
+        getCurrentSpan().setAllAttributes(allAttributes);
+    }
 }