[Tracing Framework] Add support for SpanKind

Signed-off-by: Gagan Juneja <[email protected]>
Gaganjuneja · Sep 19, 2023 · 82c9732 · 82c9732
1 parent 2a5b124
commit 82c9732
Show file tree

Hide file tree

Showing 18 changed files with 243 additions and 76 deletions.
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
@@ -44,26 +44,26 @@ public DefaultTracer(TracingTelemetry tracingTelemetry, TracerContextStorage<Str
 
     @Override
     public Span startSpan(SpanCreationContext context) {
-        return startSpan(context.getSpanName(), context.getAttributes());
+        return startSpan(context.getSpanName(), context.getAttributes(), context.getSpanKind());
     }
 
     @Override
-    public Span startSpan(String spanName) {
-        return startSpan(spanName, Attributes.EMPTY);
+    public Span startSpan(String spanName, SpanKind spanKind) {
+        return startSpan(spanName, Attributes.EMPTY, spanKind);
     }
 
     @Override
-    public Span startSpan(String spanName, Attributes attributes) {
-        return startSpan(spanName, (SpanContext) null, attributes);
+    public Span startSpan(String spanName, Attributes attributes, SpanKind spanKind) {
+        return startSpan(spanName, (SpanContext) null, attributes, spanKind);
     }
 
     @Override
-    public Span startSpan(String spanName, SpanContext parentSpan, Attributes attributes) {
+    public Span startSpan(String spanName, SpanContext parentSpan, Attributes attributes, SpanKind spanKind) {
         Span span = null;
         if (parentSpan != null) {
-            span = createSpan(spanName, parentSpan.getSpan(), attributes);
+            span = createSpan(spanName, parentSpan.getSpan(), attributes, spanKind);
         } else {
-            span = createSpan(spanName, getCurrentSpanInternal(), attributes);
+            span = createSpan(spanName, getCurrentSpanInternal(), attributes, spanKind);
         }
         setCurrentSpanInContext(span);
         addDefaultAttributes(span);
@@ -92,7 +92,12 @@ public ScopedSpan startScopedSpan(SpanCreationContext spanCreationContext) {
 
     @Override
     public ScopedSpan startScopedSpan(SpanCreationContext spanCreationContext, SpanContext parentSpan) {
-        Span span = startSpan(spanCreationContext.getSpanName(), parentSpan, spanCreationContext.getAttributes());
+        Span span = startSpan(
+            spanCreationContext.getSpanName(),
+            parentSpan,
+            spanCreationContext.getAttributes(),
+            spanCreationContext.getSpanKind()
+        );
         SpanScope spanScope = withSpanInScope(span);
         return new DefaultScopedSpan(span, spanScope);
     }
@@ -102,8 +107,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(String spanName, Span parentSpan, Attributes attributes, SpanKind spanKind) {
+        return tracingTelemetry.createSpan(spanName, parentSpan, attributes, spanKind);
     }
 
     private void setCurrentSpanInContext(Span span) {
@@ -124,7 +129,8 @@ public Span startSpan(SpanCreationContext spanCreationContext, Map<String, List<
         return startSpan(
             spanCreationContext.getSpanName(),
             propagatedSpan.map(SpanContext::new).orElse(null),
-            spanCreationContext.getAttributes()
+            spanCreationContext.getAttributes(),
+            spanCreationContext.getSpanKind()
         );
     }
 

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
@@ -20,15 +20,27 @@
 public final class SpanCreationContext {
     private final String spanName;
     private final Attributes attributes;
+    private final SpanKind spanKind;
 
     /**
      * Constructor.
      * @param spanName span name.
      * @param attributes attributes.
      */
     public SpanCreationContext(String spanName, Attributes attributes) {
+        this(spanName, attributes, SpanKind.INTERNAL);
+    }
+
+    /**
+     * Constructor
+     * @param spanName span name.
+     * @param attributes attributes.
+     * @param spanKind span type.
+     */
+    public SpanCreationContext(String spanName, Attributes attributes, SpanKind spanKind) {
         this.spanName = spanName;
         this.attributes = attributes;
+        this.spanKind = spanKind;
     }
 
     /**
@@ -46,4 +58,12 @@ public String getSpanName() {
     public Attributes getAttributes() {
         return attributes;
     }
+
+    /**
+     * Returns the span kind.
+     * @return spankind.
+     */
+    public SpanKind getSpanKind() {
+        return spanKind;
+    }
 }
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,28 @@
+/*
+ * 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;
+
+/**
+ * Type of Span.
+ */
+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
@@ -36,29 +36,32 @@ public interface Tracer extends HttpTracer, Closeable {
      * Starts the {@link Span} with given name
      *
      * @param spanName span name
+     * @param spanKind type of span.
      * @return span, must be closed.
      */
-    Span startSpan(String spanName);
+    Span startSpan(String spanName, SpanKind spanKind);
 
     /**
      * 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.
+     * @param spanKind type of span.
      * @return span, must be closed.
      */
-    Span startSpan(String spanName, Attributes attributes);
+    Span startSpan(String spanName, Attributes attributes, SpanKind spanKind);
 
     /**
      * 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.
+     * @param spanKind type of span.
      * @return span, must be closed.
      */
-    Span startSpan(String spanName, SpanContext parentSpan, Attributes attributes);
+    Span startSpan(String spanName, SpanContext parentSpan, Attributes attributes, SpanKind spanKind);
 
     /**
      * Returns the current span.

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
@@ -27,9 +27,10 @@ public interface TracingTelemetry extends Closeable {
      * @param spanName   name of the span
      * @param parentSpan span's parent span
      * @param attributes attributes to be added.
+     * @param spanKind type of span.
      * @return span instance
      */
-    Span createSpan(String spanName, Span parentSpan, Attributes attributes);
+    Span createSpan(String spanName, Span parentSpan, Attributes attributes, SpanKind spanKind);
 
     /**
      * 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
@@ -13,6 +13,7 @@
 import org.opensearch.telemetry.tracing.Span;
 import org.opensearch.telemetry.tracing.SpanContext;
 import org.opensearch.telemetry.tracing.SpanCreationContext;
+import org.opensearch.telemetry.tracing.SpanKind;
 import org.opensearch.telemetry.tracing.SpanScope;
 import org.opensearch.telemetry.tracing.Tracer;
 import org.opensearch.telemetry.tracing.attributes.Attributes;
@@ -41,17 +42,17 @@ public Span startSpan(SpanCreationContext context) {
     }
 
     @Override
-    public Span startSpan(String spanName) {
+    public Span startSpan(String spanName, SpanKind spanKind) {
         return NoopSpan.INSTANCE;
     }
 
     @Override
-    public Span startSpan(String spanName, Attributes attributes) {
+    public Span startSpan(String spanName, Attributes attributes, SpanKind spanKind) {
         return NoopSpan.INSTANCE;
     }
 
     @Override
-    public Span startSpan(String spanName, SpanContext parentSpan, Attributes attributes) {
+    public Span startSpan(String spanName, SpanContext parentSpan, Attributes attributes, SpanKind spanKind) {
         return NoopSpan.INSTANCE;
     }
 

diff --git a/libs/telemetry/src/test/java/org/opensearch/telemetry/tracing/DefaultTracerTests.java b/libs/telemetry/src/test/java/org/opensearch/telemetry/tracing/DefaultTracerTests.java
@@ -58,7 +58,7 @@ public void tearDown() throws Exception {
     public void testCreateSpan() {
         DefaultTracer defaultTracer = new DefaultTracer(mockTracingTelemetry, mockTracerContextStorage);
 
-        defaultTracer.startSpan("span_name");
+        defaultTracer.startSpan("span_name", SpanKind.INTERNAL);
 
         assertEquals("span_name", defaultTracer.getCurrentSpan().getSpan().getSpanName());
     }
@@ -67,18 +67,22 @@ public void testCreateSpan() {
     public void testCreateSpanWithAttributesWithMock() {
         DefaultTracer defaultTracer = new DefaultTracer(mockTracingTelemetry, mockTracerContextStorage);
         Attributes attributes = Attributes.create().addAttribute("name", "value");
-        when(mockTracingTelemetry.createSpan(eq("span_name"), eq(mockParentSpan), eq(attributes))).thenReturn(mockSpan);
-        defaultTracer.startSpan("span_name", attributes);
-        verify(mockTracingTelemetry).createSpan(eq("span_name"), eq(mockParentSpan), eq(attributes));
+        when(mockTracingTelemetry.createSpan(eq("span_name"), eq(mockParentSpan), eq(attributes), eq(SpanKind.INTERNAL))).thenReturn(
+            mockSpan
+        );
+        defaultTracer.startSpan("span_name", attributes, SpanKind.INTERNAL);
+        verify(mockTracingTelemetry).createSpan(eq("span_name"), eq(mockParentSpan), eq(attributes), eq(SpanKind.INTERNAL));
     }
 
     @SuppressWarnings("unchecked")
     public void testCreateSpanWithAttributesWithParentMock() {
         DefaultTracer defaultTracer = new DefaultTracer(mockTracingTelemetry, mockTracerContextStorage);
         Attributes attributes = Attributes.create().addAttribute("name", "value");
-        when(mockTracingTelemetry.createSpan(eq("span_name"), eq(mockParentSpan), eq(attributes))).thenReturn(mockSpan);
-        defaultTracer.startSpan("span_name", new SpanContext(mockParentSpan), attributes);
-        verify(mockTracingTelemetry).createSpan(eq("span_name"), eq(mockParentSpan), eq(attributes));
+        when(mockTracingTelemetry.createSpan(eq("span_name"), eq(mockParentSpan), eq(attributes), eq(SpanKind.INTERNAL))).thenReturn(
+            mockSpan
+        );
+        defaultTracer.startSpan("span_name", new SpanContext(mockParentSpan), attributes, SpanKind.INTERNAL);
+        verify(mockTracingTelemetry).createSpan(eq("span_name"), eq(mockParentSpan), eq(attributes), eq(SpanKind.INTERNAL));
         verify(mockTracerContextStorage, never()).get(TracerContextStorage.CURRENT_SPAN);
     }
 
@@ -92,7 +96,8 @@ public void testCreateSpanWithAttributes() {
 
         Span span = defaultTracer.startSpan(
             "span_name",
-            Attributes.create().addAttribute("key1", 1.0).addAttribute("key2", 2l).addAttribute("key3", true).addAttribute("key4", "key4")
+            Attributes.create().addAttribute("key1", 1.0).addAttribute("key2", 2l).addAttribute("key3", true).addAttribute("key4", "key4"),
+            SpanKind.INTERNAL
         );
 
         assertEquals("span_name", defaultTracer.getCurrentSpan().getSpan().getSpanName());
@@ -114,7 +119,7 @@ public void testCreateSpanWithParent() {
 
         SpanContext parentSpan = defaultTracer.getCurrentSpan();
 
-        Span span1 = defaultTracer.startSpan("span_name_1", parentSpan, Attributes.EMPTY);
+        Span span1 = defaultTracer.startSpan("span_name_1", parentSpan, Attributes.EMPTY, SpanKind.INTERNAL);
 
         assertEquals("span_name_1", defaultTracer.getCurrentSpan().getSpan().getSpanName());
         assertEquals(parentSpan.getSpan(), defaultTracer.getCurrentSpan().getSpan().getParentSpan());
@@ -126,9 +131,11 @@ public void testCreateSpanWithParent() {
     public void testCreateSpanWithContext() {
         DefaultTracer defaultTracer = new DefaultTracer(mockTracingTelemetry, mockTracerContextStorage);
         Attributes attributes = Attributes.create().addAttribute("name", "value");
-        when(mockTracingTelemetry.createSpan(eq("span_name"), eq(mockParentSpan), eq(attributes))).thenReturn(mockSpan);
+        when(mockTracingTelemetry.createSpan(eq("span_name"), eq(mockParentSpan), eq(attributes), eq(SpanKind.INTERNAL))).thenReturn(
+            mockSpan
+        );
         defaultTracer.startSpan(new SpanCreationContext("span_name", attributes));
-        verify(mockTracingTelemetry).createSpan(eq("span_name"), eq(mockParentSpan), eq(attributes));
+        verify(mockTracingTelemetry).createSpan(eq("span_name"), eq(mockParentSpan), eq(attributes), eq(SpanKind.INTERNAL));
     }
 
     public void testCreateSpanWithNullParent() {
@@ -139,7 +146,7 @@ public void testCreateSpanWithNullParent() {
             new ThreadContextBasedTracerContextStorage(threadContext, tracingTelemetry)
         );
 
-        Span span = defaultTracer.startSpan("span_name", (SpanContext) null, Attributes.EMPTY);
+        Span span = defaultTracer.startSpan("span_name", (SpanContext) null, Attributes.EMPTY, SpanKind.INTERNAL);
 
         assertEquals("span_name", defaultTracer.getCurrentSpan().getSpan().getSpanName());
         assertEquals(null, defaultTracer.getCurrentSpan().getSpan().getParentSpan());
@@ -391,6 +398,8 @@ private void setupMocks() {
         when(mockParentSpan.getSpanId()).thenReturn("parent_span_id");
         when(mockParentSpan.getTraceId()).thenReturn("trace_id");
         when(mockTracerContextStorage.get(TracerContextStorage.CURRENT_SPAN)).thenReturn(mockParentSpan, mockSpan);
-        when(mockTracingTelemetry.createSpan(eq("span_name"), eq(mockParentSpan), any(Attributes.class))).thenReturn(mockSpan);
+        when(mockTracingTelemetry.createSpan(eq("span_name"), eq(mockParentSpan), any(Attributes.class), eq(SpanKind.INTERNAL))).thenReturn(
+            mockSpan
+        );
     }
 }
diff --git a/.../telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelSpanKindConverter.java b/.../telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelSpanKindConverter.java
@@ -0,0 +1,43 @@
+/*
+ * 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 io.opentelemetry.api.trace.SpanKind;
+
+/**
+ * Converts {@link org.opensearch.telemetry.tracing.SpanKind} to OTel {@link SpanKind}
+ */
+final class OTelSpanKindConverter {
+
+    /**
+     * Constructor.
+     */
+    private OTelSpanKindConverter() {}
+
+    /**
+     * SpanKind converter.
+     * @param spanKind span kind.
+     * @return otel attributes.
+     */
+    static SpanKind convert(org.opensearch.telemetry.tracing.SpanKind spanKind) {
+        if (spanKind == null) {
+            return SpanKind.INTERNAL;
+        } else {
+            switch (spanKind) {
+                case CLIENT:
+                    return SpanKind.CLIENT;
+                case SERVER:
+                    return SpanKind.SERVER;
+                case INTERNAL:
+                default:
+                    return SpanKind.INTERNAL;
+            }
+        }
+    }
+}
diff --git a/...s/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelTracingTelemetry.java b/...s/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelTracingTelemetry.java
@@ -47,27 +47,38 @@ public void close() {
     }
 
     @Override
-    public Span createSpan(String spanName, Span parentSpan, Attributes attributes) {
-        return createOtelSpan(spanName, parentSpan, attributes);
+    public Span createSpan(String spanName, Span parentSpan, Attributes attributes, SpanKind spanKind) {
+        return createOtelSpan(spanName, parentSpan, attributes, spanKind);
     }
 
     @Override
     public TracingContextPropagator getContextPropagator() {
         return new OTelTracingContextPropagator(openTelemetry);
     }
 
-    private Span createOtelSpan(String spanName, Span parentSpan, Attributes attributes) {
-        io.opentelemetry.api.trace.Span otelSpan = otelSpan(spanName, parentSpan, OTelAttributesConverter.convert(attributes));
+    private Span createOtelSpan(String spanName, Span parentSpan, Attributes attributes, SpanKind spanKind) {
+        io.opentelemetry.api.trace.Span otelSpan = otelSpan(
+            spanName,
+            parentSpan,
+            OTelAttributesConverter.convert(attributes),
+            OTelSpanKindConverter.convert(spanKind)
+        );
         Span newSpan = new OTelSpan(spanName, otelSpan, parentSpan);
         return newSpan;
     }
 
-    io.opentelemetry.api.trace.Span otelSpan(String spanName, Span parentOTelSpan, io.opentelemetry.api.common.Attributes attributes) {
+    io.opentelemetry.api.trace.Span otelSpan(
+        String spanName,
+        Span parentOTelSpan,
+        io.opentelemetry.api.common.Attributes attributes,
+        io.opentelemetry.api.trace.SpanKind spanKind
+    ) {
         return parentOTelSpan == null || !(parentOTelSpan instanceof OTelSpan)
             ? otelTracer.spanBuilder(spanName).setAllAttributes(attributes).startSpan()
             : otelTracer.spanBuilder(spanName)
                 .setParent(Context.current().with(((OTelSpan) parentOTelSpan).getDelegateSpan()))
                 .setAllAttributes(attributes)
+                .setSpanKind(spanKind)
                 .startSpan();
     }
 }