lettuce-5.0

Author: SylvainJuge

instrumentation/lettuce/lettuce-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/lettuce/v5_0/LettuceAsyncCommandInstrumentation.java

@@ -5,14 +5,14 @@
 
 package io.opentelemetry.javaagent.instrumentation.lettuce.v5_0;
 
+import static io.opentelemetry.javaagent.instrumentation.lettuce.v5_0.LettuceSingletons.CONTEXT;
 import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
 import static net.bytebuddy.matcher.ElementMatchers.named;
 import static net.bytebuddy.matcher.ElementMatchers.namedOneOf;
 
 import io.lettuce.core.protocol.AsyncCommand;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
-import io.opentelemetry.instrumentation.api.util.VirtualField;
 import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
@@ -44,22 +44,22 @@ public static void saveContext(@Advice.This AsyncCommand<?, ?, ?> asyncCommand)
       Context context = Java8BytecodeBridge.currentContext();
       // get the context that submitted this command and attach it, it will be used to run callbacks
       context = context.get(LettuceSingletons.COMMAND_CONTEXT_KEY);
-      VirtualField.find(AsyncCommand.class, Context.class).set(asyncCommand, context);
+      CONTEXT.set(asyncCommand, context);
     }
   }
 
   @SuppressWarnings("unused")
   public static class RestoreContextAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void onEnter(
-        @Advice.This AsyncCommand<?, ?, ?> asyncCommand, @Advice.Local("otelScope") Scope scope) {
-      Context context = VirtualField.find(AsyncCommand.class, Context.class).get(asyncCommand);
-      scope = context.makeCurrent();
+    public static Scope onEnter(@Advice.This AsyncCommand<?, ?, ?> asyncCommand) {
+
+      Context context = CONTEXT.get(asyncCommand);
+      return context.makeCurrent();
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
-    public static void onExit(@Advice.Local("otelScope") Scope scope) {
+    public static void onExit(@Advice.Enter Scope scope) {
       scope.close();
     }
   }

instrumentation/lettuce/lettuce-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/lettuce/v5_0/LettuceAsyncCommandsInstrumentation.java

@@ -18,6 +18,7 @@
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import javax.annotation.Nullable;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -41,39 +42,53 @@ public void transform(TypeTransformer transformer) {
   @SuppressWarnings("unused")
   public static class DispatchAdvice {
 
+    public static class AdviceScope {
+      private final Context context;
+      private final Scope scope;
+
+      public AdviceScope(Context context, Scope scope) {
+        this.context = context;
+        this.scope = scope;
+      }
+
+      public void end(
+          @Nullable Throwable throwable,
+          RedisCommand<?, ?, ?> command,
+          AsyncCommand<?, ?, ?> asyncCommand) {
+        scope.close();
+
+        if (throwable != null) {
+          instrumenter().end(context, command, null, throwable);
+          return;
+        }
+
+        // close spans on error or normal completion
+        if (expectsResponse(command)) {
+          asyncCommand.handleAsync(new EndCommandAsyncBiFunction<>(context, command));
+        } else {
+          instrumenter().end(context, command, null, null);
+        }
+      }
+    }
+
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void onEnter(
-        @Advice.Argument(0) RedisCommand<?, ?, ?> command,
-        @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelScope") Scope scope) {
+    public static AdviceScope onEnter(@Advice.Argument(0) RedisCommand<?, ?, ?> command) {
 
       Context parentContext = currentContext();
-      context = instrumenter().start(parentContext, command);
+      Context context = instrumenter().start(parentContext, command);
       // remember the context that called dispatch, it is used in LettuceAsyncCommandInstrumentation
       context = context.with(LettuceSingletons.COMMAND_CONTEXT_KEY, parentContext);
-      scope = context.makeCurrent();
+      return new AdviceScope(context, context.makeCurrent());
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void stopSpan(
         @Advice.Argument(0) RedisCommand<?, ?, ?> command,
-        @Advice.Thrown Throwable throwable,
+        @Advice.Thrown @Nullable Throwable throwable,
         @Advice.Return AsyncCommand<?, ?, ?> asyncCommand,
-        @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelScope") Scope scope) {
-      scope.close();
+        @Advice.Enter AdviceScope adviceScope) {
 
-      if (throwable != null) {
-        instrumenter().end(context, command, null, throwable);
-        return;
-      }
-
-      // close spans on error or normal completion
-      if (expectsResponse(command)) {
-        asyncCommand.handleAsync(new EndCommandAsyncBiFunction<>(context, command));
-      } else {
-        instrumenter().end(context, command, null, null);
-      }
+      adviceScope.end(throwable, command, asyncCommand);
     }
   }
 }

instrumentation/lettuce/lettuce-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/lettuce/v5_0/LettuceClientInstrumentation.java

@@ -21,6 +21,7 @@
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import javax.annotation.Nullable;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -47,29 +48,41 @@ public void transform(TypeTransformer transformer) {
   @SuppressWarnings("unused")
   public static class ConnectAdvice {
 
+    public static class AdviceScope {
+      private final Context context;
+      private final Scope scope;
+
+      public AdviceScope(Context context, Scope scope) {
+        this.context = context;
+        this.scope = scope;
+      }
+
+      public void end(
+          @Nullable Throwable throwable, RedisURI redisUri, ConnectionFuture<?> connectionFuture) {
+
+        scope.close();
+
+        if (throwable != null) {
+          connectInstrumenter().end(context, redisUri, null, throwable);
+          return;
+        }
+        connectionFuture.handleAsync(new EndConnectAsyncBiFunction<>(context, redisUri));
+      }
+    }
+
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void onEnter(
-        @Advice.Argument(1) RedisURI redisUri,
-        @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelScope") Scope scope) {
-      context = connectInstrumenter().start(currentContext(), redisUri);
-      scope = context.makeCurrent();
+    public static AdviceScope onEnter(@Advice.Argument(1) RedisURI redisUri) {
+      Context context = connectInstrumenter().start(currentContext(), redisUri);
+      return new AdviceScope(context, context.makeCurrent());
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void stopSpan(
         @Advice.Argument(1) RedisURI redisUri,
-        @Advice.Thrown Throwable throwable,
+        @Advice.Thrown @Nullable Throwable throwable,
         @Advice.Return ConnectionFuture<?> connectionFuture,
-        @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelScope") Scope scope) {
-      scope.close();
-
-      if (throwable != null) {
-        connectInstrumenter().end(context, redisUri, null, throwable);
-        return;
-      }
-      connectionFuture.handleAsync(new EndConnectAsyncBiFunction<>(context, redisUri));
+        @Advice.Enter AdviceScope adviceScope) {
+      adviceScope.end(throwable, redisUri, connectionFuture);
     }
   }
 }

instrumentation/lettuce/lettuce-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/lettuce/v5_0/LettuceInstrumentationModule.java

@@ -12,12 +12,14 @@
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.ExperimentalInstrumentationModule;
 import io.opentelemetry.javaagent.instrumentation.lettuce.v5_0.rx.LettuceReactiveCommandsInstrumentation;
 import java.util.List;
 import net.bytebuddy.matcher.ElementMatcher;
 
 @AutoService(InstrumentationModule.class)
-public class LettuceInstrumentationModule extends InstrumentationModule {
+public class LettuceInstrumentationModule extends InstrumentationModule
+    implements ExperimentalInstrumentationModule {
   public LettuceInstrumentationModule() {
     super("lettuce", "lettuce-5.0");
   }
@@ -35,4 +37,9 @@ public List<TypeInstrumentation> typeInstrumentations() {
         new LettuceClientInstrumentation(),
         new LettuceReactiveCommandsInstrumentation());
   }
+
+  @Override
+  public boolean isIndyReady() {
+    return true;
+  }
 }

instrumentation/lettuce/lettuce-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/lettuce/v5_0/LettuceSingletons.java

@@ -6,6 +6,7 @@
 package io.opentelemetry.javaagent.instrumentation.lettuce.v5_0;
 
 import io.lettuce.core.RedisURI;
+import io.lettuce.core.protocol.AsyncCommand;
 import io.lettuce.core.protocol.RedisCommand;
 import io.opentelemetry.api.GlobalOpenTelemetry;
 import io.opentelemetry.context.Context;
@@ -17,6 +18,7 @@
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
 import io.opentelemetry.instrumentation.api.instrumenter.SpanKindExtractor;
 import io.opentelemetry.instrumentation.api.semconv.network.ServerAttributesExtractor;
+import io.opentelemetry.instrumentation.api.util.VirtualField;
 import io.opentelemetry.javaagent.bootstrap.internal.AgentCommonConfig;
 import io.opentelemetry.javaagent.bootstrap.internal.AgentInstrumentationConfig;
 
@@ -29,6 +31,9 @@ public final class LettuceSingletons {
   public static final ContextKey<Context> COMMAND_CONTEXT_KEY =
       ContextKey.named("opentelemetry-lettuce-v5_0-context-key");
 
+  public static final VirtualField<AsyncCommand<?, ?, ?>, Context> CONTEXT =
+      VirtualField.find(AsyncCommand.class, Context.class);
+
   static {
     LettuceDbAttributesGetter dbAttributesGetter = new LettuceDbAttributesGetter();
 

instrumentation/lettuce/lettuce-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/lettuce/v5_0/rx/LettuceReactiveCommandsInstrumentation.java

@@ -19,6 +19,7 @@
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import java.util.function.Supplier;
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.asm.Advice.AssignReturned;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 import org.reactivestreams.Subscription;
@@ -61,10 +62,11 @@ public static <K, V, T> RedisCommand<K, V, T> extractCommandName(
 
     // throwables wouldn't matter here, because no spans have been started due to redis command not
     // being run until the user subscribes to the Mono publisher
+    @AssignReturned.ToReturned
     @Advice.OnMethodExit(suppress = Throwable.class)
-    public static <K, V, T> void monitorSpan(
-        @Advice.Enter RedisCommand<K, V, T> command,
-        @Advice.Return(readOnly = false) Mono<T> publisher) {
+    public static <K, V, T> Mono<T> monitorSpan(
+        @Advice.Return Mono<T> originalPublisher, @Advice.Enter RedisCommand<K, V, T> command) {
+      Mono<T> publisher = originalPublisher;
       boolean finishSpanOnClose = !expectsResponse(command);
       LettuceMonoDualConsumer<? super Subscription, T> mdc =
           new LettuceMonoDualConsumer<>(command, finishSpanOnClose);
@@ -73,6 +75,7 @@ public static <K, V, T> void monitorSpan(
       if (!finishSpanOnClose) {
         publisher = publisher.doOnSuccessOrError(mdc);
       }
+      return publisher;
     }
   }
 
@@ -86,10 +89,11 @@ public static <K, V, T> RedisCommand<K, V, T> extractCommandName(
     }
 
     // if there is an exception thrown, then don't make spans
+    @AssignReturned.ToReturned
     @Advice.OnMethodExit(suppress = Throwable.class)
-    public static <K, V, T> void monitorSpan(
-        @Advice.Enter RedisCommand<K, V, T> command,
-        @Advice.Return(readOnly = false) Flux<T> publisher) {
+    public static <K, V, T> Flux<T> monitorSpan(
+        @Advice.Return Flux<T> originalPublisher, @Advice.Enter RedisCommand<K, V, T> command) {
+      Flux<T> publisher = originalPublisher;
 
       boolean expectsResponse = expectsResponse(command);
       LettuceFluxTerminationRunnable handler =
@@ -102,6 +106,7 @@ public static <K, V, T> void monitorSpan(
         publisher = publisher.doOnEach(handler);
         publisher = publisher.doOnCancel(handler);
       }
+      return publisher;
     }
   }
 }