Fixing comments

Signed-off-by: Nishchay Malhotra <nishcha@amazon.com>

Adding new sampled attributes

Adding new inferred sampler attributes

Signed-off-by: Nishchay Malhotra <nishcha@amazon.com>

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

@@ -36,6 +36,9 @@
 import org.opensearch.core.common.io.stream.Writeable;
 import org.opensearch.core.common.transport.TransportAddress;
 
+import java.util.Collections;
+import java.util.Map;
+
 /**
  * Message over the transport interface
  *
@@ -45,22 +48,22 @@ public abstract class TransportMessage implements Writeable {
 
     private TransportAddress remoteAddress;
 
-    private boolean sampled;
+    private Map<String, String> header = Collections.emptyMap();
 
     public void remoteAddress(TransportAddress remoteAddress) {
         this.remoteAddress = remoteAddress;
     }
 
-    public void markResponseAsSampled() {
-        this.sampled = true;
+    public void setResponseHeaders(Map<String, String> header) {
+        this.header = header;
     }
 
     public TransportAddress remoteAddress() {
         return remoteAddress;
     }
 
-    public boolean sampled() {
-        return sampled;
+    public Map<String, String> getResponseHeaders() {
+        return header;
     }
 
     /**

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

@@ -9,6 +9,7 @@
 package org.opensearch.telemetry.tracing;
 
 import org.opensearch.common.annotation.InternalApi;
+import org.opensearch.telemetry.tracing.attributes.SamplingAttributes;
 
 import java.io.Closeable;
 import java.io.IOException;
@@ -99,7 +100,7 @@ private Span createSpan(SpanCreationContext spanCreationContext, Span parentSpan
      * @param span the current active span
      */
     protected void addDefaultAttributes(Span parentSpan, Span span) {
-        addCommonParentAttributes(parentSpan, span);
+        copyInheritableParentAttributes(parentSpan, span);
         span.addAttribute(THREAD_NAME, Thread.currentThread().getName());
     }
 
@@ -111,19 +112,18 @@ public Span startSpan(SpanCreationContext spanCreationContext, Map<String, Colle
     }
 
     private void addRequestAttributeToContext(SpanCreationContext spanCreationContext, Map<String, Collection<String>> headers) {
-        if (headers != null && headers.containsKey(TracerContextStorage.INFERRED_SAMPLER)) {
-            spanCreationContext.getAttributes().addAttribute(TracerContextStorage.INFERRED_SAMPLER, true);
+        if (headers != null && headers.containsKey(SamplingAttributes.SAMPLER.getValue())) {
+            spanCreationContext.getAttributes()
+                .addAttribute(SamplingAttributes.SAMPLER.getValue(), SamplingAttributes.INFERRED_SAMPLER.getValue());
         }
     }
 
-    private void addCommonParentAttributes(Span parentSpan, Span currentSpan) {
+    private void copyInheritableParentAttributes(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)
-            );
+            Optional<String> inferredAttribute = Optional.ofNullable(parentSpan.getAttributeString(SamplingAttributes.SAMPLER.getValue()));
             if (inferredAttribute.isPresent()) {
-                currentSpan.addAttribute(TracerContextStorage.INFERRED_SAMPLER, true);
+                currentSpan.addAttribute(SamplingAttributes.SAMPLER.getValue(), SamplingAttributes.INFERRED_SAMPLER.getValue());
             }
         }
     }

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

