Review fixes for reactor-3.1:library

Automated code review of instrumentation/reactor/reactor-3.1/library.

Author: otelbot[bot]

instrumentation/reactor/reactor-3.1/library/src/main/java/io/opentelemetry/instrumentation/reactor/v3_1/ContextPropagationOperator.java

@@ -52,6 +52,8 @@ public final class ContextPropagationOperator {
 
   private static final Object VALUE = new Object();
 
+  private static final String SCHEDULERS_HOOK_KEY = RunnableWrapper.class.getName();
+
   @Nullable
   private static final MethodHandle MONO_CONTEXT_WRITE_METHOD = getContextWriteMethod(Mono.class);
 
@@ -60,6 +62,9 @@ public final class ContextPropagationOperator {
 
   @Nullable private static final MethodHandle SCHEDULERS_HOOK_METHOD = getSchedulersHookMethod();
 
+  @Nullable
+  private static final MethodHandle SCHEDULERS_RESET_HOOK_METHOD = getSchedulersResetHookMethod();
+
   @Nullable
   private static MethodHandle getContextWriteMethod(Class<?> type) {
     MethodHandles.Lookup lookup = MethodHandles.publicLookup();
@@ -88,6 +93,18 @@ private static MethodHandle getSchedulersHookMethod() {
     return null;
   }
 
+  @Nullable
+  private static MethodHandle getSchedulersResetHookMethod() {
+    MethodHandles.Lookup lookup = MethodHandles.publicLookup();
+    try {
+      return lookup.findStatic(
+          Schedulers.class, "resetOnScheduleHook", methodType(void.class, String.class));
+    } catch (NoSuchMethodException | IllegalAccessException e) {
+      // ignore
+    }
+    return null;
+  }
+
   public static ContextPropagationOperator create() {
     return builder().build();
   }
@@ -171,7 +188,7 @@ public void registerOnEachOperator() {
       Hooks.onEachOperator(
           TracingSubscriber.class.getName(), tracingLift(asyncOperationEndStrategy));
       AsyncOperationEndStrategies.instance().registerStrategy(asyncOperationEndStrategy);
-      registerScheduleHook(RunnableWrapper.class.getName(), RunnableWrapper::new);
+      registerScheduleHook(SCHEDULERS_HOOK_KEY, RunnableWrapper::new);
       enabled = true;
     }
   }
@@ -187,6 +204,17 @@ private static void registerScheduleHook(String key, Function<Runnable, Runnable
     }
   }
 
+  private static void resetScheduleHook(String key) {
+    if (SCHEDULERS_RESET_HOOK_METHOD == null) {
+      return;
+    }
+    try {
+      SCHEDULERS_RESET_HOOK_METHOD.invoke(key);
+    } catch (Throwable t) {
+      logger.log(WARNING, "Failed to remove scheduler hook", t);
+    }
+  }
+
   /** Unregisters the hook registered by {@link #registerOnEachOperator()}. */
   public void resetOnEachOperator() {
     synchronized (lock) {
@@ -195,6 +223,7 @@ public void resetOnEachOperator() {
       }
       Hooks.resetOnEachOperator(TracingSubscriber.class.getName());
       AsyncOperationEndStrategies.instance().unregisterStrategy(asyncOperationEndStrategy);
+      resetScheduleHook(SCHEDULERS_HOOK_KEY);
       enabled = false;
     }
   }

instrumentation/reactor/reactor-3.1/library/src/test/java/io/opentelemetry/instrumentation/reactor/v3_1/HooksTest.java

@@ -6,18 +6,47 @@
 package io.opentelemetry.instrumentation.reactor.v3_1;
 
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
+import static java.util.concurrent.TimeUnit.SECONDS;
 import static org.assertj.core.api.Assertions.assertThat;
-
+import static org.junit.jupiter.api.Assumptions.assumeTrue;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.api.trace.TraceFlags;
+import io.opentelemetry.api.trace.TraceState;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
 import java.util.concurrent.Callable;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Function;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Test;
 import reactor.core.CoreSubscriber;
 import reactor.core.Disposable;
+import reactor.core.publisher.Hooks;
 import reactor.core.publisher.Mono;
 import reactor.core.scheduler.Schedulers;
 
 class HooksTest {
 
+  private static final Span PARENT_SPAN =
+      Span.wrap(
+          SpanContext.create(
+              "11111111111111111111111111111111",
+              "1111111111111111",
+              TraceFlags.getSampled(),
+              TraceState.getDefault()));
+
+  @AfterEach
+  void resetHooks() {
+    Hooks.resetOnEachOperator(TracingSubscriber.class.getName());
+    if (schedulerHooksSupported()) {
+      Schedulers.resetOnScheduleHook(RunnableWrapper.class.getName());
+    }
+  }
+
   @Test
   void canResetOurHooks() {
     ContextPropagationOperator operator = ContextPropagationOperator.create();
@@ -35,6 +64,23 @@ void canResetOurHooks() {
     assertThat(subscriber.get()).extracting("actual").isNotInstanceOf(TracingSubscriber.class);
   }
 
+  @Test
+  void canResetSchedulerHook() throws InterruptedException {
+    assumeTrue(schedulerHooksSupported());
+
+    ContextPropagationOperator operator = ContextPropagationOperator.create();
+
+    assertThat(scheduledSpanIsCurrent()).isFalse();
+
+    operator.registerOnEachOperator();
+    try {
+      assertThat(scheduledSpanIsCurrent()).isTrue();
+    } finally {
+      operator.resetOnEachOperator();
+    }
+    assertThat(scheduledSpanIsCurrent()).isFalse();
+  }
+
   private static class CapturingMono extends Mono<Integer> {
     private final AtomicReference<CoreSubscriber<? super Integer>> subscriber;
 
@@ -71,4 +117,32 @@ void testInvalidBlockUsage() throws InterruptedException {
     disposable.dispose();
     operator.resetOnEachOperator();
   }
+
+  private static boolean schedulerHooksSupported() {
+    try {
+      Schedulers.class.getMethod("onScheduleHook", String.class, Function.class);
+      Schedulers.class.getMethod("resetOnScheduleHook", String.class);
+      return true;
+    } catch (NoSuchMethodException e) {
+      return false;
+    }
+  }
+
+  private static boolean scheduledSpanIsCurrent() throws InterruptedException {
+    CountDownLatch latch = new CountDownLatch(1);
+    AtomicBoolean currentSpanValid = new AtomicBoolean(false);
+    // The scope is closed before the scheduled task runs; that is the point of this test:
+    // the schedulers hook must capture the current context at schedule time, not at execution
+    // time.
+    try (Scope ignored = Context.root().with(PARENT_SPAN).makeCurrent()) {
+      Schedulers.single()
+          .schedule(
+              () -> {
+                currentSpanValid.set(Span.current().getSpanContext().isValid());
+                latch.countDown();
+              });
+    }
+    assertThat(latch.await(5, SECONDS)).isTrue();
+    return currentSpanValid.get();
+  }
 }