wip: tracing

Signed-off-by: Javier Aliaga <[email protected]>

Author: javier-aliaga

durabletask-client/pom.xml

@@ -74,6 +74,18 @@
       <artifactId>testcontainers</artifactId>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>io.micrometer</groupId>
+      <artifactId>micrometer-observation</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-api</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-context</artifactId>
+    </dependency>
   </dependencies>
   <build>
     <plugins>
@@ -109,24 +121,32 @@
         </executions>
       </plugin>
       <plugin>
-        <groupId>org.xolstice.maven.plugins</groupId>
+        <groupId>io.github.ascopes</groupId>
         <artifactId>protobuf-maven-plugin</artifactId>
         <version>${protobuf-maven-plugin.version}</version>
         <configuration>
-          <protocArtifact>com.google.protobuf:protoc:${protobuf.version}:exe:${os.detected.classifier}</protocArtifact>
-          <pluginId>grpc-java</pluginId>
-          <pluginArtifact>io.grpc:protoc-gen-grpc-java:${grpc.version}:exe:${os.detected.classifier}</pluginArtifact>
-          <protoSourceRoot>${protobuf.input.directory}</protoSourceRoot>
+          <protocVersion>${protobuf.version}</protocVersion>
+          <embedSourcesInClassOutputs>true</embedSourcesInClassOutputs>
+          <sourceDirectories>
+            <sourceDirectory>${protobuf.input.directory}</sourceDirectory>
+          </sourceDirectories>
+          <binaryMavenPlugins>
+            <binaryMavenPlugin>
+              <groupId>io.grpc</groupId>
+              <artifactId>protoc-gen-grpc-java</artifactId>
+              <version>${grpc.version}</version>
+            </binaryMavenPlugin>
+          </binaryMavenPlugins>
         </configuration>
         <executions>
           <execution>
             <goals>
-              <goal>compile</goal>
-              <goal>compile-custom</goal>
+              <goal>generate</goal>
             </goals>
           </execution>
         </executions>
       </plugin>
+
       <plugin>
         <groupId>org.apache.maven.plugins</groupId>
         <artifactId>maven-source-plugin</artifactId>

durabletask-client/src/main/java/io/dapr/durabletask/DurableTaskGrpcClient.java

@@ -16,6 +16,7 @@
 import com.google.protobuf.StringValue;
 import com.google.protobuf.Timestamp;
 import io.dapr.durabletask.implementation.protobuf.OrchestratorService;
+import io.dapr.durabletask.implementation.protobuf.OrchestratorService.TraceContext;
 import io.dapr.durabletask.implementation.protobuf.TaskHubSidecarServiceGrpc;
 import io.grpc.Channel;
 import io.grpc.ChannelCredentials;
@@ -28,6 +29,9 @@
 import io.grpc.netty.GrpcSslContexts;
 import io.grpc.netty.NettyChannelBuilder;
 import io.netty.handler.ssl.util.InsecureTrustManagerFactory;
+import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.Tracer;
 
 import javax.annotation.Nullable;
 import java.io.FileInputStream;
