Merge pull request #1445 from DataDog/tyler/migrate-lettuce

DataDog · May 12, 2020 · 9fc7d79 · 9fc7d79
2 parents 19daca5 + be766f2
commit 9fc7d79
Show file tree

Hide file tree

Showing 6 changed files with 25 additions and 16 deletions.
diff --git a/...lettuce-4/src/main/java8/datadog/trace/instrumentation/lettuce/InstrumentationPoints.java b/...lettuce-4/src/main/java8/datadog/trace/instrumentation/lettuce/InstrumentationPoints.java
@@ -36,7 +36,7 @@ public static AgentScope beforeCommand(final RedisCommand<?, ?, ?> command) {
     final AgentSpan span = startSpan("redis.query");
     DECORATE.afterStart(span);
     DECORATE.onCommand(span, command);
-    return activateSpan(span, finishSpanEarly(command));
+    return activateSpan(span);
   }
 
   public static void afterCommand(
@@ -49,7 +49,7 @@ public static void afterCommand(
       DECORATE.onError(span, throwable);
       DECORATE.beforeFinish(span);
       span.finish();
-    } else if (!finishSpanEarly(command)) {
+    } else if (expectsResponse(command)) {
       asyncCommand.handleAsync(
           (value, ex) -> {
             if (ex instanceof CancellationException) {
@@ -61,6 +61,10 @@ public static void afterCommand(
             span.finish();
             return null;
           });
+    } else {
+      // No response is expected, so we must finish the span now.
+      DECORATE.beforeFinish(span);
+      span.finish();
     }
     scope.close();
     // span may be finished by handleAsync call above.
@@ -89,11 +93,11 @@ public static void afterConnect(final AgentScope scope, final Throwable throwabl
    * we must close the span early in order to provide info for the users
    *
    * @param command
-   * @return true if finish the span early (the command will not have a return value)
+   * @return false if the span should finish early (the command will not have a return value)
    */
-  public static boolean finishSpanEarly(final RedisCommand<?, ?, ?> command) {
+  public static boolean expectsResponse(final RedisCommand<?, ?, ?> command) {
     final ProtocolKeyword keyword = command.getType();
-    return isNonInstrumentingCommand(keyword) || isNonInstrumentingKeyword(keyword);
+    return !(isNonInstrumentingCommand(keyword) || isNonInstrumentingKeyword(keyword));
   }
 
   private static boolean isNonInstrumentingCommand(final ProtocolKeyword keyword) {

diff --git a/...ce-5/src/main/java8/datadog/trace/instrumentation/lettuce/LettuceAsyncCommandsAdvice.java b/...ce-5/src/main/java8/datadog/trace/instrumentation/lettuce/LettuceAsyncCommandsAdvice.java
@@ -3,7 +3,7 @@
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
 import static datadog.trace.instrumentation.lettuce.LettuceClientDecorator.DECORATE;
-import static datadog.trace.instrumentation.lettuce.LettuceInstrumentationUtil.doFinishSpanEarly;
+import static datadog.trace.instrumentation.lettuce.LettuceInstrumentationUtil.expectsResponse;
 
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -20,7 +20,7 @@ public static AgentScope onEnter(@Advice.Argument(0) final RedisCommand command)
     DECORATE.afterStart(span);
     DECORATE.onCommand(span, command);
 
-    return activateSpan(span, doFinishSpanEarly(command));
+    return activateSpan(span);
   }
 
   @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
@@ -40,10 +40,13 @@ public static void stopSpan(
     }
 
     // close spans on error or normal completion
-    if (!doFinishSpanEarly(command)) {
+    if (expectsResponse(command)) {
       asyncCommand.handleAsync(new LettuceAsyncBiFunction<>(span));
+    } else {
+      DECORATE.beforeFinish(span);
+      span.finish();
     }
     scope.close();
-    // span finished by LettuceAsyncBiFunction
+    // span may be finished by LettuceAsyncBiFunction
   }
 }
diff --git a/...ce-5/src/main/java8/datadog/trace/instrumentation/lettuce/LettuceInstrumentationUtil.java b/...ce-5/src/main/java8/datadog/trace/instrumentation/lettuce/LettuceInstrumentationUtil.java
@@ -27,11 +27,11 @@ public class LettuceInstrumentationUtil {
    * we must close the span early in order to provide info for the users
    *
    * @param command
-   * @return true if finish the span early (the command will not have a return value)
+   * @return false if the span should finish early (the command will not have a return value)
    */
-  public static boolean doFinishSpanEarly(final RedisCommand command) {
+  public static boolean expectsResponse(final RedisCommand command) {
     final String commandName = LettuceInstrumentationUtil.getCommandName(command);
-    return nonInstrumentingCommands.contains(commandName);
+    return !nonInstrumentingCommands.contains(commandName);
   }
 
   // Workaround to keep trace agent from crashing

diff --git a/...-5/src/main/java8/datadog/trace/instrumentation/lettuce/rx/LettuceFluxCreationAdvice.java b/...-5/src/main/java8/datadog/trace/instrumentation/lettuce/rx/LettuceFluxCreationAdvice.java
@@ -1,6 +1,6 @@
 package datadog.trace.instrumentation.lettuce.rx;
 
-import static datadog.trace.instrumentation.lettuce.LettuceInstrumentationUtil.doFinishSpanEarly;
+import static datadog.trace.instrumentation.lettuce.LettuceInstrumentationUtil.expectsResponse;
 
 import io.lettuce.core.protocol.RedisCommand;
 import java.util.function.Supplier;
@@ -21,7 +21,7 @@ public static void monitorSpan(
       @Advice.Enter final RedisCommand command,
       @Advice.Return(readOnly = false) Flux<?> publisher) {
 
-    final boolean finishSpanOnClose = doFinishSpanEarly(command);
+    final boolean finishSpanOnClose = !expectsResponse(command);
     final LettuceFluxTerminationRunnable handler =
         new LettuceFluxTerminationRunnable(command, finishSpanOnClose);
     publisher = publisher.doOnSubscribe(handler.getOnSubscribeConsumer());

diff --git a/...-5/src/main/java8/datadog/trace/instrumentation/lettuce/rx/LettuceMonoCreationAdvice.java b/...-5/src/main/java8/datadog/trace/instrumentation/lettuce/rx/LettuceMonoCreationAdvice.java
@@ -1,6 +1,7 @@
 package datadog.trace.instrumentation.lettuce.rx;
 
-import datadog.trace.instrumentation.lettuce.LettuceInstrumentationUtil;
+import static datadog.trace.instrumentation.lettuce.LettuceInstrumentationUtil.expectsResponse;
+
 import io.lettuce.core.protocol.RedisCommand;
 import java.util.function.Supplier;
 import net.bytebuddy.asm.Advice;
@@ -20,7 +21,7 @@ public static RedisCommand extractCommandName(
   public static void monitorSpan(
       @Advice.Enter final RedisCommand command,
       @Advice.Return(readOnly = false) Mono<?> publisher) {
-    final boolean finishSpanOnClose = LettuceInstrumentationUtil.doFinishSpanEarly(command);
+    final boolean finishSpanOnClose = !expectsResponse(command);
     final LettuceMonoDualConsumer mdc = new LettuceMonoDualConsumer(command, finishSpanOnClose);
     publisher = publisher.doOnSubscribe(mdc);
     // register the call back to close the span only if necessary

diff --git a/...ce-5/src/main/java8/datadog/trace/instrumentation/lettuce/rx/LettuceMonoDualConsumer.java b/...ce-5/src/main/java8/datadog/trace/instrumentation/lettuce/rx/LettuceMonoDualConsumer.java
@@ -28,6 +28,7 @@ public void accept(final R r) {
     DECORATE.afterStart(span);
     DECORATE.onCommand(span, command);
     if (finishSpanOnClose) {
+      DECORATE.beforeFinish(span);
       span.finish();
     }
   }