Adds support for tracing runnable scenarios

CHANGELOG.md

@@ -46,6 +46,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Change http code on create index API with bad input raising NotXContentException from 500 to 400 ([#4773](https://github.com/opensearch-project/OpenSearch/pull/4773))
 - Improve summary error message for invalid setting updates ([#4792](https://github.com/opensearch-project/OpenSearch/pull/4792))
 - Remote Segment Store Repository setting moved from `index.remote_store.repository` to `index.remote_store.segment.repository` and `cluster.remote_store.repository` to `cluster.remote_store.segment.repository` respectively for Index and Cluster level settings ([#8719](https://github.com/opensearch-project/OpenSearch/pull/8719))
+- Add support for tracer to create new TracerContextStorage ([#8831](https://github.com/opensearch-project/OpenSearch/pull/8831))
 
 ### Deprecated
 
@@ -88,6 +89,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Make Span exporter configurable ([#8620](https://github.com/opensearch-project/OpenSearch/issues/8620))
 - Change InternalSignificantTerms to sum shard-level superset counts only in final reduce ([#8735](https://github.com/opensearch-project/OpenSearch/pull/8735))
 - Exclude 'benchmarks' from codecov report ([#8805](https://github.com/opensearch-project/OpenSearch/pull/8805))
+- Add support for tracer to create new TracerContextStorage ([#8831](https://github.com/opensearch-project/OpenSearch/pull/8831))
 
 ### Deprecated
 

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}"

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 @@
         span.addAttribute(THREAD_NAME, Thread.currentThread().getName());
     }
 
+    @Override
+    public Releasable newTracerContextStorage() {
+        return tracerContextStorage.newTracerContextStorage();
+    }
+
 }

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();
+
 }

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();
 }

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 void close() {
 
     }
+
+    @Override
+    public Releasable newTracerContextStorage() {
+        return () -> {};
+    }
 }

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<String, Object> transients(Map<String, Object> source) {
         final Map<String, Object> transients = new HashMap<>();

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 @@
         return delegateTracer.startSpan(spanName);
     }
 
+    @Override
+    public Releasable newTracerContextStorage() {
+        return defaultTracer.newTracerContextStorage();
+    }
+
     @Override
     public void close() throws IOException {
         defaultTracer.close();

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());
+    }
+}