add OpenTelemetry baggage endpoints

DataDog · Dec 5, 2024 · ac256d8 · ac256d8
1 parent ba70912
commit ac256d8
Show file tree

Hide file tree

Showing 6 changed files with 78 additions and 16 deletions.
diff --git a/...c/src/main/java/com/datadoghq/trace/opentelemetry/controller/OpenTelemetryController.java b/...c/src/main/java/com/datadoghq/trace/opentelemetry/controller/OpenTelemetryController.java
@@ -8,27 +8,13 @@
 import static io.opentelemetry.api.trace.SpanKind.SERVER;
 import static java.util.concurrent.TimeUnit.MICROSECONDS;
 
-import com.datadoghq.trace.opentelemetry.dto.AddEventArgs;
-import com.datadoghq.trace.opentelemetry.dto.EndSpanArgs;
-import com.datadoghq.trace.opentelemetry.dto.FlushArgs;
-import com.datadoghq.trace.opentelemetry.dto.FlushResult;
-import com.datadoghq.trace.opentelemetry.dto.IsRecordingArgs;
-import com.datadoghq.trace.opentelemetry.dto.IsRecordingResult;
-import com.datadoghq.trace.opentelemetry.dto.KeyValue;
-import com.datadoghq.trace.opentelemetry.dto.RecordExceptionArgs;
-import com.datadoghq.trace.opentelemetry.dto.SetAttributesArgs;
-import com.datadoghq.trace.opentelemetry.dto.SetNameArgs;
-import com.datadoghq.trace.opentelemetry.dto.SetStatusArgs;
-import com.datadoghq.trace.opentelemetry.dto.SpanContextArgs;
-import com.datadoghq.trace.opentelemetry.dto.SpanContextResult;
-import com.datadoghq.trace.opentelemetry.dto.SpanLink;
-import com.datadoghq.trace.opentelemetry.dto.StartSpanArgs;
-import com.datadoghq.trace.opentelemetry.dto.StartSpanResult;
+import com.datadoghq.trace.opentelemetry.dto.*;
 import datadog.trace.api.DDSpanId;
 import datadog.trace.api.DDTraceId;
 import datadog.trace.api.GlobalTracer;
 import datadog.trace.api.internal.InternalTracer;
 import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.api.baggage.Baggage;
 import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.common.AttributesBuilder;
 import io.opentelemetry.api.trace.Span;
@@ -58,11 +44,13 @@ public class OpenTelemetryController {
   private final Tracer tracer;
   private final TextMapPropagator propagator;
   private final Map<Long, Span> spans;
+  private Baggage baggage;
 
   public OpenTelemetryController() {
     this.tracer = GlobalOpenTelemetry.getTracer("java-client");
     this.propagator = GlobalOpenTelemetry.getPropagators().getTextMapPropagator();
     this.spans = new HashMap<>();
+    this.baggage = Baggage.empty();
   }
 
   private static SpanKind parseSpanKindNumber(int spanKindNumber) {
@@ -318,6 +306,50 @@ public FlushResult flush(@RequestBody FlushArgs args) {
     }
   }
 
+  // set_baggage
+  @PostMapping("set_baggage")
+  public void setBaggage(@RequestBody SetBaggageArgs args) {
+    LOGGER.info("Setting OTel baggage: {}", args);
+    this.baggage = this.baggage
+            .toBuilder()
+            .put(args.key(), args.value())
+            .build();
+  }
+
+  // get_baggage
+  @GetMapping("get_baggage")
+  public GetBaggageResult getBaggage(@RequestBody GetBaggageArgs args) {
+    LOGGER.info("Getting an OTel baggage entry");
+    var value = this.baggage.getEntryValue(args.key());
+    return new GetBaggageResult(value);
+  }
+
+  // get_all_baggage
+  @GetMapping("get_all_baggage")
+  public GetAllBaggageResult getAllBaggage() {
+    LOGGER.info("Getting all OTel baggage entries");
+    Map<String, String> baggageMap = new HashMap<>();
+    this.baggage.forEach((key, entry) -> baggageMap.put(key, entry.getValue()));
+    return new GetAllBaggageResult(baggageMap);
+  }
+
+  // remove_baggage
+  @PostMapping("remove_baggage")
+  public void removeBaggage(@RequestBody RemoveBaggageArgs args) {
+    LOGGER.info("Removing OTel baggage entry: {}", args);
+    this.baggage = this.baggage
+            .toBuilder()
+            .remove(args.key())
+            .build();
+  }
+
+  // remove_all_baggage
+  @PostMapping("remove_all_baggage")
+  public void removeAllBaggage() {
+    LOGGER.info("Removing all OTel baggage entries");
+    this.baggage = Baggage.empty();
+  }
+
   private Span getSpan(long spanId) {
     Span span = this.spans.get(spanId);
     if (span == null) {

diff --git a/...a/parametric/src/main/java/com/datadoghq/trace/opentelemetry/dto/GetAllBaggageResult.java b/...a/parametric/src/main/java/com/datadoghq/trace/opentelemetry/dto/GetAllBaggageResult.java
@@ -0,0 +1,7 @@
+package com.datadoghq.trace.opentelemetry.dto;
+
+import java.util.Map;
+
+public record GetAllBaggageResult(
+        Map<String, String> baggage) {
+}
diff --git a/...r/java/parametric/src/main/java/com/datadoghq/trace/opentelemetry/dto/GetBaggageArgs.java b/...r/java/parametric/src/main/java/com/datadoghq/trace/opentelemetry/dto/GetBaggageArgs.java
@@ -0,0 +1,5 @@
+package com.datadoghq.trace.opentelemetry.dto;
+
+public record GetBaggageArgs(
+        String key) {
+}
diff --git a/...java/parametric/src/main/java/com/datadoghq/trace/opentelemetry/dto/GetBaggageResult.java b/...java/parametric/src/main/java/com/datadoghq/trace/opentelemetry/dto/GetBaggageResult.java
@@ -0,0 +1,6 @@
+package com.datadoghq.trace.opentelemetry.dto;
+
+public record GetBaggageResult(
+        String baggage) {
+}
+
diff --git a/...ava/parametric/src/main/java/com/datadoghq/trace/opentelemetry/dto/RemoveBaggageArgs.java b/...ava/parametric/src/main/java/com/datadoghq/trace/opentelemetry/dto/RemoveBaggageArgs.java
@@ -0,0 +1,5 @@
+package com.datadoghq.trace.opentelemetry.dto;
+
+public record RemoveBaggageArgs(
+        String key) {
+}
diff --git a/...r/java/parametric/src/main/java/com/datadoghq/trace/opentelemetry/dto/SetBaggageArgs.java b/...r/java/parametric/src/main/java/com/datadoghq/trace/opentelemetry/dto/SetBaggageArgs.java
@@ -0,0 +1,7 @@
+package com.datadoghq.trace.opentelemetry.dto;
+
+public record SetBaggageArgs(
+        String key,
+        String value) {
+}
+