Instrument ForkJoinTask.fork (#3849)

* Instrument ForkJoinTask.fork

* spotless

instrumentation/executors/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/javaconcurrent/JavaForkJoinTaskInstrumentation.java

@@ -11,11 +11,14 @@ import static net.bytebuddy.matcher.ElementMatchers.named;
 import static net.bytebuddy.matcher.ElementMatchers.not;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
+import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
 import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
+import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
+import io.opentelemetry.javaagent.instrumentation.api.concurrent.ExecutorAdviceHelper;
 import io.opentelemetry.javaagent.instrumentation.api.concurrent.PropagatedContext;
 import io.opentelemetry.javaagent.instrumentation.api.concurrent.TaskAdviceHelper;
 import java.util.concurrent.Callable;
@@ -43,6 +46,9 @@ public class JavaForkJoinTaskInstrumentation implements TypeInstrumentation {
     transformer.applyAdviceToMethod(
         named("exec").and(takesArguments(0)).and(not(isAbstract())),
         JavaForkJoinTaskInstrumentation.class.getName() + "$ForkJoinTaskAdvice");
+    transformer.applyAdviceToMethod(
+        named("fork").and(takesArguments(0)),
+        JavaForkJoinTaskInstrumentation.class.getName() + "$ForkAdvice");
   }
 
   @SuppressWarnings("unused")
@@ -95,4 +101,25 @@ public class JavaForkJoinTaskInstrumentation implements TypeInstrumentation {
       }
     }
   }
+
+  @SuppressWarnings("unused")
+  public static class ForkAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static PropagatedContext enterFork(@Advice.This ForkJoinTask<?> task) {
+      Context context = Java8BytecodeBridge.currentContext();
+      if (ExecutorAdviceHelper.shouldPropagateContext(context, task)) {
+        ContextStore<ForkJoinTask<?>, PropagatedContext> contextStore =
+            InstrumentationContext.get(ForkJoinTask.class, PropagatedContext.class);
+        return ExecutorAdviceHelper.attachContextToTask(context, contextStore, task);
+      }
+      return null;
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void exitFork(
+        @Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
+      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
+    }
+  }
 }

instrumentation/executors/javaagent/src/test/groovy/ForkJoinTaskTest.groovy

@@ -0,0 +1,35 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import static io.opentelemetry.instrumentation.test.utils.TraceUtils.runUnderTrace
+
+import io.opentelemetry.instrumentation.test.AgentInstrumentationSpecification
+import java.util.stream.IntStream
+
+class ForkJoinTaskTest extends AgentInstrumentationSpecification {
+
+  def "test parallel"() {
+    when:
+    runUnderTrace("parent") {
+      IntStream.range(0, 20)
+        .parallel()
+        .forEach({ runUnderTrace("child") {} })
+    }
+
+    then:
+    assertTraces(1) {
+      trace(0, 21) {
+        span(0) {
+          name "parent"
+        }
+        (1..20).each { index ->
+          span(index) {
+            childOf(span(0))
+          }
+        }
+      }
+    }
+  }
+}