From 46b28eb45e617a94ad8a1e322b7ca74206d52870 Mon Sep 17 00:00:00 2001
From: viveksing <viveksing.patil@gmail.com>
Date: Mon, 30 Sep 2024 21:13:09 +0200
Subject: [PATCH] Fix ottracepropagation for short span ids (#6734)

---
 .../extension/trace/propagation/Common.java   |  4 +-
 .../trace/propagation/OtTracePropagator.java  | 10 ++++-
 .../propagation/OtTracePropagatorTest.java    | 41 +++++++++++++++++++
 3 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/extensions/trace-propagators/src/main/java/io/opentelemetry/extension/trace/propagation/Common.java b/extensions/trace-propagators/src/main/java/io/opentelemetry/extension/trace/propagation/Common.java
index aabbc190e03..9d04262e7d1 100644
--- a/extensions/trace-propagators/src/main/java/io/opentelemetry/extension/trace/propagation/Common.java
+++ b/extensions/trace-propagators/src/main/java/io/opentelemetry/extension/trace/propagation/Common.java
@@ -29,6 +29,8 @@ final class Common {
   static final int MAX_TRACE_ID_LENGTH = TraceId.getLength();
   static final int MIN_TRACE_ID_LENGTH = MAX_TRACE_ID_LENGTH / 2;
 
+  static final int MAX_SPAN_ID_LENGTH = SpanId.getLength();
+
   private Common() {}
 
   static SpanContext buildSpanContext(
@@ -44,7 +46,7 @@ static SpanContext buildSpanContext(
 
       return SpanContext.createFromRemoteParent(
           StringUtils.padLeft(traceId, MAX_TRACE_ID_LENGTH),
-          spanId,
+          StringUtils.padLeft(spanId, MAX_SPAN_ID_LENGTH),
           traceFlags,
           TraceState.getDefault());
     } catch (RuntimeException e) {
diff --git a/extensions/trace-propagators/src/main/java/io/opentelemetry/extension/trace/propagation/OtTracePropagator.java b/extensions/trace-propagators/src/main/java/io/opentelemetry/extension/trace/propagation/OtTracePropagator.java
index 105c3830c7e..3106382abd4 100644
--- a/extensions/trace-propagators/src/main/java/io/opentelemetry/extension/trace/propagation/OtTracePropagator.java
+++ b/extensions/trace-propagators/src/main/java/io/opentelemetry/extension/trace/propagation/OtTracePropagator.java
@@ -5,6 +5,7 @@
 
 package io.opentelemetry.extension.trace.propagation;
 
+import static io.opentelemetry.extension.trace.propagation.Common.MAX_SPAN_ID_LENGTH;
 import static io.opentelemetry.extension.trace.propagation.Common.MAX_TRACE_ID_LENGTH;
 
 import io.opentelemetry.api.baggage.Baggage;
@@ -12,6 +13,7 @@
 import io.opentelemetry.api.internal.StringUtils;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.api.trace.SpanId;
 import io.opentelemetry.api.trace.TraceId;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.propagation.TextMapGetter;
@@ -96,7 +98,13 @@ public <C> Context extract(Context context, @Nullable C carrier, TextMapGetter<C
         incomingTraceId == null
             ? TraceId.getInvalid()
             : StringUtils.padLeft(incomingTraceId, MAX_TRACE_ID_LENGTH);
-    String spanId = getter.get(carrier, SPAN_ID_HEADER);
+
+    String incomingSpanId = getter.get(carrier, SPAN_ID_HEADER);
+    String spanId =
+        incomingSpanId == null
+            ? SpanId.getInvalid()
+            : StringUtils.padLeft(incomingSpanId, MAX_SPAN_ID_LENGTH);
+
     String sampled = getter.get(carrier, SAMPLED_HEADER);
     SpanContext spanContext = buildSpanContext(traceId, spanId, sampled);
     if (!spanContext.isValid()) {
diff --git a/extensions/trace-propagators/src/test/java/io/opentelemetry/extension/trace/propagation/OtTracePropagatorTest.java b/extensions/trace-propagators/src/test/java/io/opentelemetry/extension/trace/propagation/OtTracePropagatorTest.java
index 7b1c7434299..22001427a94 100644
--- a/extensions/trace-propagators/src/test/java/io/opentelemetry/extension/trace/propagation/OtTracePropagatorTest.java
+++ b/extensions/trace-propagators/src/test/java/io/opentelemetry/extension/trace/propagation/OtTracePropagatorTest.java
@@ -32,6 +32,8 @@ class OtTracePropagatorTest {
   private static final String SHORT_TRACE_ID = "ff00000000000000";
   private static final String SHORT_TRACE_ID_FULL = "0000000000000000ff00000000000000";
   private static final String SPAN_ID = "ff00000000000041";
+  private static final String SHORT_SPAN_ID = "f00000000000041";
+  private static final String SHORT_SPAN_ID_FULL = "0f00000000000041";
   private static final TextMapSetter<Map<String, String>> setter = Map::put;
   private static final TextMapGetter<Map<String, String>> getter =
       new TextMapGetter<Map<String, String>>() {
@@ -261,6 +263,45 @@ void extract_NotSampledContext_Short_TraceId() {
                 SHORT_TRACE_ID_FULL, SPAN_ID, TraceFlags.getDefault(), TraceState.getDefault()));
   }
 
+  @Test
+  void extract_SampledContext_Int_Short_SPanId() {
+    Map<String, String> carrier = new LinkedHashMap<>();
+    carrier.put(OtTracePropagator.TRACE_ID_HEADER, TRACE_ID);
+    carrier.put(OtTracePropagator.SPAN_ID_HEADER, SHORT_SPAN_ID);
+    carrier.put(OtTracePropagator.SAMPLED_HEADER, Common.TRUE_INT);
+
+    assertThat(getSpanContext(propagator.extract(Context.current(), carrier, getter)))
+        .isEqualTo(
+            SpanContext.createFromRemoteParent(
+                TRACE_ID, SHORT_SPAN_ID_FULL, TraceFlags.getSampled(), TraceState.getDefault()));
+  }
+
+  @Test
+  void extract_SampledContext_Bool_Short_SpanId() {
+    Map<String, String> carrier = new LinkedHashMap<>();
+    carrier.put(OtTracePropagator.TRACE_ID_HEADER, TRACE_ID);
+    carrier.put(OtTracePropagator.SPAN_ID_HEADER, SHORT_SPAN_ID);
+    carrier.put(OtTracePropagator.SAMPLED_HEADER, "true");
+
+    assertThat(getSpanContext(propagator.extract(Context.current(), carrier, getter)))
+        .isEqualTo(
+            SpanContext.createFromRemoteParent(
+                TRACE_ID, SHORT_SPAN_ID_FULL, TraceFlags.getSampled(), TraceState.getDefault()));
+  }
+
+  @Test
+  void extract_NotSampledContext_Short_SpanId() {
+    Map<String, String> carrier = new LinkedHashMap<>();
+    carrier.put(OtTracePropagator.TRACE_ID_HEADER, TRACE_ID);
+    carrier.put(OtTracePropagator.SPAN_ID_HEADER, SHORT_SPAN_ID);
+    carrier.put(OtTracePropagator.SAMPLED_HEADER, Common.FALSE_INT);
+
+    assertThat(getSpanContext(propagator.extract(Context.current(), carrier, getter)))
+        .isEqualTo(
+            SpanContext.createFromRemoteParent(
+                TRACE_ID, SHORT_SPAN_ID_FULL, TraceFlags.getDefault(), TraceState.getDefault()));
+  }
+
   @Test
   void extract_InvalidTraceId() {
     Map<String, String> invalidHeaders = new LinkedHashMap<>();