Add support for tracer to create new TracerContextStorage

Signed-off-by: Gagan Juneja <gjjuneja@amazon.com>
opensearch-project · Jul 23, 2023 · b6f0917 · b6f0917
1 parent 7769682
commit b6f0917
Show file tree

Hide file tree

Showing 8 changed files with 104 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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
@@ -8,6 +8,8 @@
 
 package org.opensearch.telemetry.tracing;
 
+import org.opensearch.common.lease.Releasable;
+
 /**
  * Storage interface used for storing tracing context
  * @param <K> key type
@@ -34,4 +36,10 @@ public interface TracerContextStorage<K, V> {
      * @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
@@ -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/...rc/main/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorage.java b/...rc/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<String, Object> transients(Map<String, Object> source) {
         final Map<String, Object> 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
@@ -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
@@ -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<SpanScope> 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<SpanScope> 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());
+    }
+}