diff --git a/server/src/main/java/org/opensearch/transport/TransportService.java b/server/src/main/java/org/opensearch/transport/TransportService.java index d8fb8205f3dc4..ecd801e6c3f76 100644 --- a/server/src/main/java/org/opensearch/transport/TransportService.java +++ b/server/src/main/java/org/opensearch/transport/TransportService.java @@ -1043,6 +1043,7 @@ public void onRejection(Exception e) { @Override public void onFailure(Exception e) { logger.error("Failing from onFailure ", e); + contextToNotify.handler().handleRejection(e); logger.warn( () -> new ParameterizedMessage( "failed to notify response handler on exception, action: {}", diff --git a/test/telemetry/src/main/java/org/opensearch/test/telemetry/tracing/MockTracingTelemetry.java b/test/telemetry/src/main/java/org/opensearch/test/telemetry/tracing/MockTracingTelemetry.java index 7f22757eecd29..64f0299aca614 100644 --- a/test/telemetry/src/main/java/org/opensearch/test/telemetry/tracing/MockTracingTelemetry.java +++ b/test/telemetry/src/main/java/org/opensearch/test/telemetry/tracing/MockTracingTelemetry.java @@ -13,12 +13,15 @@ import org.opensearch.telemetry.tracing.TracingTelemetry; import org.opensearch.telemetry.tracing.attributes.Attributes; +import java.util.concurrent.atomic.AtomicBoolean; + /** * Mock {@link TracingTelemetry} implementation for testing. */ public class MockTracingTelemetry implements TracingTelemetry { private final SpanProcessor spanProcessor = new StrictCheckSpanProcessor(); + private final AtomicBoolean shutdown = new AtomicBoolean(); /** * Base constructor. @@ -28,7 +31,9 @@ public MockTracingTelemetry() {} @Override public Span createSpan(String spanName, Span parentSpan, Attributes attributes) { Span span = new MockSpan(spanName, parentSpan, spanProcessor, attributes); - spanProcessor.onStart(span); + if (shutdown.get() == false) { + spanProcessor.onStart(span); + } return span; } @@ -39,7 +44,8 @@ public TracingContextPropagator getContextPropagator() { @Override public void close() { - StrictCheckSpanProcessor.shutdown(); + // StrictCheckSpanProcessor.shutdown(); + shutdown.set(true); } } diff --git a/test/telemetry/src/main/java/org/opensearch/test/telemetry/tracing/StrictCheckSpanProcessor.java b/test/telemetry/src/main/java/org/opensearch/test/telemetry/tracing/StrictCheckSpanProcessor.java index ac26336de6f08..756ab183542dc 100644 --- a/test/telemetry/src/main/java/org/opensearch/test/telemetry/tracing/StrictCheckSpanProcessor.java +++ b/test/telemetry/src/main/java/org/opensearch/test/telemetry/tracing/StrictCheckSpanProcessor.java @@ -31,7 +31,7 @@ public StrictCheckSpanProcessor() { } private void initProcessor() { - restartTheProcessor(); + // restartTheProcessor(); } private static Map spanMap = new ConcurrentHashMap<>(); @@ -81,7 +81,7 @@ private MockSpanData toMockSpanData(Span span) { */ public static synchronized void shutdown() { if (shutdown.get() == false) { - shutdown.set(true); + // shutdown.set(true); } } @@ -109,7 +109,7 @@ public static void validateTracingStateOnShutdown() { spanMap.clear(); throw e; } finally { - restartTheProcessor(); + // restartTheProcessor(); } }