@@ -24,16 +24,6 @@ public interface TracerContextStorage<K, V> {
      */
     String CURRENT_SPAN = "current_span";
 
-    /**
-     * Key for storing sample information
-     */
-    String SAMPLED = "sampled";
-
-    /**
-     * Key for storing inferred Sampling information
-     */
-    String INFERRED_SAMPLER = "Inferred_sampler";
-
     /**
      * Fetches value corresponding to key
      * @param key of the tracing context

libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/attributes/SamplingAttributes.java

@@ -0,0 +1,40 @@
+/*
+ * 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 org.opensearch.common.annotation.InternalApi;
+
+import java.util.Locale;
+
+/**
+ * Enum for Inferred Sampling*
+ * @opensearch.internal*
+ */
+@InternalApi
+public enum SamplingAttributes {
+
+    /**
+     * Attribute added if the span is sampled by inferred sampler*
+     */
+    INFERRED_SAMPLER,
+
+    /**
+     * Attribute Added if the span is an outlier*
+     */
+    SAMPLED,
+
+    /**
+     * Sampler Used in the framework*
+     */
+    SAMPLER;
+
+    public String getValue() {
+        return name().toLowerCase(Locale.ROOT);
+    }
+}

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

@@ -8,6 +8,8 @@
 
 package org.opensearch.telemetry.tracing;
 
+import org.opensearch.telemetry.tracing.attributes.SamplingAttributes;
+
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.StatusCode;
@@ -34,17 +36,25 @@ public OTelSpan(String spanName, Span span, org.opensearch.telemetry.tracing.Spa
 
     @Override
     public void endSpan() {
-        if (getAttributeBoolean(TracerContextStorage.SAMPLED) != null && getAttributeBoolean(TracerContextStorage.SAMPLED)) {
+        if (isSpanOutlier()) {
             markParentForSampling();
         }
         delegateSpan.end();
     }
 
+    /*
+     * This is added temporarily will remove this after the evaluation framework PR.
+     * This Framework will be used to evaluate a span if that is an outlier or not.
+     */
+    private boolean isSpanOutlier() {
+        return false;
+    }
+
     private void markParentForSampling() {
-        org.opensearch.telemetry.tracing.Span current_parent = getParentSpan();
-        while (current_parent != null && current_parent.getAttributeBoolean(TracerContextStorage.SAMPLED) == null) {
-            current_parent.addAttribute(TracerContextStorage.SAMPLED, true);
-            current_parent = current_parent.getParentSpan();
+        org.opensearch.telemetry.tracing.Span currentParent = getParentSpan();
+        while (currentParent != null && currentParent.getAttributeBoolean(SamplingAttributes.SAMPLED.getValue()) == null) {
+            currentParent.addAttribute(SamplingAttributes.SAMPLED.getValue(), true);
+            currentParent = currentParent.getParentSpan();
         }
     }
 

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

@@ -62,9 +62,7 @@ private Span createOtelSpan(SpanCreationContext spanCreationContext, Span parent
             OTelAttributesConverter.convert(spanCreationContext.getAttributes()),
             OTelSpanKindConverter.convert(spanCreationContext.getSpanKind())
         );
-
         Span newSpan = new OTelSpan(spanCreationContext.getSpanName(), otelSpan, parentSpan);
-
         return newSpan;
     }
 

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

@@ -8,7 +8,9 @@
 
 package org.opensearch.telemetry.tracing.processor;
 
-import org.opensearch.telemetry.tracing.TracerContextStorage;
+import org.opensearch.telemetry.tracing.attributes.SamplingAttributes;
+
+import java.util.Objects;
 
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.trace.Span;
@@ -71,8 +73,11 @@ public boolean isStartRequired() {
     public void onEnd(ReadableSpan span) {
         if (span != null
             && span.getSpanContext().isSampled()
-            && Boolean.TRUE.equals(span.getAttribute(AttributeKey.booleanKey(TracerContextStorage.INFERRED_SAMPLER)))) {
-            if (span.getAttribute(AttributeKey.booleanKey(TracerContextStorage.SAMPLED)) != null) {
+            && Objects.equals(
+                span.getAttribute(AttributeKey.stringKey(SamplingAttributes.SAMPLER.getValue())),
+                SamplingAttributes.INFERRED_SAMPLER.getValue()
+            )) {
+            if (span.getAttribute(AttributeKey.booleanKey(SamplingAttributes.SAMPLED.getValue())) != null) {
                 this.delegateProcessor.onEnd(span);
             }
         } else {

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

@@ -10,7 +10,7 @@
 
 import org.opensearch.common.settings.Settings;
 import org.opensearch.telemetry.TelemetrySettings;
-import org.opensearch.telemetry.tracing.TracerContextStorage;
+import org.opensearch.telemetry.tracing.attributes.SamplingAttributes;
 import org.opensearch.telemetry.tracing.samplingResult.OTelSamplingResult;
 
 import java.util.List;
@@ -68,7 +68,7 @@ public SamplingResult shouldSample(
         boolean inferredSamplingAllowListed = telemetrySettings.getInferredSamplingAllowListed();
         if (inferredSamplingAllowListed) {
             Attributes customSampleAttributes = Attributes.builder()
-                .put(TracerContextStorage.INFERRED_SAMPLER, true)
+                .put(SamplingAttributes.SAMPLER.getValue(), SamplingAttributes.INFERRED_SAMPLER.getValue())
                 .putAll(attributes)
                 .build();
             SamplingResult result = SamplingResult.recordAndSample();

server/src/main/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorage.java

@@ -11,11 +11,13 @@
 import org.opensearch.common.annotation.InternalApi;
 import org.opensearch.common.util.concurrent.ThreadContext;
 import org.opensearch.common.util.concurrent.ThreadContextStatePropagator;
+import org.opensearch.telemetry.tracing.attributes.SamplingAttributes;
 
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 
 /**
  * Core's ThreadContext based TracerContextStorage implementation
@@ -83,12 +85,18 @@ public Map<String, String> headers(Map<String, Object> source) {
                 tracingTelemetry.getContextPropagator().inject(current.getSpan(), headers::put);
 
                 // We will be sending one more header with the response if the request is marked for sampling
-                if (current.getSpan().getAttributeBoolean(SAMPLED) != null && current.getSpan().getAttributeBoolean(SAMPLED)) {
-                    headers.put(SAMPLED, "true");
+                Optional<Boolean> isSpanSampled = Optional.ofNullable(
+                    current.getSpan().getAttributeBoolean(SamplingAttributes.SAMPLED.getValue())
+                );
+                if (isSpanSampled.isPresent()) {
+                    headers.put(SamplingAttributes.SAMPLED.getValue(), "true");
                 }
-                if (current.getSpan().getAttributeBoolean(INFERRED_SAMPLER) != null
-                    && current.getSpan().getAttributeBoolean(INFERRED_SAMPLER)) {
-                    headers.put(INFERRED_SAMPLER, "true");
+                Optional<String> isSpanInferredSampled = Optional.ofNullable(
+                    current.getSpan().getAttributeString(SamplingAttributes.SAMPLER.getValue())
+                );
+                if (isSpanInferredSampled.isPresent()
+                    && isSpanInferredSampled.get().equals(SamplingAttributes.INFERRED_SAMPLER.getValue())) {
+                    headers.put(SamplingAttributes.SAMPLER.getValue(), isSpanInferredSampled.get());
                 }
             }
         }

server/src/main/java/org/opensearch/telemetry/tracing/handler/TraceableTransportResponseHandler.java

@@ -13,7 +13,7 @@
 import org.opensearch.telemetry.tracing.Span;
 import org.opensearch.telemetry.tracing.SpanScope;
 import org.opensearch.telemetry.tracing.Tracer;
-import org.opensearch.telemetry.tracing.TracerContextStorage;
+import org.opensearch.telemetry.tracing.attributes.SamplingAttributes;
 import org.opensearch.transport.TransportException;
 import org.opensearch.transport.TransportResponseHandler;
 
@@ -70,8 +70,9 @@ public T read(StreamInput in) throws IOException {
     @Override
     public void handleResponse(T response) {
         try (SpanScope scope = tracer.withSpanInScope(span)) {
-            if (response.sampled()) {
-                span.addAttribute(TracerContextStorage.SAMPLED, true);
+            String sampleInformation = response.getResponseHeaders().getOrDefault(SamplingAttributes.SAMPLED.getValue(), "");
+            if (sampleInformation.equals("true")) {
+                span.addAttribute(SamplingAttributes.SAMPLED.getValue(), true);
             }
             span.endSpan();
         } finally {
@@ -82,7 +83,6 @@ public void handleResponse(T response) {
     @Override
     public void handleException(TransportException exp) {
         try (SpanScope scope = tracer.withSpanInScope(span)) {
-            span.addAttribute(TracerContextStorage.SAMPLED, true);
             span.setError(exp);
             span.endSpan();
         } finally {
@@ -103,7 +103,6 @@ public String toString() {
     @Override
     public void handleRejection(Exception exp) {
         try (SpanScope scope = tracer.withSpanInScope(span)) {
-            span.addAttribute(TracerContextStorage.SAMPLED, true);
             span.setError(exp);
             span.endSpan();
         } finally {

server/src/main/java/org/opensearch/transport/InboundHandler.java

@@ -50,7 +50,6 @@
 import org.opensearch.telemetry.tracing.SpanBuilder;
 import org.opensearch.telemetry.tracing.SpanScope;
 import org.opensearch.telemetry.tracing.Tracer;
-import org.opensearch.telemetry.tracing.TracerContextStorage;
 import org.opensearch.telemetry.tracing.channels.TraceableTcpTransportChannel;
 import org.opensearch.threadpool.ThreadPool;
 
@@ -137,15 +136,12 @@ private void messageReceived(TcpChannel channel, InboundMessage message, long st
         final Header header = message.getHeader();
         assert header.needsToReadVariableHeader() == false;
         ThreadContext threadContext = threadPool.getThreadContext();
-        boolean responseSampled = false;
+        Map<String, String> responseHeader;
         try (ThreadContext.StoredContext existing = threadContext.stashContext()) {
             // Place the context with the headers from the message
             threadContext.setHeaders(header.getHeaders());
             threadContext.putTransient("_remote_address", remoteAddress);
-            String sampleInformation = message.getHeader().getHeaders().v1().getOrDefault(TracerContextStorage.SAMPLED, "");
-            if (sampleInformation.equals("true")) {
-                responseSampled = true;
-            }
+            responseHeader = message.getHeader().getHeaders().v1();
             if (header.isRequest()) {
                 handleRequest(channel, header, message);
             } else {
@@ -175,11 +171,11 @@ private void messageReceived(TcpChannel channel, InboundMessage message, long st
                         if (header.isError()) {
                             handlerResponseError(requestId, streamInput, handler);
                         } else {
-                            handleResponse(requestId, remoteAddress, streamInput, handler, responseSampled);
+                            handleResponse(requestId, remoteAddress, streamInput, handler, responseHeader);
                         }
                     } else {
                         assert header.isError() == false;
-                        handleResponse(requestId, remoteAddress, EMPTY_STREAM_INPUT, handler, responseSampled);
+                        handleResponse(requestId, remoteAddress, EMPTY_STREAM_INPUT, handler, responseHeader);
                     }
                 }
 
@@ -398,15 +394,13 @@ private <T extends TransportResponse> void handleResponse(
         InetSocketAddress remoteAddress,
         final StreamInput stream,
         final TransportResponseHandler<T> handler,
-        boolean responseSampled
+        Map<String, String> responseHeader
     ) {
         final T response;
         try {
             response = handler.read(stream);
             response.remoteAddress(new TransportAddress(remoteAddress));
-            if (responseSampled) {
-                response.markResponseAsSampled();
-            }
+            response.setResponseHeaders(responseHeader);
             checkStreamIsFullyConsumed(requestId, handler, stream, false);
         } catch (Exception e) {
             final Exception serializationException = new TransportSerializationException(