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 efebfc1ad..ceccfbe9d 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
@@ -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/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 f5954363e..2d37bb892 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
@@ -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/TrafficCapture/trafficReplayer/src/testFixtures/java/org/opensearch/migrations/tracing/TestContext.java b/TrafficCapture/trafficReplayer/src/testFixtures/java/org/opensearch/migrations/tracing/TestContext.java
index 145eaf57e..3c7155504 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
@@ -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