From d166f2471ccab79d1c92aebd647eee65d72a844b Mon Sep 17 00:00:00 2001 From: Greg Schohn Date: Fri, 5 Apr 2024 11:40:21 -0400 Subject: [PATCH 1/3] Add the notion of context tracking when a context is created to the time that it is closed as a first-class feature of the coreUtilities tracing library. A number of changes were made to existing interfaces to reduce the scope of 'naked' Spans so that it's easier to audit if they're being used directly to close spans. The getCurrentSpan() needs to be public because it's in an interface and used internally, but it should really be private and require friend access only within its package. However, I'm not sure how to do better guardrails in Java. Signed-off-by: Greg Schohn --- .../kafkaoffloader/KafkaCaptureFactory.java | 4 +- .../tracing/KafkaRecordContext.java | 2 +- .../TestRootKafkaOffloaderContext.java | 2 +- .../tracing/ConnectionContext.java | 2 +- .../tracing/ActiveContextTracker.java | 34 +++++++++++++ .../ActiveContextTrackerByActivityType.java | 48 +++++++++++++++++++ .../tracing/BaseNestedSpanContext.java | 14 ++---- .../migrations/tracing/BaseSpanContext.java | 26 +++++----- .../tracing/CompositeContextTracker.java | 32 +++++++++++++ .../migrations/tracing/IContextTracker.java | 16 +++++++ .../tracing/IInstrumentConstructor.java | 13 +---- .../tracing/IInstrumentationAttributes.java | 7 +-- .../IScopedInstrumentationAttributes.java | 35 ++++++++++---- .../tracing/LoggingContextMonitor.java | 18 +++++++ .../migrations/tracing/RootOtelContext.java | 18 ++++--- .../IInstrumentationAttributesTest.java | 2 +- ...er.java => BacktracingContextTracker.java} | 8 +++- .../netty/tracing/RootWireLoggingContext.java | 9 ++-- ...ionallyReliableLoggingHttpHandlerTest.java | 12 +++-- .../proxyserver/CaptureProxy.java | 6 ++- .../proxyserver/RootCaptureContext.java | 9 ++-- .../netty/NettyScanningHttpProxyTest.java | 4 +- .../migrations/replay/TrafficReplayer.java | 16 +++++-- .../NettyPacketToHttpConsumer.java | 2 +- .../datatypes/ConnectionReplaySession.java | 3 +- .../replay/tracing/ReplayContexts.java | 6 +-- .../replay/tracing/RootReplayerContext.java | 5 +- .../replay/tracing/TracingTest.java | 2 +- .../migrations/tracing/TestContext.java | 29 ++++------- 29 files changed, 277 insertions(+), 107 deletions(-) create mode 100644 TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/ActiveContextTracker.java create mode 100644 TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/ActiveContextTrackerByActivityType.java create mode 100644 TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/CompositeContextTracker.java create mode 100644 TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/IContextTracker.java create mode 100644 TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/LoggingContextMonitor.java rename TrafficCapture/coreUtilities/src/testFixtures/java/org/opensearch/migrations/tracing/{ContextTracker.java => BacktracingContextTracker.java} (88%) diff --git a/TrafficCapture/captureKafkaOffloader/src/main/java/org/opensearch/migrations/trafficcapture/kafkaoffloader/KafkaCaptureFactory.java b/TrafficCapture/captureKafkaOffloader/src/main/java/org/opensearch/migrations/trafficcapture/kafkaoffloader/KafkaCaptureFactory.java index f78ac146d..26402b4b7 100644 --- a/TrafficCapture/captureKafkaOffloader/src/main/java/org/opensearch/migrations/trafficcapture/kafkaoffloader/KafkaCaptureFactory.java +++ b/TrafficCapture/captureKafkaOffloader/src/main/java/org/opensearch/migrations/trafficcapture/kafkaoffloader/KafkaCaptureFactory.java @@ -85,7 +85,7 @@ public StreamManager(IRootKafkaOffloaderContext rootScope, IConnectionContext ct @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 @@ public CodedOutputStreamWrapper createStream() { 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 {}", diff --git a/TrafficCapture/captureKafkaOffloader/src/main/java/org/opensearch/migrations/trafficcapture/kafkaoffloader/tracing/KafkaRecordContext.java b/TrafficCapture/captureKafkaOffloader/src/main/java/org/opensearch/migrations/trafficcapture/kafkaoffloader/tracing/KafkaRecordContext.java index 98233036f..af5ec5ed3 100644 --- a/TrafficCapture/captureKafkaOffloader/src/main/java/org/opensearch/migrations/trafficcapture/kafkaoffloader/tracing/KafkaRecordContext.java +++ b/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 { diff --git a/TrafficCapture/captureKafkaOffloader/src/test/java/org/opensearch/migrations/trafficcapture/kafkaoffloader/tracing/TestRootKafkaOffloaderContext.java b/TrafficCapture/captureKafkaOffloader/src/test/java/org/opensearch/migrations/trafficcapture/kafkaoffloader/tracing/TestRootKafkaOffloaderContext.java index 00029d717..d71cb3ca4 100644 --- a/TrafficCapture/captureKafkaOffloader/src/test/java/org/opensearch/migrations/trafficcapture/kafkaoffloader/tracing/TestRootKafkaOffloaderContext.java +++ b/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"); diff --git a/TrafficCapture/captureOffloader/src/main/java/org/opensearch/migrations/trafficcapture/tracing/ConnectionContext.java b/TrafficCapture/captureOffloader/src/main/java/org/opensearch/migrations/trafficcapture/tracing/ConnectionContext.java index d4ce9906c..5b011fce8 100644 --- a/TrafficCapture/captureOffloader/src/main/java/org/opensearch/migrations/trafficcapture/tracing/ConnectionContext.java +++ b/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); } diff --git a/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/ActiveContextTracker.java b/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/ActiveContextTracker.java new file mode 100644 index 000000000..435dac784 --- /dev/null +++ b/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 orderedScopes; + + public ActiveContextTracker() { + orderedScopes = makeScopeSkipList(); + } + + static ConcurrentSkipListSet 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 getOldestActiveScopes() { + return orderedScopes.stream(); + } +} diff --git a/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/ActiveContextTrackerByActivityType.java b/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/ActiveContextTrackerByActivityType.java new file mode 100644 index 000000000..f577ef1d0 --- /dev/null +++ b/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, + ConcurrentSkipListSet> orderedScopesByScopeType; + + public ActiveContextTrackerByActivityType() { + orderedScopesByScopeType = new ConcurrentHashMap<>(); + } + + @Override + @SuppressWarnings("unchecked") + public void onContextCreated(IScopedInstrumentationAttributes scopedContext) { + orderedScopesByScopeType + .computeIfAbsent((Class)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 + getOldestActiveScopes(Class activityType) { + return Optional.ofNullable(orderedScopesByScopeType.getOrDefault(activityType, null)) + .stream() + .flatMap(Collection::stream); + } + + public Stream> getActiveScopeTypes() { + return orderedScopesByScopeType.entrySet().stream() + .filter(kvp->!kvp.getValue().isEmpty()) + .map(Map.Entry::getKey); + } +} diff --git a/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/BaseNestedSpanContext.java b/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/BaseNestedSpanContext.java index a7f35dc36..a7f49e727 100644 --- a/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/BaseNestedSpanContext.java +++ b/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 extends BaseSpanContext @@ -21,6 +11,10 @@ protected BaseNestedSpanContext(S rootScope, T enclosingScope) { this.enclosingScope = enclosingScope; } + protected void initializeSpan() { + initializeSpan(rootInstrumentationScope); + } + @Override public IScopedInstrumentationAttributes getEnclosingScope() { return enclosingScope; diff --git a/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/BaseSpanContext.java b/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/BaseSpanContext.java index 7d4ef4056..33aacb49c 100644 --- a/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/BaseSpanContext.java +++ b/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/BaseSpanContext.java @@ -10,7 +10,7 @@ import java.util.stream.Stream; public abstract class BaseSpanContext - implements IScopedInstrumentationAttributes, IWithStartTimeAndAttributes, IHasRootInstrumentationScope, AutoCloseable { + implements IScopedInstrumentationAttributes, IHasRootInstrumentationScope, AutoCloseable { @Getter protected final S rootInstrumentationScope; @Getter @@ -20,10 +20,9 @@ public abstract class BaseSpanContext @Getter private Span currentSpan; - public BaseSpanContext(S rootScope) { + protected BaseSpanContext(S rootScope) { this.startNanoTime = System.nanoTime(); this.rootInstrumentationScope = rootScope; - rootScope.onContextCreated(this); } protected static AttributesBuilder addAttributeIfPresent(AttributesBuilder attributesBuilder, @@ -32,27 +31,28 @@ protected static 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 linkedSpans) { - initializeSpan(rootInstrumentationScope.buildSpan(this, getActivityName(), linkedSpans)); + protected void initializeSpanWithLinkedSpans(@NonNull IInstrumentConstructor constructor, + Stream 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; } diff --git a/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/CompositeContextTracker.java b/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/CompositeContextTracker.java new file mode 100644 index 000000000..2dc978fa1 --- /dev/null +++ b/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 trackers; + + public CompositeContextTracker(IContextTracker...trackers) { + this.trackers = Arrays.stream(trackers).collect(Collectors.toUnmodifiableList()); + } + public CompositeContextTracker(List 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 getTrackers() { + return trackers.stream(); + } +} diff --git a/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/IContextTracker.java b/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/IContextTracker.java new file mode 100644 index 000000000..193e8f7d3 --- /dev/null +++ b/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() {}; +} diff --git a/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/IInstrumentConstructor.java b/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/IInstrumentConstructor.java index 4be0a3b65..ee6af9e2d 100644 --- a/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/IInstrumentConstructor.java +++ b/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 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) {} - } diff --git a/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/IInstrumentationAttributes.java b/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/IInstrumentationAttributes.java index b74e1057d..4c5654970 100644 --- a/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/IInstrumentationAttributes.java +++ b/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 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); } diff --git a/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/IScopedInstrumentationAttributes.java b/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/IScopedInstrumentationAttributes.java index d81a73b1a..4dc7e5bd7 100644 --- a/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/IScopedInstrumentationAttributes.java +++ b/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(); 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 attributeKey, long attributeValue) { + getCurrentSpan().setAttribute(attributeKey, attributeValue); + } + + default void setAttribute(AttributeKey attributeKey, String attributeValue) { + getCurrentSpan().setAttribute(attributeKey, attributeValue); + } + + default void setAllAttributes(Attributes allAttributes) { + getCurrentSpan().setAllAttributes(allAttributes); + } } diff --git a/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/LoggingContextMonitor.java b/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/LoggingContextMonitor.java new file mode 100644 index 000000000..06aa02b7a --- /dev/null +++ b/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/LoggingContextMonitor.java @@ -0,0 +1,18 @@ +package org.opensearch.migrations.tracing; + +import org.slf4j.Logger; + +import java.time.Duration; +import java.util.concurrent.ExecutorService; + +public class LoggingContextMonitor implements AutoCloseable { + ExecutorService executorService; + public LoggingContextMonitor(Logger logger, Duration period) { + + } + + @Override + public void close() throws Exception { + + } +} diff --git a/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/RootOtelContext.java b/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/RootOtelContext.java index 55137db41..7d7490f51 100644 --- a/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/RootOtelContext.java +++ b/TrafficCapture/coreUtilities/src/main/java/org/opensearch/migrations/tracing/RootOtelContext.java @@ -34,6 +34,8 @@ public class RootOtelContext implements IRootOtelContext { private final String scopeName; @Getter private final MetricInstruments metrics; + @Getter + private final IContextTracker contextTracker; public static OpenTelemetry initializeOpenTelemetryForCollector(@NonNull String collectorEndpoint, @NonNull String serviceName) { @@ -97,18 +99,20 @@ public MetricInstruments(Meter meter, String activityName) { } } - public RootOtelContext(String scopeName) { - this(scopeName, null); + public RootOtelContext(String scopeName, IContextTracker contextTracker) { + this(scopeName, contextTracker, null); } - public RootOtelContext(String scopeName, String collectorEndpoint, String serviceName) { - this(scopeName, initializeOpenTelemetryWithCollectorOrAsNoop(collectorEndpoint, serviceName)); + public RootOtelContext(String scopeName, IContextTracker contextTracker, + String collectorEndpoint, String serviceName) { + this(scopeName, contextTracker, initializeOpenTelemetryWithCollectorOrAsNoop(collectorEndpoint, serviceName)); } - public RootOtelContext(String scopeName, OpenTelemetry sdk) { + public RootOtelContext(String scopeName, IContextTracker contextTracker, OpenTelemetry sdk) { openTelemetryImpl = sdk != null ? sdk : initializeOpenTelemetryWithCollectorOrAsNoop(null, null); this.scopeName = scopeName; - metrics = new MetricInstruments(this.getMeterProvider().get(scopeName), "root"); + this.metrics = new MetricInstruments(this.getMeterProvider().get(scopeName), "root"); + this.contextTracker = contextTracker; } @Override @@ -121,7 +125,7 @@ public RootOtelContext getEnclosingScope() { return null; } - OpenTelemetry getOpenTelemetry() { + private OpenTelemetry getOpenTelemetry() { return openTelemetryImpl; } diff --git a/TrafficCapture/coreUtilities/src/test/java/org/opensearch/migrations/tracing/IInstrumentationAttributesTest.java b/TrafficCapture/coreUtilities/src/test/java/org/opensearch/migrations/tracing/IInstrumentationAttributesTest.java index 45afa2769..87444114a 100644 --- a/TrafficCapture/coreUtilities/src/test/java/org/opensearch/migrations/tracing/IInstrumentationAttributesTest.java +++ b/TrafficCapture/coreUtilities/src/test/java/org/opensearch/migrations/tracing/IInstrumentationAttributesTest.java @@ -64,7 +64,7 @@ public AttributesBuilder fillAttributesForSpansBelow(AttributesBuilder builder) @Test public void getPopulatedAttributesAreOverrideCorrectly() { - var rootCtx = new RootOtelContext("test"); + var rootCtx = new RootOtelContext("test", IContextTracker.DO_NOTHING_TRACKER); var aCtx = new AContext(rootCtx); var bCtx = new BContext(rootCtx, aCtx); diff --git a/TrafficCapture/coreUtilities/src/testFixtures/java/org/opensearch/migrations/tracing/ContextTracker.java b/TrafficCapture/coreUtilities/src/testFixtures/java/org/opensearch/migrations/tracing/BacktracingContextTracker.java similarity index 88% rename from TrafficCapture/coreUtilities/src/testFixtures/java/org/opensearch/migrations/tracing/ContextTracker.java rename to TrafficCapture/coreUtilities/src/testFixtures/java/org/opensearch/migrations/tracing/BacktracingContextTracker.java index a3214ce3c..922d7bece 100644 --- a/TrafficCapture/coreUtilities/src/testFixtures/java/org/opensearch/migrations/tracing/ContextTracker.java +++ b/TrafficCapture/coreUtilities/src/testFixtures/java/org/opensearch/migrations/tracing/BacktracingContextTracker.java @@ -8,7 +8,7 @@ import java.util.stream.Collectors; @Slf4j -public class ContextTracker implements AutoCloseable { +public class BacktracingContextTracker implements IContextTracker, AutoCloseable { private static class ExceptionForStackTracingOnly extends Exception { } @@ -33,6 +33,12 @@ public void onCreated(IScopedInstrumentationAttributes ctx) { return; } var oldItem = scopedContextToCallDetails.putIfAbsent(ctx, new CallDetails()); + if (oldItem != null) { + var oldCtx = scopedContextToCallDetails.entrySet().stream().findFirst().get().getKey(); + if (oldCtx.equals(ctx)) { + log.error("check"); + } + } assert oldItem == null; } } diff --git a/TrafficCapture/nettyWireLogging/src/main/java/org/opensearch/migrations/trafficcapture/netty/tracing/RootWireLoggingContext.java b/TrafficCapture/nettyWireLogging/src/main/java/org/opensearch/migrations/trafficcapture/netty/tracing/RootWireLoggingContext.java index 503861b1f..2d7acc751 100644 --- a/TrafficCapture/nettyWireLogging/src/main/java/org/opensearch/migrations/trafficcapture/netty/tracing/RootWireLoggingContext.java +++ b/TrafficCapture/nettyWireLogging/src/main/java/org/opensearch/migrations/trafficcapture/netty/tracing/RootWireLoggingContext.java @@ -2,6 +2,7 @@ import io.opentelemetry.api.OpenTelemetry; import lombok.Getter; +import org.opensearch.migrations.tracing.IContextTracker; import org.opensearch.migrations.tracing.RootOtelContext; @Getter @@ -14,12 +15,12 @@ public class RootWireLoggingContext extends RootOtelContext implements IRootWire public final WireCaptureContexts.WaitingForResponseContext.MetricInstruments waitingForResponseInstruments; public final WireCaptureContexts.ResponseContext.MetricInstruments responseInstruments; - public RootWireLoggingContext(OpenTelemetry openTelemetry) { - this(openTelemetry, SCOPE_NAME); + public RootWireLoggingContext(OpenTelemetry openTelemetry, IContextTracker contextTracker) { + this(openTelemetry, contextTracker, SCOPE_NAME); } - public RootWireLoggingContext(OpenTelemetry openTelemetry, String scopeName) { - super(scopeName, openTelemetry); + public RootWireLoggingContext(OpenTelemetry openTelemetry, IContextTracker contextTracker, String scopeName) { + super(scopeName, contextTracker, openTelemetry); var meter = this.getMeterProvider().get(scopeName); connectionInstruments = WireCaptureContexts.ConnectionContext.makeMetrics(meter); requestInstruments = WireCaptureContexts.RequestContext.makeMetrics(meter); diff --git a/TrafficCapture/nettyWireLogging/src/test/java/org/opensearch/migrations/trafficcapture/netty/ConditionallyReliableLoggingHttpHandlerTest.java b/TrafficCapture/nettyWireLogging/src/test/java/org/opensearch/migrations/trafficcapture/netty/ConditionallyReliableLoggingHttpHandlerTest.java index 4d17893fe..26d4d3100 100644 --- a/TrafficCapture/nettyWireLogging/src/test/java/org/opensearch/migrations/trafficcapture/netty/ConditionallyReliableLoggingHttpHandlerTest.java +++ b/TrafficCapture/nettyWireLogging/src/test/java/org/opensearch/migrations/trafficcapture/netty/ConditionallyReliableLoggingHttpHandlerTest.java @@ -5,6 +5,7 @@ import io.netty.buffer.Unpooled; import io.netty.channel.embedded.EmbeddedChannel; import lombok.Getter; +import lombok.NonNull; import lombok.SneakyThrows; import lombok.extern.slf4j.Slf4j; import org.junit.jupiter.api.Assertions; @@ -14,6 +15,7 @@ import org.junit.jupiter.params.provider.ValueSource; import org.opensearch.migrations.testutils.TestUtilities; import org.opensearch.migrations.testutils.WrapWithNettyLeakDetection; +import org.opensearch.migrations.tracing.IContextTracker; import org.opensearch.migrations.tracing.InMemoryInstrumentationBundle; import org.opensearch.migrations.trafficcapture.CodedOutputStreamAndByteBufferWrapper; import org.opensearch.migrations.trafficcapture.CodedOutputStreamHolder; @@ -51,11 +53,15 @@ public TestRootContext() { this(false, false); } public TestRootContext(boolean trackMetrics, boolean trackTraces) { - this(new InMemoryInstrumentationBundle(trackTraces, trackTraces)); + this(trackMetrics, trackTraces, DO_NOTHING_TRACKER); + } + public TestRootContext(boolean trackMetrics, boolean trackTraces, @NonNull IContextTracker contextTracker) { + this(new InMemoryInstrumentationBundle(trackMetrics, trackTraces), contextTracker); } - public TestRootContext(InMemoryInstrumentationBundle inMemoryInstrumentationBundle) { - super(inMemoryInstrumentationBundle.openTelemetrySdk); + public TestRootContext(InMemoryInstrumentationBundle inMemoryInstrumentationBundle, + IContextTracker contextTracker) { + super(inMemoryInstrumentationBundle.openTelemetrySdk, contextTracker); this.instrumentationBundle = inMemoryInstrumentationBundle; } diff --git a/TrafficCapture/trafficCaptureProxyServer/src/main/java/org/opensearch/migrations/trafficcapture/proxyserver/CaptureProxy.java b/TrafficCapture/trafficCaptureProxyServer/src/main/java/org/opensearch/migrations/trafficcapture/proxyserver/CaptureProxy.java index 64aa22008..8f14eae36 100644 --- a/TrafficCapture/trafficCaptureProxyServer/src/main/java/org/opensearch/migrations/trafficcapture/proxyserver/CaptureProxy.java +++ b/TrafficCapture/trafficCaptureProxyServer/src/main/java/org/opensearch/migrations/trafficcapture/proxyserver/CaptureProxy.java @@ -17,6 +17,9 @@ import org.apache.kafka.clients.producer.ProducerConfig; import org.apache.kafka.common.config.SaslConfigs; import org.opensearch.common.settings.Settings; +import org.opensearch.migrations.tracing.ActiveContextTracker; +import org.opensearch.migrations.tracing.ActiveContextTrackerByActivityType; +import org.opensearch.migrations.tracing.CompositeContextTracker; import org.opensearch.migrations.tracing.RootOtelContext; import org.opensearch.migrations.trafficcapture.CodedOutputStreamHolder; import org.opensearch.migrations.trafficcapture.FileConnectionCaptureFactory; @@ -315,7 +318,8 @@ public static void main(String[] args) throws InterruptedException, IOException var backsideUri = convertStringToUri(params.backsideUriString); var rootContext = new RootCaptureContext( - RootOtelContext.initializeOpenTelemetryWithCollectorOrAsNoop(params.otelCollectorEndpoint, "capture")); + RootOtelContext.initializeOpenTelemetryWithCollectorOrAsNoop(params.otelCollectorEndpoint, "capture"), + new CompositeContextTracker(new ActiveContextTracker(), new ActiveContextTrackerByActivityType())); var sksOp = Optional.ofNullable(params.sslConfigFilePath) .map(sslConfigFile->new DefaultSecurityKeyStore(getSettings(sslConfigFile), diff --git a/TrafficCapture/trafficCaptureProxyServer/src/main/java/org/opensearch/migrations/trafficcapture/proxyserver/RootCaptureContext.java b/TrafficCapture/trafficCaptureProxyServer/src/main/java/org/opensearch/migrations/trafficcapture/proxyserver/RootCaptureContext.java index af2be2a0c..6731a7998 100644 --- a/TrafficCapture/trafficCaptureProxyServer/src/main/java/org/opensearch/migrations/trafficcapture/proxyserver/RootCaptureContext.java +++ b/TrafficCapture/trafficCaptureProxyServer/src/main/java/org/opensearch/migrations/trafficcapture/proxyserver/RootCaptureContext.java @@ -2,6 +2,7 @@ import io.opentelemetry.api.OpenTelemetry; import lombok.Getter; +import org.opensearch.migrations.tracing.IContextTracker; import org.opensearch.migrations.trafficcapture.kafkaoffloader.tracing.IRootKafkaOffloaderContext; import org.opensearch.migrations.trafficcapture.kafkaoffloader.tracing.KafkaRecordContext; import org.opensearch.migrations.trafficcapture.netty.tracing.RootWireLoggingContext; @@ -13,12 +14,12 @@ public class RootCaptureContext extends RootWireLoggingContext implements IRootK @Getter public final KafkaRecordContext.MetricInstruments kafkaOffloadingInstruments; - public RootCaptureContext(OpenTelemetry openTelemetry) { - this(openTelemetry, SCOPE_NAME); + public RootCaptureContext(OpenTelemetry openTelemetry, IContextTracker contextTracker) { + this(openTelemetry, contextTracker, SCOPE_NAME); } - public RootCaptureContext(OpenTelemetry openTelemetry, String scopeName) { - super(openTelemetry, scopeName); + public RootCaptureContext(OpenTelemetry openTelemetry, IContextTracker contextTracker, String scopeName) { + super(openTelemetry, contextTracker, scopeName); var meter = this.getMeterProvider().get(scopeName); kafkaOffloadingInstruments = KafkaRecordContext.makeMetrics(meter); } diff --git a/TrafficCapture/trafficCaptureProxyServer/src/test/java/org/opensearch/migrations/trafficcapture/proxyserver/netty/NettyScanningHttpProxyTest.java b/TrafficCapture/trafficCaptureProxyServer/src/test/java/org/opensearch/migrations/trafficcapture/proxyserver/netty/NettyScanningHttpProxyTest.java index 0db298df5..4d02b8fb9 100644 --- a/TrafficCapture/trafficCaptureProxyServer/src/test/java/org/opensearch/migrations/trafficcapture/proxyserver/netty/NettyScanningHttpProxyTest.java +++ b/TrafficCapture/trafficCaptureProxyServer/src/test/java/org/opensearch/migrations/trafficcapture/proxyserver/netty/NettyScanningHttpProxyTest.java @@ -10,6 +10,7 @@ import org.opensearch.migrations.testutils.SimpleHttpClientForTesting; import org.opensearch.migrations.testutils.SimpleHttpResponse; import org.opensearch.migrations.testutils.SimpleHttpServer; +import org.opensearch.migrations.tracing.IContextTracker; import org.opensearch.migrations.tracing.InMemoryInstrumentationBundle; import org.opensearch.migrations.trafficcapture.IConnectionCaptureFactory; import org.opensearch.migrations.trafficcapture.InMemoryConnectionCaptureFactory; @@ -90,7 +91,8 @@ public void testRoundTrip() throws var captureFactory = new InMemoryConnectionCaptureFactory(TEST_NODE_ID_STRING, 1024*1024, () -> interactionsCapturedCountdown.countDown()); var inMemoryInstrumentationBundle = new InMemoryInstrumentationBundle(true, true); - var rootCtx = new RootWireLoggingContext(inMemoryInstrumentationBundle.openTelemetrySdk); + var rootCtx = new RootWireLoggingContext(inMemoryInstrumentationBundle.openTelemetrySdk, + IContextTracker.DO_NOTHING_TRACKER); var servers = startServers(rootCtx, captureFactory); try (var client = new SimpleHttpClientForTesting()) { diff --git a/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/TrafficReplayer.java b/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/TrafficReplayer.java index 1e00c1d96..f3453caa2 100644 --- a/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/TrafficReplayer.java +++ b/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/TrafficReplayer.java @@ -27,6 +27,10 @@ import org.opensearch.migrations.replay.traffic.source.TrafficStreamLimiter; import org.opensearch.migrations.replay.util.DiagnosticTrackableCompletableFuture; import org.opensearch.migrations.replay.util.StringTrackableCompletableFuture; +import org.opensearch.migrations.tracing.ActiveContextTracker; +import org.opensearch.migrations.tracing.ActiveContextTrackerByActivityType; +import org.opensearch.migrations.tracing.CompositeContextTracker; +import org.opensearch.migrations.tracing.LoggingContextMonitor; import org.opensearch.migrations.tracing.RootOtelContext; import org.opensearch.migrations.trafficcapture.protos.TrafficStreamUtils; import org.opensearch.migrations.transform.IAuthTransformer; @@ -35,6 +39,8 @@ import org.opensearch.migrations.transform.IJsonTransformer; import org.opensearch.migrations.transform.RemovingAuthTransformerFactory; import org.opensearch.migrations.transform.StaticAuthTransformerFactory; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import org.slf4j.event.Level; import org.slf4j.spi.LoggingEventBuilder; import software.amazon.awssdk.arns.Arn; @@ -386,10 +392,14 @@ public static void main(String[] args) throws Exception { System.exit(3); return; } - var topContext = new RootReplayerContext(RootOtelContext.initializeOpenTelemetryWithCollectorOrAsNoop(params.otelCollectorEndpoint, - "replay")); + var contextTrackers = new CompositeContextTracker( + new ActiveContextTracker(), + new ActiveContextTrackerByActivityType()); + var topContext = new RootReplayerContext( + RootOtelContext.initializeOpenTelemetryWithCollectorOrAsNoop(params.otelCollectorEndpoint, "replay"), + contextTrackers); try (var blockingTrafficSource = TrafficCaptureSourceFactory.createTrafficCaptureSource(topContext, params, - Duration.ofSeconds(params.lookaheadTimeSeconds)); + Duration.ofSeconds(params.lookaheadTimeSeconds)); var authTransformer = buildAuthTransformerFactory(params)) { String transformerConfig = getTransformerConfig(params); diff --git a/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/datahandlers/NettyPacketToHttpConsumer.java b/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/datahandlers/NettyPacketToHttpConsumer.java index 769d30dd2..dc4e373ff 100644 --- a/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/datahandlers/NettyPacketToHttpConsumer.java +++ b/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/datahandlers/NettyPacketToHttpConsumer.java @@ -110,7 +110,7 @@ private void activateLiveChannel(DiagnosticTrackableCompletableFuture"error creating channel, not retrying") .setCause(connectFuture.cause()).log(); initialFuture.future.completeExceptionally(connectFuture.cause()); diff --git a/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/datatypes/ConnectionReplaySession.java b/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/datatypes/ConnectionReplaySession.java index e7995f05f..c591a86c4 100644 --- a/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/datatypes/ConnectionReplaySession.java +++ b/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/datatypes/ConnectionReplaySession.java @@ -1,7 +1,6 @@ package org.opensearch.migrations.replay.datatypes; import io.netty.channel.ChannelFuture; -import io.netty.channel.ConnectTimeoutException; import io.netty.channel.EventLoop; import lombok.Getter; import lombok.NonNull; @@ -91,7 +90,7 @@ private void createNewChannelFuture(boolean requireActiveChannel, int retries, eventLoopFuture.future.complete(v); } } else if (t != null) { - channelKeyContext.addException(t, true); + channelKeyContext.addTraceException(t, true); eventLoopFuture.future.completeExceptionally(t); } else { eventLoopFuture.future.complete(v); diff --git a/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/tracing/ReplayContexts.java b/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/tracing/ReplayContexts.java index e39ba747c..3d5fc6f99 100644 --- a/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/tracing/ReplayContexts.java +++ b/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/tracing/ReplayContexts.java @@ -748,7 +748,7 @@ public AttributesBuilder fillExtraAttributesForThisSpan(AttributesBuilder builde public void sendMeterEventsForEnd() { super.sendMeterEventsForEnd(); AttributesBuilder attributesBuilderForAggregate = getSharedAttributes(Attributes.builder()); - getCurrentSpan().setAllAttributes(attributesBuilderForAggregate.build()); + setAllAttributes(attributesBuilderForAggregate.build()); meterIncrementEvent(getMetrics().resultCounter, 1, attributesBuilderForAggregate); } @@ -769,7 +769,7 @@ public static long categorizeStatus(int status) { */ @Override public void setEndpoint(String endpointUrl) { - getCurrentSpan().setAttribute(ENDPOINT_KEY, endpointUrl); + setAttribute(ENDPOINT_KEY, endpointUrl); } /** @@ -779,7 +779,7 @@ public void setEndpoint(String endpointUrl) { */ @Override public void setHttpVersion(String httpVersion) { - getCurrentSpan().setAttribute(HTTP_VERSION_KEY, httpVersion); + setAttribute(HTTP_VERSION_KEY, httpVersion); } @Override diff --git a/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/tracing/RootReplayerContext.java b/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/tracing/RootReplayerContext.java index 8fe51cacf..8aacf243b 100644 --- a/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/tracing/RootReplayerContext.java +++ b/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/tracing/RootReplayerContext.java @@ -5,6 +5,7 @@ import org.opensearch.migrations.replay.datatypes.ISourceTrafficChannelKey; import org.opensearch.migrations.replay.datatypes.ITrafficStreamKey; import org.opensearch.migrations.replay.traffic.source.InputStreamOfTraffic; +import org.opensearch.migrations.tracing.IContextTracker; import org.opensearch.migrations.tracing.RootOtelContext; @Getter @@ -36,8 +37,8 @@ public class RootReplayerContext extends RootOtelContext implements IRootReplaye public final ReplayContexts.TupleHandlingContext.MetricInstruments tupleHandlingInstruments; public final ReplayContexts.SocketContext.MetricInstruments socketInstruments; - public RootReplayerContext(OpenTelemetry sdk) { - super(SCOPE_NAME, sdk); + public RootReplayerContext(OpenTelemetry sdk, IContextTracker contextTracker) { + super(SCOPE_NAME, contextTracker, sdk); var meter = this.getMeterProvider().get(SCOPE_NAME); asyncListeningInstruments = KafkaConsumerContexts.AsyncListeningContext.makeMetrics(meter); diff --git a/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/tracing/TracingTest.java b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/tracing/TracingTest.java index 941fe0bfe..f8753fbb3 100644 --- a/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/tracing/TracingTest.java +++ b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/tracing/TracingTest.java @@ -66,7 +66,7 @@ public void tracingWorks() { checkSpans(recordedSpans); checkMetrics(recordedMetrics); - Assertions.assertTrue(rootContext.contextTracker.getAllRemainingActiveScopes().isEmpty()); + Assertions.assertTrue(rootContext.getBacktracingContextTracker().getAllRemainingActiveScopes().isEmpty()); } private void checkMetrics(Collection recordedMetrics) { diff --git a/TrafficCapture/trafficReplayer/src/testFixtures/java/org/opensearch/migrations/tracing/TestContext.java b/TrafficCapture/trafficReplayer/src/testFixtures/java/org/opensearch/migrations/tracing/TestContext.java index 1c5927e46..73882dd8a 100644 --- a/TrafficCapture/trafficReplayer/src/testFixtures/java/org/opensearch/migrations/tracing/TestContext.java +++ b/TrafficCapture/trafficReplayer/src/testFixtures/java/org/opensearch/migrations/tracing/TestContext.java @@ -1,7 +1,5 @@ package org.opensearch.migrations.tracing; -import io.opentelemetry.sdk.testing.exporter.InMemoryMetricExporter; -import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; import org.opensearch.migrations.replay.datatypes.ITrafficStreamKey; import org.opensearch.migrations.replay.datatypes.PojoTrafficStreamKeyAndContext; import org.opensearch.migrations.replay.datatypes.UniqueReplayerRequestKey; @@ -16,12 +14,11 @@ public class TestContext extends RootReplayerContext implements AutoCloseable { public static final String TEST_NODE_ID = "testNodeId"; public static final String DEFAULT_TEST_CONNECTION = "testConnection"; public final InMemoryInstrumentationBundle inMemoryInstrumentationBundle; - public final ContextTracker contextTracker = new ContextTracker(); public final ChannelContextManager channelContextManager = new ChannelContextManager(this); private final Object channelContextManagerLock = new Object(); public static TestContext withTracking(boolean tracing, boolean metrics) { - return new TestContext(new InMemoryInstrumentationBundle(tracing, metrics)); + return new TestContext(new InMemoryInstrumentationBundle(tracing, metrics), new BacktracingContextTracker()); } public static TestContext withAllTracking() { @@ -29,36 +26,30 @@ public static TestContext withAllTracking() { } public static TestContext noOtelTracking() { - return new TestContext(new InMemoryInstrumentationBundle(null, null)); + return new TestContext(new InMemoryInstrumentationBundle(null, null), new BacktracingContextTracker()); } - public TestContext(InMemoryInstrumentationBundle inMemoryInstrumentationBundle) { - super(inMemoryInstrumentationBundle.openTelemetrySdk); + public TestContext(InMemoryInstrumentationBundle inMemoryInstrumentationBundle, IContextTracker contextTracker) { + super(inMemoryInstrumentationBundle.openTelemetrySdk, contextTracker); this.inMemoryInstrumentationBundle = inMemoryInstrumentationBundle; } - @Override - public void onContextCreated(IScopedInstrumentationAttributes newScopedContext) { - contextTracker.onCreated(newScopedContext); - } - - @Override - public void onContextClosed(IScopedInstrumentationAttributes newScopedContext) { - contextTracker.onClosed(newScopedContext); - } - public IReplayContexts.ITrafficStreamsLifecycleContext createTrafficStreamContextForTest(ITrafficStreamKey tsk) { synchronized (channelContextManagerLock) { return createTrafficStreamContextForStreamSource(channelContextManager.retainOrCreateContext(tsk), tsk); } } + public BacktracingContextTracker getBacktracingContextTracker() { + return (BacktracingContextTracker) getContextTracker(); + } + @Override public void close() { - contextTracker.close(); - inMemoryInstrumentationBundle.close(); // Assertions.assertEquals("", contextTracker.getAllRemainingActiveScopes().entrySet().stream() // .map(kvp->kvp.getKey().toString()).collect(Collectors.joining())); + getBacktracingContextTracker().close(); + inMemoryInstrumentationBundle.close(); } From eb9d957fb70f441149a4b9fac9dfa0bdd530de2d Mon Sep 17 00:00:00 2001 From: Greg Schohn Date: Sat, 6 Apr 2024 09:02:22 -0400 Subject: [PATCH 2/3] Refactor TrafficReplayer to separate its main/command line parameters from the business logic. Some of the new class's constructors have been simplified/streamlined a bit at the expense of not creating a JsonTransfomer automatically. Signed-off-by: Greg Schohn --- .../migrations/replay/TrafficReplayer.java | 680 +---------------- .../replay/TrafficReplayerTopLevel.java | 690 ++++++++++++++++++ .../replay/TrafficReplayerTest.java | 9 +- .../replay/TransformationLoaderTest.java | 2 +- .../NettyPacketToHttpConsumerTest.java | 3 +- .../FullReplayerWithTracingChecksTest.java | 10 +- .../e2etests/FullTrafficReplayerTest.java | 40 +- ...ficStreamBecomesTwoTargetChannelsTest.java | 7 +- .../e2etests/TrafficReplayerRunner.java | 13 +- 9 files changed, 751 insertions(+), 703 deletions(-) create mode 100644 TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/TrafficReplayerTopLevel.java diff --git a/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/TrafficReplayer.java b/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/TrafficReplayer.java index 3bd04c2a4..ea572820f 100644 --- a/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/TrafficReplayer.java +++ b/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/TrafficReplayer.java @@ -3,52 +3,23 @@ import com.beust.jcommander.JCommander; import com.beust.jcommander.Parameter; import com.beust.jcommander.ParameterException; -import io.netty.buffer.Unpooled; -import io.netty.handler.ssl.SslContext; -import io.netty.handler.ssl.SslContextBuilder; -import io.netty.handler.ssl.util.InsecureTrustManagerFactory; -import io.netty.util.concurrent.Future; -import lombok.AllArgsConstructor; -import lombok.Lombok; -import lombok.NonNull; -import lombok.SneakyThrows; import lombok.extern.slf4j.Slf4j; -import org.opensearch.migrations.replay.datahandlers.IPacketFinalizingConsumer; -import org.opensearch.migrations.replay.datatypes.HttpRequestTransformationStatus; -import org.opensearch.migrations.replay.datatypes.ITrafficStreamKey; -import org.opensearch.migrations.replay.datatypes.TransformedPackets; -import org.opensearch.migrations.replay.datatypes.UniqueReplayerRequestKey; -import org.opensearch.migrations.replay.tracing.IReplayContexts; -import org.opensearch.migrations.replay.tracing.IRootReplayerContext; import org.opensearch.migrations.replay.tracing.RootReplayerContext; -import org.opensearch.migrations.replay.traffic.source.BlockingTrafficSource; -import org.opensearch.migrations.replay.traffic.source.ITrafficCaptureSource; -import org.opensearch.migrations.replay.traffic.source.ITrafficStreamWithKey; -import org.opensearch.migrations.replay.traffic.source.TrafficStreamLimiter; -import org.opensearch.migrations.replay.util.DiagnosticTrackableCompletableFuture; -import org.opensearch.migrations.replay.util.StringTrackableCompletableFuture; import org.opensearch.migrations.tracing.ActiveContextTracker; import org.opensearch.migrations.tracing.ActiveContextTrackerByActivityType; import org.opensearch.migrations.tracing.CompositeContextTracker; import org.opensearch.migrations.tracing.LoggingContextMonitor; import org.opensearch.migrations.tracing.RootOtelContext; -import org.opensearch.migrations.trafficcapture.protos.TrafficStreamUtils; import org.opensearch.migrations.transform.IAuthTransformer; import org.opensearch.migrations.transform.IAuthTransformerFactory; import org.opensearch.migrations.transform.IHttpMessage; -import org.opensearch.migrations.transform.IJsonTransformer; import org.opensearch.migrations.transform.RemovingAuthTransformerFactory; import org.opensearch.migrations.transform.StaticAuthTransformerFactory; -import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import org.slf4j.event.Level; -import org.slf4j.spi.LoggingEventBuilder; import software.amazon.awssdk.arns.Arn; import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider; import software.amazon.awssdk.regions.Region; -import javax.net.ssl.SSLException; -import java.io.EOFException; import java.io.IOException; import java.lang.ref.WeakReference; import java.net.URI; @@ -56,54 +27,20 @@ import java.nio.file.Files; import java.nio.file.Paths; import java.time.Duration; -import java.time.Instant; -import java.util.ArrayList; -import java.util.Arrays; import java.util.List; -import java.util.Map; import java.util.Optional; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.CompletionException; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.ExecutionException; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.TimeoutException; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicReference; -import java.util.function.Consumer; -import java.util.function.Supplier; -import java.util.stream.Collectors; -import java.util.stream.Stream; @Slf4j -public class TrafficReplayer implements AutoCloseable { +public class TrafficReplayer { + private static final String ALL_ACTIVE_CONTEXTS_MONITOR = "AllActiveContextsMonitor"; public static final String SIGV_4_AUTH_HEADER_SERVICE_REGION_ARG = "--sigv4-auth-header-service-region"; public static final String AUTH_HEADER_VALUE_ARG = "--auth-header-value"; public static final String REMOVE_AUTH_HEADER_VALUE_ARG = "--remove-auth-header"; public static final String AWS_AUTH_HEADER_USER_AND_SECRET_ARG = "--auth-header-user-and-secret"; public static final String PACKET_TIMEOUT_SECONDS_PARAMETER_NAME = "--packet-timeout-seconds"; - public static final int MAX_ITEMS_TO_SHOW_FOR_LEFTOVER_WORK_AT_INFO_LEVEL = 10; - public static final String TARGET_CONNECTION_POOL_NAME = "targetConnectionPool"; public static final String LOOKAHEAD_TIME_WINDOW_PARAMETER_NAME = "--lookahead-time-window"; - public static AtomicInteger targetConnectionPoolUniqueCounter = new AtomicInteger(); - - private final PacketToTransformingHttpHandlerFactory inputRequestTransformerFactory; - private final ClientConnectionPool clientConnectionPool; - private final TrafficStreamLimiter liveTrafficStreamLimiter; - private final AtomicInteger successfulRequestCount; - private final AtomicInteger exceptionRequestCount; - public final IRootReplayerContext topLevelContext; - private final ConcurrentHashMap> requestToFinalWorkFuturesMap; - - private final AtomicBoolean stopReadingRef; - private final AtomicReference> allRemainingWorkFutureOrShutdownSignalRef; - private final AtomicReference shutdownReasonRef; - private final AtomicReference> shutdownFutureRef; - private final AtomicReference>> nextChunkFutureRef; public static class DualException extends Exception { public final Throwable originalCause; @@ -126,80 +63,6 @@ public TerminationException(Throwable originalCause, Throwable immediateCause) { } } - public TrafficReplayer(IRootReplayerContext context, - URI serverUri, - String fullTransformerConfig, - IAuthTransformerFactory authTransformerFactory, - boolean allowInsecureConnections) - throws SSLException { - this(context, serverUri, fullTransformerConfig, authTransformerFactory, null, allowInsecureConnections, - 0, 1024, - TARGET_CONNECTION_POOL_NAME + targetConnectionPoolUniqueCounter.incrementAndGet()); - } - - - public TrafficReplayer(IRootReplayerContext context, - URI serverUri, - String fullTransformerConfig, - IAuthTransformerFactory authTransformerFactory, - String userAgent, - boolean allowInsecureConnections, - int numSendingThreads, - int maxConcurrentOutstandingRequests, - String targetConnectionPoolName) - throws SSLException { - this(context, serverUri, authTransformerFactory, allowInsecureConnections, - numSendingThreads, maxConcurrentOutstandingRequests, targetConnectionPoolName, - new TransformationLoader() - .getTransformerFactoryLoader(serverUri.getHost(), userAgent, fullTransformerConfig) - ); - } - - public TrafficReplayer(IRootReplayerContext context, - URI serverUri, - IAuthTransformerFactory authTransformer, - boolean allowInsecureConnections, - int numSendingThreads, int maxConcurrentOutstandingRequests, - String clientThreadNamePrefix, - IJsonTransformer jsonTransformer) - throws SSLException - { - this.topLevelContext = context; - if (serverUri.getPort() < 0) { - throw new IllegalArgumentException("Port not present for URI: "+serverUri); - } - if (serverUri.getHost() == null) { - throw new IllegalArgumentException("Hostname not present for URI: "+serverUri); - } - if (serverUri.getScheme() == null) { - throw new IllegalArgumentException("Scheme (http|https) is not present for URI: "+serverUri); - } - inputRequestTransformerFactory = new PacketToTransformingHttpHandlerFactory(jsonTransformer, authTransformer); - clientConnectionPool = new ClientConnectionPool(serverUri, - loadSslContext(serverUri, allowInsecureConnections), clientThreadNamePrefix, numSendingThreads); - requestToFinalWorkFuturesMap = new ConcurrentHashMap<>(); - successfulRequestCount = new AtomicInteger(); - exceptionRequestCount = new AtomicInteger(); - liveTrafficStreamLimiter = new TrafficStreamLimiter(maxConcurrentOutstandingRequests); - allRemainingWorkFutureOrShutdownSignalRef = new AtomicReference<>(); - shutdownReasonRef = new AtomicReference<>(); - shutdownFutureRef = new AtomicReference<>(); - nextChunkFutureRef = new AtomicReference<>(); - stopReadingRef = new AtomicBoolean(); - } - - private static SslContext loadSslContext(URI serverUri, boolean allowInsecureConnections) throws SSLException { - if (serverUri.getScheme().equalsIgnoreCase("https")) { - var sslContextBuilder = SslContextBuilder.forClient(); - if (allowInsecureConnections) { - sslContextBuilder.trustManager(InsecureTrustManagerFactory.INSTANCE); - } - return sslContextBuilder.build(); - } else { - return null; - } - } - public static boolean validateRequiredKafkaParams(String brokers, String topic, String groupId) { if (brokers == null && topic == null && groupId == null) { return false; @@ -412,14 +275,17 @@ public static void main(String[] args) throws Exception { contextTrackers); try (var blockingTrafficSource = TrafficCaptureSourceFactory.createTrafficCaptureSource(topContext, params, Duration.ofSeconds(params.lookaheadTimeSeconds)); - var authTransformer = buildAuthTransformerFactory(params)) + var authTransformer = buildAuthTransformerFactory(params); + var globalContextMonitor = new LoggingContextMonitor(LoggerFactory.getLogger(ALL_ACTIVE_CONTEXTS_MONITOR), + Duration.ofSeconds(30))) { String transformerConfig = getTransformerConfig(params); if (transformerConfig != null) { log.atInfo().setMessage(()->"Transformations config string: " + transformerConfig).log(); } - var tr = new TrafficReplayer(topContext, uri, transformerConfig, authTransformer, params.userAgent, - params.allowInsecureConnections, params.numClientThreads, params.maxConcurrentRequests, TARGET_CONNECTION_POOL_NAME); + var tr = new TrafficReplayerTopLevel(topContext, uri, authTransformer, + params.allowInsecureConnections, params.numClientThreads, params.maxConcurrentRequests, + new TransformationLoader().getTransformerFactoryLoader(uri.getHost(), params.userAgent, transformerConfig)); setupShutdownHookForReplayer(tr); var tupleWriter = new TupleParserChainConsumer(new ResultsToLogsConsumer()); @@ -430,7 +296,7 @@ public static void main(String[] args) throws Exception { } } - private static void setupShutdownHookForReplayer(TrafficReplayer tr) { + private static void setupShutdownHookForReplayer(TrafficReplayerTopLevel tr) { var weakTrafficReplayer = new WeakReference<>(tr); Runtime.getRuntime().addShutdownHook(new Thread(() -> { // both Log4J and the java builtin loggers add shutdown hooks. @@ -522,532 +388,4 @@ public void close() { return null; // default is to do nothing to auth headers } } - - public void setupRunAndWaitForReplayToFinish(Duration observedPacketConnectionTimeout, - BlockingTrafficSource trafficSource, - TimeShifter timeShifter, - Consumer resultTupleConsumer) - throws InterruptedException, ExecutionException { - - var senderOrchestrator = new RequestSenderOrchestrator(clientConnectionPool); - var replayEngine = new ReplayEngine(senderOrchestrator, trafficSource, timeShifter); - - CapturedTrafficToHttpTransactionAccumulator trafficToHttpTransactionAccumulator = - new CapturedTrafficToHttpTransactionAccumulator(observedPacketConnectionTimeout, - "(see " + PACKET_TIMEOUT_SECONDS_PARAMETER_NAME + ")", - new TrafficReplayerAccumulationCallbacks(replayEngine, resultTupleConsumer, trafficSource)); - try { - pullCaptureFromSourceToAccumulator(trafficSource, trafficToHttpTransactionAccumulator); - } catch (InterruptedException ex) { - throw ex; - } catch (Exception e) { - log.atWarn().setCause(e).setMessage("Terminating runReplay due to exception").log(); - throw e; - } finally { - trafficToHttpTransactionAccumulator.close(); - wrapUpWorkAndEmitSummary(replayEngine, trafficToHttpTransactionAccumulator); - assert shutdownFutureRef.get() != null || requestToFinalWorkFuturesMap.isEmpty() : - "expected to wait for all the in flight requests to fully flush and self destruct themselves"; - } - } - - protected void wrapUpWorkAndEmitSummary(ReplayEngine replayEngine, - CapturedTrafficToHttpTransactionAccumulator trafficToHttpTransactionAccumulator) - throws ExecutionException, InterruptedException { - final var primaryLogLevel = Level.INFO; - final var secondaryLogLevel = Level.WARN; - var logLevel = primaryLogLevel; - for (var timeout = Duration.ofSeconds(60); ; timeout = timeout.multipliedBy(2)) { - if (shutdownFutureRef.get() != null) { - log.warn("Not waiting for work because the TrafficReplayer is shutting down."); - break; - } - try { - waitForRemainingWork(logLevel, timeout); - break; - } catch (TimeoutException e) { - log.atLevel(logLevel).log("Timed out while waiting for the remaining " + - "requests to be finalized..."); - logLevel = secondaryLogLevel; - } - } - if (!requestToFinalWorkFuturesMap.isEmpty() || exceptionRequestCount.get() > 0) { - log.atWarn().setMessage("{} in-flight requests being dropped due to pending shutdown; " + - "{} requests to the target threw an exception; " + - "{} requests were successfully processed.") - .addArgument(requestToFinalWorkFuturesMap.size()) - .addArgument(exceptionRequestCount.get()) - .addArgument(successfulRequestCount.get()) - .log(); - } else { - log.info(successfulRequestCount.get() + " requests were successfully processed."); - } - log.info("# of connections created: {}; # of requests on reused keep-alive connections: {}; " + - "# of expired connections: {}; # of connections closed: {}; " + - "# of connections terminated upon accumulator termination: {}", - trafficToHttpTransactionAccumulator.numberOfConnectionsCreated(), - trafficToHttpTransactionAccumulator.numberOfRequestsOnReusedConnections(), - trafficToHttpTransactionAccumulator.numberOfConnectionsExpired(), - trafficToHttpTransactionAccumulator.numberOfConnectionsClosed(), - trafficToHttpTransactionAccumulator.numberOfRequestsTerminatedUponAccumulatorClose() - ); - } - - public void setupRunAndWaitForReplayWithShutdownChecks(Duration observedPacketConnectionTimeout, - BlockingTrafficSource trafficSource, - TimeShifter timeShifter, - Consumer resultTupleConsumer) - throws TerminationException, ExecutionException, InterruptedException { - try { - setupRunAndWaitForReplayToFinish(observedPacketConnectionTimeout, trafficSource, - timeShifter, resultTupleConsumer); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new TerminationException(shutdownReasonRef.get(), e); - } catch (Throwable t) { - throw new TerminationException(shutdownReasonRef.get(), t); - } - if (shutdownReasonRef.get() != null) { - throw new TerminationException(shutdownReasonRef.get(), null); - } - // if nobody has run shutdown yet, do so now so that we can tear down the netty resources - shutdown(null).get(); // if somebody already HAD run shutdown, it will return the future already created - } - - @AllArgsConstructor - class TrafficReplayerAccumulationCallbacks implements AccumulationCallbacks { - private final ReplayEngine replayEngine; - private Consumer resultTupleConsumer; - private ITrafficCaptureSource trafficCaptureSource; - - @Override - public Consumer - onRequestReceived(@NonNull IReplayContexts.IReplayerHttpTransactionContext ctx, - @NonNull HttpMessageAndTimestamp request) { - replayEngine.setFirstTimestamp(request.getFirstPacketTimestamp()); - - var allWorkFinishedForTransaction = - new StringTrackableCompletableFuture(new CompletableFuture<>(), - ()->"waiting for work to be queued and run through TrafficStreamLimiter"); - var requestPushFuture = new StringTrackableCompletableFuture( - new CompletableFuture<>(), () -> "Waiting to get response from target"); - var requestKey = ctx.getReplayerRequestKey(); - var workItem = liveTrafficStreamLimiter.queueWork(1, ctx, wi -> { - transformAndSendRequest(replayEngine, request, ctx).future.whenComplete((v,t)->{ - liveTrafficStreamLimiter.doneProcessing(wi); - if (t != null) { - requestPushFuture.future.completeExceptionally(t); - } else { - requestPushFuture.future.complete(v); - } - }); - }); - if (!allWorkFinishedForTransaction.future.isDone()) { - log.trace("Adding " + requestKey + " to targetTransactionInProgressMap"); - requestToFinalWorkFuturesMap.put(requestKey, allWorkFinishedForTransaction); - if (allWorkFinishedForTransaction.future.isDone()) { - requestToFinalWorkFuturesMap.remove(requestKey); - } - } - - return rrPair -> - requestPushFuture.map(f -> f.handle((v, t) -> { - log.atInfo().setMessage(() -> "Done receiving captured stream for " + ctx + - ":" + rrPair.requestData).log(); - log.atTrace().setMessage(() -> - "Summary response value for " + requestKey + " returned=" + v).log(); - return handleCompletedTransaction(ctx, rrPair, v, t); - }), () -> "logging summary") - .whenComplete((v,t)->{ - if (t != null) { - allWorkFinishedForTransaction.future.completeExceptionally(t); - } else { - allWorkFinishedForTransaction.future.complete(null); - } - }, ()->""); - } - - Void handleCompletedTransaction(@NonNull IReplayContexts.IReplayerHttpTransactionContext context, - RequestResponsePacketPair rrPair, - TransformedTargetRequestAndResponse summary, Throwable t) { - try (var httpContext = rrPair.getHttpTransactionContext()) { - // if this comes in with a serious Throwable (not an Exception), don't bother - // packaging it up and calling the callback. - // Escalate it up out handling stack and shutdown. - if (t == null || t instanceof Exception) { - try (var tupleHandlingContext = httpContext.createTupleContext()) { - packageAndWriteResponse(tupleHandlingContext, resultTupleConsumer, - rrPair, summary, (Exception) t); - } - commitTrafficStreams(rrPair.completionStatus, rrPair.trafficStreamKeysBeingHeld); - return null; - } else { - log.atError().setCause(t).setMessage(() -> "Throwable passed to handle() for " + context + - ". Rethrowing.").log(); - throw Lombok.sneakyThrow(t); - } - } catch (Error error) { - log.atError() - .setCause(error) - .setMessage(() -> "Caught error and initiating TrafficReplayer shutdown") - .log(); - shutdown(error); - throw error; - } catch (Exception e) { - log.atError() - .setMessage("Unexpected exception while sending the " + - "aggregated response and context for {} to the callback. " + - "Proceeding, but the tuple receiver context may be compromised.") - .addArgument(context) - .setCause(e) - .log(); - throw e; - } finally { - var requestKey = context.getReplayerRequestKey(); - requestToFinalWorkFuturesMap.remove(requestKey); - log.trace("removed rrPair.requestData to " + - "targetTransactionInProgressMap for " + - requestKey); - } - } - - @Override - public void onTrafficStreamsExpired(RequestResponsePacketPair.ReconstructionStatus status, - @NonNull IReplayContexts.IChannelKeyContext ctx, - @NonNull List trafficStreamKeysBeingHeld) { - commitTrafficStreams(status, trafficStreamKeysBeingHeld); - } - - @SneakyThrows - private void commitTrafficStreams(RequestResponsePacketPair.ReconstructionStatus status, - List trafficStreamKeysBeingHeld) { - commitTrafficStreams(status != RequestResponsePacketPair.ReconstructionStatus.CLOSED_PREMATURELY, - trafficStreamKeysBeingHeld); - } - - @SneakyThrows - private void commitTrafficStreams(boolean shouldCommit, - List trafficStreamKeysBeingHeld) { - if (shouldCommit && trafficStreamKeysBeingHeld != null) { - for (var tsk : trafficStreamKeysBeingHeld) { - tsk.getTrafficStreamsContext().close(); - trafficCaptureSource.commitTrafficStream(tsk); - } - } - } - - @Override - public void onConnectionClose(int channelInteractionNum, - @NonNull IReplayContexts.IChannelKeyContext ctx, int channelSessionNumber, - RequestResponsePacketPair.ReconstructionStatus status, - @NonNull Instant timestamp, @NonNull List trafficStreamKeysBeingHeld) { - replayEngine.setFirstTimestamp(timestamp); - var cf = replayEngine.closeConnection(channelInteractionNum, ctx, channelSessionNumber, timestamp); - cf.map(f->f.whenComplete((v,t)->{ - commitTrafficStreams(status, trafficStreamKeysBeingHeld); - }), ()->"closing the channel in the ReplayEngine"); - } - - @Override - public void onTrafficStreamIgnored(@NonNull IReplayContexts.ITrafficStreamsLifecycleContext ctx) { - commitTrafficStreams(true, List.of(ctx.getTrafficStreamKey())); - } - - private TransformedTargetRequestAndResponse - packageAndWriteResponse(IReplayContexts.ITupleHandlingContext tupleHandlingContext, - Consumer tupleWriter, - RequestResponsePacketPair rrPair, - TransformedTargetRequestAndResponse summary, - Exception t) { - log.trace("done sending and finalizing data to the packet handler"); - - try (var requestResponseTuple = getSourceTargetCaptureTuple(tupleHandlingContext, rrPair, summary, t)) { - log.atInfo().setMessage(()->"Source/Target Request/Response tuple: " + requestResponseTuple).log(); - tupleWriter.accept(requestResponseTuple); - } - - if (t != null) { throw new CompletionException(t); } - if (summary.getError() != null) { - log.atInfo().setCause(summary.getError()).setMessage("Exception for {}: ") - .addArgument(tupleHandlingContext).log(); - exceptionRequestCount.incrementAndGet(); - } else if (summary.getTransformationStatus() == HttpRequestTransformationStatus.ERROR) { - log.atInfo().setCause(summary.getError()).setMessage("Unknown error transforming {}: ") - .addArgument(tupleHandlingContext).log(); - exceptionRequestCount.incrementAndGet(); - } else { - successfulRequestCount.incrementAndGet(); - } - return summary; - } - } - - protected void waitForRemainingWork(Level logLevel, @NonNull Duration timeout) - throws ExecutionException, InterruptedException, TimeoutException { - - if (!liveTrafficStreamLimiter.isStopped()) { - var streamLimiterHasRunEverything = new CompletableFuture(); - liveTrafficStreamLimiter.queueWork(1, null, wi -> { - streamLimiterHasRunEverything.complete(null); - liveTrafficStreamLimiter.doneProcessing(wi); - }); - streamLimiterHasRunEverything.get(timeout.toMillis(), TimeUnit.MILLISECONDS); - } - - Map.Entry>[] - allRemainingWorkArray = requestToFinalWorkFuturesMap.entrySet().toArray(Map.Entry[]::new); - writeStatusLogsForRemainingWork(logLevel, allRemainingWorkArray); - - // remember, this block is ONLY for the leftover items. Lots of other items have been processed - // and were removed from the live map (hopefully) - DiagnosticTrackableCompletableFuture[] allCompletableFuturesArray = - Arrays.stream(allRemainingWorkArray) - .map(Map.Entry::getValue).toArray(DiagnosticTrackableCompletableFuture[]::new); - var allWorkFuture = StringTrackableCompletableFuture.allOf(allCompletableFuturesArray, - () -> "TrafficReplayer.AllWorkFinished"); - try { - if (allRemainingWorkFutureOrShutdownSignalRef.compareAndSet(null, allWorkFuture)) { - allWorkFuture.get(timeout); - } else { - handleAlreadySetFinishedSignal(); - } - } catch (TimeoutException e) { - var didCancel = allWorkFuture.future.cancel(true); - if (!didCancel) { - assert allWorkFuture.future.isDone() : "expected future to have finished if cancel didn't succeed"; - // continue with the rest of the function - } else { - throw e; - } - } finally { - allRemainingWorkFutureOrShutdownSignalRef.set(null); - } - } - - private void handleAlreadySetFinishedSignal() throws InterruptedException, ExecutionException { - try { - var finishedSignal = allRemainingWorkFutureOrShutdownSignalRef.get().future; - assert finishedSignal.isDone() : "Expected this reference to be EITHER the current work futures " + - "or a sentinel value indicating a shutdown has commenced. The signal, when set, should " + - "have been completed at the time that the reference was set"; - finishedSignal.get(); - log.debug("Did shutdown cleanly"); - } catch (ExecutionException e) { - var c = e.getCause(); - if (c instanceof Error) { - throw (Error) c; - } else { - throw e; - } - } catch (Error t) { - log.atError().setCause(t).setMessage(() -> "Not waiting for all work to finish. " + - "The TrafficReplayer is shutting down").log(); - throw t; - } - } - - private static void writeStatusLogsForRemainingWork(Level logLevel, - Map.Entry>[] - allRemainingWorkArray) { - log.atLevel(logLevel).log("All remaining work to wait on " + allRemainingWorkArray.length); - if (log.isInfoEnabled()) { - LoggingEventBuilder loggingEventBuilderToUse = log.isTraceEnabled() ? log.atTrace() : log.atInfo(); - long itemLimit = log.isTraceEnabled() ? Long.MAX_VALUE : MAX_ITEMS_TO_SHOW_FOR_LEFTOVER_WORK_AT_INFO_LEVEL; - loggingEventBuilderToUse.setMessage(() -> " items: " + - Arrays.stream(allRemainingWorkArray) - .map(kvp -> kvp.getKey() + " --> " + - kvp.getValue().formatAsString(TrafficReplayer::formatWorkItem)) - .limit(itemLimit) - .collect(Collectors.joining("\n"))) - .log(); - } - } - - private static String formatWorkItem(DiagnosticTrackableCompletableFuture cf) { - try { - var resultValue = cf.get(); - if (resultValue instanceof TransformedTargetRequestAndResponse) { - return "" + ((TransformedTargetRequestAndResponse) resultValue).getTransformationStatus(); - } - return null; - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - return "Exception: " + e.getMessage(); - } catch (ExecutionException e) { - return e.getMessage(); - } - } - - private static SourceTargetCaptureTuple - getSourceTargetCaptureTuple(@NonNull IReplayContexts.ITupleHandlingContext tupleHandlingContext, - RequestResponsePacketPair rrPair, - TransformedTargetRequestAndResponse summary, - Exception t) - { - SourceTargetCaptureTuple requestResponseTuple; - if (t != null) { - log.error("Got exception in CompletableFuture callback: ", t); - requestResponseTuple = new SourceTargetCaptureTuple(tupleHandlingContext, rrPair, - new TransformedPackets(), new ArrayList<>(), - HttpRequestTransformationStatus.ERROR, t, Duration.ZERO); - } else { - requestResponseTuple = new SourceTargetCaptureTuple(tupleHandlingContext, rrPair, - summary.requestPackets, - summary.getReceiptTimeAndResponsePackets() - .map(Map.Entry::getValue).collect(Collectors.toList()), - summary.getTransformationStatus(), - summary.getError(), - summary.getResponseDuration() - ); - } - return requestResponseTuple; - } - - public DiagnosticTrackableCompletableFuture - transformAndSendRequest(ReplayEngine replayEngine, HttpMessageAndTimestamp request, - IReplayContexts.IReplayerHttpTransactionContext ctx) { - return transformAndSendRequest(inputRequestTransformerFactory, replayEngine, ctx, - request.getFirstPacketTimestamp(), request.getLastPacketTimestamp(), - request.packetBytes::stream); - } - - public static DiagnosticTrackableCompletableFuture - transformAndSendRequest(PacketToTransformingHttpHandlerFactory inputRequestTransformerFactory, - ReplayEngine replayEngine, - IReplayContexts.IReplayerHttpTransactionContext ctx, - @NonNull Instant start, @NonNull Instant end, - Supplier> packetsSupplier) - { - try { - var transformationCompleteFuture = replayEngine.scheduleTransformationWork(ctx, start, ()-> - transformAllData(inputRequestTransformerFactory.create(ctx), packetsSupplier)); - log.atDebug().setMessage(()->"finalizeRequest future for transformation of " + ctx + - " = " + transformationCompleteFuture).log(); - // It might be safer to chain this work directly inside the scheduleWork call above so that the - // read buffer horizons aren't set after the transformation work finishes, but after the packets - // are fully handled - return transformationCompleteFuture.thenCompose(transformedResult -> - replayEngine.scheduleRequest(ctx, start, end, - transformedResult.transformedOutput.size(), - transformedResult.transformedOutput.streamRetained()) - .map(future->future.thenApply(t -> - new TransformedTargetRequestAndResponse(transformedResult.transformedOutput, - t, transformedResult.transformationStatus, t.error)), - ()->"(if applicable) packaging transformed result into a completed TransformedTargetRequestAndResponse object") - .map(future->future.exceptionally(t -> - new TransformedTargetRequestAndResponse(transformedResult.transformedOutput, - transformedResult.transformationStatus, t)), - ()->"(if applicable) packaging transformed result into a failed TransformedTargetRequestAndResponse object"), - () -> "transitioning transformed packets onto the wire") - .map(future->future.exceptionally(t->new TransformedTargetRequestAndResponse(null, null, t)), - ()->"Checking for exception out of sending data to the target server"); - } catch (Exception e) { - log.debug("Caught exception in writeToSocket, so failing future"); - return StringTrackableCompletableFuture.failedFuture(e, ()->"TrafficReplayer.writeToSocketAndClose"); - } - } - - private static DiagnosticTrackableCompletableFuture - transformAllData(IPacketFinalizingConsumer packetHandler, Supplier> packetSupplier) { - try { - var logLabel = packetHandler.getClass().getSimpleName(); - var packets = packetSupplier.get().map(Unpooled::wrappedBuffer); - packets.forEach(packetData -> { - log.atDebug().setMessage(() -> logLabel + " sending " + packetData.readableBytes() + - " bytes to the packetHandler").log(); - var consumeFuture = packetHandler.consumeBytes(packetData); - log.atDebug().setMessage(() -> logLabel + " consumeFuture = " + consumeFuture).log(); - }); - log.atDebug().setMessage(() -> logLabel + " done sending bytes, now finalizing the request").log(); - return packetHandler.finalizeRequest(); - } catch (Exception e) { - log.atInfo().setCause(e).setMessage("Encountered an exception while transforming the http request. " + - "The base64 gzipped traffic stream, for later diagnostic purposes, is: " + - Utils.packetsToCompressedTrafficStream(packetSupplier.get())).log(); - throw e; - } - } - - @SneakyThrows - public @NonNull CompletableFuture shutdown(Error error) { - log.atWarn().setCause(error).setMessage(()->"Shutting down " + this).log(); - shutdownReasonRef.compareAndSet(null, error); - if (!shutdownFutureRef.compareAndSet(null, new CompletableFuture<>())) { - log.atError().setMessage(()->"Shutdown was already signaled by {}. " + - "Ignoring this shutdown request due to {}.") - .addArgument(shutdownReasonRef.get()) - .addArgument(error) - .log(); - return shutdownFutureRef.get(); - } - stopReadingRef.set(true); - liveTrafficStreamLimiter.close(); - - var nettyShutdownFuture = clientConnectionPool.shutdownNow(); - nettyShutdownFuture.whenComplete((v,t) -> { - if (t != null) { - shutdownFutureRef.get().completeExceptionally(t); - } else { - shutdownFutureRef.get().complete(null); - } - }); - Optional.ofNullable(this.nextChunkFutureRef.get()).ifPresent(f->f.cancel(true)); - var shutdownWasSignalledFuture = error == null ? - StringTrackableCompletableFuture.completedFuture(null, ()->"TrafficReplayer shutdown") : - StringTrackableCompletableFuture.failedFuture(error, ()->"TrafficReplayer shutdown"); - while (!allRemainingWorkFutureOrShutdownSignalRef.compareAndSet(null, shutdownWasSignalledFuture)) { - var otherRemainingWorkObj = allRemainingWorkFutureOrShutdownSignalRef.get(); - if (otherRemainingWorkObj != null) { - otherRemainingWorkObj.future.cancel(true); - break; - } - } - var shutdownFuture = shutdownFutureRef.get(); - log.atWarn().setMessage(()->"Shutdown setup has been initiated").log(); - return shutdownFuture; - } - - - @Override - public void close() throws Exception { - shutdown(null).get(); - } - - @SneakyThrows - public void pullCaptureFromSourceToAccumulator( - ITrafficCaptureSource trafficChunkStream, - CapturedTrafficToHttpTransactionAccumulator trafficToHttpTransactionAccumulator) - throws InterruptedException { - while (true) { - log.trace("Reading next chunk from TrafficStream supplier"); - if (stopReadingRef.get()) { - break; - } - this.nextChunkFutureRef.set(trafficChunkStream - .readNextTrafficStreamChunk(topLevelContext::createReadChunkContext)); - List trafficStreams = null; - try { - trafficStreams = this.nextChunkFutureRef.get().get(); - } catch (ExecutionException ex) { - if (ex.getCause() instanceof EOFException) { - log.atWarn().setCause(ex.getCause()).setMessage("Got an EOF on the stream. " + - "Done reading traffic streams.").log(); - break; - } else { - log.atWarn().setCause(ex).setMessage("Done reading traffic streams due to exception.").log(); - throw ex.getCause(); - } - } - if (log.isInfoEnabled()) { - Optional.of(trafficStreams.stream() - .map(ts -> TrafficStreamUtils.summarizeTrafficStream(ts.getStream())) - .collect(Collectors.joining(";"))) - .filter(s -> !s.isEmpty()) - .ifPresent(s -> log.atInfo().log("TrafficStream Summary: {" + s + "}")); - } - trafficStreams.forEach(trafficToHttpTransactionAccumulator::accept); - } - } } diff --git a/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/TrafficReplayerTopLevel.java b/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/TrafficReplayerTopLevel.java new file mode 100644 index 000000000..f202cfc58 --- /dev/null +++ b/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/TrafficReplayerTopLevel.java @@ -0,0 +1,690 @@ +package org.opensearch.migrations.replay; + +import io.netty.buffer.Unpooled; +import io.netty.handler.ssl.SslContext; +import io.netty.handler.ssl.SslContextBuilder; +import io.netty.handler.ssl.util.InsecureTrustManagerFactory; +import lombok.AllArgsConstructor; +import lombok.Lombok; +import lombok.NonNull; +import lombok.SneakyThrows; +import lombok.extern.slf4j.Slf4j; +import org.opensearch.migrations.replay.datahandlers.IPacketFinalizingConsumer; +import org.opensearch.migrations.replay.datatypes.HttpRequestTransformationStatus; +import org.opensearch.migrations.replay.datatypes.ITrafficStreamKey; +import org.opensearch.migrations.replay.datatypes.TransformedPackets; +import org.opensearch.migrations.replay.datatypes.UniqueReplayerRequestKey; +import org.opensearch.migrations.replay.tracing.IReplayContexts; +import org.opensearch.migrations.replay.tracing.IRootReplayerContext; +import org.opensearch.migrations.replay.traffic.source.BlockingTrafficSource; +import org.opensearch.migrations.replay.traffic.source.ITrafficCaptureSource; +import org.opensearch.migrations.replay.traffic.source.ITrafficStreamWithKey; +import org.opensearch.migrations.replay.traffic.source.TrafficStreamLimiter; +import org.opensearch.migrations.replay.util.DiagnosticTrackableCompletableFuture; +import org.opensearch.migrations.replay.util.StringTrackableCompletableFuture; +import org.opensearch.migrations.trafficcapture.protos.TrafficStreamUtils; +import org.opensearch.migrations.transform.IAuthTransformerFactory; +import org.opensearch.migrations.transform.IJsonTransformer; +import org.slf4j.event.Level; +import org.slf4j.spi.LoggingEventBuilder; + +import javax.net.ssl.SSLException; +import java.io.EOFException; +import java.net.URI; +import java.time.Duration; +import java.time.Instant; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; +import java.util.function.Supplier; +import java.util.stream.Collectors; +import java.util.stream.Stream; + +@Slf4j +public class TrafficReplayerTopLevel implements AutoCloseable { + public static final String TARGET_CONNECTION_POOL_NAME = "targetConnectionPool"; + public static final int MAX_ITEMS_TO_SHOW_FOR_LEFTOVER_WORK_AT_INFO_LEVEL = 10; + + public static AtomicInteger targetConnectionPoolUniqueCounter = new AtomicInteger(); + + private final PacketToTransformingHttpHandlerFactory inputRequestTransformerFactory; + private final ClientConnectionPool clientConnectionPool; + private final TrafficStreamLimiter liveTrafficStreamLimiter; + private final AtomicInteger successfulRequestCount; + private final AtomicInteger exceptionRequestCount; + public final IRootReplayerContext topLevelContext; + private final ConcurrentHashMap> requestToFinalWorkFuturesMap; + + private final AtomicBoolean stopReadingRef; + private final AtomicReference> allRemainingWorkFutureOrShutdownSignalRef; + private final AtomicReference shutdownReasonRef; + private final AtomicReference> shutdownFutureRef; + private final AtomicReference>> nextChunkFutureRef; + + + public TrafficReplayerTopLevel(IRootReplayerContext context, + URI serverUri, + IAuthTransformerFactory authTransformerFactory, + boolean allowInsecureConnections) + throws SSLException { + this(context, serverUri, authTransformerFactory, allowInsecureConnections, + new TransformationLoader().getTransformerFactoryLoader(serverUri.getHost())); + } + + public TrafficReplayerTopLevel(IRootReplayerContext context, + URI serverUri, + IAuthTransformerFactory authTransformerFactory, + boolean allowInsecureConnections, + IJsonTransformer jsonTransformer) + throws SSLException { + this(context, serverUri, authTransformerFactory, allowInsecureConnections, 0, + 1024, + jsonTransformer); + } + + public TrafficReplayerTopLevel(IRootReplayerContext context, + URI serverUri, + IAuthTransformerFactory authTransformerFactory, + boolean allowInsecureConnections, + int numSendingThreads, + int maxConcurrentOutstandingRequests, + IJsonTransformer jsonTransformer) + throws SSLException { + this(context, serverUri, authTransformerFactory, allowInsecureConnections, + numSendingThreads, maxConcurrentOutstandingRequests, + jsonTransformer, + getTargetConnectionPoolName(targetConnectionPoolUniqueCounter.getAndIncrement())); + } + + private static String getTargetConnectionPoolName(int i) { + return TARGET_CONNECTION_POOL_NAME + (i == 0 ? "" : Integer.toString(i)); + } + + public TrafficReplayerTopLevel(IRootReplayerContext context, + URI serverUri, + IAuthTransformerFactory authTransformer, + boolean allowInsecureConnections, + int numSendingThreads, + int maxConcurrentOutstandingRequests, + IJsonTransformer jsonTransformer, + String clientThreadNamePrefix) + throws SSLException + { + this.topLevelContext = context; + if (serverUri.getPort() < 0) { + throw new IllegalArgumentException("Port not present for URI: "+serverUri); + } + if (serverUri.getHost() == null) { + throw new IllegalArgumentException("Hostname not present for URI: "+serverUri); + } + if (serverUri.getScheme() == null) { + throw new IllegalArgumentException("Scheme (http|https) is not present for URI: "+serverUri); + } + inputRequestTransformerFactory = new PacketToTransformingHttpHandlerFactory(jsonTransformer, authTransformer); + clientConnectionPool = new ClientConnectionPool(serverUri, + loadSslContext(serverUri, allowInsecureConnections), clientThreadNamePrefix, numSendingThreads); + requestToFinalWorkFuturesMap = new ConcurrentHashMap<>(); + successfulRequestCount = new AtomicInteger(); + exceptionRequestCount = new AtomicInteger(); + liveTrafficStreamLimiter = new TrafficStreamLimiter(maxConcurrentOutstandingRequests); + allRemainingWorkFutureOrShutdownSignalRef = new AtomicReference<>(); + shutdownReasonRef = new AtomicReference<>(); + shutdownFutureRef = new AtomicReference<>(); + nextChunkFutureRef = new AtomicReference<>(); + stopReadingRef = new AtomicBoolean(); + } + + private static SslContext loadSslContext(URI serverUri, boolean allowInsecureConnections) throws SSLException { + if (serverUri.getScheme().equalsIgnoreCase("https")) { + var sslContextBuilder = SslContextBuilder.forClient(); + if (allowInsecureConnections) { + sslContextBuilder.trustManager(InsecureTrustManagerFactory.INSTANCE); + } + return sslContextBuilder.build(); + } else { + return null; + } + } + + public void setupRunAndWaitForReplayToFinish(Duration observedPacketConnectionTimeout, + BlockingTrafficSource trafficSource, + TimeShifter timeShifter, + Consumer resultTupleConsumer) + throws InterruptedException, ExecutionException { + + var senderOrchestrator = new RequestSenderOrchestrator(clientConnectionPool); + var replayEngine = new ReplayEngine(senderOrchestrator, trafficSource, timeShifter); + + CapturedTrafficToHttpTransactionAccumulator trafficToHttpTransactionAccumulator = + new CapturedTrafficToHttpTransactionAccumulator(observedPacketConnectionTimeout, + "(see command line option " + + TrafficReplayer.PACKET_TIMEOUT_SECONDS_PARAMETER_NAME + ")", + new TrafficReplayerAccumulationCallbacks(replayEngine, resultTupleConsumer, trafficSource)); + try { + pullCaptureFromSourceToAccumulator(trafficSource, trafficToHttpTransactionAccumulator); + } catch (InterruptedException ex) { + throw ex; + } catch (Exception e) { + log.atWarn().setCause(e).setMessage("Terminating runReplay due to exception").log(); + throw e; + } finally { + trafficToHttpTransactionAccumulator.close(); + wrapUpWorkAndEmitSummary(replayEngine, trafficToHttpTransactionAccumulator); + assert shutdownFutureRef.get() != null || requestToFinalWorkFuturesMap.isEmpty() : + "expected to wait for all the in flight requests to fully flush and self destruct themselves"; + } + } + + protected void wrapUpWorkAndEmitSummary(ReplayEngine replayEngine, + CapturedTrafficToHttpTransactionAccumulator trafficToHttpTransactionAccumulator) + throws ExecutionException, InterruptedException { + final var primaryLogLevel = Level.INFO; + final var secondaryLogLevel = Level.WARN; + var logLevel = primaryLogLevel; + for (var timeout = Duration.ofSeconds(60); ; timeout = timeout.multipliedBy(2)) { + if (shutdownFutureRef.get() != null) { + log.warn("Not waiting for work because the TrafficReplayer is shutting down."); + break; + } + try { + waitForRemainingWork(logLevel, timeout); + break; + } catch (TimeoutException e) { + log.atLevel(logLevel).log("Timed out while waiting for the remaining " + + "requests to be finalized..."); + logLevel = secondaryLogLevel; + } + } + if (!requestToFinalWorkFuturesMap.isEmpty() || exceptionRequestCount.get() > 0) { + log.atWarn().setMessage("{} in-flight requests being dropped due to pending shutdown; " + + "{} requests to the target threw an exception; " + + "{} requests were successfully processed.") + .addArgument(requestToFinalWorkFuturesMap.size()) + .addArgument(exceptionRequestCount.get()) + .addArgument(successfulRequestCount.get()) + .log(); + } else { + log.info(successfulRequestCount.get() + " requests were successfully processed."); + } + log.info("# of connections created: {}; # of requests on reused keep-alive connections: {}; " + + "# of expired connections: {}; # of connections closed: {}; " + + "# of connections terminated upon accumulator termination: {}", + trafficToHttpTransactionAccumulator.numberOfConnectionsCreated(), + trafficToHttpTransactionAccumulator.numberOfRequestsOnReusedConnections(), + trafficToHttpTransactionAccumulator.numberOfConnectionsExpired(), + trafficToHttpTransactionAccumulator.numberOfConnectionsClosed(), + trafficToHttpTransactionAccumulator.numberOfRequestsTerminatedUponAccumulatorClose() + ); + } + + public void setupRunAndWaitForReplayWithShutdownChecks(Duration observedPacketConnectionTimeout, + BlockingTrafficSource trafficSource, + TimeShifter timeShifter, + Consumer resultTupleConsumer) + throws TrafficReplayer.TerminationException, ExecutionException, InterruptedException { + try { + setupRunAndWaitForReplayToFinish(observedPacketConnectionTimeout, trafficSource, + timeShifter, resultTupleConsumer); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new TrafficReplayer.TerminationException(shutdownReasonRef.get(), e); + } catch (Throwable t) { + throw new TrafficReplayer.TerminationException(shutdownReasonRef.get(), t); + } + if (shutdownReasonRef.get() != null) { + throw new TrafficReplayer.TerminationException(shutdownReasonRef.get(), null); + } + // if nobody has run shutdown yet, do so now so that we can tear down the netty resources + shutdown(null).get(); // if somebody already HAD run shutdown, it will return the future already created + } + + @AllArgsConstructor + class TrafficReplayerAccumulationCallbacks implements AccumulationCallbacks { + private final ReplayEngine replayEngine; + private Consumer resultTupleConsumer; + private ITrafficCaptureSource trafficCaptureSource; + + @Override + public Consumer + onRequestReceived(@NonNull IReplayContexts.IReplayerHttpTransactionContext ctx, + @NonNull HttpMessageAndTimestamp request) { + replayEngine.setFirstTimestamp(request.getFirstPacketTimestamp()); + + var allWorkFinishedForTransaction = + new StringTrackableCompletableFuture(new CompletableFuture<>(), + ()->"waiting for work to be queued and run through TrafficStreamLimiter"); + var requestPushFuture = new StringTrackableCompletableFuture( + new CompletableFuture<>(), () -> "Waiting to get response from target"); + var requestKey = ctx.getReplayerRequestKey(); + liveTrafficStreamLimiter.queueWork(1, ctx, wi -> { + transformAndSendRequest(replayEngine, request, ctx).future.whenComplete((v,t)->{ + liveTrafficStreamLimiter.doneProcessing(wi); + if (t != null) { + requestPushFuture.future.completeExceptionally(t); + } else { + requestPushFuture.future.complete(v); + } + }); + }); + if (!allWorkFinishedForTransaction.future.isDone()) { + log.trace("Adding " + requestKey + " to targetTransactionInProgressMap"); + requestToFinalWorkFuturesMap.put(requestKey, allWorkFinishedForTransaction); + if (allWorkFinishedForTransaction.future.isDone()) { + requestToFinalWorkFuturesMap.remove(requestKey); + } + } + + return rrPair -> + requestPushFuture.map(f -> f.handle((v, t) -> { + log.atInfo().setMessage(() -> "Done receiving captured stream for " + ctx + + ":" + rrPair.requestData).log(); + log.atTrace().setMessage(() -> + "Summary response value for " + requestKey + " returned=" + v).log(); + return handleCompletedTransaction(ctx, rrPair, v, t); + }), () -> "logging summary") + .whenComplete((v,t)->{ + if (t != null) { + allWorkFinishedForTransaction.future.completeExceptionally(t); + } else { + allWorkFinishedForTransaction.future.complete(null); + } + }, ()->""); + } + + Void handleCompletedTransaction(@NonNull IReplayContexts.IReplayerHttpTransactionContext context, + RequestResponsePacketPair rrPair, + TransformedTargetRequestAndResponse summary, Throwable t) { + try (var httpContext = rrPair.getHttpTransactionContext()) { + // if this comes in with a serious Throwable (not an Exception), don't bother + // packaging it up and calling the callback. + // Escalate it up out handling stack and shutdown. + if (t == null || t instanceof Exception) { + try (var tupleHandlingContext = httpContext.createTupleContext()) { + packageAndWriteResponse(tupleHandlingContext, resultTupleConsumer, + rrPair, summary, (Exception) t); + } + commitTrafficStreams(rrPair.completionStatus, rrPair.trafficStreamKeysBeingHeld); + return null; + } else { + log.atError().setCause(t).setMessage(() -> "Throwable passed to handle() for " + context + + ". Rethrowing.").log(); + throw Lombok.sneakyThrow(t); + } + } catch (Error error) { + log.atError() + .setCause(error) + .setMessage(() -> "Caught error and initiating TrafficReplayer shutdown") + .log(); + shutdown(error); + throw error; + } catch (Exception e) { + log.atError() + .setMessage("Unexpected exception while sending the " + + "aggregated response and context for {} to the callback. " + + "Proceeding, but the tuple receiver context may be compromised.") + .addArgument(context) + .setCause(e) + .log(); + throw e; + } finally { + var requestKey = context.getReplayerRequestKey(); + requestToFinalWorkFuturesMap.remove(requestKey); + log.trace("removed rrPair.requestData to " + + "targetTransactionInProgressMap for " + + requestKey); + } + } + + @Override + public void onTrafficStreamsExpired(RequestResponsePacketPair.ReconstructionStatus status, + @NonNull IReplayContexts.IChannelKeyContext ctx, + @NonNull List trafficStreamKeysBeingHeld) { + commitTrafficStreams(status, trafficStreamKeysBeingHeld); + } + + @SneakyThrows + private void commitTrafficStreams(RequestResponsePacketPair.ReconstructionStatus status, + List trafficStreamKeysBeingHeld) { + commitTrafficStreams(status != RequestResponsePacketPair.ReconstructionStatus.CLOSED_PREMATURELY, + trafficStreamKeysBeingHeld); + } + + @SneakyThrows + private void commitTrafficStreams(boolean shouldCommit, + List trafficStreamKeysBeingHeld) { + if (shouldCommit && trafficStreamKeysBeingHeld != null) { + for (var tsk : trafficStreamKeysBeingHeld) { + tsk.getTrafficStreamsContext().close(); + trafficCaptureSource.commitTrafficStream(tsk); + } + } + } + + @Override + public void onConnectionClose(int channelInteractionNum, + @NonNull IReplayContexts.IChannelKeyContext ctx, int channelSessionNumber, + RequestResponsePacketPair.ReconstructionStatus status, + @NonNull Instant timestamp, @NonNull List trafficStreamKeysBeingHeld) { + replayEngine.setFirstTimestamp(timestamp); + var cf = replayEngine.closeConnection(channelInteractionNum, ctx, channelSessionNumber, timestamp); + cf.map(f->f.whenComplete((v,t)->{ + commitTrafficStreams(status, trafficStreamKeysBeingHeld); + }), ()->"closing the channel in the ReplayEngine"); + } + + @Override + public void onTrafficStreamIgnored(@NonNull IReplayContexts.ITrafficStreamsLifecycleContext ctx) { + commitTrafficStreams(true, List.of(ctx.getTrafficStreamKey())); + } + + private TransformedTargetRequestAndResponse + packageAndWriteResponse(IReplayContexts.ITupleHandlingContext tupleHandlingContext, + Consumer tupleWriter, + RequestResponsePacketPair rrPair, + TransformedTargetRequestAndResponse summary, + Exception t) { + log.trace("done sending and finalizing data to the packet handler"); + + try (var requestResponseTuple = getSourceTargetCaptureTuple(tupleHandlingContext, rrPair, summary, t)) { + log.atInfo().setMessage(()->"Source/Target Request/Response tuple: " + requestResponseTuple).log(); + tupleWriter.accept(requestResponseTuple); + } + + if (t != null) { throw new CompletionException(t); } + if (summary.getError() != null) { + log.atInfo().setCause(summary.getError()).setMessage("Exception for {}: ") + .addArgument(tupleHandlingContext).log(); + exceptionRequestCount.incrementAndGet(); + } else if (summary.getTransformationStatus() == HttpRequestTransformationStatus.ERROR) { + log.atInfo().setCause(summary.getError()).setMessage("Unknown error transforming {}: ") + .addArgument(tupleHandlingContext).log(); + exceptionRequestCount.incrementAndGet(); + } else { + successfulRequestCount.incrementAndGet(); + } + return summary; + } + } + + protected void waitForRemainingWork(Level logLevel, @NonNull Duration timeout) + throws ExecutionException, InterruptedException, TimeoutException { + + if (!liveTrafficStreamLimiter.isStopped()) { + var streamLimiterHasRunEverything = new CompletableFuture(); + liveTrafficStreamLimiter.queueWork(1, null, wi -> { + streamLimiterHasRunEverything.complete(null); + liveTrafficStreamLimiter.doneProcessing(wi); + }); + streamLimiterHasRunEverything.get(timeout.toMillis(), TimeUnit.MILLISECONDS); + } + + Map.Entry>[] + allRemainingWorkArray = requestToFinalWorkFuturesMap.entrySet().toArray(Map.Entry[]::new); + writeStatusLogsForRemainingWork(logLevel, allRemainingWorkArray); + + // remember, this block is ONLY for the leftover items. Lots of other items have been processed + // and were removed from the live map (hopefully) + DiagnosticTrackableCompletableFuture[] allCompletableFuturesArray = + Arrays.stream(allRemainingWorkArray) + .map(Map.Entry::getValue).toArray(DiagnosticTrackableCompletableFuture[]::new); + var allWorkFuture = StringTrackableCompletableFuture.allOf(allCompletableFuturesArray, + () -> "TrafficReplayer.AllWorkFinished"); + try { + if (allRemainingWorkFutureOrShutdownSignalRef.compareAndSet(null, allWorkFuture)) { + allWorkFuture.get(timeout); + } else { + handleAlreadySetFinishedSignal(); + } + } catch (TimeoutException e) { + var didCancel = allWorkFuture.future.cancel(true); + if (!didCancel) { + assert allWorkFuture.future.isDone() : "expected future to have finished if cancel didn't succeed"; + // continue with the rest of the function + } else { + throw e; + } + } finally { + allRemainingWorkFutureOrShutdownSignalRef.set(null); + } + } + + private void handleAlreadySetFinishedSignal() throws InterruptedException, ExecutionException { + try { + var finishedSignal = allRemainingWorkFutureOrShutdownSignalRef.get().future; + assert finishedSignal.isDone() : "Expected this reference to be EITHER the current work futures " + + "or a sentinel value indicating a shutdown has commenced. The signal, when set, should " + + "have been completed at the time that the reference was set"; + finishedSignal.get(); + log.debug("Did shutdown cleanly"); + } catch (ExecutionException e) { + var c = e.getCause(); + if (c instanceof Error) { + throw (Error) c; + } else { + throw e; + } + } catch (Error t) { + log.atError().setCause(t).setMessage(() -> "Not waiting for all work to finish. " + + "The TrafficReplayer is shutting down").log(); + throw t; + } + } + + private static void writeStatusLogsForRemainingWork(Level logLevel, + Map.Entry>[] + allRemainingWorkArray) { + log.atLevel(logLevel).log("All remaining work to wait on " + allRemainingWorkArray.length); + if (log.isInfoEnabled()) { + LoggingEventBuilder loggingEventBuilderToUse = log.isTraceEnabled() ? log.atTrace() : log.atInfo(); + long itemLimit = log.isTraceEnabled() ? Long.MAX_VALUE : MAX_ITEMS_TO_SHOW_FOR_LEFTOVER_WORK_AT_INFO_LEVEL; + loggingEventBuilderToUse.setMessage(() -> " items: " + + Arrays.stream(allRemainingWorkArray) + .map(kvp -> kvp.getKey() + " --> " + + kvp.getValue().formatAsString(TrafficReplayerTopLevel::formatWorkItem)) + .limit(itemLimit) + .collect(Collectors.joining("\n"))) + .log(); + } + } + + private static String formatWorkItem(DiagnosticTrackableCompletableFuture cf) { + try { + var resultValue = cf.get(); + if (resultValue instanceof TransformedTargetRequestAndResponse) { + return "" + ((TransformedTargetRequestAndResponse) resultValue).getTransformationStatus(); + } + return null; + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return "Exception: " + e.getMessage(); + } catch (ExecutionException e) { + return e.getMessage(); + } + } + + private static SourceTargetCaptureTuple + getSourceTargetCaptureTuple(@NonNull IReplayContexts.ITupleHandlingContext tupleHandlingContext, + RequestResponsePacketPair rrPair, + TransformedTargetRequestAndResponse summary, + Exception t) + { + SourceTargetCaptureTuple requestResponseTuple; + if (t != null) { + log.error("Got exception in CompletableFuture callback: ", t); + requestResponseTuple = new SourceTargetCaptureTuple(tupleHandlingContext, rrPair, + new TransformedPackets(), new ArrayList<>(), + HttpRequestTransformationStatus.ERROR, t, Duration.ZERO); + } else { + requestResponseTuple = new SourceTargetCaptureTuple(tupleHandlingContext, rrPair, + summary.requestPackets, + summary.getReceiptTimeAndResponsePackets() + .map(Map.Entry::getValue).collect(Collectors.toList()), + summary.getTransformationStatus(), + summary.getError(), + summary.getResponseDuration() + ); + } + return requestResponseTuple; + } + + public DiagnosticTrackableCompletableFuture + transformAndSendRequest(ReplayEngine replayEngine, HttpMessageAndTimestamp request, + IReplayContexts.IReplayerHttpTransactionContext ctx) { + return transformAndSendRequest(inputRequestTransformerFactory, replayEngine, ctx, + request.getFirstPacketTimestamp(), request.getLastPacketTimestamp(), + request.packetBytes::stream); + } + + public static DiagnosticTrackableCompletableFuture + transformAndSendRequest(PacketToTransformingHttpHandlerFactory inputRequestTransformerFactory, + ReplayEngine replayEngine, + IReplayContexts.IReplayerHttpTransactionContext ctx, + @NonNull Instant start, @NonNull Instant end, + Supplier> packetsSupplier) + { + try { + var transformationCompleteFuture = replayEngine.scheduleTransformationWork(ctx, start, ()-> + transformAllData(inputRequestTransformerFactory.create(ctx), packetsSupplier)); + log.atDebug().setMessage(()->"finalizeRequest future for transformation of " + ctx + + " = " + transformationCompleteFuture).log(); + // It might be safer to chain this work directly inside the scheduleWork call above so that the + // read buffer horizons aren't set after the transformation work finishes, but after the packets + // are fully handled + return transformationCompleteFuture.thenCompose(transformedResult -> + replayEngine.scheduleRequest(ctx, start, end, + transformedResult.transformedOutput.size(), + transformedResult.transformedOutput.streamRetained()) + .map(future->future.thenApply(t -> + new TransformedTargetRequestAndResponse(transformedResult.transformedOutput, + t, transformedResult.transformationStatus, t.error)), + ()->"(if applicable) packaging transformed result into a completed TransformedTargetRequestAndResponse object") + .map(future->future.exceptionally(t -> + new TransformedTargetRequestAndResponse(transformedResult.transformedOutput, + transformedResult.transformationStatus, t)), + ()->"(if applicable) packaging transformed result into a failed TransformedTargetRequestAndResponse object"), + () -> "transitioning transformed packets onto the wire") + .map(future->future.exceptionally(t->new TransformedTargetRequestAndResponse(null, null, t)), + ()->"Checking for exception out of sending data to the target server"); + } catch (Exception e) { + log.debug("Caught exception in writeToSocket, so failing future"); + return StringTrackableCompletableFuture.failedFuture(e, ()->"TrafficReplayer.writeToSocketAndClose"); + } + } + + private static DiagnosticTrackableCompletableFuture + transformAllData(IPacketFinalizingConsumer packetHandler, Supplier> packetSupplier) { + try { + var logLabel = packetHandler.getClass().getSimpleName(); + var packets = packetSupplier.get().map(Unpooled::wrappedBuffer); + packets.forEach(packetData -> { + log.atDebug().setMessage(() -> logLabel + " sending " + packetData.readableBytes() + + " bytes to the packetHandler").log(); + var consumeFuture = packetHandler.consumeBytes(packetData); + log.atDebug().setMessage(() -> logLabel + " consumeFuture = " + consumeFuture).log(); + }); + log.atDebug().setMessage(() -> logLabel + " done sending bytes, now finalizing the request").log(); + return packetHandler.finalizeRequest(); + } catch (Exception e) { + log.atInfo().setCause(e).setMessage("Encountered an exception while transforming the http request. " + + "The base64 gzipped traffic stream, for later diagnostic purposes, is: " + + Utils.packetsToCompressedTrafficStream(packetSupplier.get())).log(); + throw e; + } + } + + @SneakyThrows + public @NonNull CompletableFuture shutdown(Error error) { + log.atWarn().setCause(error).setMessage(()->"Shutting down " + this).log(); + shutdownReasonRef.compareAndSet(null, error); + if (!shutdownFutureRef.compareAndSet(null, new CompletableFuture<>())) { + log.atError().setMessage(()->"Shutdown was already signaled by {}. " + + "Ignoring this shutdown request due to {}.") + .addArgument(shutdownReasonRef.get()) + .addArgument(error) + .log(); + return shutdownFutureRef.get(); + } + stopReadingRef.set(true); + liveTrafficStreamLimiter.close(); + + var nettyShutdownFuture = clientConnectionPool.shutdownNow(); + nettyShutdownFuture.whenComplete((v,t) -> { + if (t != null) { + shutdownFutureRef.get().completeExceptionally(t); + } else { + shutdownFutureRef.get().complete(null); + } + }); + Optional.ofNullable(this.nextChunkFutureRef.get()).ifPresent(f->f.cancel(true)); + var shutdownWasSignalledFuture = error == null ? + StringTrackableCompletableFuture.completedFuture(null, ()->"TrafficReplayer shutdown") : + StringTrackableCompletableFuture.failedFuture(error, ()->"TrafficReplayer shutdown"); + while (!allRemainingWorkFutureOrShutdownSignalRef.compareAndSet(null, shutdownWasSignalledFuture)) { + var otherRemainingWorkObj = allRemainingWorkFutureOrShutdownSignalRef.get(); + if (otherRemainingWorkObj != null) { + otherRemainingWorkObj.future.cancel(true); + break; + } + } + var shutdownFuture = shutdownFutureRef.get(); + log.atWarn().setMessage(()->"Shutdown setup has been initiated").log(); + return shutdownFuture; + } + + + @Override + public void close() throws Exception { + shutdown(null).get(); + } + + @SneakyThrows + public void pullCaptureFromSourceToAccumulator( + ITrafficCaptureSource trafficChunkStream, + CapturedTrafficToHttpTransactionAccumulator trafficToHttpTransactionAccumulator) + throws InterruptedException { + while (true) { + log.trace("Reading next chunk from TrafficStream supplier"); + if (stopReadingRef.get()) { + break; + } + this.nextChunkFutureRef.set(trafficChunkStream + .readNextTrafficStreamChunk(topLevelContext::createReadChunkContext)); + List trafficStreams = null; + try { + trafficStreams = this.nextChunkFutureRef.get().get(); + } catch (ExecutionException ex) { + if (ex.getCause() instanceof EOFException) { + log.atWarn().setCause(ex.getCause()).setMessage("Got an EOF on the stream. " + + "Done reading traffic streams.").log(); + break; + } else { + log.atWarn().setCause(ex).setMessage("Done reading traffic streams due to exception.").log(); + throw ex.getCause(); + } + } + if (log.isInfoEnabled()) { + Optional.of(trafficStreams.stream() + .map(ts -> TrafficStreamUtils.summarizeTrafficStream(ts.getStream())) + .collect(Collectors.joining(";"))) + .filter(s -> !s.isEmpty()) + .ifPresent(s -> log.atInfo().log("TrafficStream Summary: {" + s + "}")); + } + trafficStreams.forEach(trafficToHttpTransactionAccumulator::accept); + } + } +} diff --git a/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/TrafficReplayerTest.java b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/TrafficReplayerTest.java index 5677e47f6..561e41275 100644 --- a/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/TrafficReplayerTest.java +++ b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/TrafficReplayerTest.java @@ -152,8 +152,8 @@ static byte[] synthesizeTrafficStreamsIntoByteArray(Instant timestamp, int numSt @Test public void testReader() throws Exception { - try (var tr = new TrafficReplayer(rootContext, - new URI("http://localhost:9200"), null, null, false)) { + try (var tr = new TrafficReplayerTopLevel(rootContext, + new URI("http://localhost:9200"), null, false)) { List> byteArrays = new ArrayList<>(); CapturedTrafficToHttpTransactionAccumulator trafficAccumulator = new CapturedTrafficToHttpTransactionAccumulator(Duration.ofSeconds(30), null, @@ -208,8 +208,9 @@ public void onTrafficStreamIgnored(@NonNull IReplayContexts.ITrafficStreamsLifec @Test public void testCapturedReadsAfterCloseAreHandledAsNew() throws Exception { - try (var tr = new TrafficReplayer(rootContext, - new URI("http://localhost:9200"), null, null, false)) { + try (var tr = new TrafficReplayerTopLevel(rootContext, + new URI("http://localhost:9200"), null, false, + new TransformationLoader().getTransformerFactoryLoader("localhost"))) { List> byteArrays = new ArrayList<>(); var remainingAccumulations = new AtomicInteger(); CapturedTrafficToHttpTransactionAccumulator trafficAccumulator = diff --git a/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/TransformationLoaderTest.java b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/TransformationLoaderTest.java index 98e25bd5a..f1967f796 100644 --- a/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/TransformationLoaderTest.java +++ b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/TransformationLoaderTest.java @@ -46,7 +46,7 @@ public void testMisconfiguration() throws Exception { @Test public void testThatNoConfigMeansNoThrow() throws Exception { var transformer = Assertions.assertDoesNotThrow(()->new TransformationLoader() - .getTransformerFactoryLoader("localhost", null, null)); + .getTransformerFactoryLoader("localhost")); Assertions.assertNotNull(transformer); var origDoc = parseAsMap(SampleContents.loadSampleJsonRequestAsString()); Assertions.assertNotNull(transformer.transformJson(origDoc)); diff --git a/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/datahandlers/NettyPacketToHttpConsumerTest.java b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/datahandlers/NettyPacketToHttpConsumerTest.java index 8fd447f82..13e764732 100644 --- a/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/datahandlers/NettyPacketToHttpConsumerTest.java +++ b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/datahandlers/NettyPacketToHttpConsumerTest.java @@ -23,6 +23,7 @@ import org.opensearch.migrations.replay.RequestSenderOrchestrator; import org.opensearch.migrations.replay.TimeShifter; import org.opensearch.migrations.replay.TrafficReplayer; +import org.opensearch.migrations.replay.TrafficReplayerTopLevel; import org.opensearch.migrations.replay.TransformationLoader; import org.opensearch.migrations.replay.datatypes.ConnectionReplaySession; import org.opensearch.migrations.replay.traffic.source.BufferedFlowController; @@ -235,7 +236,7 @@ public void testThatConnectionsAreKeptAliveAndShared(boolean useTls, boolean lar for (int j = 0; j < 2; ++j) { for (int i = 0; i < 2; ++i) { var ctx = rootContext.getTestConnectionRequestContext("TEST_" + i, j); - var requestFinishFuture = TrafficReplayer.transformAndSendRequest(transformingHttpHandlerFactory, + var requestFinishFuture = TrafficReplayerTopLevel.transformAndSendRequest(transformingHttpHandlerFactory, sendingFactory, ctx, Instant.now(), Instant.now(), () -> Stream.of(EXPECTED_REQUEST_STRING.getBytes(StandardCharsets.UTF_8))); log.info("requestFinishFuture=" + requestFinishFuture); diff --git a/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/FullReplayerWithTracingChecksTest.java b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/FullReplayerWithTracingChecksTest.java index c763482e7..68a94338a 100644 --- a/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/FullReplayerWithTracingChecksTest.java +++ b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/FullReplayerWithTracingChecksTest.java @@ -14,6 +14,8 @@ import org.opensearch.migrations.replay.TestHttpServerContext; import org.opensearch.migrations.replay.TimeShifter; import org.opensearch.migrations.replay.TrafficReplayer; +import org.opensearch.migrations.replay.TrafficReplayerTopLevel; +import org.opensearch.migrations.replay.TransformationLoader; import org.opensearch.migrations.replay.traffic.source.ArrayCursorTrafficCaptureSource; import org.opensearch.migrations.replay.traffic.source.ArrayCursorTrafficSourceContext; import org.opensearch.migrations.replay.traffic.source.BlockingTrafficSource; @@ -89,10 +91,12 @@ public void testStreamWithRequestsWithCloseIsCommittedOnce(int numRequests) thro new ArrayCursorTrafficSourceContext(List.of(trafficStream))); var tuplesReceived = new HashSet(); - try (var tr = new TrafficReplayer(rootContext, httpServer.localhostEndpoint(), null, - new StaticAuthTransformerFactory("TEST"), null, + var serverUri = httpServer.localhostEndpoint(); + try (var tr = new TrafficReplayerTopLevel(rootContext, serverUri, + new StaticAuthTransformerFactory("TEST"), true, 10, 10 * 1024, - "targetConnectionPool for testStreamWithRequestsWithCloseIsCommittedOnce"); + new TransformationLoader() + .getTransformerFactoryLoader(serverUri.getHost())); var blockingTrafficSource = new BlockingTrafficSource(trafficSource, Duration.ofMinutes(2))) { tr.setupRunAndWaitForReplayToFinish(Duration.ofSeconds(70), blockingTrafficSource, new TimeShifter(10 * 1000), diff --git a/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/FullTrafficReplayerTest.java b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/FullTrafficReplayerTest.java index 643a0b5d2..1afdcec7b 100644 --- a/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/FullTrafficReplayerTest.java +++ b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/FullTrafficReplayerTest.java @@ -15,6 +15,8 @@ import org.opensearch.migrations.replay.TestHttpServerContext; import org.opensearch.migrations.replay.TimeShifter; import org.opensearch.migrations.replay.TrafficReplayer; +import org.opensearch.migrations.replay.TrafficReplayerTopLevel; +import org.opensearch.migrations.replay.TransformationLoader; import org.opensearch.migrations.replay.tracing.IRootReplayerContext; import org.opensearch.migrations.replay.traffic.generator.ExhaustiveTrafficStreamGenerator; import org.opensearch.migrations.replay.datatypes.ITrafficStreamKey; @@ -68,23 +70,22 @@ public class FullTrafficReplayerTest extends InstrumentationTest { public static final String TEST_CONNECTION_ID = "testConnectionId"; public static final String DUMMY_URL_THAT_WILL_NEVER_BE_CONTACTED = "http://localhost:9999/"; - protected static class TrafficReplayerWithWaitOnClose extends TrafficReplayer { + protected static class TrafficReplayerWithWaitOnClose extends TrafficReplayerTopLevel { private final Duration maxWaitTime; public TrafficReplayerWithWaitOnClose(Duration maxWaitTime, IRootReplayerContext context, URI serverUri, - String fullTransformerConfig, IAuthTransformerFactory authTransformerFactory, - String userAgent, boolean allowInsecureConnections, int numSendingThreads, int maxConcurrentOutstandingRequests, + IJsonTransformer jsonTransformer, String targetConnectionPoolName) throws SSLException { - super(context, serverUri, fullTransformerConfig, authTransformerFactory, userAgent, + super(context, serverUri, authTransformerFactory, allowInsecureConnections, numSendingThreads, maxConcurrentOutstandingRequests, - targetConnectionPoolName); + jsonTransformer, targetConnectionPoolName); this.maxWaitTime = maxWaitTime; } @@ -145,9 +146,11 @@ public void testLongRequestEndingAfterEOFStillCountsCorrectly() throws Throwable (rc, threadPrefix) -> { try { return new TrafficReplayerWithWaitOnClose(Duration.ofSeconds(600), - rc, httpServer.localhostEndpoint(), null, - new StaticAuthTransformerFactory("TEST"), null, - true, 1, 1, threadPrefix); + rc, httpServer.localhostEndpoint(), + new StaticAuthTransformerFactory("TEST"), + true, 1, 1, + new TransformationLoader().getTransformerFactoryLoader("localhost"), + threadPrefix); } catch (SSLException e) { throw new RuntimeException(e); } @@ -179,10 +182,12 @@ public void testSingleStreamWithCloseIsCommitted() throws Throwable { 0, (rc, threadPrefix) -> { try { - return new TrafficReplayerWithWaitOnClose(Duration.ofSeconds(10), - rc, httpServer.localhostEndpoint(), null, - new StaticAuthTransformerFactory("TEST"), null, - true, 1, 1, threadPrefix); + return new TrafficReplayerWithWaitOnClose(Duration.ofSeconds(600), + rc, httpServer.localhostEndpoint(), + new StaticAuthTransformerFactory("TEST"), + true, 1, 1, + new TransformationLoader().getTransformerFactoryLoader("localhost"), + threadPrefix); } catch (SSLException e) { throw new RuntimeException(e); } @@ -238,9 +243,12 @@ public CommitResult commitTrafficStream(ITrafficStreamKey trafficStreamKey) thro numExpectedRequests, (rc, threadPrefix) -> { try { - return new TrafficReplayer(rc, httpServer.localhostEndpoint(), null, - new StaticAuthTransformerFactory("TEST"), null, - true, 1, 1, threadPrefix); + return new TrafficReplayerWithWaitOnClose(Duration.ofSeconds(600), + rc, httpServer.localhostEndpoint(), + new StaticAuthTransformerFactory("TEST"), + true, 1, 1, + new TransformationLoader().getTransformerFactoryLoader("localhost"), + threadPrefix); } catch (SSLException e) { throw new RuntimeException(e); } @@ -264,7 +272,7 @@ public void makeSureThatCollateralDamageDoesntFreezeTests() throws Throwable { throw new RuntimeException(e); } } - }, TrafficReplayer.TARGET_CONNECTION_POOL_NAME + " Just to break a test"); + }, TrafficReplayerTopLevel.TARGET_CONNECTION_POOL_NAME + " Just to break a test"); imposterThread.start(); try { diff --git a/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/SlowAndExpiredTrafficStreamBecomesTwoTargetChannelsTest.java b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/SlowAndExpiredTrafficStreamBecomesTwoTargetChannelsTest.java index 23ac90199..dba92f499 100644 --- a/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/SlowAndExpiredTrafficStreamBecomesTwoTargetChannelsTest.java +++ b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/SlowAndExpiredTrafficStreamBecomesTwoTargetChannelsTest.java @@ -9,6 +9,8 @@ import org.junit.jupiter.api.Test; import org.opensearch.migrations.replay.TimeShifter; import org.opensearch.migrations.replay.TrafficReplayer; +import org.opensearch.migrations.replay.TrafficReplayerTopLevel; +import org.opensearch.migrations.replay.TransformationLoader; import org.opensearch.migrations.replay.traffic.source.ArrayCursorTrafficCaptureSource; import org.opensearch.migrations.replay.traffic.source.ArrayCursorTrafficSourceContext; import org.opensearch.migrations.replay.traffic.source.BlockingTrafficSource; @@ -96,9 +98,10 @@ public void test() throws Exception { var trafficSource = new BlockingTrafficSource(arraySource, Duration.ofSeconds(SPACING_SECONDS)); try (var httpServer = SimpleNettyHttpServer.makeServer(false, Duration.ofMillis(200), responseTracker); - var replayer = new TrafficReplayer(rc, httpServer.localhostEndpoint(), null, - new StaticAuthTransformerFactory("TEST"), null, + var replayer = new TrafficReplayerTopLevel(rc, httpServer.localhostEndpoint(), + new StaticAuthTransformerFactory("TEST"), true, 1, 1, + new TransformationLoader().getTransformerFactoryLoader("localhost"), "targetConnectionPool for SlowAndExpiredTrafficStreamBecomesTwoTargetChannelsTest")) { new Thread(()->responseTracker.onCountDownFinished(Duration.ofSeconds(10), ()->replayer.shutdown(null).join())); diff --git a/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/TrafficReplayerRunner.java b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/TrafficReplayerRunner.java index cabdec074..07cfd2689 100644 --- a/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/TrafficReplayerRunner.java +++ b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/TrafficReplayerRunner.java @@ -7,6 +7,8 @@ import org.opensearch.migrations.replay.SourceTargetCaptureTuple; import org.opensearch.migrations.replay.TimeShifter; import org.opensearch.migrations.replay.TrafficReplayer; +import org.opensearch.migrations.replay.TrafficReplayerTopLevel; +import org.opensearch.migrations.replay.TransformationLoader; import org.opensearch.migrations.replay.datatypes.ISourceTrafficChannelKey; import org.opensearch.migrations.replay.tracing.IRootReplayerContext; import org.opensearch.migrations.replay.traffic.source.BlockingTrafficSource; @@ -63,9 +65,10 @@ public static void runReplayer(int numExpectedRequests, runReplayer(numExpectedRequests, (rootContext, targetConnectionPoolPrefix) -> { try { - return new TrafficReplayer(rootContext, endpoint, null, - new StaticAuthTransformerFactory("TEST"), null, + return new TrafficReplayerTopLevel(rootContext, endpoint, + new StaticAuthTransformerFactory("TEST"), true, 10, 10*1024, + new TransformationLoader().getTransformerFactoryLoader(endpoint.getHost()), targetConnectionPoolPrefix); } catch (SSLException e) { throw new RuntimeException(e); @@ -75,7 +78,7 @@ public static void runReplayer(int numExpectedRequests, } public static void runReplayer(int numExpectedRequests, - BiFunction trafficReplayerFactory, + BiFunction trafficReplayerFactory, Supplier> tupleListenerSupplier, Supplier rootContextSupplier, Function trafficSourceFactory, @@ -93,7 +96,7 @@ public static void runReplayer(int numExpectedRequests, int runNumber = runNumberRef.get(); var counter = new AtomicInteger(); var tupleReceiver = tupleListenerSupplier.get(); - String targetConnectionPoolPrefix = TrafficReplayer.TARGET_CONNECTION_POOL_NAME + " run: " + runNumber; + String targetConnectionPoolPrefix = TrafficReplayerTopLevel.TARGET_CONNECTION_POOL_NAME + " run: " + runNumber; try (var rootContext = rootContextSupplier.get(); var trafficReplayer = trafficReplayerFactory.apply(rootContext, targetConnectionPoolPrefix)) { runTrafficReplayer(trafficReplayer, ()->trafficSourceFactory.apply(rootContext), @@ -183,7 +186,7 @@ public static void runReplayer(int numExpectedRequests, Assertions.assertEquals(numExpectedRequests, totalUniqueEverReceived.get()); } - private static void runTrafficReplayer(TrafficReplayer trafficReplayer, + private static void runTrafficReplayer(TrafficReplayerTopLevel trafficReplayer, Supplier captureSourceSupplier, Consumer tupleReceiver, TimeShifter timeShifter) throws Exception { From 02027a1ce6ae11b5463b8d3f032665315d8c9ae4 Mon Sep 17 00:00:00 2001 From: Greg Schohn Date: Sat, 6 Apr 2024 15:41:05 -0400 Subject: [PATCH 3/3] Refactor "Core" parts of TrafficReplayerTopLevel out into a separate class. TopLevel deals with the total lifecycle (shutdown) of critical resources while Core does the orchestration between various top-level components. Signed-off-by: Greg Schohn --- .../migrations/replay/TrafficReplayer.java | 4 +- .../replay/TrafficReplayerCore.java | 382 ++++++++++++++++ .../replay/TrafficReplayerTopLevel.java | 415 ++---------------- .../replay/TrafficReplayerTest.java | 6 +- .../FullReplayerWithTracingChecksTest.java | 8 +- .../e2etests/FullTrafficReplayerTest.java | 7 +- ...ficStreamBecomesTwoTargetChannelsTest.java | 4 +- .../e2etests/TrafficReplayerRunner.java | 3 +- 8 files changed, 434 insertions(+), 395 deletions(-) create mode 100644 TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/TrafficReplayerCore.java diff --git a/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/TrafficReplayer.java b/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/TrafficReplayer.java index ea572820f..958d1b182 100644 --- a/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/TrafficReplayer.java +++ b/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/TrafficReplayer.java @@ -284,8 +284,8 @@ public static void main(String[] args) throws Exception { log.atInfo().setMessage(()->"Transformations config string: " + transformerConfig).log(); } var tr = new TrafficReplayerTopLevel(topContext, uri, authTransformer, - params.allowInsecureConnections, params.numClientThreads, params.maxConcurrentRequests, - new TransformationLoader().getTransformerFactoryLoader(uri.getHost(), params.userAgent, transformerConfig)); + new TransformationLoader().getTransformerFactoryLoader(uri.getHost(), params.userAgent, transformerConfig), params.allowInsecureConnections, params.numClientThreads, params.maxConcurrentRequests + ); setupShutdownHookForReplayer(tr); var tupleWriter = new TupleParserChainConsumer(new ResultsToLogsConsumer()); diff --git a/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/TrafficReplayerCore.java b/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/TrafficReplayerCore.java new file mode 100644 index 000000000..c3b8ae999 --- /dev/null +++ b/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/TrafficReplayerCore.java @@ -0,0 +1,382 @@ +package org.opensearch.migrations.replay; + +import io.netty.buffer.Unpooled; +import lombok.AllArgsConstructor; +import lombok.Lombok; +import lombok.NonNull; +import lombok.SneakyThrows; +import lombok.extern.slf4j.Slf4j; +import org.opensearch.migrations.replay.datahandlers.IPacketFinalizingConsumer; +import org.opensearch.migrations.replay.datatypes.HttpRequestTransformationStatus; +import org.opensearch.migrations.replay.datatypes.ITrafficStreamKey; +import org.opensearch.migrations.replay.datatypes.TransformedPackets; +import org.opensearch.migrations.replay.datatypes.UniqueReplayerRequestKey; +import org.opensearch.migrations.replay.tracing.IReplayContexts; +import org.opensearch.migrations.replay.tracing.IRootReplayerContext; +import org.opensearch.migrations.replay.traffic.source.ITrafficCaptureSource; +import org.opensearch.migrations.replay.traffic.source.ITrafficStreamWithKey; +import org.opensearch.migrations.replay.traffic.source.TrafficStreamLimiter; +import org.opensearch.migrations.replay.util.DiagnosticTrackableCompletableFuture; +import org.opensearch.migrations.replay.util.StringTrackableCompletableFuture; +import org.opensearch.migrations.trafficcapture.protos.TrafficStreamUtils; +import org.opensearch.migrations.transform.IAuthTransformerFactory; +import org.opensearch.migrations.transform.IJsonTransformer; + +import javax.net.ssl.SSLException; +import java.io.EOFException; +import java.net.URI; +import java.time.Duration; +import java.time.Instant; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; +import java.util.function.Supplier; +import java.util.stream.Collectors; +import java.util.stream.Stream; + +@Slf4j +public abstract class TrafficReplayerCore { + + private final PacketToTransformingHttpHandlerFactory inputRequestTransformerFactory; + protected final ClientConnectionPool clientConnectionPool; + protected final TrafficStreamLimiter liveTrafficStreamLimiter; + protected final AtomicInteger successfulRequestCount; + protected final AtomicInteger exceptionRequestCount; + public final IRootReplayerContext topLevelContext; + protected final ConcurrentHashMap> requestToFinalWorkFuturesMap; + + protected final AtomicBoolean stopReadingRef; + protected final AtomicReference>> nextChunkFutureRef; + + public TrafficReplayerCore(IRootReplayerContext context, + URI serverUri, + IAuthTransformerFactory authTransformer, + IJsonTransformer jsonTransformer, + ClientConnectionPool clientConnectionPool, + TrafficStreamLimiter trafficStreamLimiter) + throws SSLException + { + this.topLevelContext = context; + if (serverUri.getPort() < 0) { + throw new IllegalArgumentException("Port not present for URI: "+serverUri); + } + if (serverUri.getHost() == null) { + throw new IllegalArgumentException("Hostname not present for URI: "+serverUri); + } + if (serverUri.getScheme() == null) { + throw new IllegalArgumentException("Scheme (http|https) is not present for URI: "+serverUri); + } + this.liveTrafficStreamLimiter = trafficStreamLimiter; + this.clientConnectionPool = clientConnectionPool; + inputRequestTransformerFactory = new PacketToTransformingHttpHandlerFactory(jsonTransformer, authTransformer); + requestToFinalWorkFuturesMap = new ConcurrentHashMap<>(); + successfulRequestCount = new AtomicInteger(); + exceptionRequestCount = new AtomicInteger(); + nextChunkFutureRef = new AtomicReference<>(); + stopReadingRef = new AtomicBoolean(); + } + + protected abstract CompletableFuture shutdown(Error error); + + @AllArgsConstructor + class TrafficReplayerAccumulationCallbacks implements AccumulationCallbacks { + private final ReplayEngine replayEngine; + private Consumer resultTupleConsumer; + private ITrafficCaptureSource trafficCaptureSource; + + @Override + public Consumer + onRequestReceived(@NonNull IReplayContexts.IReplayerHttpTransactionContext ctx, + @NonNull HttpMessageAndTimestamp request) { + replayEngine.setFirstTimestamp(request.getFirstPacketTimestamp()); + + var allWorkFinishedForTransaction = + new StringTrackableCompletableFuture(new CompletableFuture<>(), + ()->"waiting for work to be queued and run through TrafficStreamLimiter"); + var requestPushFuture = new StringTrackableCompletableFuture( + new CompletableFuture<>(), () -> "Waiting to get response from target"); + var requestKey = ctx.getReplayerRequestKey(); + liveTrafficStreamLimiter.queueWork(1, ctx, wi -> { + transformAndSendRequest(replayEngine, request, ctx).future.whenComplete((v,t)->{ + liveTrafficStreamLimiter.doneProcessing(wi); + if (t != null) { + requestPushFuture.future.completeExceptionally(t); + } else { + requestPushFuture.future.complete(v); + } + }); + }); + if (!allWorkFinishedForTransaction.future.isDone()) { + log.trace("Adding " + requestKey + " to targetTransactionInProgressMap"); + requestToFinalWorkFuturesMap.put(requestKey, allWorkFinishedForTransaction); + if (allWorkFinishedForTransaction.future.isDone()) { + requestToFinalWorkFuturesMap.remove(requestKey); + } + } + + return rrPair -> + requestPushFuture.map(f -> f.handle((v, t) -> { + log.atInfo().setMessage(() -> "Done receiving captured stream for " + ctx + + ":" + rrPair.requestData).log(); + log.atTrace().setMessage(() -> + "Summary response value for " + requestKey + " returned=" + v).log(); + return handleCompletedTransaction(ctx, rrPair, v, t); + }), () -> "logging summary") + .whenComplete((v,t)->{ + if (t != null) { + allWorkFinishedForTransaction.future.completeExceptionally(t); + } else { + allWorkFinishedForTransaction.future.complete(null); + } + }, ()->""); + } + + Void handleCompletedTransaction(@NonNull IReplayContexts.IReplayerHttpTransactionContext context, + RequestResponsePacketPair rrPair, + TransformedTargetRequestAndResponse summary, Throwable t) { + try (var httpContext = rrPair.getHttpTransactionContext()) { + // if this comes in with a serious Throwable (not an Exception), don't bother + // packaging it up and calling the callback. + // Escalate it up out handling stack and shutdown. + if (t == null || t instanceof Exception) { + try (var tupleHandlingContext = httpContext.createTupleContext()) { + packageAndWriteResponse(tupleHandlingContext, resultTupleConsumer, + rrPair, summary, (Exception) t); + } + commitTrafficStreams(rrPair.completionStatus, rrPair.trafficStreamKeysBeingHeld); + return null; + } else { + log.atError().setCause(t).setMessage(() -> "Throwable passed to handle() for " + context + + ". Rethrowing.").log(); + throw Lombok.sneakyThrow(t); + } + } catch (Error error) { + log.atError() + .setCause(error) + .setMessage(() -> "Caught error and initiating TrafficReplayer shutdown") + .log(); + shutdown(error); + throw error; + } catch (Exception e) { + log.atError() + .setMessage("Unexpected exception while sending the " + + "aggregated response and context for {} to the callback. " + + "Proceeding, but the tuple receiver context may be compromised.") + .addArgument(context) + .setCause(e) + .log(); + throw e; + } finally { + var requestKey = context.getReplayerRequestKey(); + requestToFinalWorkFuturesMap.remove(requestKey); + log.trace("removed rrPair.requestData to " + + "targetTransactionInProgressMap for " + + requestKey); + } + } + + @Override + public void onTrafficStreamsExpired(RequestResponsePacketPair.ReconstructionStatus status, + @NonNull IReplayContexts.IChannelKeyContext ctx, + @NonNull List trafficStreamKeysBeingHeld) { + commitTrafficStreams(status, trafficStreamKeysBeingHeld); + } + + @SneakyThrows + private void commitTrafficStreams(RequestResponsePacketPair.ReconstructionStatus status, + List trafficStreamKeysBeingHeld) { + commitTrafficStreams(status != RequestResponsePacketPair.ReconstructionStatus.CLOSED_PREMATURELY, + trafficStreamKeysBeingHeld); + } + + @SneakyThrows + private void commitTrafficStreams(boolean shouldCommit, + List trafficStreamKeysBeingHeld) { + if (shouldCommit && trafficStreamKeysBeingHeld != null) { + for (var tsk : trafficStreamKeysBeingHeld) { + tsk.getTrafficStreamsContext().close(); + trafficCaptureSource.commitTrafficStream(tsk); + } + } + } + + @Override + public void onConnectionClose(int channelInteractionNum, + @NonNull IReplayContexts.IChannelKeyContext ctx, int channelSessionNumber, + RequestResponsePacketPair.ReconstructionStatus status, + @NonNull Instant timestamp, @NonNull List trafficStreamKeysBeingHeld) { + replayEngine.setFirstTimestamp(timestamp); + var cf = replayEngine.closeConnection(channelInteractionNum, ctx, channelSessionNumber, timestamp); + cf.map(f->f.whenComplete((v,t)->{ + commitTrafficStreams(status, trafficStreamKeysBeingHeld); + }), ()->"closing the channel in the ReplayEngine"); + } + + @Override + public void onTrafficStreamIgnored(@NonNull IReplayContexts.ITrafficStreamsLifecycleContext ctx) { + commitTrafficStreams(true, List.of(ctx.getTrafficStreamKey())); + } + + private void packageAndWriteResponse(IReplayContexts.ITupleHandlingContext tupleHandlingContext, + Consumer tupleWriter, + RequestResponsePacketPair rrPair, + TransformedTargetRequestAndResponse summary, + Exception t) { + log.trace("done sending and finalizing data to the packet handler"); + + try (var requestResponseTuple = getSourceTargetCaptureTuple(tupleHandlingContext, rrPair, summary, t)) { + log.atInfo().setMessage(()->"Source/Target Request/Response tuple: " + requestResponseTuple).log(); + tupleWriter.accept(requestResponseTuple); + } + + if (t != null) { throw new CompletionException(t); } + if (summary.getError() != null) { + log.atInfo().setCause(summary.getError()).setMessage("Exception for {}: ") + .addArgument(tupleHandlingContext).log(); + exceptionRequestCount.incrementAndGet(); + } else if (summary.getTransformationStatus() == HttpRequestTransformationStatus.ERROR) { + log.atInfo().setCause(summary.getError()).setMessage("Unknown error transforming {}: ") + .addArgument(tupleHandlingContext).log(); + exceptionRequestCount.incrementAndGet(); + } else { + successfulRequestCount.incrementAndGet(); + } + } + } + + private static SourceTargetCaptureTuple + getSourceTargetCaptureTuple(@NonNull IReplayContexts.ITupleHandlingContext tupleHandlingContext, + RequestResponsePacketPair rrPair, + TransformedTargetRequestAndResponse summary, + Exception t) + { + SourceTargetCaptureTuple requestResponseTuple; + if (t != null) { + log.error("Got exception in CompletableFuture callback: ", t); + requestResponseTuple = new SourceTargetCaptureTuple(tupleHandlingContext, rrPair, + new TransformedPackets(), new ArrayList<>(), + HttpRequestTransformationStatus.ERROR, t, Duration.ZERO); + } else { + requestResponseTuple = new SourceTargetCaptureTuple(tupleHandlingContext, rrPair, + summary.requestPackets, + summary.getReceiptTimeAndResponsePackets() + .map(Map.Entry::getValue).collect(Collectors.toList()), + summary.getTransformationStatus(), + summary.getError(), + summary.getResponseDuration() + ); + } + return requestResponseTuple; + } + + public DiagnosticTrackableCompletableFuture + transformAndSendRequest(ReplayEngine replayEngine, HttpMessageAndTimestamp request, + IReplayContexts.IReplayerHttpTransactionContext ctx) { + return transformAndSendRequest(inputRequestTransformerFactory, replayEngine, ctx, + request.getFirstPacketTimestamp(), request.getLastPacketTimestamp(), + request.packetBytes::stream); + } + + public static DiagnosticTrackableCompletableFuture + transformAndSendRequest(PacketToTransformingHttpHandlerFactory inputRequestTransformerFactory, + ReplayEngine replayEngine, + IReplayContexts.IReplayerHttpTransactionContext ctx, + @NonNull Instant start, @NonNull Instant end, + Supplier> packetsSupplier) + { + try { + var transformationCompleteFuture = replayEngine.scheduleTransformationWork(ctx, start, ()-> + transformAllData(inputRequestTransformerFactory.create(ctx), packetsSupplier)); + log.atDebug().setMessage(()->"finalizeRequest future for transformation of " + ctx + + " = " + transformationCompleteFuture).log(); + // It might be safer to chain this work directly inside the scheduleWork call above so that the + // read buffer horizons aren't set after the transformation work finishes, but after the packets + // are fully handled + return transformationCompleteFuture.thenCompose(transformedResult -> + replayEngine.scheduleRequest(ctx, start, end, + transformedResult.transformedOutput.size(), + transformedResult.transformedOutput.streamRetained()) + .map(future->future.thenApply(t -> + new TransformedTargetRequestAndResponse(transformedResult.transformedOutput, + t, transformedResult.transformationStatus, t.error)), + ()->"(if applicable) packaging transformed result into a completed TransformedTargetRequestAndResponse object") + .map(future->future.exceptionally(t -> + new TransformedTargetRequestAndResponse(transformedResult.transformedOutput, + transformedResult.transformationStatus, t)), + ()->"(if applicable) packaging transformed result into a failed TransformedTargetRequestAndResponse object"), + () -> "transitioning transformed packets onto the wire") + .map(future->future.exceptionally(t->new TransformedTargetRequestAndResponse(null, null, t)), + ()->"Checking for exception out of sending data to the target server"); + } catch (Exception e) { + log.debug("Caught exception in writeToSocket, so failing future"); + return StringTrackableCompletableFuture.failedFuture(e, ()->"TrafficReplayer.writeToSocketAndClose"); + } + } + + private static DiagnosticTrackableCompletableFuture + transformAllData(IPacketFinalizingConsumer packetHandler, Supplier> packetSupplier) { + try { + var logLabel = packetHandler.getClass().getSimpleName(); + var packets = packetSupplier.get().map(Unpooled::wrappedBuffer); + packets.forEach(packetData -> { + log.atDebug().setMessage(() -> logLabel + " sending " + packetData.readableBytes() + + " bytes to the packetHandler").log(); + var consumeFuture = packetHandler.consumeBytes(packetData); + log.atDebug().setMessage(() -> logLabel + " consumeFuture = " + consumeFuture).log(); + }); + log.atDebug().setMessage(() -> logLabel + " done sending bytes, now finalizing the request").log(); + return packetHandler.finalizeRequest(); + } catch (Exception e) { + log.atInfo().setCause(e).setMessage("Encountered an exception while transforming the http request. " + + "The base64 gzipped traffic stream, for later diagnostic purposes, is: " + + Utils.packetsToCompressedTrafficStream(packetSupplier.get())).log(); + throw e; + } + } + + @SneakyThrows + public void pullCaptureFromSourceToAccumulator( + ITrafficCaptureSource trafficChunkStream, + CapturedTrafficToHttpTransactionAccumulator trafficToHttpTransactionAccumulator) + throws InterruptedException { + while (true) { + log.trace("Reading next chunk from TrafficStream supplier"); + if (stopReadingRef.get()) { + break; + } + this.nextChunkFutureRef.set(trafficChunkStream + .readNextTrafficStreamChunk(topLevelContext::createReadChunkContext)); + List trafficStreams = null; + try { + trafficStreams = this.nextChunkFutureRef.get().get(); + } catch (ExecutionException ex) { + if (ex.getCause() instanceof EOFException) { + log.atWarn().setCause(ex.getCause()).setMessage("Got an EOF on the stream. " + + "Done reading traffic streams.").log(); + break; + } else { + log.atWarn().setCause(ex).setMessage("Done reading traffic streams due to exception.").log(); + throw ex.getCause(); + } + } + if (log.isInfoEnabled()) { + Optional.of(trafficStreams.stream() + .map(ts -> TrafficStreamUtils.summarizeTrafficStream(ts.getStream())) + .collect(Collectors.joining(";"))) + .filter(s -> !s.isEmpty()) + .ifPresent(s -> log.atInfo().log("TrafficStream Summary: {" + s + "}")); + } + trafficStreams.forEach(trafficToHttpTransactionAccumulator::accept); + } + } +} diff --git a/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/TrafficReplayerTopLevel.java b/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/TrafficReplayerTopLevel.java index f202cfc58..48a703def 100644 --- a/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/TrafficReplayerTopLevel.java +++ b/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/TrafficReplayerTopLevel.java @@ -1,150 +1,103 @@ package org.opensearch.migrations.replay; -import io.netty.buffer.Unpooled; import io.netty.handler.ssl.SslContext; import io.netty.handler.ssl.SslContextBuilder; import io.netty.handler.ssl.util.InsecureTrustManagerFactory; -import lombok.AllArgsConstructor; -import lombok.Lombok; import lombok.NonNull; import lombok.SneakyThrows; import lombok.extern.slf4j.Slf4j; -import org.opensearch.migrations.replay.datahandlers.IPacketFinalizingConsumer; -import org.opensearch.migrations.replay.datatypes.HttpRequestTransformationStatus; -import org.opensearch.migrations.replay.datatypes.ITrafficStreamKey; -import org.opensearch.migrations.replay.datatypes.TransformedPackets; import org.opensearch.migrations.replay.datatypes.UniqueReplayerRequestKey; -import org.opensearch.migrations.replay.tracing.IReplayContexts; import org.opensearch.migrations.replay.tracing.IRootReplayerContext; import org.opensearch.migrations.replay.traffic.source.BlockingTrafficSource; -import org.opensearch.migrations.replay.traffic.source.ITrafficCaptureSource; -import org.opensearch.migrations.replay.traffic.source.ITrafficStreamWithKey; import org.opensearch.migrations.replay.traffic.source.TrafficStreamLimiter; import org.opensearch.migrations.replay.util.DiagnosticTrackableCompletableFuture; import org.opensearch.migrations.replay.util.StringTrackableCompletableFuture; -import org.opensearch.migrations.trafficcapture.protos.TrafficStreamUtils; import org.opensearch.migrations.transform.IAuthTransformerFactory; import org.opensearch.migrations.transform.IJsonTransformer; import org.slf4j.event.Level; import org.slf4j.spi.LoggingEventBuilder; import javax.net.ssl.SSLException; -import java.io.EOFException; import java.net.URI; import java.time.Duration; -import java.time.Instant; -import java.util.ArrayList; import java.util.Arrays; -import java.util.List; import java.util.Map; import java.util.Optional; import java.util.concurrent.CompletableFuture; -import java.util.concurrent.CompletionException; -import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; -import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; -import java.util.function.Supplier; import java.util.stream.Collectors; -import java.util.stream.Stream; @Slf4j -public class TrafficReplayerTopLevel implements AutoCloseable { +public class TrafficReplayerTopLevel extends TrafficReplayerCore implements AutoCloseable { public static final String TARGET_CONNECTION_POOL_NAME = "targetConnectionPool"; public static final int MAX_ITEMS_TO_SHOW_FOR_LEFTOVER_WORK_AT_INFO_LEVEL = 10; - public static AtomicInteger targetConnectionPoolUniqueCounter = new AtomicInteger(); + public static final AtomicInteger targetConnectionPoolUniqueCounter = new AtomicInteger(); - private final PacketToTransformingHttpHandlerFactory inputRequestTransformerFactory; - private final ClientConnectionPool clientConnectionPool; - private final TrafficStreamLimiter liveTrafficStreamLimiter; - private final AtomicInteger successfulRequestCount; - private final AtomicInteger exceptionRequestCount; - public final IRootReplayerContext topLevelContext; - private final ConcurrentHashMap> requestToFinalWorkFuturesMap; - - private final AtomicBoolean stopReadingRef; private final AtomicReference> allRemainingWorkFutureOrShutdownSignalRef; private final AtomicReference shutdownReasonRef; private final AtomicReference> shutdownFutureRef; - private final AtomicReference>> nextChunkFutureRef; - public TrafficReplayerTopLevel(IRootReplayerContext context, URI serverUri, IAuthTransformerFactory authTransformerFactory, boolean allowInsecureConnections) throws SSLException { - this(context, serverUri, authTransformerFactory, allowInsecureConnections, - new TransformationLoader().getTransformerFactoryLoader(serverUri.getHost())); + this(context, serverUri, authTransformerFactory, new TransformationLoader().getTransformerFactoryLoader(serverUri.getHost()), allowInsecureConnections + ); } public TrafficReplayerTopLevel(IRootReplayerContext context, URI serverUri, IAuthTransformerFactory authTransformerFactory, - boolean allowInsecureConnections, - IJsonTransformer jsonTransformer) + IJsonTransformer jsonTransformer, + boolean allowInsecureConnections) throws SSLException { - this(context, serverUri, authTransformerFactory, allowInsecureConnections, 0, - 1024, - jsonTransformer); + this(context, serverUri, authTransformerFactory, jsonTransformer, allowInsecureConnections, 0, + 1024); + } + + public TrafficReplayerTopLevel(IRootReplayerContext topContext, + URI uri, + IAuthTransformerFactory authTransformer, + IJsonTransformer transformerFactoryLoader, + boolean allowInsecureConnections, + int numClientThreads, + int maxConcurrentRequests) throws SSLException { + this(topContext, uri, authTransformer, transformerFactoryLoader, allowInsecureConnections, + numClientThreads, maxConcurrentRequests, + getTargetConnectionPoolName(targetConnectionPoolUniqueCounter.getAndIncrement())); } public TrafficReplayerTopLevel(IRootReplayerContext context, URI serverUri, IAuthTransformerFactory authTransformerFactory, + IJsonTransformer jsonTransformer, boolean allowInsecureConnections, int numSendingThreads, int maxConcurrentOutstandingRequests, - IJsonTransformer jsonTransformer) - throws SSLException { - this(context, serverUri, authTransformerFactory, allowInsecureConnections, - numSendingThreads, maxConcurrentOutstandingRequests, - jsonTransformer, - getTargetConnectionPoolName(targetConnectionPoolUniqueCounter.getAndIncrement())); - } - - private static String getTargetConnectionPoolName(int i) { - return TARGET_CONNECTION_POOL_NAME + (i == 0 ? "" : Integer.toString(i)); + String connectionPoolName) throws SSLException { + this(context, serverUri, authTransformerFactory, jsonTransformer, + new ClientConnectionPool(serverUri, + loadSslContext(serverUri, allowInsecureConnections), connectionPoolName, numSendingThreads), + new TrafficStreamLimiter(maxConcurrentOutstandingRequests)); } public TrafficReplayerTopLevel(IRootReplayerContext context, URI serverUri, - IAuthTransformerFactory authTransformer, - boolean allowInsecureConnections, - int numSendingThreads, - int maxConcurrentOutstandingRequests, + IAuthTransformerFactory authTransformerFactory, IJsonTransformer jsonTransformer, - String clientThreadNamePrefix) - throws SSLException - { - this.topLevelContext = context; - if (serverUri.getPort() < 0) { - throw new IllegalArgumentException("Port not present for URI: "+serverUri); - } - if (serverUri.getHost() == null) { - throw new IllegalArgumentException("Hostname not present for URI: "+serverUri); - } - if (serverUri.getScheme() == null) { - throw new IllegalArgumentException("Scheme (http|https) is not present for URI: "+serverUri); - } - inputRequestTransformerFactory = new PacketToTransformingHttpHandlerFactory(jsonTransformer, authTransformer); - clientConnectionPool = new ClientConnectionPool(serverUri, - loadSslContext(serverUri, allowInsecureConnections), clientThreadNamePrefix, numSendingThreads); - requestToFinalWorkFuturesMap = new ConcurrentHashMap<>(); - successfulRequestCount = new AtomicInteger(); - exceptionRequestCount = new AtomicInteger(); - liveTrafficStreamLimiter = new TrafficStreamLimiter(maxConcurrentOutstandingRequests); + ClientConnectionPool clientConnectionPool, + TrafficStreamLimiter trafficStreamLimiter) throws SSLException { + super(context, serverUri, authTransformerFactory, jsonTransformer, clientConnectionPool, trafficStreamLimiter); allRemainingWorkFutureOrShutdownSignalRef = new AtomicReference<>(); shutdownReasonRef = new AtomicReference<>(); shutdownFutureRef = new AtomicReference<>(); - nextChunkFutureRef = new AtomicReference<>(); - stopReadingRef = new AtomicBoolean(); } private static SslContext loadSslContext(URI serverUri, boolean allowInsecureConnections) throws SSLException { @@ -159,6 +112,10 @@ private static SslContext loadSslContext(URI serverUri, boolean allowInsecureCon } } + private static String getTargetConnectionPoolName(int i) { + return TARGET_CONNECTION_POOL_NAME + (i == 0 ? "" : Integer.toString(i)); + } + public void setupRunAndWaitForReplayToFinish(Duration observedPacketConnectionTimeout, BlockingTrafficSource trafficSource, TimeShifter timeShifter, @@ -188,6 +145,10 @@ public void setupRunAndWaitForReplayToFinish(Duration observedPacketConnectionTi } } + /** + * + * @param replayEngine is not used here but might be of use to extensions of this class + */ protected void wrapUpWorkAndEmitSummary(ReplayEngine replayEngine, CapturedTrafficToHttpTransactionAccumulator trafficToHttpTransactionAccumulator) throws ExecutionException, InterruptedException { @@ -251,174 +212,6 @@ public void setupRunAndWaitForReplayWithShutdownChecks(Duration observedPacketCo shutdown(null).get(); // if somebody already HAD run shutdown, it will return the future already created } - @AllArgsConstructor - class TrafficReplayerAccumulationCallbacks implements AccumulationCallbacks { - private final ReplayEngine replayEngine; - private Consumer resultTupleConsumer; - private ITrafficCaptureSource trafficCaptureSource; - - @Override - public Consumer - onRequestReceived(@NonNull IReplayContexts.IReplayerHttpTransactionContext ctx, - @NonNull HttpMessageAndTimestamp request) { - replayEngine.setFirstTimestamp(request.getFirstPacketTimestamp()); - - var allWorkFinishedForTransaction = - new StringTrackableCompletableFuture(new CompletableFuture<>(), - ()->"waiting for work to be queued and run through TrafficStreamLimiter"); - var requestPushFuture = new StringTrackableCompletableFuture( - new CompletableFuture<>(), () -> "Waiting to get response from target"); - var requestKey = ctx.getReplayerRequestKey(); - liveTrafficStreamLimiter.queueWork(1, ctx, wi -> { - transformAndSendRequest(replayEngine, request, ctx).future.whenComplete((v,t)->{ - liveTrafficStreamLimiter.doneProcessing(wi); - if (t != null) { - requestPushFuture.future.completeExceptionally(t); - } else { - requestPushFuture.future.complete(v); - } - }); - }); - if (!allWorkFinishedForTransaction.future.isDone()) { - log.trace("Adding " + requestKey + " to targetTransactionInProgressMap"); - requestToFinalWorkFuturesMap.put(requestKey, allWorkFinishedForTransaction); - if (allWorkFinishedForTransaction.future.isDone()) { - requestToFinalWorkFuturesMap.remove(requestKey); - } - } - - return rrPair -> - requestPushFuture.map(f -> f.handle((v, t) -> { - log.atInfo().setMessage(() -> "Done receiving captured stream for " + ctx + - ":" + rrPair.requestData).log(); - log.atTrace().setMessage(() -> - "Summary response value for " + requestKey + " returned=" + v).log(); - return handleCompletedTransaction(ctx, rrPair, v, t); - }), () -> "logging summary") - .whenComplete((v,t)->{ - if (t != null) { - allWorkFinishedForTransaction.future.completeExceptionally(t); - } else { - allWorkFinishedForTransaction.future.complete(null); - } - }, ()->""); - } - - Void handleCompletedTransaction(@NonNull IReplayContexts.IReplayerHttpTransactionContext context, - RequestResponsePacketPair rrPair, - TransformedTargetRequestAndResponse summary, Throwable t) { - try (var httpContext = rrPair.getHttpTransactionContext()) { - // if this comes in with a serious Throwable (not an Exception), don't bother - // packaging it up and calling the callback. - // Escalate it up out handling stack and shutdown. - if (t == null || t instanceof Exception) { - try (var tupleHandlingContext = httpContext.createTupleContext()) { - packageAndWriteResponse(tupleHandlingContext, resultTupleConsumer, - rrPair, summary, (Exception) t); - } - commitTrafficStreams(rrPair.completionStatus, rrPair.trafficStreamKeysBeingHeld); - return null; - } else { - log.atError().setCause(t).setMessage(() -> "Throwable passed to handle() for " + context + - ". Rethrowing.").log(); - throw Lombok.sneakyThrow(t); - } - } catch (Error error) { - log.atError() - .setCause(error) - .setMessage(() -> "Caught error and initiating TrafficReplayer shutdown") - .log(); - shutdown(error); - throw error; - } catch (Exception e) { - log.atError() - .setMessage("Unexpected exception while sending the " + - "aggregated response and context for {} to the callback. " + - "Proceeding, but the tuple receiver context may be compromised.") - .addArgument(context) - .setCause(e) - .log(); - throw e; - } finally { - var requestKey = context.getReplayerRequestKey(); - requestToFinalWorkFuturesMap.remove(requestKey); - log.trace("removed rrPair.requestData to " + - "targetTransactionInProgressMap for " + - requestKey); - } - } - - @Override - public void onTrafficStreamsExpired(RequestResponsePacketPair.ReconstructionStatus status, - @NonNull IReplayContexts.IChannelKeyContext ctx, - @NonNull List trafficStreamKeysBeingHeld) { - commitTrafficStreams(status, trafficStreamKeysBeingHeld); - } - - @SneakyThrows - private void commitTrafficStreams(RequestResponsePacketPair.ReconstructionStatus status, - List trafficStreamKeysBeingHeld) { - commitTrafficStreams(status != RequestResponsePacketPair.ReconstructionStatus.CLOSED_PREMATURELY, - trafficStreamKeysBeingHeld); - } - - @SneakyThrows - private void commitTrafficStreams(boolean shouldCommit, - List trafficStreamKeysBeingHeld) { - if (shouldCommit && trafficStreamKeysBeingHeld != null) { - for (var tsk : trafficStreamKeysBeingHeld) { - tsk.getTrafficStreamsContext().close(); - trafficCaptureSource.commitTrafficStream(tsk); - } - } - } - - @Override - public void onConnectionClose(int channelInteractionNum, - @NonNull IReplayContexts.IChannelKeyContext ctx, int channelSessionNumber, - RequestResponsePacketPair.ReconstructionStatus status, - @NonNull Instant timestamp, @NonNull List trafficStreamKeysBeingHeld) { - replayEngine.setFirstTimestamp(timestamp); - var cf = replayEngine.closeConnection(channelInteractionNum, ctx, channelSessionNumber, timestamp); - cf.map(f->f.whenComplete((v,t)->{ - commitTrafficStreams(status, trafficStreamKeysBeingHeld); - }), ()->"closing the channel in the ReplayEngine"); - } - - @Override - public void onTrafficStreamIgnored(@NonNull IReplayContexts.ITrafficStreamsLifecycleContext ctx) { - commitTrafficStreams(true, List.of(ctx.getTrafficStreamKey())); - } - - private TransformedTargetRequestAndResponse - packageAndWriteResponse(IReplayContexts.ITupleHandlingContext tupleHandlingContext, - Consumer tupleWriter, - RequestResponsePacketPair rrPair, - TransformedTargetRequestAndResponse summary, - Exception t) { - log.trace("done sending and finalizing data to the packet handler"); - - try (var requestResponseTuple = getSourceTargetCaptureTuple(tupleHandlingContext, rrPair, summary, t)) { - log.atInfo().setMessage(()->"Source/Target Request/Response tuple: " + requestResponseTuple).log(); - tupleWriter.accept(requestResponseTuple); - } - - if (t != null) { throw new CompletionException(t); } - if (summary.getError() != null) { - log.atInfo().setCause(summary.getError()).setMessage("Exception for {}: ") - .addArgument(tupleHandlingContext).log(); - exceptionRequestCount.incrementAndGet(); - } else if (summary.getTransformationStatus() == HttpRequestTransformationStatus.ERROR) { - log.atInfo().setCause(summary.getError()).setMessage("Unknown error transforming {}: ") - .addArgument(tupleHandlingContext).log(); - exceptionRequestCount.incrementAndGet(); - } else { - successfulRequestCount.incrementAndGet(); - } - return summary; - } - } - protected void waitForRemainingWork(Level logLevel, @NonNull Duration timeout) throws ExecutionException, InterruptedException, TimeoutException { @@ -483,7 +276,7 @@ private void handleAlreadySetFinishedSignal() throws InterruptedException, Execu } } - private static void writeStatusLogsForRemainingWork(Level logLevel, + protected static void writeStatusLogsForRemainingWork(Level logLevel, Map.Entry>[] @@ -502,7 +295,7 @@ private static void writeStatusLogsForRemainingWork(Level logLevel, } } - private static String formatWorkItem(DiagnosticTrackableCompletableFuture cf) { + protected static String formatWorkItem(DiagnosticTrackableCompletableFuture cf) { try { var resultValue = cf.get(); if (resultValue instanceof TransformedTargetRequestAndResponse) { @@ -517,97 +310,8 @@ private static String formatWorkItem(DiagnosticTrackableCompletableFuture(), - HttpRequestTransformationStatus.ERROR, t, Duration.ZERO); - } else { - requestResponseTuple = new SourceTargetCaptureTuple(tupleHandlingContext, rrPair, - summary.requestPackets, - summary.getReceiptTimeAndResponsePackets() - .map(Map.Entry::getValue).collect(Collectors.toList()), - summary.getTransformationStatus(), - summary.getError(), - summary.getResponseDuration() - ); - } - return requestResponseTuple; - } - - public DiagnosticTrackableCompletableFuture - transformAndSendRequest(ReplayEngine replayEngine, HttpMessageAndTimestamp request, - IReplayContexts.IReplayerHttpTransactionContext ctx) { - return transformAndSendRequest(inputRequestTransformerFactory, replayEngine, ctx, - request.getFirstPacketTimestamp(), request.getLastPacketTimestamp(), - request.packetBytes::stream); - } - - public static DiagnosticTrackableCompletableFuture - transformAndSendRequest(PacketToTransformingHttpHandlerFactory inputRequestTransformerFactory, - ReplayEngine replayEngine, - IReplayContexts.IReplayerHttpTransactionContext ctx, - @NonNull Instant start, @NonNull Instant end, - Supplier> packetsSupplier) - { - try { - var transformationCompleteFuture = replayEngine.scheduleTransformationWork(ctx, start, ()-> - transformAllData(inputRequestTransformerFactory.create(ctx), packetsSupplier)); - log.atDebug().setMessage(()->"finalizeRequest future for transformation of " + ctx + - " = " + transformationCompleteFuture).log(); - // It might be safer to chain this work directly inside the scheduleWork call above so that the - // read buffer horizons aren't set after the transformation work finishes, but after the packets - // are fully handled - return transformationCompleteFuture.thenCompose(transformedResult -> - replayEngine.scheduleRequest(ctx, start, end, - transformedResult.transformedOutput.size(), - transformedResult.transformedOutput.streamRetained()) - .map(future->future.thenApply(t -> - new TransformedTargetRequestAndResponse(transformedResult.transformedOutput, - t, transformedResult.transformationStatus, t.error)), - ()->"(if applicable) packaging transformed result into a completed TransformedTargetRequestAndResponse object") - .map(future->future.exceptionally(t -> - new TransformedTargetRequestAndResponse(transformedResult.transformedOutput, - transformedResult.transformationStatus, t)), - ()->"(if applicable) packaging transformed result into a failed TransformedTargetRequestAndResponse object"), - () -> "transitioning transformed packets onto the wire") - .map(future->future.exceptionally(t->new TransformedTargetRequestAndResponse(null, null, t)), - ()->"Checking for exception out of sending data to the target server"); - } catch (Exception e) { - log.debug("Caught exception in writeToSocket, so failing future"); - return StringTrackableCompletableFuture.failedFuture(e, ()->"TrafficReplayer.writeToSocketAndClose"); - } - } - - private static DiagnosticTrackableCompletableFuture - transformAllData(IPacketFinalizingConsumer packetHandler, Supplier> packetSupplier) { - try { - var logLabel = packetHandler.getClass().getSimpleName(); - var packets = packetSupplier.get().map(Unpooled::wrappedBuffer); - packets.forEach(packetData -> { - log.atDebug().setMessage(() -> logLabel + " sending " + packetData.readableBytes() + - " bytes to the packetHandler").log(); - var consumeFuture = packetHandler.consumeBytes(packetData); - log.atDebug().setMessage(() -> logLabel + " consumeFuture = " + consumeFuture).log(); - }); - log.atDebug().setMessage(() -> logLabel + " done sending bytes, now finalizing the request").log(); - return packetHandler.finalizeRequest(); - } catch (Exception e) { - log.atInfo().setCause(e).setMessage("Encountered an exception while transforming the http request. " + - "The base64 gzipped traffic stream, for later diagnostic purposes, is: " + - Utils.packetsToCompressedTrafficStream(packetSupplier.get())).log(); - throw e; - } - } - @SneakyThrows + @Override public @NonNull CompletableFuture shutdown(Error error) { log.atWarn().setCause(error).setMessage(()->"Shutting down " + this).log(); shutdownReasonRef.compareAndSet(null, error); @@ -646,45 +350,8 @@ private static String formatWorkItem(DiagnosticTrackableCompletableFuture trafficStreams = null; - try { - trafficStreams = this.nextChunkFutureRef.get().get(); - } catch (ExecutionException ex) { - if (ex.getCause() instanceof EOFException) { - log.atWarn().setCause(ex.getCause()).setMessage("Got an EOF on the stream. " + - "Done reading traffic streams.").log(); - break; - } else { - log.atWarn().setCause(ex).setMessage("Done reading traffic streams due to exception.").log(); - throw ex.getCause(); - } - } - if (log.isInfoEnabled()) { - Optional.of(trafficStreams.stream() - .map(ts -> TrafficStreamUtils.summarizeTrafficStream(ts.getStream())) - .collect(Collectors.joining(";"))) - .filter(s -> !s.isEmpty()) - .ifPresent(s -> log.atInfo().log("TrafficStream Summary: {" + s + "}")); - } - trafficStreams.forEach(trafficToHttpTransactionAccumulator::accept); - } - } } diff --git a/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/TrafficReplayerTest.java b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/TrafficReplayerTest.java index 561e41275..ef8ba36d9 100644 --- a/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/TrafficReplayerTest.java +++ b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/TrafficReplayerTest.java @@ -2,7 +2,6 @@ import com.google.protobuf.ByteString; import com.google.protobuf.Timestamp; -import io.netty.handler.logging.LogLevel; import lombok.NonNull; import lombok.extern.slf4j.Slf4j; import org.junit.jupiter.api.Assertions; @@ -209,8 +208,9 @@ public void onTrafficStreamIgnored(@NonNull IReplayContexts.ITrafficStreamsLifec @Test public void testCapturedReadsAfterCloseAreHandledAsNew() throws Exception { try (var tr = new TrafficReplayerTopLevel(rootContext, - new URI("http://localhost:9200"), null, false, - new TransformationLoader().getTransformerFactoryLoader("localhost"))) { + new URI("http://localhost:9200"), null, + new TransformationLoader().getTransformerFactoryLoader("localhost"), false + )) { List> byteArrays = new ArrayList<>(); var remainingAccumulations = new AtomicInteger(); CapturedTrafficToHttpTransactionAccumulator trafficAccumulator = diff --git a/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/FullReplayerWithTracingChecksTest.java b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/FullReplayerWithTracingChecksTest.java index 68a94338a..26864435b 100644 --- a/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/FullReplayerWithTracingChecksTest.java +++ b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/FullReplayerWithTracingChecksTest.java @@ -3,17 +3,13 @@ import com.google.protobuf.ByteString; import com.google.protobuf.Timestamp; import io.opentelemetry.sdk.trace.data.SpanData; -import lombok.Lombok; import lombok.extern.slf4j.Slf4j; import org.junit.jupiter.api.Assertions; -import org.junit.jupiter.api.Tag; -import org.junit.jupiter.api.Test; import org.junit.jupiter.api.parallel.ResourceLock; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.ValueSource; import org.opensearch.migrations.replay.TestHttpServerContext; import org.opensearch.migrations.replay.TimeShifter; -import org.opensearch.migrations.replay.TrafficReplayer; import org.opensearch.migrations.replay.TrafficReplayerTopLevel; import org.opensearch.migrations.replay.TransformationLoader; import org.opensearch.migrations.replay.traffic.source.ArrayCursorTrafficCaptureSource; @@ -94,9 +90,9 @@ public void testStreamWithRequestsWithCloseIsCommittedOnce(int numRequests) thro var serverUri = httpServer.localhostEndpoint(); try (var tr = new TrafficReplayerTopLevel(rootContext, serverUri, new StaticAuthTransformerFactory("TEST"), - true, 10, 10 * 1024, new TransformationLoader() - .getTransformerFactoryLoader(serverUri.getHost())); + .getTransformerFactoryLoader(serverUri.getHost()), true, 10, 10 * 1024 + ); var blockingTrafficSource = new BlockingTrafficSource(trafficSource, Duration.ofMinutes(2))) { tr.setupRunAndWaitForReplayToFinish(Duration.ofSeconds(70), blockingTrafficSource, new TimeShifter(10 * 1000), diff --git a/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/FullTrafficReplayerTest.java b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/FullTrafficReplayerTest.java index 1afdcec7b..47889b3bf 100644 --- a/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/FullTrafficReplayerTest.java +++ b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/FullTrafficReplayerTest.java @@ -14,7 +14,6 @@ import org.opensearch.migrations.replay.SourceTargetCaptureTuple; import org.opensearch.migrations.replay.TestHttpServerContext; import org.opensearch.migrations.replay.TimeShifter; -import org.opensearch.migrations.replay.TrafficReplayer; import org.opensearch.migrations.replay.TrafficReplayerTopLevel; import org.opensearch.migrations.replay.TransformationLoader; import org.opensearch.migrations.replay.tracing.IRootReplayerContext; @@ -44,11 +43,9 @@ import java.io.IOException; import java.net.URI; import java.time.Duration; -import java.time.Instant; import java.util.List; import java.util.Random; import java.util.concurrent.CompletableFuture; -import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Consumer; @@ -84,8 +81,8 @@ public TrafficReplayerWithWaitOnClose(Duration maxWaitTime, IJsonTransformer jsonTransformer, String targetConnectionPoolName) throws SSLException { super(context, serverUri, authTransformerFactory, - allowInsecureConnections, numSendingThreads, maxConcurrentOutstandingRequests, - jsonTransformer, targetConnectionPoolName); + jsonTransformer, allowInsecureConnections, numSendingThreads, maxConcurrentOutstandingRequests, + targetConnectionPoolName); this.maxWaitTime = maxWaitTime; } diff --git a/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/SlowAndExpiredTrafficStreamBecomesTwoTargetChannelsTest.java b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/SlowAndExpiredTrafficStreamBecomesTwoTargetChannelsTest.java index dba92f499..da0f44d98 100644 --- a/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/SlowAndExpiredTrafficStreamBecomesTwoTargetChannelsTest.java +++ b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/SlowAndExpiredTrafficStreamBecomesTwoTargetChannelsTest.java @@ -8,7 +8,6 @@ import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; import org.opensearch.migrations.replay.TimeShifter; -import org.opensearch.migrations.replay.TrafficReplayer; import org.opensearch.migrations.replay.TrafficReplayerTopLevel; import org.opensearch.migrations.replay.TransformationLoader; import org.opensearch.migrations.replay.traffic.source.ArrayCursorTrafficCaptureSource; @@ -100,8 +99,7 @@ public void test() throws Exception { try (var httpServer = SimpleNettyHttpServer.makeServer(false, Duration.ofMillis(200), responseTracker); var replayer = new TrafficReplayerTopLevel(rc, httpServer.localhostEndpoint(), new StaticAuthTransformerFactory("TEST"), - true, 1, 1, - new TransformationLoader().getTransformerFactoryLoader("localhost"), + new TransformationLoader().getTransformerFactoryLoader("localhost"), true, 1, 1, "targetConnectionPool for SlowAndExpiredTrafficStreamBecomesTwoTargetChannelsTest")) { new Thread(()->responseTracker.onCountDownFinished(Duration.ofSeconds(10), ()->replayer.shutdown(null).join())); diff --git a/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/TrafficReplayerRunner.java b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/TrafficReplayerRunner.java index 07cfd2689..08c94fe1f 100644 --- a/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/TrafficReplayerRunner.java +++ b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/e2etests/TrafficReplayerRunner.java @@ -67,8 +67,7 @@ public static void runReplayer(int numExpectedRequests, try { return new TrafficReplayerTopLevel(rootContext, endpoint, new StaticAuthTransformerFactory("TEST"), - true, 10, 10*1024, - new TransformationLoader().getTransformerFactoryLoader(endpoint.getHost()), + new TransformationLoader().getTransformerFactoryLoader(endpoint.getHost()), true, 10, 10*1024, targetConnectionPoolPrefix); } catch (SSLException e) { throw new RuntimeException(e);