diff --git a/test/framework/src/main/java/org/opensearch/test/OpenSearchIntegTestCase.java b/test/framework/src/main/java/org/opensearch/test/OpenSearchIntegTestCase.java index 63fe5ec03b059..f06a6e26defeb 100644 --- a/test/framework/src/main/java/org/opensearch/test/OpenSearchIntegTestCase.java +++ b/test/framework/src/main/java/org/opensearch/test/OpenSearchIntegTestCase.java @@ -157,6 +157,7 @@ import org.opensearch.test.disruption.ServiceDisruptionScheme; import org.opensearch.test.store.MockFSIndexStore; import org.opensearch.test.telemetry.MockTelemetryPlugin; +import org.opensearch.test.telemetry.tracing.StrictCheckSpanProcessor; import org.opensearch.test.transport.MockTransportService; import org.opensearch.transport.TransportInterceptor; import org.opensearch.transport.TransportRequest; @@ -2295,7 +2296,9 @@ public static void afterClass() throws Exception { INSTANCE.printTestMessage("cleaning up after"); INSTANCE.afterInternal(true); checkStaticState(true); + StrictCheckSpanProcessor.validateTracingStateOnShutdown(); } + } finally { SUITE_SEED = null; currentCluster = null; diff --git a/test/framework/src/main/java/org/opensearch/test/OpenSearchSingleNodeTestCase.java b/test/framework/src/main/java/org/opensearch/test/OpenSearchSingleNodeTestCase.java index 2de4c8fdbdfb8..f14fe3bf3961c 100644 --- a/test/framework/src/main/java/org/opensearch/test/OpenSearchSingleNodeTestCase.java +++ b/test/framework/src/main/java/org/opensearch/test/OpenSearchSingleNodeTestCase.java @@ -73,6 +73,7 @@ import org.opensearch.search.internal.SearchContext; import org.opensearch.telemetry.TelemetrySettings; import org.opensearch.test.telemetry.MockTelemetryPlugin; +import org.opensearch.test.telemetry.tracing.StrictCheckSpanProcessor; import org.opensearch.transport.TransportSettings; import org.junit.AfterClass; import org.junit.BeforeClass; @@ -190,6 +191,7 @@ public static void setUpClass() throws Exception { @AfterClass public static void tearDownClass() throws Exception { stopNode(); + StrictCheckSpanProcessor.validateTracingStateOnShutdown(); } /** diff --git a/test/framework/src/main/java/org/opensearch/test/telemetry/MockTelemetry.java b/test/framework/src/main/java/org/opensearch/test/telemetry/MockTelemetry.java index 894e8a67cea1f..6b428a7f65594 100644 --- a/test/framework/src/main/java/org/opensearch/test/telemetry/MockTelemetry.java +++ b/test/framework/src/main/java/org/opensearch/test/telemetry/MockTelemetry.java @@ -13,46 +13,22 @@ import org.opensearch.telemetry.metrics.MetricsTelemetry; import org.opensearch.telemetry.tracing.TracingTelemetry; import org.opensearch.test.telemetry.tracing.MockTracingTelemetry; -import org.opensearch.threadpool.ThreadPool; - -import java.util.concurrent.TimeUnit; /** * Mock {@link Telemetry} implementation for testing. */ public class MockTelemetry implements Telemetry { - private final ThreadPool threadPool; - /** * Constructor with settings. * @param settings telemetry settings. */ public MockTelemetry(TelemetrySettings settings) { - this(settings, null); - } - /** - * Constructor with settings. - * @param settings telemetry settings. - * @param threadPool thread pool to watch for termination - */ - public MockTelemetry(TelemetrySettings settings, ThreadPool threadPool) { - this.threadPool = threadPool; } @Override public TracingTelemetry getTracingTelemetry() { - return new MockTracingTelemetry(() -> { - // There could be some asynchronous tasks running that we should await for before the closing - // up the tracer instance. - if (threadPool != null) { - try { - threadPool.awaitTermination(10, TimeUnit.SECONDS); - } catch (final InterruptedException ex) { - Thread.currentThread().interrupt(); - } - } - }); + return new MockTracingTelemetry(); } @Override diff --git a/test/framework/src/main/java/org/opensearch/test/telemetry/MockTelemetryPlugin.java b/test/framework/src/main/java/org/opensearch/test/telemetry/MockTelemetryPlugin.java index ebba9857aa8f1..4f483098caf82 100644 --- a/test/framework/src/main/java/org/opensearch/test/telemetry/MockTelemetryPlugin.java +++ b/test/framework/src/main/java/org/opensearch/test/telemetry/MockTelemetryPlugin.java @@ -8,34 +8,18 @@ package org.opensearch.test.telemetry; -import org.opensearch.client.Client; -import org.opensearch.cluster.metadata.IndexNameExpressionResolver; -import org.opensearch.cluster.service.ClusterService; -import org.opensearch.common.SetOnce; -import org.opensearch.core.common.io.stream.NamedWriteableRegistry; -import org.opensearch.core.xcontent.NamedXContentRegistry; -import org.opensearch.env.Environment; -import org.opensearch.env.NodeEnvironment; import org.opensearch.plugins.Plugin; import org.opensearch.plugins.TelemetryPlugin; -import org.opensearch.repositories.RepositoriesService; -import org.opensearch.script.ScriptService; import org.opensearch.telemetry.Telemetry; import org.opensearch.telemetry.TelemetrySettings; -import org.opensearch.threadpool.ThreadPool; -import org.opensearch.watcher.ResourceWatcherService; -import java.util.Collection; -import java.util.Collections; import java.util.Optional; -import java.util.function.Supplier; /** * Mock {@link TelemetryPlugin} implementation for testing. */ public class MockTelemetryPlugin extends Plugin implements TelemetryPlugin { private static final String MOCK_TRACER_NAME = "mock"; - private final SetOnce threadPool = new SetOnce<>(); /** * Base constructor. @@ -44,27 +28,9 @@ public MockTelemetryPlugin() { } - @Override - public Collection createComponents( - Client client, - ClusterService clusterService, - ThreadPool threadPool, - ResourceWatcherService resourceWatcherService, - ScriptService scriptService, - NamedXContentRegistry xContentRegistry, - Environment environment, - NodeEnvironment nodeEnvironment, - NamedWriteableRegistry namedWriteableRegistry, - IndexNameExpressionResolver indexNameExpressionResolver, - Supplier repositoriesServiceSupplier - ) { - this.threadPool.set(threadPool); - return Collections.emptyList(); - } - @Override public Optional getTelemetry(TelemetrySettings settings) { - return Optional.of(new MockTelemetry(settings, threadPool.get())); + return Optional.of(new MockTelemetry(settings)); } @Override 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 a5e51dd27541b..c7f5943719230 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 @@ -12,11 +12,8 @@ import org.opensearch.telemetry.tracing.TracingContextPropagator; import org.opensearch.telemetry.tracing.TracingTelemetry; import org.opensearch.telemetry.tracing.attributes.Attributes; -import org.opensearch.test.telemetry.tracing.validators.AllSpansAreEndedProperly; -import org.opensearch.test.telemetry.tracing.validators.AllSpansHaveUniqueId; -import java.util.Arrays; -import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; /** * Mock {@link TracingTelemetry} implementation for testing. @@ -24,28 +21,19 @@ public class MockTracingTelemetry implements TracingTelemetry { private final SpanProcessor spanProcessor = new StrictCheckSpanProcessor(); - private final Runnable onClose; + private final AtomicBoolean shutdown = new AtomicBoolean(false); /** * Base constructor. */ - public MockTracingTelemetry() { - this(() -> {}); - } - - /** - * Base constructor. - * - * @param onClose on close hook - */ - public MockTracingTelemetry(final Runnable onClose) { - this.onClose = onClose; - } + 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; } @@ -56,15 +44,7 @@ public TracingContextPropagator getContextPropagator() { @Override public void close() { - // Run onClose hook - onClose.run(); - - List spanData = ((StrictCheckSpanProcessor) spanProcessor).getFinishedSpanItems(); - if (spanData.size() != 0) { - TelemetryValidators validators = new TelemetryValidators( - Arrays.asList(new AllSpansAreEndedProperly(), new AllSpansHaveUniqueId()) - ); - validators.validate(spanData, 1); - } + 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 e3fca8813b696..c6e57531d23df 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 @@ -9,8 +9,11 @@ package org.opensearch.test.telemetry.tracing; import org.opensearch.telemetry.tracing.Span; +import org.opensearch.test.telemetry.tracing.validators.AllSpansAreEndedProperly; +import org.opensearch.test.telemetry.tracing.validators.AllSpansHaveUniqueId; import java.util.ArrayList; +import java.util.Arrays; import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; @@ -61,4 +64,23 @@ private MockSpanData toMockSpanData(Span span) { ); return spanData; } + + /** + * Ensures the strict check succeeds for all the spans. + */ + public static void validateTracingStateOnShutdown() { + List spanData = new ArrayList<>(spanMap.values()); + if (spanData.size() != 0) { + TelemetryValidators validators = new TelemetryValidators( + Arrays.asList(new AllSpansAreEndedProperly(), new AllSpansHaveUniqueId()) + ); + try { + validators.validate(spanData, 1); + } catch (Error e) { + spanMap.clear(); + throw e; + } + } + + } }