@@ -41,6 +45,8 @@
 import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Consumer;
 import java.util.logging.Logger;
 
 /**
@@ -57,6 +63,7 @@ public final class DurableTaskGrpcClient extends DurableTaskClient {
   private final DataConverter dataConverter;
   private final ManagedChannel managedSidecarChannel;
   private final TaskHubSidecarServiceGrpc.TaskHubSidecarServiceBlockingStub sidecarClient;
+  private final Tracer tracer;
 
   DurableTaskGrpcClient(DurableTaskGrpcClientBuilder builder) {
     this.dataConverter = builder.dataConverter != null ? builder.dataConverter : new JacksonDataConverter();
@@ -130,6 +137,12 @@ public final class DurableTaskGrpcClient extends DurableTaskClient {
       sidecarGrpcChannel = this.managedSidecarChannel;
     }
 
+    if (builder.tracer != null) {
+      this.tracer = builder.tracer;
+    } else {
+      this.tracer = OpenTelemetry.noop().getTracer("DurableTaskGrpcClient");
+    }
+
     this.sidecarClient = TaskHubSidecarServiceGrpc.newBlockingStub(sidecarGrpcChannel);
   }
 
@@ -188,9 +201,28 @@ public String scheduleNewOrchestrationInstance(
       builder.setScheduledStartTimestamp(ts);
     }
 
-    OrchestratorService.CreateInstanceRequest request = builder.build();
-    OrchestratorService.CreateInstanceResponse response = this.sidecarClient.startInstance(request);
-    return response.getInstanceId();
+    AtomicReference<OrchestratorService.CreateInstanceResponse> response = new AtomicReference<>();
+
+    this.observe("dapr.workflow.grpc.startInstance", (TraceContext context) -> {
+      builder.setParentTraceContext(context);
+      OrchestratorService.CreateInstanceRequest request = builder.build();
+      response.set(this.sidecarClient.startInstance(request));
+    });
+
+    return response.get().getInstanceId();
+  }
+
+  private void observe(String spanName, Consumer<TraceContext> fn) {
+    Span span = tracer.spanBuilder(spanName).startSpan();
+    try {
+      TraceContext.Builder traceContextBuilder = TraceContext.newBuilder();
+      traceContextBuilder.setTraceParent(span.getSpanContext().getTraceId());
+      var context = traceContextBuilder.build();
+
+      fn.accept(context);
+    } finally {
+      span.end();
+    }
   }
 
   @Override
@@ -206,8 +238,11 @@ public void raiseEvent(String instanceId, String eventName, Object eventPayload)
       builder.setInput(StringValue.of(serializedPayload));
     }
 
-    OrchestratorService.RaiseEventRequest request = builder.build();
-    this.sidecarClient.raiseEvent(request);
+    this.observe("dapr.workflow.grpc.raiseEvent", (TraceContext context) -> {
+      OrchestratorService.RaiseEventRequest request = builder.build();
+      this.sidecarClient.raiseEvent(request);
+    });
+
   }
 
   @Override

durabletask-client/src/main/java/io/dapr/durabletask/DurableTaskGrpcClientBuilder.java

@@ -14,6 +14,7 @@
 package io.dapr.durabletask;
 
 import io.grpc.Channel;
+import io.opentelemetry.api.trace.Tracer;
 
 /**
  * Builder class for constructing new {@link DurableTaskClient} objects that communicate with a sidecar process
@@ -27,6 +28,7 @@ public final class DurableTaskGrpcClientBuilder {
   String tlsCertPath;
   String tlsKeyPath;
   boolean insecure;
+  Tracer tracer;
 
   /**
    * Sets the {@link DataConverter} to use for converting serializable data payloads.
@@ -57,6 +59,17 @@ public DurableTaskGrpcClientBuilder grpcChannel(Channel channel) {
     return this;
   }
 
+  /**
+   * Sets the Tracer object to be used by DurableTaskClient to emit traces.
+   *
+   * @param tracer to  be used by the DurableTaskClient
+   * @return this builder object
+   */
+  public DurableTaskGrpcClientBuilder tracer(Tracer tracer) {
+    this.tracer = tracer;
+    return this;
+  }
+
   /**
    * Sets the gRPC endpoint port to connect to. If not specified, the default Durable Task port number will be used.
    *

durabletask-client/src/main/java/io/dapr/durabletask/DurableTaskGrpcWorker.java

@@ -22,10 +22,14 @@
 import io.grpc.ManagedChannelBuilder;
 import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
+import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.propagation.TextMapGetter;
 
 import java.time.Duration;
 import java.util.HashMap;
 import java.util.Iterator;
+import java.util.Map;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
@@ -85,7 +89,11 @@ public final class DurableTaskGrpcWorker implements AutoCloseable {
     this.dataConverter = builder.dataConverter != null ? builder.dataConverter : new JacksonDataConverter();
     this.maximumTimerInterval = builder.maximumTimerInterval != null ? builder.maximumTimerInterval
         : DEFAULT_MAXIMUM_TIMER_INTERVAL;
-    this.workerPool = builder.executorService != null ? builder.executorService : Executors.newCachedThreadPool();
+
+    ExecutorService rawExecutor = builder.executorService != null
+        ? builder.executorService : Executors.newCachedThreadPool();
+    this.workerPool = Context.taskWrapping(rawExecutor);
+
     this.isExecutorServiceManaged = builder.executorService == null;
   }
 
@@ -201,10 +209,16 @@ public void startAndBlock() {
             });
           } else if (requestType == OrchestratorService.WorkItem.RequestCase.ACTIVITYREQUEST) {
             OrchestratorService.ActivityRequest activityRequest = workItem.getActivityRequest();
-            logger.log(Level.FINEST,
-                String.format("Processing activity request: %s for instance: %s}",
+
+
+            logger.log(Level.INFO,
+                String.format("Processing activity request: %s for instance: %s, gRPC thread context: %s",
                     activityRequest.getName(),
-                    activityRequest.getOrchestrationInstance().getInstanceId()));
+                    activityRequest.getOrchestrationInstance().getInstanceId(),
+                    Context.current()));
+
+            // Extract trace context from the ActivityRequest and set it as current
+            Context traceContext = extractTraceContext(activityRequest);
 
             // TODO: Error handling
             this.workerPool.submit(() -> {
@@ -215,7 +229,8 @@ public void startAndBlock() {
                     activityRequest.getName(),
                     activityRequest.getInput().getValue(),
                     activityRequest.getTaskExecutionId(),
-                    activityRequest.getTaskId());
+                    activityRequest.getTaskId(),
+                    activityRequest.getParentTraceContext().getTraceParent());
               } catch (Throwable e) {
                 failureDetails = TaskFailureDetails.newBuilder()
                     .setErrorType(e.getClass().getName())
@@ -325,4 +340,57 @@ private void shutDownWorkerPool() {
   private String getSidecarAddress() {
     return this.sidecarClient.getChannel().authority();
   }
+
+  /**
+   * Extracts trace context from the ActivityRequest's ParentTraceContext field
+   * and creates an OpenTelemetry Context with the parent span set.
+   *
+   * @param activityRequest The activity request containing the parent trace context
+   * @return A Context with the parent span set, or the current context if no trace context is present
+   */
+  private Context extractTraceContext(OrchestratorService.ActivityRequest activityRequest) {
+    if (!activityRequest.hasParentTraceContext()) {
+      logger.log(Level.FINE, "No parent trace context in activity request");
+      return Context.current();
+    }
+
+    OrchestratorService.TraceContext traceContext = activityRequest.getParentTraceContext();
+    String traceParent = traceContext.getTraceParent();
+
+    if (traceParent.isEmpty()) {
+      logger.log(Level.FINE, "Empty traceparent in activity request");
+      return Context.current();
+    }
+
+    logger.log(Level.INFO,
+        String.format("Extracting trace context from ActivityRequest: traceparent=%s", traceParent));
+
+    // Use W3CTraceContextPropagator to extract the trace context
+    Map<String, String> carrier = new HashMap<>();
+    carrier.put("traceparent", traceParent);
+    if (traceContext.hasTraceState()) {
+      carrier.put("tracestate", traceContext.getTraceState().getValue());
+    }
+
+    TextMapGetter<Map<String, String>> getter = new TextMapGetter<Map<String, String>>() {
+      @Override
+      public Iterable<String> keys(Map<String, String> carrier) {
+        return carrier.keySet();
+      }
+
+      @Override
+      public String get(Map<String, String> carrier, String key) {
+        return carrier.get(key);
+      }
+    };
+
+
+    Context extractedContext = W3CTraceContextPropagator.getInstance()
+        .extract(Context.current(), carrier, getter);
+
+    logger.log(Level.INFO,
+        String.format("Extracted trace context: %s", extractedContext));
+
+    return extractedContext;
+  }
 }

