Add tracing support for ThreadPoolExecutor invokeAll and invokeAny

CHANGES.md

@@ -11,6 +11,7 @@ Release Notes.
 * Add a Jetty 12 server plugin (`jetty-server-12.x`). Jetty 12 removed the `HttpChannel` handle target and moved request handling to the async `Server#handle(Request, Response, Callback)` core API, so it needs a separate plugin from the merged `jetty-server`.
 * Add a Struts 7 plugin (`struts2-7.x`) for Jakarta Struts, whose `DefaultActionInvocation` moved to `org.apache.struts2`.
 * Added support for Lettuce reactive Redis commands.
+* Add tracing support for `invokeAll` and `invokeAny` in the JDK thread pool plugin.
 * Add Spring AI 1.x plugin and GenAI layer.
 * Fix httpclient-5.x plugin injecting sw8 propagation headers into ClickHouse HTTP requests (port 8123), causing HTTP 400. Add `PROPAGATION_EXCLUDE_PORTS` config to skip tracing (including header injection) for specified ports in the classic client interceptor.
 * Add Spring RabbitMQ 2.x - 4.x plugin.

apm-sniffer/bootstrap-plugins/jdk-threadpool-plugin/src/main/java/org/apache/skywalking/apm/plugin/ThreadPoolInvokeMethodInterceptor.java

@@ -0,0 +1,96 @@
+/*
+ * 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;
+
+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.context.trace.AbstractSpan;
+import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.EnhancedInstance;
+import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.v2.InstanceMethodsAroundInterceptorV2;
+import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.v2.MethodInvocationContext;
+import org.apache.skywalking.apm.network.trace.component.ComponentsDefine;
+import org.apache.skywalking.apm.plugin.wrapper.SwCallableWrapper;
+
+public class ThreadPoolInvokeMethodInterceptor implements InstanceMethodsAroundInterceptorV2 {
+
+    private static final String OPERATION_NAME_PREFIX = "ThreadPoolExecutor/";
+
+    @Override
+    public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
+                             MethodInvocationContext context) throws Throwable {
+        if (!shouldEnhance(allArguments)) {
+            return;
+        }
+
+        AbstractSpan span = ContextManager.createLocalSpan(OPERATION_NAME_PREFIX + method.getName());
+        span.setComponent(ComponentsDefine.JDK_THREADING);
+        context.setContext(span);
+
+        ContextSnapshot contextSnapshot = ContextManager.capture();
+        Collection<?> callables = (Collection<?>) allArguments[0];
+        List<Object> wrappedCallables = new ArrayList<>(callables.size());
+        for (Object callable : callables) {
+            wrappedCallables.add(wrap(callable, contextSnapshot));
+        }
+        allArguments[0] = wrappedCallables;
+    }
+
+    @Override
+    public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
+                              Object ret, MethodInvocationContext context) throws Throwable {
+        AbstractSpan span = (AbstractSpan) context.getContext();
+        if (span != null) {
+            ContextManager.stopSpan(span);
+        }
+        return ret;
+    }
+
+    @Override
+    public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments,
+                                      Class<?>[] argumentsTypes, Throwable t, MethodInvocationContext context) {
+        AbstractSpan span = (AbstractSpan) context.getContext();
+        if (span != null) {
+            span.log(t);
+        }
+    }
+
+    private boolean shouldEnhance(Object[] allArguments) {
+        return ContextManager.isActive()
+            && allArguments != null
+            && allArguments.length > 0
+            && allArguments[0] instanceof Collection;
+    }
+
+    private Object wrap(Object callable, ContextSnapshot contextSnapshot) {
+        if (!(callable instanceof Callable) || callable instanceof SwCallableWrapper || hasCapturedContext(callable)) {
+            return callable;
+        }
+        return new SwCallableWrapper((Callable) callable, contextSnapshot);
+    }
+
+    private boolean hasCapturedContext(Object callable) {
+        return callable instanceof EnhancedInstance
+            && ((EnhancedInstance) callable).getSkyWalkingDynamicField() instanceof ContextSnapshot;
+    }
+}

apm-sniffer/bootstrap-plugins/jdk-threadpool-plugin/src/main/java/org/apache/skywalking/apm/plugin/define/ThreadPoolExecutorInvokeInstrumentation.java

@@ -0,0 +1,83 @@
+/*
+ * 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.define;
+
+import java.util.Collection;
+import java.util.concurrent.TimeUnit;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import net.bytebuddy.matcher.ElementMatchers;
+import org.apache.skywalking.apm.agent.core.plugin.interceptor.ConstructorInterceptPoint;
+import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.v2.ClassInstanceMethodsEnhancePluginDefineV2;
+import org.apache.skywalking.apm.agent.core.plugin.interceptor.v2.InstanceMethodsInterceptV2Point;
+import org.apache.skywalking.apm.agent.core.plugin.match.ClassMatch;
+import org.apache.skywalking.apm.agent.core.plugin.match.NameMatch;
+
+public class ThreadPoolExecutorInvokeInstrumentation extends ClassInstanceMethodsEnhancePluginDefineV2 {
+
+    private static final String ENHANCE_CLASS = "java.util.concurrent.ThreadPoolExecutor";
+
+    private static final String INTERCEPT_INVOKE_ALL_METHOD = "invokeAll";
+
+    private static final String INTERCEPT_INVOKE_ANY_METHOD = "invokeAny";
+
+    private static final String INTERCEPT_INVOKE_METHOD_HANDLE =
+        "org.apache.skywalking.apm.plugin.ThreadPoolInvokeMethodInterceptor";
+
+    @Override
+    public boolean isBootstrapInstrumentation() {
+        return true;
+    }
+
+    @Override
+    protected ClassMatch enhanceClass() {
+        return NameMatch.byName(ENHANCE_CLASS);
+    }
+
+    @Override
+    public ConstructorInterceptPoint[] getConstructorsInterceptPoints() {
+        return new ConstructorInterceptPoint[0];
+    }
+
+    @Override
+    public InstanceMethodsInterceptV2Point[] getInstanceMethodsInterceptV2Points() {
+        return new InstanceMethodsInterceptV2Point[] {
+            new InstanceMethodsInterceptV2Point() {
+                @Override
+                public ElementMatcher<MethodDescription> getMethodsMatcher() {
+                    return ElementMatchers.named(INTERCEPT_INVOKE_ALL_METHOD)
+                                          .or(ElementMatchers.named(INTERCEPT_INVOKE_ANY_METHOD))
+                                          .and(ElementMatchers.takesArguments(Collection.class)
+                                                              .or(ElementMatchers.takesArguments(
+                                                                  Collection.class, long.class, TimeUnit.class)));
+                }
+
+                @Override
+                public String getMethodsInterceptorV2() {
+                    return INTERCEPT_INVOKE_METHOD_HANDLE;
+                }
+
+                @Override
+                public boolean isOverrideArgs() {
+                    return true;
+                }
+            }
+        };
+    }
+}

apm-sniffer/bootstrap-plugins/jdk-threadpool-plugin/src/main/resources/skywalking-plugin.def

@@ -15,3 +15,4 @@
 # limitations under the License.
 
 jdk-threadpool-plugin=org.apache.skywalking.apm.plugin.define.ThreadPoolExecutorInstrumentation
+jdk-threadpool-plugin=org.apache.skywalking.apm.plugin.define.ThreadPoolExecutorInvokeInstrumentation