From b6f0917314320ce635621afd505646c0ec45143d Mon Sep 17 00:00:00 2001 From: Gagan Juneja Date: Sun, 23 Jul 2023 20:13:58 +0530 Subject: [PATCH] Add support for tracer to create new TracerContextStorage Signed-off-by: Gagan Juneja --- libs/telemetry/build.gradle | 1 + .../telemetry/tracing/DefaultTracer.java | 6 ++ .../opensearch/telemetry/tracing/Tracer.java | 7 +++ .../tracing/TracerContextStorage.java | 8 +++ .../telemetry/tracing/noop/NoopTracer.java | 6 ++ ...hreadContextBasedTracerContextStorage.java | 7 +++ .../telemetry/tracing/WrappedTracer.java | 6 ++ .../telemetry/tracing/TracerTests.java | 63 +++++++++++++++++++ 8 files changed, 104 insertions(+) create mode 100644 server/src/test/java/org/opensearch/telemetry/tracing/TracerTests.java diff --git a/libs/telemetry/build.gradle b/libs/telemetry/build.gradle index ce94698836b4f..9e00feb6bedde 100644 --- a/libs/telemetry/build.gradle +++ b/libs/telemetry/build.gradle @@ -10,6 +10,7 @@ */ dependencies { + api project(':libs:opensearch-common') testImplementation "com.carrotsearch.randomizedtesting:randomizedtesting-runner:${versions.randomizedrunner}" testImplementation "junit:junit:${versions.junit}" testImplementation "org.hamcrest:hamcrest:${versions.hamcrest}" diff --git a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultTracer.java b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultTracer.java index ea59eec645420..30ca6de22e3d1 100644 --- a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultTracer.java +++ b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultTracer.java @@ -10,6 +10,7 @@ import java.io.Closeable; import java.io.IOException; +import org.opensearch.common.lease.Releasable; /** * @@ -76,4 +77,9 @@ protected void addDefaultAttributes(Span span) { span.addAttribute(THREAD_NAME, Thread.currentThread().getName()); } + @Override + public Releasable newTracerContextStorage() { + return tracerContextStorage.newTracerContextStorage(); + } + } diff --git a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/Tracer.java b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/Tracer.java index d422b58aa0a9f..79e6c6b0a3304 100644 --- a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/Tracer.java +++ b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/Tracer.java @@ -9,6 +9,7 @@ package org.opensearch.telemetry.tracing; import java.io.Closeable; +import org.opensearch.common.lease.Releasable; /** * Tracer is the interface used to create a {@link Span} @@ -26,4 +27,10 @@ public interface Tracer extends Closeable { */ SpanScope startSpan(String spanName); + /** + * Creates new TracerContextStorage. + * @return returns {@link Releasable} which should be able to release this TracerContextStorage. + */ + Releasable newTracerContextStorage(); + } diff --git a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/TracerContextStorage.java b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/TracerContextStorage.java index d85b404b0ce41..d46412bdbab4a 100644 --- a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/TracerContextStorage.java +++ b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/TracerContextStorage.java @@ -8,6 +8,8 @@ package org.opensearch.telemetry.tracing; +import org.opensearch.common.lease.Releasable; + /** * Storage interface used for storing tracing context * @param key type @@ -34,4 +36,10 @@ public interface TracerContextStorage { * @param value of the tracing context */ void put(K key, V value); + + /** + * Creates new TracerContextStorage. + * @return returns {@link Releasable} which should be able to release this TracerContextStorage. + */ + Releasable newTracerContextStorage(); } diff --git a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/noop/NoopTracer.java b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/noop/NoopTracer.java index a1768d7d59116..3bde90a7020e6 100644 --- a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/noop/NoopTracer.java +++ b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/noop/NoopTracer.java @@ -8,6 +8,7 @@ package org.opensearch.telemetry.tracing.noop; +import org.opensearch.common.lease.Releasable; import org.opensearch.telemetry.tracing.SpanScope; import org.opensearch.telemetry.tracing.Tracer; @@ -34,4 +35,9 @@ public SpanScope startSpan(String spanName) { public void close() { } + + @Override + public Releasable newTracerContextStorage() { + return () -> {}; + } } diff --git a/server/src/main/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorage.java b/server/src/main/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorage.java index c009ab2391aab..09fc614a6f3d3 100644 --- a/server/src/main/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorage.java +++ b/server/src/main/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorage.java @@ -8,6 +8,7 @@ package org.opensearch.telemetry.tracing; +import org.opensearch.common.lease.Releasable; import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.common.util.concurrent.ThreadContextStatePropagator; @@ -51,6 +52,12 @@ public void put(String key, Span span) { } } + @Override + public Releasable newTracerContextStorage() { + ThreadContext.StoredContext newContext = threadContext.newStoredContext(true); + return () -> newContext.close(); + } + @Override public Map transients(Map source) { final Map transients = new HashMap<>(); diff --git a/server/src/main/java/org/opensearch/telemetry/tracing/WrappedTracer.java b/server/src/main/java/org/opensearch/telemetry/tracing/WrappedTracer.java index 0ba9a8ea5fd88..5607c963c4b31 100644 --- a/server/src/main/java/org/opensearch/telemetry/tracing/WrappedTracer.java +++ b/server/src/main/java/org/opensearch/telemetry/tracing/WrappedTracer.java @@ -8,6 +8,7 @@ package org.opensearch.telemetry.tracing; +import org.opensearch.common.lease.Releasable; import org.opensearch.telemetry.TelemetrySettings; import org.opensearch.telemetry.tracing.noop.NoopTracer; @@ -40,6 +41,11 @@ public SpanScope startSpan(String spanName) { return delegateTracer.startSpan(spanName); } + @Override + public Releasable newTracerContextStorage() { + return defaultTracer.newTracerContextStorage(); + } + @Override public void close() throws IOException { defaultTracer.close(); diff --git a/server/src/test/java/org/opensearch/telemetry/tracing/TracerTests.java b/server/src/test/java/org/opensearch/telemetry/tracing/TracerTests.java new file mode 100644 index 0000000000000..d7f0d42fd27c1 --- /dev/null +++ b/server/src/test/java/org/opensearch/telemetry/tracing/TracerTests.java @@ -0,0 +1,63 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.telemetry.tracing; + +import java.util.ArrayList; +import java.util.List; +import org.opensearch.common.lease.Releasable; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.test.telemetry.tracing.MockTracingTelemetry; + +public class TracerTests extends OpenSearchTestCase { + + private final ThreadContextBasedTracerContextStorage contextStorage = new ThreadContextBasedTracerContextStorage( + new ThreadContext(Settings.EMPTY), + new MockTracingTelemetry() + ); + + public void testIterationScenarioFail() throws Exception { + List spansToBeClosed = new ArrayList<>(); + DefaultTracer defaultTracer = new DefaultTracer(new MockTracingTelemetry(), contextStorage); + try (SpanScope parentSpanScope = defaultTracer.startSpan("parentSpan")) { + Span parentSpan = defaultTracer.getCurrentSpan(); + Span newWrongParent = null; + for (int i = 0; i < 3; i++) { + String spanName = "childSpan_" + i; + SpanScope child = defaultTracer.startSpan(spanName); + spansToBeClosed.add(child); + if (i == 0) { + assertEquals(parentSpan, defaultTracer.getCurrentSpan().getParentSpan()); + } else { + assertEquals(newWrongParent, defaultTracer.getCurrentSpan().getParentSpan()); + } + newWrongParent = defaultTracer.getCurrentSpan(); + } + } + spansToBeClosed.forEach(a -> a.close()); + } + + public void testIterationScenarioSuccess() throws Exception { + List spansToBeClosed = new ArrayList<>(); + DefaultTracer defaultTracer = new DefaultTracer(new MockTracingTelemetry(), contextStorage); + try (SpanScope parentSpanScope = defaultTracer.startSpan("parentSpan")) { + Span parentSpan = defaultTracer.getCurrentSpan(); + for (int i = 0; i < 3; i++) { + String spanName = "childSpan_" + i; + try (Releasable releasable = contextStorage.newTracerContextStorage()) { + SpanScope child = defaultTracer.startSpan(spanName); + spansToBeClosed.add(child); + assertEquals(parentSpan, defaultTracer.getCurrentSpan().getParentSpan()); + } + } + } + spansToBeClosed.forEach(a -> a.close()); + } +}