diff --git a/apm-sniffer/bootstrap-plugins/jdk-threading-plugin/src/main/java/org/apache/skywalking/apm/plugin/jdk/threading/ThreadPoolExecutorInvokeInterceptor.java b/apm-sniffer/bootstrap-plugins/jdk-threading-plugin/src/main/java/org/apache/skywalking/apm/plugin/jdk/threading/ThreadPoolExecutorInvokeInterceptor.java new file mode 100644 index 0000000000..b666cf04f9 --- /dev/null +++ b/apm-sniffer/bootstrap-plugins/jdk-threading-plugin/src/main/java/org/apache/skywalking/apm/plugin/jdk/threading/ThreadPoolExecutorInvokeInterceptor.java @@ -0,0 +1,112 @@ +/* + * 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.jdk.threading; + +import java.lang.reflect.Method; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.concurrent.Callable; + +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.plugin.interceptor.enhance.EnhancedInstance; +import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.InstanceMethodsAroundInterceptor; +import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.MethodInterceptResult; + +/** + * Interceptor for ThreadPoolExecutor.invokeAll() and invokeAny() methods. + * + *

These methods wrap the original Callable tasks in RunnableFuture, which causes + * the tracing context to be lost. This interceptor captures the active context before + * the method call and wraps each Callable to restore the context when executed.

+ */ +public class ThreadPoolExecutorInvokeInterceptor implements InstanceMethodsAroundInterceptor { + + @Override + public void beforeMethod(final EnhancedInstance objInst, final Method method, final Object[] allArguments, + final Class[] argumentsTypes, final MethodInterceptResult result) { + + if (allArguments.length == 0) { + return; + } + + // Get the Callable collection (first argument) + Object arg = allArguments[0]; + if (!(arg instanceof Collection)) { + return; + } + + Collection callables = (Collection) arg; + if (callables.isEmpty()) { + return; + } + + // Capture the current tracing context + ContextSnapshot snapshot = ContextManager.capture(); + + // Wrap each Callable to restore the tracing context + List> wrappedCallables = new ArrayList<>(callables.size()); + for (Object callable : callables) { + if (callable instanceof Callable) { + wrappedCallables.add(new ContextRestoringCallable<>((Callable) callable, snapshot)); + } else { + wrappedCallables.add((Callable) callable); + } + } + + // Replace the original argument with wrapped callables + allArguments[0] = wrappedCallables; + } + + @Override + public Object afterMethod(final EnhancedInstance objInst, final Method method, final Object[] allArguments, + final Class[] argumentsTypes, final Object ret) { + return ret; + } + + @Override + public void handleMethodException(final EnhancedInstance objInst, final Method method, final Object[] allArguments, + final Class[] argumentsTypes, final Throwable t) { + // No special handling needed + } + + /** + * A Callable wrapper that restores the tracing context before execution. + */ + private static class ContextRestoringCallable implements Callable { + private final Callable delegate; + private final ContextSnapshot snapshot; + + ContextRestoringCallable(Callable delegate, ContextSnapshot snapshot) { + this.delegate = delegate; + this.snapshot = snapshot; + } + + @Override + public T call() throws Exception { + ContextManager.continued(snapshot); + try { + return delegate.call(); + } finally { + ContextManager.stopSpan(); + } + } + } +} diff --git a/apm-sniffer/bootstrap-plugins/jdk-threading-plugin/src/main/java/org/apache/skywalking/apm/plugin/jdk/threading/define/ThreadPoolExecutorInstrumentation.java b/apm-sniffer/bootstrap-plugins/jdk-threading-plugin/src/main/java/org/apache/skywalking/apm/plugin/jdk/threading/define/ThreadPoolExecutorInstrumentation.java new file mode 100644 index 0000000000..3828ce652d --- /dev/null +++ b/apm-sniffer/bootstrap-plugins/jdk-threading-plugin/src/main/java/org/apache/skywalking/apm/plugin/jdk/threading/define/ThreadPoolExecutorInstrumentation.java @@ -0,0 +1,141 @@ +/* + * 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.jdk.threading.define; + +import net.bytebuddy.description.method.MethodDescription; +import net.bytebuddy.matcher.ElementMatcher; +import org.apache.skywalking.apm.agent.core.plugin.interceptor.ConstructorInterceptPoint; +import org.apache.skywalking.apm.agent.core.plugin.interceptor.InstanceMethodsInterceptPoint; +import org.apache.skywalking.apm.agent.core.plugin.interceptor.StaticMethodsInterceptPoint; +import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.ClassEnhancePluginDefine; +import org.apache.skywalking.apm.agent.core.plugin.match.ClassMatch; + +import static net.bytebuddy.matcher.ElementMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; +import static org.apache.skywalking.apm.agent.core.plugin.match.NameMatch.byName; + +/** + * Instrumentation for ThreadPoolExecutor to intercept invokeAll and invokeAny methods. + * + *

These methods wrap original Callable tasks in RunnableFuture, causing tracing context + * to be lost. This interceptor captures the active context and wraps each Callable to restore + * the context when executed.

+ */ +public class ThreadPoolExecutorInstrumentation extends ClassEnhancePluginDefine { + + private static final String THREAD_POOL_EXECUTOR_CLASS = "java.util.concurrent.ThreadPoolExecutor"; + + private static final String INVOKE_ALL_METHOD = "invokeAll"; + private static final String INVOKE_ANY_METHOD = "invokeAny"; + private static final String INTERCEPTOR_CLASS = "org.apache.skywalking.apm.plugin.jdk.threading.ThreadPoolExecutorInvokeInterceptor"; + + @Override + protected ClassMatch enhanceClass() { + return byName(THREAD_POOL_EXECUTOR_CLASS); + } + + @Override + public ConstructorInterceptPoint[] getConstructorsInterceptPoints() { + return new ConstructorInterceptPoint[0]; + } + + @Override + public InstanceMethodsInterceptPoint[] getInstanceMethodsInterceptPoints() { + return new InstanceMethodsInterceptPoint[] { + // invokeAll(Collection> tasks) + new InstanceMethodsInterceptPoint() { + @Override + public ElementMatcher getMethodsMatcher() { + return named(INVOKE_ALL_METHOD).and(takesArguments(1)); + } + + @Override + public String getMethodsInterceptor() { + return INTERCEPTOR_CLASS; + } + + @Override + public boolean isOverrideArgs() { + return true; + } + }, + // invokeAll(Collection> tasks, long timeout, TimeUnit unit) + new InstanceMethodsInterceptPoint() { + @Override + public ElementMatcher getMethodsMatcher() { + return named(INVOKE_ALL_METHOD).and(takesArguments(3)); + } + + @Override + public String getMethodsInterceptor() { + return INTERCEPTOR_CLASS; + } + + @Override + public boolean isOverrideArgs() { + return true; + } + }, + // invokeAny(Collection> tasks) + new InstanceMethodsInterceptPoint() { + @Override + public ElementMatcher getMethodsMatcher() { + return named(INVOKE_ANY_METHOD).and(takesArguments(1)); + } + + @Override + public String getMethodsInterceptor() { + return INTERCEPTOR_CLASS; + } + + @Override + public boolean isOverrideArgs() { + return true; + } + }, + // invokeAny(Collection> tasks, long timeout, TimeUnit unit) + new InstanceMethodsInterceptPoint() { + @Override + public ElementMatcher getMethodsMatcher() { + return named(INVOKE_ANY_METHOD).and(takesArguments(3)); + } + + @Override + public String getMethodsInterceptor() { + return INTERCEPTOR_CLASS; + } + + @Override + public boolean isOverrideArgs() { + return true; + } + } + }; + } + + @Override + public StaticMethodsInterceptPoint[] getStaticMethodsInterceptPoints() { + return new StaticMethodsInterceptPoint[0]; + } + + @Override + public boolean isBootstrapInstrumentation() { + return true; + } +} diff --git a/apm-sniffer/bootstrap-plugins/jdk-threading-plugin/src/main/resources/skywalking-plugin.def b/apm-sniffer/bootstrap-plugins/jdk-threading-plugin/src/main/resources/skywalking-plugin.def index 06aaee35e1..7c8aa03838 100644 --- a/apm-sniffer/bootstrap-plugins/jdk-threading-plugin/src/main/resources/skywalking-plugin.def +++ b/apm-sniffer/bootstrap-plugins/jdk-threading-plugin/src/main/resources/skywalking-plugin.def @@ -15,4 +15,5 @@ # limitations under the License. jdk-threading-plugin=org.apache.skywalking.apm.plugin.jdk.threading.define.RunnableInstrumentation -jdk-threading-plugin=org.apache.skywalking.apm.plugin.jdk.threading.define.CallableInstrumentation \ No newline at end of file +jdk-threading-plugin=org.apache.skywalking.apm.plugin.jdk.threading.define.CallableInstrumentation +jdk-threading-plugin=org.apache.skywalking.apm.plugin.jdk.threading.define.ThreadPoolExecutorInstrumentation \ No newline at end of file