Adding get attribute as readable span attribute

Signed-off-by: Nishchay Malhotra <[email protected]>

libs/core/src/main/java/org/opensearch/core/transport/TransportMessage.java

@@ -51,7 +51,7 @@ public void remoteAddress(TransportAddress remoteAddress) {
         this.remoteAddress = remoteAddress;
     }
 
-    public void markSampled() {
+    public void markResponseAsSampled() {
         this.sampled = true;
     }
 

libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultSpanScope.java

@@ -9,10 +9,7 @@
 package org.opensearch.telemetry.tracing;
 
 import org.opensearch.common.annotation.InternalApi;
-import org.opensearch.telemetry.tracing.attributes.AttributeType;
 
-import java.util.HashMap;
-import java.util.Map;
 import java.util.Objects;
 
 /**
@@ -27,7 +24,6 @@ class DefaultSpanScope implements SpanScope {
     private final Span beforeSpan;
     private static final ThreadLocal<SpanScope> spanScopeThreadLocal = new ThreadLocal<>();
     private final TracerContextStorage<String, Span> tracerContextStorage;
-    private final Map<String, AttributeType> commonAttributeMap = new HashMap<>();
 
     /**
      * Constructor
@@ -44,7 +40,6 @@ private DefaultSpanScope(
         this.beforeSpan = beforeSpan;
         this.previousSpanScope = previousSpanScope;
         this.tracerContextStorage = tracerContextStorage;
-        initializeCommonPropagationAttributes();
     }
 
     /**
@@ -60,13 +55,6 @@ public static SpanScope create(Span span, TracerContextStorage<String, Span> tra
         return newSpanScope;
     }
 
-    /**
-     * Common attributes need to be taken from parent and propagated to child span*
-     */
-    private void initializeCommonPropagationAttributes() {
-        commonAttributeMap.put(TracerContextStorage.INFERRED_SAMPLER, AttributeType.BOOLEAN);
-    }
-
     @Override
     public void close() {
         detach();
@@ -76,38 +64,9 @@ public void close() {
     public SpanScope attach() {
         spanScopeThreadLocal.set(this);
         tracerContextStorage.put(TracerContextStorage.CURRENT_SPAN, this.span);
-        addCommonParentAttributes();
         return this;
     }
 
-    private void addCommonParentAttributes() {
-        // This work as common attribute propagator from parent to child
-        for (String attribute : commonAttributeMap.keySet()) {
-            if (beforeSpan != null && beforeSpan.getAttributes().containsKey(attribute)) {
-                AttributeType attributeValue = commonAttributeMap.get(attribute);
-                this.storeAttributeValue(attribute, attributeValue);
-            }
-        }
-    }
-
-    private void storeAttributeValue(String attribute, AttributeType attributeType) {
-        switch (attributeType) {
-            case BOOLEAN:
-                span.addAttribute(attribute, (Boolean) beforeSpan.getAttribute(attribute));
-                break;
-            case LONG:
-                span.addAttribute(attribute, (Long) beforeSpan.getAttribute(attribute));
-                break;
-            case STRING:
-                span.addAttribute(attribute, (String) beforeSpan.getAttribute(attribute));
-                break;
-            case DOUBLE:
-                span.addAttribute(attribute, (Double) beforeSpan.getAttribute(attribute));
-                break;
-            // Add more cases for other types if needed
-        }
-    }
-
     private void detach() {
         spanScopeThreadLocal.set(previousSpanScope);
         if (beforeSpan != null) {

libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultTracer.java

@@ -52,8 +52,9 @@ public Span startSpan(SpanCreationContext context) {
         } else {
             parentSpan = getCurrentSpanInternal();
         }
+
         Span span = createSpan(context, parentSpan);
-        addDefaultAttributes(span);
+        addDefaultAttributes(parentSpan, span);
         return span;
     }
 
@@ -97,7 +98,8 @@ private Span createSpan(SpanCreationContext spanCreationContext, Span parentSpan
      * Adds default attributes in the span
      * @param span the current active span
      */
-    protected void addDefaultAttributes(Span span) {
+    protected void addDefaultAttributes(Span parentSpan, Span span) {
+        addCommonParentAttributes(parentSpan, span);
         span.addAttribute(THREAD_NAME, Thread.currentThread().getName());
     }
 
@@ -113,4 +115,16 @@ private void addRequestAttributeToContext(SpanCreationContext spanCreationContex
             spanCreationContext.getAttributes().addAttribute(TracerContextStorage.INFERRED_SAMPLER, true);
         }
     }
+
+    private void addCommonParentAttributes(Span parentSpan, Span currentSpan) {
+        // This work as common attribute propagator from parent to child
+        if (parentSpan != null) {
+            Optional<Boolean> inferredAttribute = Optional.ofNullable(
+                parentSpan.getAttributeBoolean(TracerContextStorage.INFERRED_SAMPLER)
+            );
+            if (inferredAttribute.isPresent()) {
+                currentSpan.addAttribute(TracerContextStorage.INFERRED_SAMPLER, true);
+            }
+        }
+    }
 }

libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/Span.java

@@ -10,8 +10,6 @@
 
 import org.opensearch.common.annotation.ExperimentalApi;
 
-import java.util.Map;
-
 /**
  * An interface that represents a tracing span.
  * Spans are created by the Tracer.startSpan method.
@@ -96,15 +94,30 @@ public interface Span {
     String getSpanId();
 
     /**
-     * Returns attribute.
-     * @param key key
-     * @return value
+     * *
+     * @param key for which we need to look for value
+     * @return string attribute value
+     */
+    String getAttributeString(String key);
+
+    /**
+     * *
+     * @param key for which we need to look for value
+     * @return Boolean attribute value
+     */
+    Boolean getAttributeBoolean(String key);
+
+    /**
+     * *
+     * @param key for which we need to look for value
+     * @return Long attribute value
      */
-    Object getAttribute(String key);
+    Long getAttributeLong(String key);
 
     /**
-     * Returns the attributes as map.
-     * @return returns the attributes map.
+     * *
+     * @param key for which we need to look for value
+     * @return Double attribute value
      */
-    Map<String, Object> getAttributes();
+    Double getAttributeDouble(String key);
 }

libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/noop/NoopSpan.java

@@ -11,8 +11,6 @@
 import org.opensearch.common.annotation.InternalApi;
 import org.opensearch.telemetry.tracing.Span;
 
-import java.util.Map;
-
 /**
  * No-op implementation of {@link org.opensearch.telemetry.tracing.Span}
  *
@@ -85,24 +83,23 @@ public String getSpanId() {
         return "noop-span-id";
     }
 
-    /**
-     * Returns attribute.
-     *
-     * @param key key
-     * @return value
-     */
     @Override
-    public Object getAttribute(String key) {
-        return null;
+    public String getAttributeString(String key) {
+        return "";
     }
 
-    /**
-     * Returns the attributes as map.
-     *
-     * @return returns the attributes map.
-     */
     @Override
-    public Map<String, Object> getAttributes() {
-        return null;
+    public Boolean getAttributeBoolean(String key) {
+        return false;
+    }
+
+    @Override
+    public Long getAttributeLong(String key) {
+        return 0L;
+    }
+
+    @Override
+    public Double getAttributeDouble(String key) {
+        return 0.0;
     }
 }

plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelSpan.java

@@ -8,11 +8,10 @@
 
 package org.opensearch.telemetry.tracing;
 
-import java.util.HashMap;
-import java.util.Map;
-
+import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.StatusCode;
+import io.opentelemetry.sdk.trace.ReadableSpan;
 
 /**
  * Default implementation of {@link Span} using Otel span. It keeps a reference of OpenTelemetry Span and handles span
@@ -22,8 +21,6 @@ class OTelSpan extends AbstractSpan {
 
     private final Span delegateSpan;
 
-    private final Map<String, Object> metadata;
-
     /**
      * Constructor
      * @param spanName span name
@@ -33,20 +30,19 @@ class OTelSpan extends AbstractSpan {
     public OTelSpan(String spanName, Span span, org.opensearch.telemetry.tracing.Span parentSpan) {
         super(spanName, parentSpan);
         this.delegateSpan = span;
-        this.metadata = new HashMap<>();
     }
 
     @Override
     public void endSpan() {
-        if (getAttributes().containsKey(TracerContextStorage.SAMPLED)) {
+        if (getAttributeBoolean(TracerContextStorage.SAMPLED) != null && getAttributeBoolean(TracerContextStorage.SAMPLED)) {
             markParentForSampling();
         }
         delegateSpan.end();
     }
 
     private void markParentForSampling() {
         org.opensearch.telemetry.tracing.Span current_parent = getParentSpan();
-        while (current_parent != null && !current_parent.getAttributes().containsKey(TracerContextStorage.SAMPLED)) {
+        while (current_parent != null && current_parent.getAttributeBoolean(TracerContextStorage.SAMPLED) == null) {
             current_parent.addAttribute(TracerContextStorage.SAMPLED, true);
             current_parent = current_parent.getParentSpan();
         }
@@ -55,25 +51,21 @@ private void markParentForSampling() {
     @Override
     public void addAttribute(String key, String value) {
         delegateSpan.setAttribute(key, value);
-        metadata.put(key, value);
     }
 
     @Override
     public void addAttribute(String key, Long value) {
         delegateSpan.setAttribute(key, value);
-        metadata.put(key, value);
     }
 
     @Override
     public void addAttribute(String key, Double value) {
         delegateSpan.setAttribute(key, value);
-        metadata.put(key, value);
     }
 
     @Override
     public void addAttribute(String key, Boolean value) {
         delegateSpan.setAttribute(key, value);
-        metadata.put(key, value);
     }
 
     @Override
@@ -99,13 +91,39 @@ public String getSpanId() {
     }
 
     @Override
-    public Object getAttribute(String key) {
-        return metadata.get(key);
+    public String getAttributeString(String key) {
+        if (delegateSpan != null && delegateSpan instanceof ReadableSpan) return ((ReadableSpan) delegateSpan).getAttribute(
+            AttributeKey.stringKey(key)
+        );
+
+        return null;
     }
 
     @Override
-    public Map<String, Object> getAttributes() {
-        return metadata;
+    public Boolean getAttributeBoolean(String key) {
+        if (delegateSpan != null && delegateSpan instanceof ReadableSpan) {
+            return ((ReadableSpan) delegateSpan).getAttribute(AttributeKey.booleanKey(key));
+        }
+
+        return null;
+    }
+
+    @Override
+    public Long getAttributeLong(String key) {
+        if (delegateSpan != null && delegateSpan instanceof ReadableSpan) return ((ReadableSpan) delegateSpan).getAttribute(
+            AttributeKey.longKey(key)
+        );
+
+        return null;
+    }
+
+    @Override
+    public Double getAttributeDouble(String key) {
+        if (delegateSpan != null && delegateSpan instanceof ReadableSpan) return ((ReadableSpan) delegateSpan).getAttribute(
+            AttributeKey.doubleKey(key)
+        );
+
+        return null;
     }
 
     io.opentelemetry.api.trace.Span getDelegateSpan() {