code cleanup

dd-java-agent/instrumentation/lettuce-4.0/src/main/java/datadog/trace/instrumentation/lettuce/LettuceAsyncCommandsInstrumentation.java

@@ -27,10 +27,7 @@ public class LettuceAsyncCommandsInstrumentation extends Instrumenter.Default {
   @Override
   public String[] helperClassNames() {
     return new String[] {
-      packageName + ".LettuceClientDecorator",
+      packageName + ".LettuceClientDecorator", packageName + ".InstrumentationPoints"
-      packageName + ".LettuceAsyncBiFunction",
-      packageName + ".LettuceInstrumentationUtil",
-      packageName + ".InstrumentationPoints"
     };
   }
 

dd-java-agent/instrumentation/lettuce-4.0/src/main/java/datadog/trace/instrumentation/lettuce/LettuceClientInstrumentation.java

@@ -26,10 +26,7 @@ public final class LettuceClientInstrumentation extends Instrumenter.Default {
   @Override
   public String[] helperClassNames() {
     return new String[] {
-      packageName + ".LettuceClientDecorator",
+      packageName + ".LettuceClientDecorator", packageName + ".InstrumentationPoints"
-      packageName + ".LettuceInstrumentationUtil",
-      packageName + ".InstrumentationPoints",
-      packageName + ".LettuceAsyncBiFunction"
     };
   }
 

dd-java-agent/instrumentation/lettuce-4.0/src/main/java8/datadog/trace/instrumentation/lettuce/InstrumentationPoints.java

@@ -1,49 +1,77 @@
 package datadog.trace.instrumentation.lettuce;
 
+import static com.lambdaworks.redis.protocol.CommandKeyword.SEGFAULT;
+import static com.lambdaworks.redis.protocol.CommandType.CLIENT;
+import static com.lambdaworks.redis.protocol.CommandType.CLUSTER;
+import static com.lambdaworks.redis.protocol.CommandType.COMMAND;
+import static com.lambdaworks.redis.protocol.CommandType.CONFIG;
+import static com.lambdaworks.redis.protocol.CommandType.DEBUG;
+import static com.lambdaworks.redis.protocol.CommandType.SCRIPT;
+import static com.lambdaworks.redis.protocol.CommandType.SHUTDOWN;
 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.finishSpanEarly;
 
 import com.lambdaworks.redis.RedisURI;
 import com.lambdaworks.redis.protocol.AsyncCommand;
+import com.lambdaworks.redis.protocol.CommandType;
+import com.lambdaworks.redis.protocol.ProtocolKeyword;
 import com.lambdaworks.redis.protocol.RedisCommand;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 
-@SuppressWarnings("rawtypes")
+import java.util.EnumSet;
+import java.util.Set;
+import java.util.concurrent.CancellationException;
+
 public final class InstrumentationPoints {
 
-  public static AgentScope onEnter(RedisCommand command) {
+  private static final Set<CommandType> NON_INSTRUMENTING_COMMANDS = EnumSet.of(SHUTDOWN, DEBUG);
+
+  private static final Set<CommandType> AGENT_CRASHING_COMMANDS =
+    EnumSet.of(CLIENT, CLUSTER, COMMAND, CONFIG, DEBUG, SCRIPT);
+
+  public static final String AGENT_CRASHING_COMMAND_PREFIX = "COMMAND-NAME:";
+
+  public static AgentScope beforeCommand(RedisCommand<?, ?, ?> command) {
     AgentSpan span = startSpan("redis.query");
     DECORATE.afterStart(span);
     DECORATE.onCommand(span, command);
     return activateSpan(span, finishSpanEarly(command));
   }
 
-  public static void onReturn(RedisCommand command,
+  public static void afterCommand(RedisCommand<?, ?, ?> command,
-                              AgentScope scope,
+                                  AgentScope scope,
-                              Throwable throwable,
+                                  Throwable throwable,
-                              AsyncCommand<?, ?, ?> asyncCommand) {
+                                  AsyncCommand<?, ?, ?> asyncCommand) {
       AgentSpan span = scope.span();
       if (throwable != null) {
         DECORATE.onError(span, throwable);
         DECORATE.beforeFinish(span);
         span.finish();
       } else if (!finishSpanEarly(command)) {
-        asyncCommand.handleAsync(new LettuceAsyncBiFunction<>(span));
+        asyncCommand.handleAsync((value, ex) -> {
+          if (ex instanceof CancellationException) {
+            span.setTag("db.command.cancelled", true);
+          } else {
+            DECORATE.onError(span, ex);
+          }
+          DECORATE.beforeFinish(span);
+          span.finish();
+          return null;
+        });
       }
       scope.close();
   }
 
-  public static AgentScope onEnter(RedisURI redisURI) {
+  public static AgentScope beforeConnect(RedisURI redisURI) {
     AgentSpan span = startSpan("redis.query");
     DECORATE.afterStart(span);
     DECORATE.onConnection(span, redisURI);
     return activateSpan(span, false);
   }
 
-  public static void onReturn(AgentScope scope, Throwable throwable) {
+  public static void afterConnect(AgentScope scope, Throwable throwable) {
     AgentSpan span = scope.span();
     if (throwable != null) {
       DECORATE.onError(span, throwable);
@@ -52,4 +80,49 @@ public final class InstrumentationPoints {
     span.finish();
     scope.close();
   }
+
+  /**
+   * Determines whether a redis command should finish its relevant span early (as soon as tags are
+   * added and the command is executed) because these commands have no return values/call backs, so
+   * 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)
+   */
+  public static boolean finishSpanEarly(RedisCommand<?, ?, ?> command) {
+    ProtocolKeyword keyword = command.getType();
+    return isNonInstrumentingCommand(keyword) || isNonInstrumentingKeyword(keyword);
+  }
+
+  private static boolean isNonInstrumentingCommand(ProtocolKeyword keyword) {
+    return keyword instanceof CommandType && NON_INSTRUMENTING_COMMANDS.contains(keyword);
+  }
+
+  private static boolean isNonInstrumentingKeyword(ProtocolKeyword keyword) {
+    return keyword == SEGFAULT;
+  }
+
+  // Workaround to keep trace agent from crashing
+  // Currently the commands in AGENT_CRASHING_COMMANDS_WORDS will crash the trace agent and
+  // traces with these commands as the resource name will not be processed by the trace agent
+  // https://github.com/DataDog/datadog-trace-agent/blob/master/quantizer/redis.go#L18 has
+  // list of commands that will currently fail at the trace agent level.
+
+  /**
+   * Workaround to keep trace agent from crashing Currently the commands in
+   * AGENT_CRASHING_COMMANDS_WORDS will crash the trace agent and traces with these commands as the
+   * resource name will not be processed by the trace agent
+   * https://github.com/DataDog/datadog-trace-agent/blob/master/quantizer/redis.go#L18 has list of
+   * commands that will currently fail at the trace agent level.
+   *
+   * @param keyword the actual redis command
+   * @return the redis command with a prefix if it is a command that will crash the trace agent,
+   *     otherwise, the original command is returned.
+   */
+  public static String getCommandResourceName(ProtocolKeyword keyword) {
+    if (keyword instanceof CommandType && AGENT_CRASHING_COMMANDS.contains(keyword)) {
+      return AGENT_CRASHING_COMMAND_PREFIX + keyword.name();
+    }
+    return keyword.name();
+  }
 }

dd-java-agent/instrumentation/lettuce-4.0/src/main/java8/datadog/trace/instrumentation/lettuce/LettuceAsyncBiFunction.java

@@ -1,39 +0,0 @@
-package datadog.trace.instrumentation.lettuce;
-
-import static datadog.trace.instrumentation.lettuce.LettuceClientDecorator.DECORATE;
-
-import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
-
-import java.util.concurrent.CancellationException;
-import java.util.function.BiFunction;
-
-/**
- * Callback class to close the span on an error or a success in the RedisFuture returned by the
- * lettuce async API
- *
- * @param <T> the normal completion result
- * @param <U> the error
- * @param <R> the return type, should be null since nothing else should happen from tracing
- *     standpoint after the span is closed
- */
-public class LettuceAsyncBiFunction<T extends Object, U extends Throwable, R extends Object>
-    implements BiFunction<T, Throwable, R> {
-
-  private final AgentSpan span;
-
-  public LettuceAsyncBiFunction(final AgentSpan span) {
-    this.span = span;
-  }
-
-  @Override
-  public R apply(final T t, final Throwable throwable) {
-    if (throwable instanceof CancellationException) {
-      span.setTag("db.command.cancelled", true);
-    } else {
-      DECORATE.onError(span, throwable);
-    }
-    DECORATE.beforeFinish(span);
-    span.finish();
-    return null;
-  }
-}

dd-java-agent/instrumentation/lettuce-4.0/src/main/java8/datadog/trace/instrumentation/lettuce/LettuceAsyncCommandsAdvice.java

@@ -1,7 +1,5 @@
 package datadog.trace.instrumentation.lettuce;
 
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
-
 import com.lambdaworks.redis.protocol.AsyncCommand;
 import com.lambdaworks.redis.protocol.RedisCommand;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
@@ -10,17 +8,17 @@ import net.bytebuddy.asm.Advice;
 public class LettuceAsyncCommandsAdvice {
 
   @Advice.OnMethodEnter(suppress = Throwable.class)
-  public static AgentScope onEnter(@Advice.Argument(0) final RedisCommand command) {
+  public static AgentScope onEnter(@Advice.Argument(0) final RedisCommand<?, ?, ?> command) {
-    return InstrumentationPoints.onEnter(command);
+    return InstrumentationPoints.beforeCommand(command);
   }
 
   @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
-  public static void stopSpan(
+  public static void onExit(
-      @Advice.Argument(0) RedisCommand command,
+      @Advice.Argument(0) final RedisCommand<?, ?, ?> command,
-      @Advice.Enter AgentScope scope,
+      @Advice.Enter final AgentScope scope,
-      @Advice.Thrown Throwable throwable,
+      @Advice.Thrown final Throwable throwable,
-      @Advice.Return AsyncCommand<?, ?, ?> asyncCommand) {
+      @Advice.Return final AsyncCommand<?, ?, ?> asyncCommand) {
-    InstrumentationPoints.onReturn(command, scope, throwable, asyncCommand);
+    InstrumentationPoints.afterCommand(command, scope, throwable, asyncCommand);
   }
 
 }

dd-java-agent/instrumentation/lettuce-4.0/src/main/java8/datadog/trace/instrumentation/lettuce/LettuceClientDecorator.java

@@ -1,6 +1,6 @@
 package datadog.trace.instrumentation.lettuce;
 
-import static datadog.trace.instrumentation.lettuce.LettuceInstrumentationUtil.getCommandResourceName;
+import static datadog.trace.instrumentation.lettuce.InstrumentationPoints.getCommandResourceName;
 
 import com.lambdaworks.redis.RedisURI;
 import com.lambdaworks.redis.protocol.RedisCommand;
@@ -40,37 +40,38 @@ public class LettuceClientDecorator extends DatabaseClientDecorator<RedisURI> {
   }
 
   @Override
-  protected String dbUser(final RedisURI connection) {
+  protected String dbUser(RedisURI connection) {
     return null;
   }
 
   @Override
-  protected String dbInstance(final RedisURI connection) {
+  protected String dbInstance(RedisURI connection) {
     return null;
   }
 
   @Override
-  public AgentSpan onConnection(final AgentSpan span, final RedisURI connection) {
+  public AgentSpan onConnection(AgentSpan span, RedisURI connection) {
     if (connection != null) {
       span.setTag(Tags.PEER_HOSTNAME, connection.getHost());
       span.setTag(Tags.PEER_PORT, connection.getPort());
       span.setTag("db.redis.dbIndex", connection.getDatabase());
-      span.setTag(
+      span.setTag(DDTags.RESOURCE_NAME, resourceName(connection));
-          DDTags.RESOURCE_NAME,
-          "CONNECT:"
-              + connection.getHost()
-              + ":"
-              + connection.getPort()
-              + "/"
-              + connection.getDatabase());
     }
     return super.onConnection(span, connection);
   }
 
-  @SuppressWarnings("rawtypes")
+  public AgentSpan onCommand(AgentSpan span, RedisCommand<?, ?, ?> command) {
-  public AgentSpan onCommand(final AgentSpan span, final RedisCommand command) {
+    span.setTag(DDTags.RESOURCE_NAME,
-      span.setTag(DDTags.RESOURCE_NAME,
         null == command ? "Redis Command" : getCommandResourceName(command.getType()));
     return span;
   }
+
+  private static String resourceName(RedisURI connection) {
+    return "CONNECT:"
+      + connection.getHost()
+      + ":"
+      + connection.getPort()
+      + "/"
+      + connection.getDatabase();
+  }
 }

dd-java-agent/instrumentation/lettuce-4.0/src/main/java8/datadog/trace/instrumentation/lettuce/LettuceInstrumentationUtil.java

@@ -1,74 +0,0 @@
-package datadog.trace.instrumentation.lettuce;
-
-import static com.lambdaworks.redis.protocol.CommandKeyword.SEGFAULT;
-import static com.lambdaworks.redis.protocol.CommandType.CLIENT;
-import static com.lambdaworks.redis.protocol.CommandType.CLUSTER;
-import static com.lambdaworks.redis.protocol.CommandType.COMMAND;
-import static com.lambdaworks.redis.protocol.CommandType.CONFIG;
-import static com.lambdaworks.redis.protocol.CommandType.DEBUG;
-import static com.lambdaworks.redis.protocol.CommandType.SCRIPT;
-import static com.lambdaworks.redis.protocol.CommandType.SHUTDOWN;
-
-import com.lambdaworks.redis.protocol.CommandKeyword;
-import com.lambdaworks.redis.protocol.CommandType;
-import com.lambdaworks.redis.protocol.ProtocolKeyword;
-import com.lambdaworks.redis.protocol.RedisCommand;
-
-import java.util.EnumSet;
-import java.util.Set;
-
-public class LettuceInstrumentationUtil {
-
-  public static final String AGENT_CRASHING_COMMAND_PREFIX = "COMMAND-NAME:";
-
-  public static final EnumSet<CommandType> NON_INSTRUMENTING_COMMANDS = EnumSet.of(SHUTDOWN, DEBUG);
-
-  public static final EnumSet<CommandKeyword> NON_INSTRUMENTING_KEYWORDS = EnumSet.of(SEGFAULT);
-
-  public static final Set<CommandType> AGENT_CRASHING_COMMANDS = EnumSet.of(CLIENT, CLUSTER, COMMAND, CONFIG, DEBUG, SCRIPT);
-
-  /**
-   * Determines whether a redis command should finish its relevant span early (as soon as tags are
-   * added and the command is executed) because these commands have no return values/call backs, so
-   * 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)
-   */
-  public static boolean finishSpanEarly(final RedisCommand<?, ?, ?> command) {
-    ProtocolKeyword keyword = command.getType();
-    return isNonInstrumentingCommand(keyword) || isNonInstrumentingKeyword(keyword);
-  }
-
-  private static boolean isNonInstrumentingCommand(ProtocolKeyword keyword) {
-    return keyword instanceof CommandType && NON_INSTRUMENTING_COMMANDS.contains(keyword);
-  }
-
-  private static boolean isNonInstrumentingKeyword(ProtocolKeyword keyword) {
-    return keyword instanceof CommandKeyword && NON_INSTRUMENTING_KEYWORDS.contains(keyword);
-  }
-
-  // Workaround to keep trace agent from crashing
-  // Currently the commands in AGENT_CRASHING_COMMANDS_WORDS will crash the trace agent and
-  // traces with these commands as the resource name will not be processed by the trace agent
-  // https://github.com/DataDog/datadog-trace-agent/blob/master/quantizer/redis.go#L18 has
-  // list of commands that will currently fail at the trace agent level.
-
-  /**
-   * Workaround to keep trace agent from crashing Currently the commands in
-   * AGENT_CRASHING_COMMANDS_WORDS will crash the trace agent and traces with these commands as the
-   * resource name will not be processed by the trace agent
-   * https://github.com/DataDog/datadog-trace-agent/blob/master/quantizer/redis.go#L18 has list of
-   * commands that will currently fail at the trace agent level.
-   *
-   * @param keyword the actual redis command
-   * @return the redis command with a prefix if it is a command that will crash the trace agent,
-   *     otherwise, the original command is returned.
-   */
-  public static String getCommandResourceName(ProtocolKeyword keyword) {
-    if (keyword instanceof CommandType && AGENT_CRASHING_COMMANDS.contains(keyword)) {
-      return AGENT_CRASHING_COMMAND_PREFIX + keyword.name();
-    }
-    return keyword.name();
-  }
-}

dd-java-agent/instrumentation/lettuce-4.0/src/main/java8/datadog/trace/instrumentation/lettuce/RedisConnectionAdvice.java

@@ -7,12 +7,13 @@ import net.bytebuddy.asm.Advice;
 public class RedisConnectionAdvice {
 
   @Advice.OnMethodEnter(suppress = Throwable.class)
-  public static AgentScope onEnter(@Advice.Argument(1) RedisURI redisURI) {
+  public static AgentScope onEnter(@Advice.Argument(1) final RedisURI redisURI) {
-    return InstrumentationPoints.onEnter(redisURI);
+    return InstrumentationPoints.beforeConnect(redisURI);
   }
 
   @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
-  public static void onReturn(@Advice.Enter AgentScope scope, @Advice.Thrown Throwable throwable) {
+  public static void onExit(@Advice.Enter final AgentScope scope,
-    InstrumentationPoints.onReturn(scope, throwable);
+                            @Advice.Thrown final Throwable throwable) {
+    InstrumentationPoints.afterConnect(scope, throwable);
   }
 }

dd-java-agent/instrumentation/lettuce-4.0/src/test/groovy/LettuceAsyncClientTest.groovy

@@ -25,7 +25,7 @@ import java.util.function.Function
 
 import com.lambdaworks.redis.RedisConnectionException
 
-import static datadog.trace.instrumentation.lettuce.LettuceInstrumentationUtil.AGENT_CRASHING_COMMAND_PREFIX
+import static datadog.trace.instrumentation.lettuce.InstrumentationPoints.AGENT_CRASHING_COMMAND_PREFIX
 
 class LettuceAsyncClientTest extends AgentTestRunner {
   public static final String HOST = "127.0.0.1"

dd-java-agent/instrumentation/lettuce-4.0/src/test/groovy/LettuceSyncClientTest.groovy

@@ -10,7 +10,7 @@ import datadog.trace.bootstrap.instrumentation.api.Tags
 import redis.embedded.RedisServer
 import spock.lang.Shared
 
-import static datadog.trace.instrumentation.lettuce.LettuceInstrumentationUtil.AGENT_CRASHING_COMMAND_PREFIX
+import static datadog.trace.instrumentation.lettuce.InstrumentationPoints.AGENT_CRASHING_COMMAND_PREFIX
 
 class LettuceSyncClientTest extends AgentTestRunner {
   public static final String HOST = "127.0.0.1"