Propagate otelContext in Request and in thread-local/coroutine context.

Author: slinkydeveloper

sdk-api-kotlin/build.gradle.kts

@@ -15,6 +15,8 @@ dependencies {
   implementation(kotlinLibs.kotlinx.serialization.json)
 
   implementation(coreLibs.log4j.api)
+  implementation(platform(coreLibs.opentelemetry.bom))
+  implementation(coreLibs.opentelemetry.kotlin)
 
   testImplementation(project(":sdk-core"))
   testImplementation(testingLibs.junit.jupiter)

sdk-api-kotlin/src/main/kotlin/dev/restate/sdk/kotlin/Service.kt

@@ -11,6 +11,7 @@ package dev.restate.sdk.kotlin
 import com.google.protobuf.ByteString
 import dev.restate.sdk.common.*
 import dev.restate.sdk.common.syscalls.*
+import io.opentelemetry.extension.kotlin.asContextElement
 import kotlin.coroutines.CoroutineContext
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
@@ -138,7 +139,8 @@ private constructor(
       val scope =
           CoroutineScope(
               options.coroutineContext +
-                  InvocationHandler.SYSCALLS_THREAD_LOCAL.asContextElement(syscalls))
+                  InvocationHandler.SYSCALLS_THREAD_LOCAL.asContextElement(syscalls) +
+                  syscalls.request().otelContext()!!.asContextElement())
       scope.launch {
         val serializedResult: ByteString
 

sdk-api/src/main/java/dev/restate/sdk/Service.java

@@ -12,6 +12,7 @@
 import dev.restate.sdk.common.*;
 import dev.restate.sdk.common.ServiceType;
 import dev.restate.sdk.common.syscalls.*;
+import io.opentelemetry.context.Scope;
 import java.util.*;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
@@ -105,7 +106,6 @@ public Service build(Service.Options options) {
     }
   }
 
-  @SuppressWarnings("unchecked")
   public static class Handler<REQ, RES> implements InvocationHandler<Service.Options> {
     private final HandlerSignature<REQ, RES> handlerSignature;
     private final HandlerType handlerType;
@@ -119,6 +119,7 @@ public Handler(
         BiFunction<? extends Context, REQ, RES> runner) {
       this.handlerSignature = handlerSignature;
       this.handlerType = handlerType;
+      //noinspection unchecked
       this.runner = (BiFunction<Context, REQ, RES>) runner;
     }
 
@@ -148,7 +149,7 @@ public void handle(
               options.executor.execute(
                   () -> {
                     SYSCALLS_THREAD_LOCAL.set(syscalls);
-                    try {
+                    try (Scope ignored = syscalls.request().otelContext().makeCurrent()) {
                       runnable.run();
                     } finally {
                       SYSCALLS_THREAD_LOCAL.remove();

sdk-common/build.gradle.kts

@@ -9,6 +9,8 @@ dependencies {
   compileOnly(coreLibs.jspecify)
 
   api(coreLibs.protobuf.java)
+  api(platform(coreLibs.opentelemetry.bom))
+  api(coreLibs.opentelemetry.api)
 
   implementation(platform(jacksonLibs.jackson.bom))
   implementation(jacksonLibs.jackson.core)

sdk-common/src/main/java/dev/restate/sdk/common/Request.java

@@ -9,17 +9,24 @@
 package dev.restate.sdk.common;
 
 import com.google.protobuf.ByteString;
+import io.opentelemetry.context.Context;
 import java.util.Map;
 import java.util.Objects;
 
 public final class Request {
 
   private final InvocationId invocationId;
+  private final Context otelContext;
   private final ByteString body;
   private final Map<String, String> headers;
 
-  public Request(InvocationId invocationId, ByteString body, Map<String, String> headers) {
+  public Request(
+      InvocationId invocationId,
+      Context otelContext,
+      ByteString body,
+      Map<String, String> headers) {
     this.invocationId = invocationId;
+    this.otelContext = otelContext;
     this.body = body;
     this.headers = headers;
   }
@@ -28,6 +35,10 @@ public InvocationId invocationId() {
     return invocationId;
   }
 
+  public Context otelContext() {
+    return otelContext;
+  }
+
   public byte[] body() {
     return body.toByteArray();
   }
@@ -46,29 +57,23 @@ public boolean equals(Object o) {
     if (o == null || getClass() != o.getClass()) return false;
 
     Request request = (Request) o;
-
-    if (!Objects.equals(invocationId, request.invocationId)) return false;
-    if (!Objects.equals(body, request.body)) return false;
-    return Objects.equals(headers, request.headers);
+    return Objects.equals(invocationId, request.invocationId)
+        && Objects.equals(otelContext, request.otelContext)
+        && Objects.equals(body, request.body)
+        && Objects.equals(headers, request.headers);
   }
 
   @Override
   public int hashCode() {
-    int result = invocationId != null ? invocationId.hashCode() : 0;
-    result = 31 * result + (body != null ? body.hashCode() : 0);
-    result = 31 * result + (headers != null ? headers.hashCode() : 0);
+    int result = Objects.hashCode(invocationId);
+    result = 31 * result + Objects.hashCode(otelContext);
+    result = 31 * result + Objects.hashCode(body);
+    result = 31 * result + Objects.hashCode(headers);
     return result;
   }
 
   @Override
   public String toString() {
-    return "Request{"
-        + "invocationId="
-        + invocationId
-        + ", body="
-        + body
-        + ", headers="
-        + headers
-        + '}';
+    return "Request{" + "invocationId=" + invocationId + ", headers=" + headers + '}';
   }
 }

sdk-core/build.gradle.kts

@@ -24,7 +24,6 @@ dependencies {
 
   implementation(platform(coreLibs.opentelemetry.bom))
   implementation(coreLibs.opentelemetry.api)
-  implementation(coreLibs.opentelemetry.semconv)
 
   testCompileOnly(coreLibs.jspecify)
   testImplementation(testingLibs.junit.jupiter)

sdk-core/src/main/java/dev/restate/sdk/core/InvocationStateMachine.java

@@ -19,6 +19,7 @@
 import dev.restate.sdk.common.syscalls.*;
 import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.context.Context;
 import java.util.*;
 import java.util.concurrent.Flow;
 import java.util.function.Consumer;
@@ -227,6 +228,7 @@ void onStartMessage(MessageLite msg) {
           Request request =
               new Request(
                   invocationId,
+                  Context.root().with(span),
                   inputEntry.getValue(),
                   inputEntry.getHeadersList().stream()
                       .collect(

sdk-core/src/main/java/dev/restate/sdk/core/RestateEndpoint.java

@@ -17,7 +17,6 @@
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.api.trace.Tracer;
-import io.opentelemetry.semconv.trace.attributes.SemanticAttributes;
 import java.util.*;
 import java.util.concurrent.Executor;
 import java.util.function.Function;
@@ -72,9 +71,6 @@ public ResolvedEndpointHandler resolve(
             .spanBuilder("Invoke method")
             .setSpanKind(SpanKind.SERVER)
             .setParent(otelContext)
-            .setAttribute(SemanticAttributes.RPC_SYSTEM, "restate")
-            .setAttribute(SemanticAttributes.RPC_SERVICE, componentName)
-            .setAttribute(SemanticAttributes.RPC_METHOD, handlerName)
             .startSpan();
 
     // Setup logging context

settings.gradle.kts

@@ -36,7 +36,7 @@ dependencyResolutionManagement {
     create("coreLibs") {
       version("protobuf", "3.24.3")
       version("log4j", "2.22.0")
-      version("opentelemetry", "1.30.1")
+      version("opentelemetry", "1.37.0")
 
       library("protoc", "com.google.protobuf", "protoc").versionRef("protobuf")
       library("protobuf-java", "com.google.protobuf", "protobuf-java").versionRef("protobuf")
@@ -48,7 +48,8 @@ dependencyResolutionManagement {
       library("opentelemetry-bom", "io.opentelemetry", "opentelemetry-bom")
           .versionRef("opentelemetry")
       library("opentelemetry-api", "io.opentelemetry", "opentelemetry-api").withoutVersion()
-      library("opentelemetry-semconv", "io.opentelemetry:opentelemetry-semconv:1.19.0-alpha")
+      library("opentelemetry-kotlin", "io.opentelemetry", "opentelemetry-extension-kotlin")
+          .withoutVersion()
 
       library("jspecify", "org.jspecify", "jspecify").version("0.3.0")
     }