Adds attributes to startSpan (#9199) (#9292)

* Adds attributes to startSpan * Update Changelog * Adds attributes to startSpan * Refactor code * Add java doc * Refactor code * Refactor code * Refactor code * Removes dependency --------- Signed-off-by: Gagan Juneja <gjjuneja@amazon.com> Co-authored-by: Gagan Juneja <gjjuneja@amazon.com>
opensearch-project · Aug 14, 2023 · a3cfccc · a3cfccc
1 parent 156118a
commit a3cfccc
Show file tree

Hide file tree

Showing 19 changed files with 439 additions and 39 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -66,6 +66,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Removed blocking wait in TransportGetSnapshotsAction which was exhausting generic threadpool ([#8377](https://github.com/opensearch-project/OpenSearch/pull/8377))
 - Adds support for tracing runnable scenarios ([#8831](https://github.com/opensearch-project/OpenSearch/pull/8831))
 - Change shard_size and shard_min_doc_count evaluation to happen in shard level reduce phase ([#9085](https://github.com/opensearch-project/OpenSearch/pull/9085))
+- Add attributes to startSpan methods ([#9199](https://github.com/opensearch-project/OpenSearch/pull/9199))
 
 ### Deprecated
 

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.telemetry.tracing.attributes.Attributes;
 
 /**
  *
@@ -37,16 +38,21 @@ public DefaultTracer(TracingTelemetry tracingTelemetry, TracerContextStorage<Str
 
     @Override
     public SpanScope startSpan(String spanName) {
-        return startSpan(spanName, null);
+        return startSpan(spanName, Attributes.EMPTY);
     }
 
     @Override
-    public SpanScope startSpan(String spanName, SpanContext parentSpan) {
+    public SpanScope startSpan(String spanName, Attributes attributes) {
+        return startSpan(spanName, null, attributes);
+    }
+
+    @Override
+    public SpanScope startSpan(String spanName, SpanContext parentSpan, Attributes attributes) {
         Span span = null;
         if (parentSpan != null) {
-            span = createSpan(spanName, parentSpan.getSpan());
+            span = createSpan(spanName, parentSpan.getSpan(), attributes);
         } else {
-            span = createSpan(spanName, getCurrentSpanInternal());
+            span = createSpan(spanName, getCurrentSpanInternal(), attributes);
         }
         setCurrentSpanInContext(span);
         addDefaultAttributes(span);
@@ -74,8 +80,8 @@ private void endSpan(Span span) {
         }
     }
 
-    private Span createSpan(String spanName, Span parentSpan) {
-        return tracingTelemetry.createSpan(spanName, parentSpan);
+    private Span createSpan(String spanName, Span parentSpan, Attributes attributes) {
+        return tracingTelemetry.createSpan(spanName, parentSpan, attributes);
     }
 
     private void setCurrentSpanInContext(Span span) {

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.telemetry.tracing.attributes.Attributes;
 
 /**
  * Tracer is the interface used to create a {@link Span}
@@ -27,12 +28,22 @@ public interface Tracer extends Closeable {
     SpanScope startSpan(String spanName);
 
     /**
-     *  Started the {@link Span} with the given name and parent.
+     * 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 scope of the span, must be closed with explicit close or with try-with-resource
+     */
+    SpanScope 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 scope of the span, must be closed with explicit close or with try-with-resource
      */
-    SpanScope startSpan(String spanName, SpanContext parentSpan);
+    SpanScope startSpan(String spanName, SpanContext parentSpan, Attributes attributes);
 
     /**
      * 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
@@ -9,6 +9,7 @@
 package org.opensearch.telemetry.tracing;
 
 import java.io.Closeable;
+import org.opensearch.telemetry.tracing.attributes.Attributes;
 
 /**
  * Interface for tracing telemetry providers
@@ -21,9 +22,10 @@ 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.
      * @return span instance
      */
-    Span createSpan(String spanName, Span parentSpan);
+    Span createSpan(String spanName, Span parentSpan, Attributes attributes);
 
     /**
      * provides tracing context propagator

diff --git a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/attributes/Attributes.java b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/attributes/Attributes.java
@@ -0,0 +1,94 @@
+/*
+ * 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.attributes;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+
+/**
+ * Class to create attributes for a span.
+ */
+public class Attributes {
+    private final Map<String, Object> attributesMap;
+    /**
+     * Empty value.
+     */
+    public final static Attributes EMPTY = new Attributes(Collections.emptyMap());
+
+    /**
+     * Factory method.
+     * @return attributes.
+     */
+    public static Attributes create() {
+        return new Attributes(new HashMap<>());
+    }
+
+    /**
+     * Constructor.
+     */
+    private Attributes(Map<String, Object> attributesMap) {
+        this.attributesMap = attributesMap;
+    }
+
+    /**
+     * Add String attribute.
+     * @param key key
+     * @param value value
+     * @return Same instance.
+     */
+    public Attributes addAttribute(String key, String value) {
+        Objects.requireNonNull(value, "value cannot be null");
+        attributesMap.put(key, value);
+        return this;
+    }
+
+    /**
+     * Add long attribute.
+     * @param key key
+     * @param value value
+     * @return Same instance.
+     */
+    public Attributes addAttribute(String key, long value) {
+        attributesMap.put(key, value);
+        return this;
+    };
+
+    /**
+     * Add double attribute.
+     * @param key key
+     * @param value value
+     * @return Same instance.
+     */
+    public Attributes addAttribute(String key, double value) {
+        attributesMap.put(key, value);
+        return this;
+    };
+
+    /**
+     * Add boolean attribute.
+     * @param key key
+     * @param value value
+     * @return Same instance.
+     */
+    public Attributes addAttribute(String key, boolean value) {
+        attributesMap.put(key, value);
+        return this;
+    };
+
+    /**
+     * Returns the attribute map.
+     * @return attributes map
+     */
+    public Map<String, ?> getAttributesMap() {
+        return Collections.unmodifiableMap(attributesMap);
+    }
+
+}
diff --git a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/attributes/package-info.java b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/attributes/package-info.java
@@ -0,0 +1,12 @@
+/*
+ * 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.
+ */
+
+/**
+ * Contains No-op implementations
+ */
+package org.opensearch.telemetry.tracing.attributes;
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
@@ -11,6 +11,7 @@
 import org.opensearch.telemetry.tracing.SpanScope;
 import org.opensearch.telemetry.tracing.Tracer;
 import org.opensearch.telemetry.tracing.SpanContext;
+import org.opensearch.telemetry.tracing.attributes.Attributes;
 
 /**
  * No-op implementation of Tracer
@@ -32,15 +33,20 @@ public SpanScope startSpan(String spanName) {
     }
 
     @Override
-    public SpanContext getCurrentSpan() {
-        return null;
+    public SpanScope startSpan(String spanName, Attributes attributes) {
+        return SpanScope.NO_OP;
     }
 
     @Override
-    public SpanScope startSpan(String spanName, SpanContext parentSpan) {
+    public SpanScope startSpan(String spanName, SpanContext parentSpan, Attributes attributes) {
         return SpanScope.NO_OP;
     }
 
+    @Override
+    public SpanContext getCurrentSpan() {
+        return null;
+    }
+
     @Override
     public void close() {
 

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
@@ -11,6 +11,7 @@
 import org.opensearch.telemetry.tracing.SpanContext;
 import org.opensearch.telemetry.tracing.SpanScope;
 import org.opensearch.telemetry.tracing.Tracer;
+import org.opensearch.telemetry.tracing.attributes.Attributes;
 
 /**
  * Wraps the runnable and add instrumentation to trace the {@link Runnable}
@@ -20,24 +21,27 @@ public class TraceableRunnable implements Runnable {
     private final SpanContext parent;
     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 runnable runnable.
      */
-    public TraceableRunnable(Tracer tracer, String spanName, SpanContext parent, Runnable runnable) {
+    public TraceableRunnable(Tracer tracer, String spanName, SpanContext parent, Attributes attributes, Runnable runnable) {
         this.tracer = tracer;
         this.spanName = spanName;
         this.parent = parent;
+        this.attributes = attributes;
         this.runnable = runnable;
     }
 
     @Override
     public void run() {
-        try (SpanScope spanScope = tracer.startSpan(spanName, parent)) {
+        try (SpanScope spanScope = tracer.startSpan(spanName, parent, attributes)) {
             runnable.run();
         }
     }

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
@@ -11,12 +11,17 @@
 import org.junit.Assert;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.util.concurrent.ThreadContext;
+import org.opensearch.telemetry.tracing.attributes.Attributes;
 import org.opensearch.test.OpenSearchTestCase;
 
 import java.io.IOException;
+import org.opensearch.test.telemetry.tracing.MockSpan;
 import org.opensearch.test.telemetry.tracing.MockTracingTelemetry;
 
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.when;
 import static org.mockito.Mockito.verify;
 
@@ -46,6 +51,42 @@ public void testCreateSpan() {
         Assert.assertEquals("span_name", defaultTracer.getCurrentSpan().getSpan().getSpanName());
     }
 
+    public void testCreateSpanWithAttributesWithMock() {
+        DefaultTracer defaultTracer = new DefaultTracer(mockTracingTelemetry, mockTracerContextStorage);
+        Attributes attributes = Attributes.create().addAttribute("name", "value");
+        when(mockTracingTelemetry.createSpan("span_name", mockParentSpan, attributes)).thenReturn(mockSpan);
+        defaultTracer.startSpan("span_name", attributes);
+        verify(mockTracingTelemetry).createSpan("span_name", mockParentSpan, attributes);
+    }
+
+    public void testCreateSpanWithAttributesWithParentMock() {
+        DefaultTracer defaultTracer = new DefaultTracer(mockTracingTelemetry, mockTracerContextStorage);
+        Attributes attributes = Attributes.create().addAttribute("name", "value");
+        when(mockTracingTelemetry.createSpan("span_name", mockParentSpan, attributes)).thenReturn(mockSpan);
+        defaultTracer.startSpan("span_name", new SpanContext(mockParentSpan), attributes);
+        verify(mockTracingTelemetry).createSpan("span_name", mockParentSpan, attributes);
+        verify(mockTracerContextStorage, never()).get(TracerContextStorage.CURRENT_SPAN);
+    }
+
+    public void testCreateSpanWithAttributes() {
+        TracingTelemetry tracingTelemetry = new MockTracingTelemetry();
+        DefaultTracer defaultTracer = new DefaultTracer(
+            tracingTelemetry,
+            new ThreadContextBasedTracerContextStorage(new ThreadContext(Settings.EMPTY), tracingTelemetry)
+        );
+
+        defaultTracer.startSpan(
+            "span_name",
+            Attributes.create().addAttribute("key1", 1.0).addAttribute("key2", 2l).addAttribute("key3", true).addAttribute("key4", "key4")
+        );
+
+        Assert.assertEquals("span_name", defaultTracer.getCurrentSpan().getSpan().getSpanName());
+        Assert.assertEquals(1.0, ((MockSpan) defaultTracer.getCurrentSpan().getSpan()).getAttribute("key1"));
+        Assert.assertEquals(2l, ((MockSpan) defaultTracer.getCurrentSpan().getSpan()).getAttribute("key2"));
+        Assert.assertEquals(true, ((MockSpan) defaultTracer.getCurrentSpan().getSpan()).getAttribute("key3"));
+        Assert.assertEquals("key4", ((MockSpan) defaultTracer.getCurrentSpan().getSpan()).getAttribute("key4"));
+    }
+
     public void testCreateSpanWithParent() {
         TracingTelemetry tracingTelemetry = new MockTracingTelemetry();
         DefaultTracer defaultTracer = new DefaultTracer(
@@ -57,7 +98,7 @@ public void testCreateSpanWithParent() {
 
         SpanContext parentSpan = defaultTracer.getCurrentSpan();
 
-        defaultTracer.startSpan("span_name_1", parentSpan);
+        defaultTracer.startSpan("span_name_1", parentSpan, Attributes.EMPTY);
 
         Assert.assertEquals("span_name_1", defaultTracer.getCurrentSpan().getSpan().getSpanName());
         Assert.assertEquals(parentSpan.getSpan(), defaultTracer.getCurrentSpan().getSpan().getParentSpan());
@@ -70,7 +111,7 @@ public void testCreateSpanWithNullParent() {
             new ThreadContextBasedTracerContextStorage(new ThreadContext(Settings.EMPTY), tracingTelemetry)
         );
 
-        defaultTracer.startSpan("span_name", null);
+        defaultTracer.startSpan("span_name", null, Attributes.EMPTY);
 
         Assert.assertEquals("span_name", defaultTracer.getCurrentSpan().getSpan().getSpanName());
         Assert.assertEquals(null, defaultTracer.getCurrentSpan().getSpan().getParentSpan());
@@ -105,6 +146,6 @@ 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("span_name", mockParentSpan)).thenReturn(mockSpan);
+        when(mockTracingTelemetry.createSpan(eq("span_name"), eq(mockParentSpan), any(Attributes.class))).thenReturn(mockSpan);
     }
 }