Fix logging context propagation (#185)

* Improve the Context propagation logic:
  * This is now more obvious, in particular It makes more clear what's the boundary between state machine and user code and how the Context is propagated
  * Renamed the methods of RestateServerCallListener to make more clear their goal
* Propagate fully qualified method name and invocation state in SyscallsInternal
* Reintroduce the GrpcContextDataProvider
* Fix the LoggingContextSetter.THREAD_LOCAL_INSTANCE
* Bump log4j2

Author: slinkydeveloper

sdk-core/build.gradle.kts

@@ -17,6 +17,9 @@ dependencies {
   implementation(coreLibs.grpc.protobuf)
   implementation(coreLibs.log4j.api)
 
+  // We don't want a hard-dependency on it
+  compileOnly(coreLibs.log4j.core)
+
   implementation(platform(coreLibs.opentelemetry.bom))
   implementation(coreLibs.opentelemetry.api)
   implementation(coreLibs.opentelemetry.semconv)

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

@@ -11,7 +11,6 @@
 import com.google.protobuf.ByteString;
 import com.google.protobuf.MessageLite;
 import dev.restate.sdk.common.TerminalException;
-import dev.restate.sdk.common.syscalls.*;
 import dev.restate.sdk.common.syscalls.DeferredResult;
 import dev.restate.sdk.common.syscalls.EnterSideEffectSyscallCallback;
 import dev.restate.sdk.common.syscalls.ExitSideEffectSyscallCallback;
@@ -48,13 +47,13 @@ private ExecutorSwitchingServerCallListener(
     }
 
     @Override
-    public void onMessageAndHalfClose(MessageLite message) {
-      userExecutor.execute(() -> listener.onMessageAndHalfClose(message));
+    public void invoke(MessageLite message) {
+      userExecutor.execute(() -> listener.invoke(message));
     }
 
     // A bit of explanation why the following methods are not executed on the user executor.
     //
-    // The listener methods onReady/onCancel/onComplete are used purely for notification reasons,
+    // The listener methods listenerReady/cancel/close are used purely for notification reasons,
     // they don't execute any user code.
     //
     // Running them in the userExecutor can also be problematic if the listener
@@ -64,18 +63,18 @@ public void onMessageAndHalfClose(MessageLite message) {
     // as thread local.
 
     @Override
-    public void onCancel() {
-      listener.onCancel();
+    public void cancel() {
+      listener.cancel();
     }
 
     @Override
-    public void onComplete() {
-      listener.onComplete();
+    public void close() {
+      listener.close();
     }
 
     @Override
-    public void onReady() {
-      listener.onReady();
+    public void listenerReady() {
+      listener.listenerReady();
     }
   }
 
@@ -182,6 +181,18 @@ public <T> void resolveDeferred(
       syscallsExecutor.execute(() -> syscalls.resolveDeferred(deferredToResolve, callback));
     }
 
+    @Override
+    public String getFullyQualifiedMethodName() {
+      // We can read this from another thread
+      return syscalls.getFullyQualifiedMethodName();
+    }
+
+    @Override
+    public InvocationState getInvocationState() {
+      // We can read this from another thread
+      return syscalls.getInvocationState();
+    }
+
     @Override
     public void close() {
       syscallsExecutor.execute(syscalls::close);

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

@@ -0,0 +1,48 @@
+// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH
+//
+// This file is part of the Restate Java SDK,
+// which is released under the MIT license.
+//
+// You can find a copy of the license in file LICENSE in the root
+// directory of this repository or package, or at
+// https://github.com/restatedev/sdk-java/blob/main/LICENSE
+package dev.restate.sdk.core;
+
+import dev.restate.sdk.common.InvocationId;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.logging.log4j.core.util.ContextDataProvider;
+
+/**
+ * Log4j2 ContextDataProvider inferring context from the Grpc context.
+ *
+ * <p>This is used to propagate the context to the user code, such that log statements from the user
+ * will contain the restate logging context variables.
+ *
+ * <p>This is based on grpc Context due to the fact that it's the only guaranteed thread
+ * local/context we can rely on that is always available in the user code.
+ */
+public class GrpcContextDataProvider implements ContextDataProvider {
+  @Override
+  public Map<String, String> supplyContextData() {
+    InvocationId invocationId = InvocationId.INVOCATION_ID_KEY.get();
+    SyscallsInternal syscalls = (SyscallsInternal) SyscallsInternal.SYSCALLS_KEY.get();
+
+    if (invocationId == null) {
+      return Collections.emptyMap();
+    }
+
+    // We can't use the immutable MapN implementation from Map.of because of
+    // https://github.com/apache/logging-log4j2/issues/2098
+    HashMap<String, String> m = new HashMap<>();
+    m.put(RestateGrpcServer.LoggingContextSetter.INVOCATION_ID_KEY, invocationId.toString());
+    m.put(
+        RestateGrpcServer.LoggingContextSetter.SERVICE_METHOD_KEY,
+        syscalls.getFullyQualifiedMethodName());
+    m.put(
+        RestateGrpcServer.LoggingContextSetter.SERVICE_INVOCATION_STATUS_KEY,
+        syscalls.getInvocationState().toString());
+    return m;
+  }
+}

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

@@ -8,9 +8,7 @@
 // https://github.com/restatedev/sdk-java/blob/main/LICENSE
 package dev.restate.sdk.core;
 
-import io.grpc.Metadata;
-import io.grpc.ServerCall;
-import io.grpc.Status;
+import io.grpc.*;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
@@ -24,50 +22,75 @@ class GrpcServerCallListenerAdaptor<ReqT, RespT> implements RestateServerCallLis
 
   private static final Logger LOG = LogManager.getLogger(GrpcServerCallListenerAdaptor.class);
 
+  private final Context context;
   private final ServerCall<ReqT, RespT> serverCall;
-
   private final ServerCall.Listener<ReqT> delegate;
 
   GrpcServerCallListenerAdaptor(
-      ServerCall.Listener<ReqT> delegate, ServerCall<ReqT, RespT> serverCall) {
-    this.delegate = delegate;
+      Context context,
+      ServerCall<ReqT, RespT> serverCall,
+      Metadata headers,
+      ServerCallHandler<ReqT, RespT> next) {
+    this.context = context;
     this.serverCall = serverCall;
+
+    // This emulates Contexts.interceptCall.
+    // We need it because some code generator depends on the fact that startCall already has the
+    // context available
+    Context previous = this.context.attach();
+    try {
+      this.delegate = next.startCall(serverCall, headers);
+    } finally {
+      this.context.detach(previous);
+    }
   }
 
   @Override
-  public void onMessageAndHalfClose(ReqT message) {
+  public void invoke(ReqT message) {
+    Context previous = context.attach();
     try {
       delegate.onMessage(message);
       delegate.onHalfClose();
     } catch (Throwable e) {
       closeWithException(e);
+    } finally {
+      context.detach(previous);
     }
   }
 
   @Override
-  public void onCancel() {
+  public void cancel() {
+    Context previous = context.attach();
     try {
       delegate.onCancel();
     } catch (Throwable e) {
       closeWithException(e);
+    } finally {
+      context.detach(previous);
     }
   }
 
   @Override
-  public void onComplete() {
+  public void close() {
+    Context previous = context.attach();
     try {
       delegate.onComplete();
     } catch (Throwable e) {
       closeWithException(e);
+    } finally {
+      context.detach(previous);
     }
   }
 
   @Override
-  public void onReady() {
+  public void listenerReady() {
+    Context previous = context.attach();
     try {
       delegate.onReady();
     } catch (Throwable e) {
       closeWithException(e);
+    } finally {
+      context.detach(previous);
     }
   }
 

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

@@ -0,0 +1,16 @@
+// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH
+//
+// This file is part of the Restate Java SDK,
+// which is released under the MIT license.
+//
+// You can find a copy of the license in file LICENSE in the root
+// directory of this repository or package, or at
+// https://github.com/restatedev/sdk-java/blob/main/LICENSE
+package dev.restate.sdk.core;
+
+public enum InvocationState {
+  WAITING_START,
+  REPLAYING,
+  PROCESSING,
+  CLOSED;
+}