Use custom Moshi JsonAdapter

DataDog · Apr 5, 2024 · ed7ebbf · ed7ebbf
1 parent 734e3c5
commit ed7ebbf
Showing 1 changed file with 57 additions and 20 deletions.
diff --git a/dd-trace-core/src/main/java/datadog/trace/core/DDSpanLink.java b/dd-trace-core/src/main/java/datadog/trace/core/DDSpanLink.java
@@ -2,19 +2,24 @@
 
 import static datadog.trace.bootstrap.instrumentation.api.SpanLinkAttributes.EMPTY;
 
-import com.squareup.moshi.FromJson;
 import com.squareup.moshi.JsonAdapter;
+import com.squareup.moshi.JsonReader;
+import com.squareup.moshi.JsonWriter;
 import com.squareup.moshi.Moshi;
-import com.squareup.moshi.ToJson;
+import com.squareup.moshi.Types;
 import datadog.trace.api.DDSpanId;
 import datadog.trace.api.DDTraceId;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink;
 import datadog.trace.bootstrap.instrumentation.api.SpanLink;
 import datadog.trace.bootstrap.instrumentation.api.SpanLinkAttributes;
 import datadog.trace.core.propagation.ExtractedContext;
 import datadog.trace.core.propagation.PropagationTags;
+import java.io.IOException;
+import java.lang.annotation.Annotation;
+import java.lang.reflect.Type;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -23,8 +28,6 @@ public class DDSpanLink extends SpanLink {
   private static final Logger LOGGER = LoggerFactory.getLogger(DDSpanLink.class);
   /** The maximum of characters a span tag value can hold. */
   private static final int TAG_MAX_LENGTH = 25_000;
-  /** JSON encoder (lazily initialized) */
-  private static JsonAdapter<AgentSpanLink> encoder;
 
   protected DDSpanLink(
       DDTraceId traceId, long spanId, byte traceFlags, String traceState, Attributes attributes) {
@@ -95,16 +98,44 @@ public static String toTag(List<AgentSpanLink> links) {
   }
 
   private static JsonAdapter<AgentSpanLink> getEncoder() {
-    if (encoder == null) {
-      Moshi moshi = new Moshi.Builder().add(new SpanLinkAdapter()).build();
-      encoder = moshi.adapter(AgentSpanLink.class);
+    return EncoderHolder.ENCODER;
+  }
+
+  /**
+   * JSON encoder (lazily initialized) This is not folded at native image build time, so it works as
+   * expected
+   */
+  private static class EncoderHolder {
+    static final JsonAdapter<AgentSpanLink> ENCODER;
+
+    static {
+      Moshi moshi = new Moshi.Builder().add(AgentSpanLinkAdapter.factory()).build();
+      ENCODER = moshi.adapter(AgentSpanLink.class);
     }
-    return encoder;
   }
 
-  private static class SpanLinkAdapter {
-    @ToJson
-    SpanLinkJson toSpanLinkJson(AgentSpanLink link) {
+  private static class AgentSpanLinkAdapter extends JsonAdapter<AgentSpanLink> {
+    private final JsonAdapter<DDSpanLink.SpanLinkJson> delegateAdapter;
+
+    private AgentSpanLinkAdapter(JsonAdapter<DDSpanLink.SpanLinkJson> delegateAdapter) {
+      this.delegateAdapter = delegateAdapter;
+    }
+
+    @Override
+    public AgentSpanLink fromJson(JsonReader reader) throws IOException {
+      SpanLinkJson json = delegateAdapter.fromJson(reader);
+      assert json != null;
+      return new DDSpanLink(
+          DDTraceId.fromHex(json.trace_id),
+          DDSpanId.fromHex(json.span_id),
+          json.flags,
+          json.tracestate,
+          SpanLinkAttributes.fromMap(json.attributes));
+    }
+
+    @Override
+    public void toJson(JsonWriter writer, AgentSpanLink link) throws IOException {
+      assert link != null;
       SpanLinkJson json = new SpanLinkJson();
       json.trace_id = link.traceId().toHexString();
       json.span_id = DDSpanId.toHexString(link.spanId());
@@ -113,17 +144,23 @@ SpanLinkJson toSpanLinkJson(AgentSpanLink link) {
       if (!link.attributes().isEmpty()) {
         json.attributes = link.attributes().asMap();
       }
-      return json;
+      delegateAdapter.toJson(writer, json);
     }
 
-    @FromJson
-    AgentSpanLink fromSpanLinkJson(SpanLinkJson json) {
-      return new DDSpanLink(
-          DDTraceId.fromHex(json.trace_id),
-          DDSpanId.fromHex(json.span_id),
-          json.flags,
-          json.tracestate,
-          SpanLinkAttributes.fromMap(json.attributes));
+    public static JsonAdapter.Factory factory() {
+      return new JsonAdapter.Factory() {
+        @Override
+        public JsonAdapter<?> create(
+            Type type, Set<? extends Annotation> annotations, Moshi moshi) {
+          if (AgentSpanLink.class.isAssignableFrom(Types.getRawType(type))) {
+            JsonAdapter<DDSpanLink.SpanLinkJson> delegateAdapter =
+                moshi.adapter(DDSpanLink.SpanLinkJson.class);
+            return new AgentSpanLinkAdapter(delegateAdapter);
+          } else {
+            return null;
+          }
+        }
+      };
     }
   }