diff --git a/apm-sniffer/optional-plugins/lettuce-5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/lettuce/v5/AsyncCommandMethodInterceptor.java b/apm-sniffer/optional-plugins/lettuce-5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/lettuce/v5/AsyncCommandMethodInterceptor.java index b713e64625..1684318b2c 100644 --- a/apm-sniffer/optional-plugins/lettuce-5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/lettuce/v5/AsyncCommandMethodInterceptor.java +++ b/apm-sniffer/optional-plugins/lettuce-5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/lettuce/v5/AsyncCommandMethodInterceptor.java @@ -27,6 +27,7 @@ import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.MethodInt import org.apache.skywalking.apm.network.trace.component.ComponentsDefine; import java.lang.reflect.Method; +import java.util.function.BiConsumer; import java.util.function.Consumer; /** @@ -42,7 +43,11 @@ public class AsyncCommandMethodInterceptor implements InstanceMethodsAroundInter String operationName = "Lettuce/" + asyncCommand.getType().name(); AbstractSpan span = ContextManager.createLocalSpan(operationName + "/onComplete"); span.setComponent(ComponentsDefine.LETTUCE); - allArguments[0] = new SWConsumer((Consumer) allArguments[0], ContextManager.capture(), operationName); + if (allArguments[0] instanceof Consumer) { + allArguments[0] = new SWConsumer((Consumer) allArguments[0], ContextManager.capture(), operationName); + } else { + allArguments[0] = new SWBiConsumer((BiConsumer) allArguments[0], ContextManager.capture(), operationName); + } } @Override diff --git a/apm-sniffer/optional-plugins/lettuce-5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/lettuce/v5/SWBiConsumer.java b/apm-sniffer/optional-plugins/lettuce-5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/lettuce/v5/SWBiConsumer.java new file mode 100644 index 0000000000..216cc415f9 --- /dev/null +++ b/apm-sniffer/optional-plugins/lettuce-5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/lettuce/v5/SWBiConsumer.java @@ -0,0 +1,56 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + +package org.apache.skywalking.apm.plugin.lettuce.v5; + +import org.apache.skywalking.apm.agent.core.context.ContextManager; +import org.apache.skywalking.apm.agent.core.context.ContextSnapshot; +import org.apache.skywalking.apm.agent.core.context.trace.AbstractSpan; +import org.apache.skywalking.apm.network.trace.component.ComponentsDefine; + +import java.util.function.BiConsumer; + +/** + * @author zhaoyuguang + */ +public class SWBiConsumer implements BiConsumer { + + private BiConsumer biConsumer; + private ContextSnapshot snapshot; + private String operationName; + + SWBiConsumer(BiConsumer biConsumer, ContextSnapshot snapshot, String operationName) { + this.biConsumer = biConsumer; + this.snapshot = snapshot; + this.operationName = operationName; + } + + @Override + public void accept(T t, U u) { + AbstractSpan span = ContextManager.createLocalSpan(operationName + "/accept"); + span.setComponent(ComponentsDefine.LETTUCE); + try { + ContextManager.continued(snapshot); + biConsumer.accept(t, u); + } catch (Throwable th) { + ContextManager.activeSpan().errorOccurred().log(th); + } finally { + ContextManager.stopSpan(); + } + } +} diff --git a/apm-sniffer/optional-plugins/lettuce-5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/lettuce/v5/define/AsyncCommandInstrumentation.java b/apm-sniffer/optional-plugins/lettuce-5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/lettuce/v5/define/AsyncCommandInstrumentation.java index 8f744504a6..52409ca801 100644 --- a/apm-sniffer/optional-plugins/lettuce-5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/lettuce/v5/define/AsyncCommandInstrumentation.java +++ b/apm-sniffer/optional-plugins/lettuce-5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/lettuce/v5/define/AsyncCommandInstrumentation.java @@ -27,6 +27,7 @@ import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.ClassInst import org.apache.skywalking.apm.agent.core.plugin.match.ClassMatch; import static net.bytebuddy.matcher.ElementMatchers.named; +import static org.apache.skywalking.apm.agent.core.plugin.bytebuddy.ArgumentTypeNameMatch.takesArgumentWithType; import static org.apache.skywalking.apm.agent.core.plugin.match.NameMatch.byName; /** @@ -49,7 +50,8 @@ public class AsyncCommandInstrumentation extends ClassInstanceMethodsEnhancePlug new InstanceMethodsInterceptPoint() { @Override public ElementMatcher getMethodsMatcher() { - return named("onComplete"); + return (named("onComplete").and(takesArgumentWithType(0,"java.util.function.Consumer"))) + .or(named("onComplete").and(takesArgumentWithType(0,"java.util.function.BiConsumer"))); } @Override