From bfdf4c858a71db0dfd95060bc7a3a890b25f7e93 Mon Sep 17 00:00:00 2001 From: Gary Huang Date: Wed, 30 May 2018 13:49:59 -0400 Subject: [PATCH] lettuce reactive api for mono publishers completed. --- .../RedisReactiveCommandsInstrumentation.java | 44 ++++++++++++++ .../lettuce/RedisAsyncCommandsAdvice.java | 1 + .../lettuce/rx/MonoCreationAdvice.java | 55 ++++++++++++++++++ .../lettuce/rx/MonoDualConsumer.java | 58 +++++++++++++++++++ 4 files changed, 158 insertions(+) create mode 100644 dd-java-agent/instrumentation/lettuce-5/src/main/java/datadog/trace/instrumentation/lettuce/rx/RedisReactiveCommandsInstrumentation.java create mode 100644 dd-java-agent/instrumentation/lettuce-5/src/main/java8/datadog/trace/instrumentation/lettuce/rx/MonoCreationAdvice.java create mode 100644 dd-java-agent/instrumentation/lettuce-5/src/main/java8/datadog/trace/instrumentation/lettuce/rx/MonoDualConsumer.java diff --git a/dd-java-agent/instrumentation/lettuce-5/src/main/java/datadog/trace/instrumentation/lettuce/rx/RedisReactiveCommandsInstrumentation.java b/dd-java-agent/instrumentation/lettuce-5/src/main/java/datadog/trace/instrumentation/lettuce/rx/RedisReactiveCommandsInstrumentation.java new file mode 100644 index 0000000000..c1c6be5060 --- /dev/null +++ b/dd-java-agent/instrumentation/lettuce-5/src/main/java/datadog/trace/instrumentation/lettuce/rx/RedisReactiveCommandsInstrumentation.java @@ -0,0 +1,44 @@ +package datadog.trace.instrumentation.lettuce.rx; + +import static net.bytebuddy.matcher.ElementMatchers.*; + +import com.google.auto.service.AutoService; +import datadog.trace.agent.tooling.DDAdvice; +import datadog.trace.agent.tooling.DDTransformers; +import datadog.trace.agent.tooling.HelperInjector; +import datadog.trace.agent.tooling.Instrumenter; +import net.bytebuddy.agent.builder.AgentBuilder; + +@AutoService(Instrumenter.class) +public class RedisReactiveCommandsInstrumentation extends Instrumenter.Configurable { + + private static final HelperInjector REDIS_ASYNC_HELPERS = + new HelperInjector( + RedisReactiveCommandsInstrumentation.class.getPackage().getName() + ".MonoCreationAdvice", + RedisReactiveCommandsInstrumentation.class.getPackage().getName() + ".MonoDualConsumer"); + + public RedisReactiveCommandsInstrumentation() { + super("redis"); + } + + @Override + protected boolean defaultEnabled() { + return false; + } + + @Override + protected AgentBuilder apply(AgentBuilder agentBuilder) { + return agentBuilder + .type(named("io.lettuce.core.AbstractRedisReactiveCommands")) + .transform(REDIS_ASYNC_HELPERS) + .transform(DDTransformers.defaultTransformers()) + .transform( + DDAdvice.create() + .advice( + isMethod() + .and(named("createMono")) + .and(takesArgument(0, named("java.util.function.Supplier"))), + MonoCreationAdvice.class.getName())) + .asDecorator(); + } +} diff --git a/dd-java-agent/instrumentation/lettuce-5/src/main/java8/datadog/trace/instrumentation/lettuce/RedisAsyncCommandsAdvice.java b/dd-java-agent/instrumentation/lettuce-5/src/main/java8/datadog/trace/instrumentation/lettuce/RedisAsyncCommandsAdvice.java index 56b48be6e0..3643a55bce 100644 --- a/dd-java-agent/instrumentation/lettuce-5/src/main/java8/datadog/trace/instrumentation/lettuce/RedisAsyncCommandsAdvice.java +++ b/dd-java-agent/instrumentation/lettuce-5/src/main/java8/datadog/trace/instrumentation/lettuce/RedisAsyncCommandsAdvice.java @@ -59,6 +59,7 @@ public class RedisAsyncCommandsAdvice { final Span span = scope.span(); Tags.ERROR.set(span, true); span.log(Collections.singletonMap("error.object", throwable)); + span.finish(); scope.close(); return; } diff --git a/dd-java-agent/instrumentation/lettuce-5/src/main/java8/datadog/trace/instrumentation/lettuce/rx/MonoCreationAdvice.java b/dd-java-agent/instrumentation/lettuce-5/src/main/java8/datadog/trace/instrumentation/lettuce/rx/MonoCreationAdvice.java new file mode 100644 index 0000000000..6601ab4131 --- /dev/null +++ b/dd-java-agent/instrumentation/lettuce-5/src/main/java8/datadog/trace/instrumentation/lettuce/rx/MonoCreationAdvice.java @@ -0,0 +1,55 @@ +package datadog.trace.instrumentation.lettuce.rx; + +import io.lettuce.core.protocol.RedisCommand; +import java.util.HashMap; +import java.util.Map; +import java.util.function.Supplier; +import net.bytebuddy.asm.Advice; +import reactor.core.publisher.Mono; + +public class MonoCreationAdvice { + + public static final String SERVICE_NAME = "redis"; + public static final String COMPONENT_NAME = SERVICE_NAME + "-client"; + public static final String MAP_KEY_CMD_NAME = "CMD_NAME"; + public static final String MAP_KEY_CMD_ARGS = "CMD_ARGS"; + + @Advice.OnMethodEnter(suppress = Throwable.class) + public static Map extractCommand( + @Advice.Argument(0) final Supplier supplier) { + RedisCommand command = supplier.get(); + + String commandName = "Redis Command"; + String commandArgs = null; + Map commandMap = new HashMap<>(); + if (command != null) { + // get the arguments passed into the redis command + if (command.getArgs() != null) { + commandArgs = command.getArgs().toCommandString(); + } + // get the redis command name (i.e. GET, SET, HMSET, etc) + if (command.getType() != null) { + commandName = command.getType().name(); + // if it is an AUTH command, then remove the extracted command arguments since it is the password + if ("AUTH".equals(commandName)) { + commandArgs = null; + } + } + } + commandMap.put(MAP_KEY_CMD_NAME, commandName); + commandMap.put(MAP_KEY_CMD_ARGS, commandArgs); + return commandMap; + } + + // 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 + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) + public static void monitorSpan( + @Advice.Enter final Map commandMap, + @Advice.Return(readOnly = false) Mono publisher) { + + MonoDualConsumer mdc = new MonoDualConsumer(commandMap); + publisher = publisher.doOnSubscribe(mdc); + publisher = publisher.doOnSuccessOrError(mdc); + } +} diff --git a/dd-java-agent/instrumentation/lettuce-5/src/main/java8/datadog/trace/instrumentation/lettuce/rx/MonoDualConsumer.java b/dd-java-agent/instrumentation/lettuce-5/src/main/java8/datadog/trace/instrumentation/lettuce/rx/MonoDualConsumer.java new file mode 100644 index 0000000000..e0fc43663e --- /dev/null +++ b/dd-java-agent/instrumentation/lettuce-5/src/main/java8/datadog/trace/instrumentation/lettuce/rx/MonoDualConsumer.java @@ -0,0 +1,58 @@ +package datadog.trace.instrumentation.lettuce.rx; + +import datadog.trace.api.DDTags; +import io.opentracing.Scope; +import io.opentracing.Span; +import io.opentracing.tag.Tags; +import io.opentracing.util.GlobalTracer; +import java.util.Collections; +import java.util.Map; +import java.util.function.BiConsumer; +import java.util.function.Consumer; +import org.slf4j.LoggerFactory; +import reactor.core.publisher.Mono; + +public class MonoDualConsumer + implements Consumer, BiConsumer { + + private Span span = null; + private final Map commandMap; + + public MonoDualConsumer(Map commandMap) { + this.commandMap = commandMap; + } + + @Override + public void accept(T t, Throwable throwable) { + if (this.span != null) { + if (throwable != null) { + Tags.ERROR.set(this.span, true); + this.span.log(Collections.singletonMap("error.object", throwable)); + } + this.span.finish(); + } else { + LoggerFactory.getLogger(Mono.class) + .error( + "Failed to finish this.span, BiConsumer cannot find this.span because " + + "it probably wasn't started."); + } + } + + @Override + public void accept(R r) { + final Scope scope = + GlobalTracer.get().buildSpan(MonoCreationAdvice.SERVICE_NAME + ".query").startActive(false); + this.span = scope.span(); + + Tags.DB_TYPE.set(this.span, MonoCreationAdvice.SERVICE_NAME); + Tags.SPAN_KIND.set(this.span, Tags.SPAN_KIND_CLIENT); + Tags.COMPONENT.set(this.span, MonoCreationAdvice.COMPONENT_NAME); + + this.span.setTag( + DDTags.RESOURCE_NAME, this.commandMap.get(MonoCreationAdvice.MAP_KEY_CMD_NAME)); + this.span.setTag("db.command.args", this.commandMap.get(MonoCreationAdvice.MAP_KEY_CMD_ARGS)); + this.span.setTag(DDTags.SERVICE_NAME, MonoCreationAdvice.SERVICE_NAME); + this.span.setTag(DDTags.SPAN_TYPE, MonoCreationAdvice.SERVICE_NAME); + scope.close(); + } +}