Fix an edge case where a socketChannel might not have been created ev…

…en though the channel context was. Also make minor fixes to other tests to make them more resilient. Signed-off-by: Greg Schohn <[email protected]>
opensearch-project · Feb 7, 2024 · 16a9010 · 16a9010
1 parent 48a45ab
commit 16a9010
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 6 deletions.
diff --git a/...rafficReplayer/src/main/java/org/opensearch/migrations/replay/tracing/ReplayContexts.java b/...rafficReplayer/src/main/java/org/opensearch/migrations/replay/tracing/ReplayContexts.java
@@ -125,9 +125,11 @@ public void onSocketConnectionCreated() {
 
         @Override
         public void onSocketConnectionClosed() {
-            assert socketContext != null;
-            socketContext.close();
-            socketContext = null;
+            if (socketContext != null) {
+                try (var toClose = socketContext) {
+                    socketContext = null;
+                }
+            }
             meterIncrementEvent(getMetrics().channelClosedCounter);
             meterDeltaEvent(getMetrics().activeSocketConnectionsCounter, -1);
         }

diff --git a/...e/trafficReplayer/src/test/java/org/opensearch/migrations/replay/tracing/TracingTest.java b/...e/trafficReplayer/src/test/java/org/opensearch/migrations/replay/tracing/TracingTest.java
@@ -71,8 +71,7 @@ private void checkSpans(List<SpanData> recordedSpans) {
             try {
                 return f.get(null);
             } catch (Exception e) {
-                Lombok.sneakyThrow(e);
-                return null;
+                throw Lombok.sneakyThrow(e);
             }
         }).toArray(String[]::new);
         Stream.of(keys).forEach(spanName -> {

diff --git a/.../trafficReplayer/src/testFixtures/java/org/opensearch/migrations/tracing/TestContext.java b/.../trafficReplayer/src/testFixtures/java/org/opensearch/migrations/tracing/TestContext.java
@@ -18,6 +18,7 @@ public class TestContext extends RootReplayerContext implements AutoCloseable {
     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));
@@ -48,7 +49,9 @@ public void onContextClosed(IScopedInstrumentationAttributes newScopedContext) {
     }
 
     public IReplayContexts.ITrafficStreamsLifecycleContext createTrafficStreamContextForTest(ITrafficStreamKey tsk) {
-        return createTrafficStreamContextForStreamSource(channelContextManager.retainOrCreateContext(tsk), tsk);
+        synchronized (channelContextManagerLock) {
+            return createTrafficStreamContextForStreamSource(channelContextManager.retainOrCreateContext(tsk), tsk);
+        }
     }
 
     @Override