Introduced basic trace post-processing

Author: ValentinZakharov

dd-trace-core/src/main/java/datadog/trace/common/writer/TraceProcessingWorker.java

@@ -7,6 +7,12 @@
 
 import datadog.communication.ddagent.DroppingPolicy;
 import datadog.trace.api.Config;
+import datadog.trace.api.gateway.CallbackProvider;
+import datadog.trace.api.gateway.Events;
+import datadog.trace.api.gateway.Flow;
+import datadog.trace.api.gateway.RequestContext;
+import datadog.trace.api.gateway.RequestContextSlot;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.common.sampling.SingleSpanSampler;
 import datadog.trace.common.writer.ddagent.FlushEvent;
 import datadog.trace.common.writer.ddagent.Prioritization;
@@ -17,6 +23,7 @@
 import java.util.List;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
 import org.jctools.queues.MessagePassingQueue;
 import org.jctools.queues.MpscBlockingConsumerArrayQueue;
 import org.slf4j.Logger;
@@ -235,6 +242,18 @@ public void onEvent(Object event) {
       try {
         if (event instanceof List) {
           List<DDSpan> trace = (List<DDSpan>) event;
+
+          // Do post-processing if required and priority > 0
+          DDSpan rootSpan = trace.get(0);
+          if (rootSpan != null
+              && rootSpan.isRequiresPostProcessing()
+              && rootSpan.samplingPriority() > 0) {
+            CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
+            Function<RequestContext, Flow<Void>> reqPostProcCb =
+                cbp.getCallback(Events.EVENTS.requestPostProcessing());
+            reqPostProcCb.apply(rootSpan.getRequestContext());
+          }
+
           // TODO populate `_sample_rate` metric in a way that accounts for lost/dropped traces
           payloadDispatcher.addTrace(trace);
         } else if (event instanceof FlushEvent) {

dd-trace-core/src/main/java/datadog/trace/core/DDSpan.java

@@ -103,6 +103,7 @@ static DDSpan create(
 
   // the request is to be blocked (AppSec)
   private volatile Flow.Action.RequestBlockingAction requestBlockingAction;
+  private boolean requiresPostProcessing;
 
   /**
    * Version of a span that can be set by the long running spans feature:
@@ -399,6 +400,16 @@ public Flow.Action.RequestBlockingAction getRequestBlockingAction() {
     return requestBlockingAction;
   }
 
+  @Override
+  public void setRequiresPostProcessing(boolean postProcessing) {
+    this.requiresPostProcessing = postProcessing;
+  }
+
+  @Override
+  public boolean isRequiresPostProcessing() {
+    return this.requiresPostProcessing;
+  }
+
   @Override
   public DDSpan setTag(final String tag, final int value) {
     // can't use tag interceptor because it might set a metric

internal-api/src/main/java/datadog/trace/api/gateway/Events.java

@@ -189,6 +189,17 @@ public EventType<BiFunction<RequestContext, Object, Flow<Void>>> grpcServerReque
     return (EventType<BiFunction<RequestContext, Object, Flow<Void>>>) GRPC_SERVER_REQUEST_MESSAGE;
   }
 
+  static final int REQUEST_POST_PROCESSING_ID = 15;
+
+  @SuppressWarnings("rawtypes")
+  private static final EventType REQUEST_POST_PROCESSING =
+      new ET<>("request.post.processing", REQUEST_POST_PROCESSING_ID);
+  /** The request has been processed and is ready to be sent to the backend */
+  @SuppressWarnings("unchecked")
+  public EventType<Function<RequestContext, Flow<Void>>> requestPostProcessing() {
+    return (EventType<Function<RequestContext, Flow<Void>>>) REQUEST_POST_PROCESSING;
+  }
+
   static final int MAX_EVENTS = nextId.get();
 
   private static final class ET<T> extends EventType<T> {

internal-api/src/main/java/datadog/trace/api/gateway/IGSpanInfo.java

@@ -16,4 +16,8 @@ public interface IGSpanInfo {
   void setRequestBlockingAction(Flow.Action.RequestBlockingAction rba);
 
   Flow.Action.RequestBlockingAction getRequestBlockingAction();
+
+  void setRequiresPostProcessing(boolean postProcessing);
+
+  boolean isRequiresPostProcessing();
 }

internal-api/src/main/java/datadog/trace/api/gateway/InstrumentationGateway.java

@@ -12,6 +12,7 @@
 import static datadog.trace.api.gateway.Events.REQUEST_INFERRED_CLIENT_ADDRESS_ID;
 import static datadog.trace.api.gateway.Events.REQUEST_METHOD_URI_RAW_ID;
 import static datadog.trace.api.gateway.Events.REQUEST_PATH_PARAMS_ID;
+import static datadog.trace.api.gateway.Events.REQUEST_POST_PROCESSING_ID;
 import static datadog.trace.api.gateway.Events.REQUEST_STARTED_ID;
 import static datadog.trace.api.gateway.Events.RESPONSE_HEADER_DONE_ID;
 import static datadog.trace.api.gateway.Events.RESPONSE_HEADER_ID;
@@ -358,6 +359,19 @@ public Flow<Void> apply(RequestContext ctx, Integer status) {
                 }
               }
             };
+      case REQUEST_POST_PROCESSING_ID:
+        return (C)
+            new Function<RequestContext, Flow<Void>>() {
+              @Override
+              public Flow<Void> apply(RequestContext ctx) {
+                try {
+                  return ((Function<RequestContext, Flow<Void>>) callback).apply(ctx);
+                } catch (Throwable t) {
+                  log.warn("Callback for {} threw.", eventType, t);
+                  return Flow.ResultFlow.empty();
+                }
+              }
+            };
       default:
         log.warn("Unwrapped callback for {}", eventType);
         return callback;

internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java

@@ -578,6 +578,14 @@ public Flow.Action.RequestBlockingAction getRequestBlockingAction() {
       return null;
     }
 
+    @Override
+    public void setRequiresPostProcessing(boolean postProcessing) {}
+
+    @Override
+    public boolean isRequiresPostProcessing() {
+      return false;
+    }
+
     @Override
     public AgentSpan setTag(final String tag, final Number value) {
       return this;

internal-api/src/test/java/datadog/trace/api/gateway/InstrumentationGatewayTest.java

@@ -194,6 +194,8 @@ public void testNormalCalls() {
     cbp.getCallback(events.responseHeader()).accept(null, null, null);
     ss.registerCallback(events.responseHeaderDone(), callback.function);
     cbp.getCallback(events.responseHeaderDone()).apply(null);
+    ss.registerCallback(events.requestPostProcessing(), callback.function);
+    cbp.getCallback(events.requestPostProcessing()).apply(null);
     assertThat(callback.count).isEqualTo(Events.MAX_EVENTS);
   }
 
@@ -240,6 +242,8 @@ public void testThrowableBlocking() {
     cbp.getCallback(events.responseHeader()).accept(null, null, null);
     ss.registerCallback(events.responseHeaderDone(), throwback.function);
     cbp.getCallback(events.responseHeaderDone()).apply(null);
+    ss.registerCallback(events.requestPostProcessing(), throwback.function);
+    cbp.getCallback(events.requestPostProcessing()).apply(null);
     assertThat(throwback.count).isEqualTo(Events.MAX_EVENTS);
   }