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,93 @@
+/*
+ * 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.InstanceMethodsAroundInterceptor;
+import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.MethodInterceptResult;
+import org.apache.skywalking.apm.network.trace.component.ComponentsDefine;
+import org.apache.skywalking.apm.plugin.wrapper.SwCallableWrapper;
+
+public class ThreadPoolInvokeMethodInterceptor implements InstanceMethodsAroundInterceptor {
+
+    private static final String OPERATION_NAME_PREFIX = "ThreadPoolExecutor/";
+
+    @Override
+    public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
+                             MethodInterceptResult result) throws Throwable {
+        if (!shouldEnhance(allArguments)) {
+            return;
+        }
+
+        AbstractSpan span = ContextManager.createLocalSpan(OPERATION_NAME_PREFIX + method.getName());
+        span.setComponent(ComponentsDefine.JDK_THREADING);
+
+        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) throws Throwable {
+        if (shouldEnhance(allArguments)) {
+            ContextManager.stopSpan();
+        }
+        return ret;
+    }
+
+    @Override
+    public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments,
+                                      Class<?>[] argumentsTypes, Throwable t) {
+        if (shouldEnhance(allArguments)) {
+            ContextManager.activeSpan().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/ThreadPoolExecutorInstrumentation.java

@@ -37,10 +37,16 @@ public class ThreadPoolExecutorInstrumentation extends ClassInstanceMethodsEnhan
 
     private static final String INTERCEPT_SUBMIT_METHOD = "submit";
 
+    private static final String INTERCEPT_INVOKE_ALL_METHOD = "invokeAll";
+
+    private static final String INTERCEPT_INVOKE_ANY_METHOD = "invokeAny";
+
     private static final String INTERCEPT_EXECUTE_METHOD_HANDLE = "org.apache.skywalking.apm.plugin.ThreadPoolExecuteMethodInterceptor";
 
     private static final String INTERCEPT_SUBMIT_METHOD_HANDLE = "org.apache.skywalking.apm.plugin.ThreadPoolSubmitMethodInterceptor";
 
+    private static final String INTERCEPT_INVOKE_METHOD_HANDLE = "org.apache.skywalking.apm.plugin.ThreadPoolInvokeMethodInterceptor";
+
     @Override
     public boolean isBootstrapInstrumentation() {
         return true;
@@ -86,6 +92,23 @@ public String getMethodsInterceptor() {
                         return INTERCEPT_SUBMIT_METHOD_HANDLE;
                     }
 
+                    @Override
+                    public boolean isOverrideArgs() {
+                        return true;
+                    }
+                },
+                new InstanceMethodsInterceptPoint() {
+                    @Override
+                    public ElementMatcher<MethodDescription> getMethodsMatcher() {
+                        return ElementMatchers.named(INTERCEPT_INVOKE_ALL_METHOD)
+                                              .or(ElementMatchers.named(INTERCEPT_INVOKE_ANY_METHOD));
+                    }
+
+                    @Override
+                    public String getMethodsInterceptor() {
+                        return INTERCEPT_INVOKE_METHOD_HANDLE;
+                    }
+
                     @Override
                     public boolean isOverrideArgs() {
                         return true;

apm-sniffer/bootstrap-plugins/jdk-threadpool-plugin/src/test/java/org/apache/skywalking/apm/plugin/ThreadPoolInvokeMethodInterceptorTest.java

@@ -0,0 +1,129 @@
+/*
+ * 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 static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+import java.lang.reflect.Method;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.AbstractExecutorService;
+import org.apache.skywalking.apm.agent.core.context.ContextCarrier;
+import org.apache.skywalking.apm.agent.core.context.ContextManager;
+import org.apache.skywalking.apm.agent.core.context.trace.AbstractTracingSpan;
+import org.apache.skywalking.apm.agent.core.context.trace.TraceSegment;
+import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.EnhancedInstance;
+import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.MethodInterceptResult;
+import org.apache.skywalking.apm.agent.test.helper.SegmentHelper;
+import org.apache.skywalking.apm.agent.test.tools.AgentServiceRule;
+import org.apache.skywalking.apm.agent.test.tools.SegmentStorage;
+import org.apache.skywalking.apm.agent.test.tools.SegmentStoragePoint;
+import org.apache.skywalking.apm.agent.test.tools.SpanAssert;
+import org.apache.skywalking.apm.agent.test.tools.TracingSegmentRunner;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+
+@RunWith(TracingSegmentRunner.class)
+public class ThreadPoolInvokeMethodInterceptorTest {
+
+    @SegmentStoragePoint
+    private SegmentStorage segmentStorage;
+
+    @Rule
+    public AgentServiceRule agentServiceRule = new AgentServiceRule();
+
+    @Mock
+    private EnhancedInstance enhancedInstance;
+
+    @Mock
+    private MethodInterceptResult result;
+
+    private final ThreadPoolInvokeMethodInterceptor interceptor = new ThreadPoolInvokeMethodInterceptor();
+
+    @Test
+    public void shouldIgnoreUnexpectedArguments() throws Throwable {
+        Object[][] unexpectedArguments = new Object[][] {
+            null,
+            new Object[0],
+            new Object[] {null},
+            new Object[] {"not-a-collection"}
+        };
+
+        for (Object[] arguments : unexpectedArguments) {
+            ContextManager.createEntrySpan("parent", new ContextCarrier());
+
+            interceptor.beforeMethod(enhancedInstance, invokeAllMethod(), arguments, null, result);
+            interceptor.handleMethodException(
+                enhancedInstance, invokeAllMethod(), arguments, null, new IllegalStateException("ignored"));
+            interceptor.afterMethod(enhancedInstance, invokeAllMethod(), arguments, null, null);
+
+            assertThat(ContextManager.isActive(), is(true));
+            ContextManager.stopSpan();
+        }
+
+        assertThat(segmentStorage.getTraceSegments().size(), is(4));
+        for (TraceSegment traceSegment : segmentStorage.getTraceSegments()) {
+            List<AbstractTracingSpan> spans = SegmentHelper.getSpans(traceSegment);
+            assertThat(spans.size(), is(1));
+            assertThat(spans.get(0).getOperationName(), is("parent"));
+            SpanAssert.assertOccurException(spans.get(0), false);
+        }
+    }
+
+    @Test
+    public void shouldTraceEmptyCollection() throws Throwable {
+        Object[] arguments = new Object[] {Collections.emptyList()};
+        ContextManager.createEntrySpan("parent", new ContextCarrier());
+
+        interceptor.beforeMethod(enhancedInstance, invokeAllMethod(), arguments, null, result);
+        interceptor.afterMethod(enhancedInstance, invokeAllMethod(), arguments, null, null);
+        ContextManager.stopSpan();
+
+        List<AbstractTracingSpan> spans = SegmentHelper.getSpans(segmentStorage.getTraceSegments().get(0));
+        assertThat(spans.size(), is(2));
+        assertThat(spans.get(0).getOperationName(), is("ThreadPoolExecutor/invokeAll"));
+    }
+
+    @Test
+    public void shouldTraceAlternativeInvokeAllSignature() throws Throwable {
+        Object[] arguments = new Object[] {Collections.emptyList(), "alternative"};
+        ContextManager.createEntrySpan("parent", new ContextCarrier());
+
+        interceptor.beforeMethod(enhancedInstance, alternativeInvokeAllMethod(), arguments, null, result);
+        interceptor.afterMethod(enhancedInstance, alternativeInvokeAllMethod(), arguments, null, null);
+        ContextManager.stopSpan();
+
+        List<AbstractTracingSpan> spans = SegmentHelper.getSpans(segmentStorage.getTraceSegments().get(0));
+        assertThat(spans.size(), is(2));
+        assertThat(spans.get(0).getOperationName(), is("ThreadPoolExecutor/invokeAll"));
+    }
+
+    private Method invokeAllMethod() throws NoSuchMethodException {
+        return AbstractExecutorService.class.getMethod("invokeAll", Collection.class);
+    }
+
+    private Method alternativeInvokeAllMethod() throws NoSuchMethodException {
+        return getClass().getDeclaredMethod("invokeAll", Collection.class, String.class);
+    }
+
+    private void invokeAll(Collection<?> callables, String alternative) {
+    }
+}