[Tracing Framework] Add support for SpanKind (opensearch-project#10122)

* [Tracing Framework] Add support for SpanKind Signed-off-by: Gagan Juneja <gjjuneja@amazon.com> * Add CHANGELOG Signed-off-by: Gagan Juneja <gjjuneja@amazon.com> * Fix test case Signed-off-by: Gagan Juneja <gjjuneja@amazon.com> * Address review comment Signed-off-by: Gagan Juneja <gjjuneja@amazon.com> * Address review comment Signed-off-by: Gagan Juneja <gjjuneja@amazon.com> * Address review comment Signed-off-by: Gagan Juneja <gjjuneja@amazon.com> --------- Signed-off-by: Gagan Juneja <gjjuneja@amazon.com> Co-authored-by: Gagan Juneja <gjjuneja@amazon.com>
samuelcostae · Sep 25, 2023 · 19fa37c · 19fa37c
1 parent 346bcfc
commit 19fa37c
Show file tree

Hide file tree

Showing 21 changed files with 353 additions and 240 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -108,6 +108,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Mute the query profile IT with concurrent execution ([#9840](https://github.com/opensearch-project/OpenSearch/pull/9840))
 - Force merge with `only_expunge_deletes` honors max segment size ([#10036](https://github.com/opensearch-project/OpenSearch/pull/10036))
 - Add instrumentation in transport service. ([#10042](https://github.com/opensearch-project/OpenSearch/pull/10042))
+- [Tracing Framework] Add support for SpanKind. ([#10122](https://github.com/opensearch-project/OpenSearch/pull/10122))
 
 ### Deprecated
 
@@ -122,4 +123,4 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Security
 
 [Unreleased 3.0]: https://github.com/opensearch-project/OpenSearch/compare/2.x...HEAD
-[Unreleased 2.x]: https://github.com/opensearch-project/OpenSearch/compare/2.11...2.x
+[Unreleased 2.x]: https://github.com/opensearch-project/OpenSearch/compare/2.11...2.x
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
@@ -9,7 +9,6 @@
 package org.opensearch.telemetry.tracing;
 
 import org.opensearch.common.annotation.InternalApi;
-import org.opensearch.telemetry.tracing.attributes.Attributes;
 
 import java.io.Closeable;
 import java.io.IOException;
@@ -44,27 +43,13 @@ public DefaultTracer(TracingTelemetry tracingTelemetry, TracerContextStorage<Str
 
     @Override
     public Span startSpan(SpanCreationContext context) {
-        return startSpan(context.getSpanName(), context.getAttributes());
-    }
-
-    @Override
-    public Span startSpan(String spanName) {
-        return startSpan(spanName, Attributes.EMPTY);
-    }
-
-    @Override
-    public Span startSpan(String spanName, Attributes attributes) {
-        return startSpan(spanName, (SpanContext) null, attributes);
-    }
-
-    @Override
-    public Span startSpan(String spanName, SpanContext parentSpan, Attributes attributes) {
-        Span span = null;
-        if (parentSpan != null) {
-            span = createSpan(spanName, parentSpan.getSpan(), attributes);
+        Span parentSpan = null;
+        if (context.getParent() != null) {
+            parentSpan = context.getParent().getSpan();
         } else {
-            span = createSpan(spanName, getCurrentSpanInternal(), attributes);
+            parentSpan = getCurrentSpanInternal();
         }
+        Span span = createSpan(context, parentSpan);
         setCurrentSpanInContext(span);
         addDefaultAttributes(span);
         return span;
@@ -87,12 +72,7 @@ public SpanContext getCurrentSpan() {
 
     @Override
     public ScopedSpan startScopedSpan(SpanCreationContext spanCreationContext) {
-        return startScopedSpan(spanCreationContext, null);
-    }
-
-    @Override
-    public ScopedSpan startScopedSpan(SpanCreationContext spanCreationContext, SpanContext parentSpan) {
-        Span span = startSpan(spanCreationContext.getSpanName(), parentSpan, spanCreationContext.getAttributes());
+        Span span = startSpan(spanCreationContext);
         SpanScope spanScope = withSpanInScope(span);
         return new DefaultScopedSpan(span, spanScope);
     }
@@ -102,8 +82,8 @@ public SpanScope withSpanInScope(Span span) {
         return DefaultSpanScope.create(span, tracerContextStorage).attach();
     }
 
-    private Span createSpan(String spanName, Span parentSpan, Attributes attributes) {
-        return tracingTelemetry.createSpan(spanName, parentSpan, attributes);
+    private Span createSpan(SpanCreationContext spanCreationContext, Span parentSpan) {
+        return tracingTelemetry.createSpan(spanCreationContext, parentSpan);
     }
 
     private void setCurrentSpanInContext(Span span) {
@@ -121,11 +101,7 @@ protected void addDefaultAttributes(Span span) {
     @Override
     public Span startSpan(SpanCreationContext spanCreationContext, Map<String, List<String>> headers) {
         Optional<Span> propagatedSpan = tracingTelemetry.getContextPropagator().extractFromHeaders(headers);
-        return startSpan(
-            spanCreationContext.getSpanName(),
-            propagatedSpan.map(SpanContext::new).orElse(null),
-            spanCreationContext.getAttributes()
-        );
+        return startSpan(spanCreationContext.parent(propagatedSpan.map(SpanContext::new).orElse(null)));
     }
 
 }
diff --git a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/SpanCreationContext.java b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/SpanCreationContext.java
@@ -18,17 +18,74 @@
  */
 @ExperimentalApi
 public final class SpanCreationContext {
-    private final String spanName;
-    private final Attributes attributes;
+    private String spanName;
+    private Attributes attributes;
+    private SpanKind spanKind = SpanKind.INTERNAL;
+    private SpanContext parent;
 
     /**
      * Constructor.
+     */
+    private SpanCreationContext() {}
+
+    /**
+     * Sets the span type to server.
+     * @return spanCreationContext
+     */
+    public static SpanCreationContext server() {
+        SpanCreationContext spanCreationContext = new SpanCreationContext();
+        spanCreationContext.spanKind = SpanKind.SERVER;
+        return spanCreationContext;
+    }
+
+    /**
+     * Sets the span type to client.
+     * @return spanCreationContext
+     */
+    public static SpanCreationContext client() {
+        SpanCreationContext spanCreationContext = new SpanCreationContext();
+        spanCreationContext.spanKind = SpanKind.CLIENT;
+        return spanCreationContext;
+    }
+
+    /**
+     * Sets the span type to internal.
+     * @return spanCreationContext
+     */
+    public static SpanCreationContext internal() {
+        SpanCreationContext spanCreationContext = new SpanCreationContext();
+        spanCreationContext.spanKind = SpanKind.INTERNAL;
+        return spanCreationContext;
+    }
+
+    /**
+     * Sets the span name.
      * @param spanName span name.
-     * @param attributes attributes.
+     * @return spanCreationContext
      */
-    public SpanCreationContext(String spanName, Attributes attributes) {
+    public SpanCreationContext name(String spanName) {
         this.spanName = spanName;
+        return this;
+    }
+
+    /**
+     * Sets the span attributes.
+     * @param attributes attributes.
+     * @return spanCreationContext
+     */
+    public SpanCreationContext attributes(Attributes attributes) {
         this.attributes = attributes;
+        return this;
+    }
+
+    /**
+     * Sets the parent for spann
+     * @param parent parent
+     * @return spanCreationContext
+     */
+    public SpanCreationContext parent(SpanContext parent) {
+        this.parent = parent;
+        return this;
     }
 
     /**
@@ -46,4 +103,20 @@ public String getSpanName() {
     public Attributes getAttributes() {
         return attributes;
     }
+
+    /**
+     * Returns the span kind.
+     * @return spankind.
+     */
+    public SpanKind getSpanKind() {
+        return spanKind;
+    }
+
+    /**
+     * Returns the parent span
+     * @return parent.
+     */
+    public SpanContext getParent() {
+        return parent;
+    }
 }
diff --git a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/SpanKind.java b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/SpanKind.java
@@ -0,0 +1,31 @@
+/*
+ * 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 org.opensearch.common.annotation.PublicApi;
+
+/**
+ * Type of Span.
+ */
+@PublicApi(since = "2.11.0")
+public enum SpanKind {
+    /**
+     * Span represents the client side code.
+     */
+    CLIENT,
+    /**
+     * Span represents the server side code.
+     */
+    SERVER,
+
+    /**
+     * Span represents the internal calls. This is the default value of a span type.
+     */
+    INTERNAL;
+}
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,7 +9,6 @@
 package org.opensearch.telemetry.tracing;
 
 import org.opensearch.common.annotation.ExperimentalApi;
-import org.opensearch.telemetry.tracing.attributes.Attributes;
 import org.opensearch.telemetry.tracing.http.HttpTracer;
 
 import java.io.Closeable;
@@ -32,34 +31,6 @@ public interface Tracer extends HttpTracer, Closeable {
      */
     Span startSpan(SpanCreationContext context);
 
-    /**
-     * Starts the {@link Span} with given name
-     *
-     * @param spanName span name
-     * @return span, must be closed.
-     */
-    Span startSpan(String spanName);
-
-    /**
-     * Starts the {@link Span} with given name and attributes. This is required in cases when some attribute based
-     * decision needs to be made before starting the span. Very useful in the case of Sampling.
-     *
-     * @param spanName   span name.
-     * @param attributes attributes to be added.
-     * @return span, must be closed.
-     */
-    Span startSpan(String spanName, Attributes attributes);
-
-    /**
-     * Starts the {@link Span} with the given name, parent and attributes.
-     *
-     * @param spanName   span name.
-     * @param parentSpan parent span.
-     * @param attributes attributes to be added.
-     * @return span, must be closed.
-     */
-    Span startSpan(String spanName, SpanContext parentSpan, Attributes attributes);
-
     /**
      * Returns the current span.
      * @return current wrapped span.
@@ -74,15 +45,6 @@ public interface Tracer extends HttpTracer, Closeable {
      */
     ScopedSpan startScopedSpan(SpanCreationContext spanCreationContext);
 
-    /**
-     * Start the span and scoped it. This must be used for scenarios where {@link SpanScope} and {@link Span} lifecycles
-     * are same and ends within the same thread where created.
-     * @param spanCreationContext span creation context
-     * @param parentSpan parent span.
-     * @return scope of the span, must be closed with explicit close or with try-with-resource
-     */
-    ScopedSpan startScopedSpan(SpanCreationContext spanCreationContext, SpanContext parentSpan);
-
     /**
      * Creates the Span Scope for a current thread. It's mandatory to scope the span just after creation so that it will
      * automatically manage the attach /detach to the current thread.

diff --git a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/TracingTelemetry.java b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/TracingTelemetry.java
@@ -9,7 +9,6 @@
 package org.opensearch.telemetry.tracing;
 
 import org.opensearch.common.annotation.ExperimentalApi;
-import org.opensearch.telemetry.tracing.attributes.Attributes;
 
 import java.io.Closeable;
 
@@ -24,12 +23,11 @@ public interface TracingTelemetry extends Closeable {
     /**
      * Creates span with provided arguments
      *
-     * @param spanName   name of the span
-     * @param parentSpan span's parent span
-     * @param attributes attributes to be added.
+     * @param spanCreationContext span creation context.
+     * @param parentSpan parent span.
      * @return span instance
      */
-    Span createSpan(String spanName, Span parentSpan, Attributes attributes);
+    Span createSpan(SpanCreationContext spanCreationContext, Span parentSpan);
 
     /**
      * provides tracing context propagator

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
@@ -15,7 +15,6 @@
 import org.opensearch.telemetry.tracing.SpanCreationContext;
 import org.opensearch.telemetry.tracing.SpanScope;
 import org.opensearch.telemetry.tracing.Tracer;
-import org.opensearch.telemetry.tracing.attributes.Attributes;
 
 import java.util.List;
 import java.util.Map;
@@ -40,21 +39,6 @@ public Span startSpan(SpanCreationContext context) {
         return NoopSpan.INSTANCE;
     }
 
-    @Override
-    public Span startSpan(String spanName) {
-        return NoopSpan.INSTANCE;
-    }
-
-    @Override
-    public Span startSpan(String spanName, Attributes attributes) {
-        return NoopSpan.INSTANCE;
-    }
-
-    @Override
-    public Span startSpan(String spanName, SpanContext parentSpan, Attributes attributes) {
-        return NoopSpan.INSTANCE;
-    }
-
     @Override
     public SpanContext getCurrentSpan() {
         return new SpanContext(NoopSpan.INSTANCE);
@@ -65,11 +49,6 @@ public ScopedSpan startScopedSpan(SpanCreationContext spanCreationContext) {
         return ScopedSpan.NO_OP;
     }
 
-    @Override
-    public ScopedSpan startScopedSpan(SpanCreationContext spanCreationContext, SpanContext parentSpan) {
-        return ScopedSpan.NO_OP;
-    }
-
     @Override
     public SpanScope withSpanInScope(Span span) {
         return SpanScope.NO_OP;

diff --git a/.../telemetry/src/main/java/org/opensearch/telemetry/tracing/runnable/TraceableRunnable.java b/.../telemetry/src/main/java/org/opensearch/telemetry/tracing/runnable/TraceableRunnable.java
@@ -9,40 +9,32 @@
 package org.opensearch.telemetry.tracing.runnable;
 
 import org.opensearch.telemetry.tracing.ScopedSpan;
-import org.opensearch.telemetry.tracing.SpanContext;
 import org.opensearch.telemetry.tracing.SpanCreationContext;
 import org.opensearch.telemetry.tracing.Tracer;
-import org.opensearch.telemetry.tracing.attributes.Attributes;
 
 /**
  * Wraps the runnable and add instrumentation to trace the {@link Runnable}
  */
 public class TraceableRunnable implements Runnable {
     private final Runnable runnable;
-    private final SpanContext parent;
+    private final SpanCreationContext spanCreationContext;
     private final Tracer tracer;
-    private final String spanName;
-    private final Attributes attributes;
 
     /**
      * Constructor.
      * @param tracer tracer
-     * @param spanName spanName
-     * @param parent parent Span.
-     * @param attributes attributes.
+     * @param spanCreationContext spanCreationContext
      * @param runnable runnable.
      */
-    public TraceableRunnable(Tracer tracer, String spanName, SpanContext parent, Attributes attributes, Runnable runnable) {
+    public TraceableRunnable(Tracer tracer, SpanCreationContext spanCreationContext, Runnable runnable) {
         this.tracer = tracer;
-        this.spanName = spanName;
-        this.parent = parent;
-        this.attributes = attributes;
+        this.spanCreationContext = spanCreationContext;
         this.runnable = runnable;
     }
 
     @Override
     public void run() {
-        try (ScopedSpan spanScope = tracer.startScopedSpan(new SpanCreationContext(spanName, attributes), parent)) {
+        try (ScopedSpan spanScope = tracer.startScopedSpan(spanCreationContext)) {
             runnable.run();
         }
     }