update InputAttributeInjector to use datadog.json component #7973

DataDog · Feb 6, 2025 · fdb01a2 · fdb01a2
1 parent 96f7f39
commit fdb01a2
Showing 1 changed file with 23 additions and 24 deletions.
diff --git a/...fn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java b/...fn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java
@@ -1,33 +1,32 @@
 package datadog.trace.instrumentation.aws.v2.sfn;
 
-import datadog.trace.bootstrap.JsonBuffer;
+import datadog.json.JsonMapper;
+import datadog.json.JsonWriter;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 
 public class InputAttributeInjector {
   public static String buildTraceContext(AgentSpan span) {
-    // Extract span tags
-    JsonBuffer spanTagsJSON = new JsonBuffer();
-    spanTagsJSON.beginObject();
-    span.getTags()
-        .forEach((tagKey, tagValue) -> spanTagsJSON.name(tagKey).value(tagValue.toString()));
-    spanTagsJSON.endObject();
+    String tagsJson = JsonMapper.toJson(span.getTags());
 
-    // Build DD trace context object
-    JsonBuffer ddTraceContextJSON = new JsonBuffer();
-    ddTraceContextJSON
-        .beginObject()
-        .name("_datadog")
-        .beginObject()
-        .name("x-datadog-trace-id")
-        .value(span.getTraceId().toString())
-        .name("x-datadog-parent-id")
-        .value(String.valueOf(span.getSpanId()))
-        .name("x-datadog-tags")
-        .value(spanTagsJSON)
-        .endObject()
-        .endObject();
+    try {
+      JsonWriter ddTraceContextJSON = new JsonWriter();
+      ddTraceContextJSON
+          .beginObject()
+          .name("_datadog")
+          .beginObject()
+          .name("x-datadog-trace-id")
+          .value(span.getTraceId().toString())
+          .name("x-datadog-parent-id")
+          .value(String.valueOf(span.getSpanId()))
+          .name("x-datadog-tags")
+          .jsonValue(tagsJson)
+          .endObject()
+          .endObject();
 
-    return ddTraceContextJSON.toString();
+      return ddTraceContextJSON.toString();
+    } catch (Exception e) {
+      return "{}";
+    }
   }
 
   public static String getModifiedInput(String request, String ddTraceContextJSON) {
@@ -38,8 +37,8 @@ public static String getModifiedInput(String request, String ddTraceContextJSON)
     if (inputContent.isEmpty()) {
       modifiedInput.insert(endPos, ddTraceContextJSON);
     } else {
-      modifiedInput.insert(
-          endPos, ",".concat(ddTraceContextJSON)); // prepend comma to existing input
+      // Prepend comma to separate from existing content
+      modifiedInput.insert(endPos, "," + ddTraceContextJSON);
     }
     return modifiedInput.toString();
   }