durabletask-client/src/main/java/io/dapr/durabletask/TaskActivityContext.java

@@ -48,4 +48,10 @@ public interface TaskActivityContext {
    * @return the task id of the current task activity
    */
   int getTaskId();
+
+  /**
+   * Gets the trace parent id for the current workflow execution.
+   * @return trace parent id
+   */
+  String getTraceParent();
 }

durabletask-client/src/main/java/io/dapr/durabletask/TaskActivityExecutor.java

@@ -30,7 +30,8 @@ public TaskActivityExecutor(
     this.logger = logger;
   }
 
-  public String execute(String taskName, String input, String taskExecutionId, int taskId) throws Throwable {
+  public String execute(String taskName, String input,
+                        String taskExecutionId, int taskId, String traceParent) throws Throwable {
     TaskActivityFactory factory = this.activityFactories.get(taskName);
     if (factory == null) {
       throw new IllegalStateException(
@@ -43,7 +44,8 @@ public String execute(String taskName, String input, String taskExecutionId, int
           String.format("The task factory '%s' returned a null TaskActivity object.", taskName));
     }
 
-    TaskActivityContextImpl context = new TaskActivityContextImpl(taskName, input, taskExecutionId, taskId);
+    TaskActivityContextImpl context = new TaskActivityContextImpl(
+        taskName, input, taskExecutionId, taskId, traceParent);
 
     // Unhandled exceptions are allowed to escape
     Object output = activity.run(context);
@@ -59,14 +61,17 @@ private class TaskActivityContextImpl implements TaskActivityContext {
     private final String rawInput;
     private final String taskExecutionId;
     private final int taskId;
+    private final String traceParent;
 
     private final DataConverter dataConverter = TaskActivityExecutor.this.dataConverter;
 
-    public TaskActivityContextImpl(String activityName, String rawInput, String taskExecutionId, int taskId) {
+    public TaskActivityContextImpl(String activityName, String rawInput,
+                                   String taskExecutionId, int taskId, String traceParent) {
       this.name = activityName;
       this.rawInput = rawInput;
       this.taskExecutionId = taskExecutionId;
       this.taskId = taskId;
+      this.traceParent = traceParent;
     }
 
     @Override
@@ -92,5 +97,10 @@ public String getTaskExecutionId() {
     public int getTaskId() {
       return this.taskId;
     }
+
+    @Override
+    public String getTraceParent() {
+      return this.traceParent;
+    }
   }
 }

pom.xml

@@ -16,7 +16,7 @@
     <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
     <grpc.version>1.69.0</grpc.version>
     <protobuf.version>3.25.5</protobuf.version>
-    <dapr.proto.baseurl>https://raw.githubusercontent.com/dapr/dapr/v1.16.0-rc.5/dapr/proto</dapr.proto.baseurl>
+    <dapr.proto.baseurl>https://raw.githubusercontent.com/dapr/dapr/master/dapr/proto</dapr.proto.baseurl>
     <durabletask.proto.url>https://raw.githubusercontent.com/dapr/durabletask-protobuf/main/protos/orchestrator_service.proto</durabletask.proto.url>
     <dapr.sdk.version>1.17.0-SNAPSHOT</dapr.sdk.version>
     <os-maven-plugin.version>1.7.1</os-maven-plugin.version>