simplify instrumentation, fix other versions of rxjava

Author: jaydeluca

instrumentation/rxjava/rxjava-2.0/library/src/main/java/io/opentelemetry/instrumentation/rxjava/v2_0/TracingAssembly.java

@@ -229,21 +229,16 @@ private static void enableObservable() {
   }
 
   @GuardedBy("TracingAssembly.class")
-  @SuppressWarnings({"rawtypes", "unchecked"})
   private static void enableObservableAssembly() {
-    oldOnObservableAssembly = RxJavaPlugins.getOnObservableAssembly();
-    RxJavaPlugins.setOnObservableAssembly(
-        compose(
-            observable -> {
-              if (observable
-                  .getClass()
-                  .getName()
-                  .equals("io.reactivex.internal.operators.observable.ObservableFromCallable")) {
-                return new TracingCallableObservable(observable, Context.current());
-              }
-              return observable;
-            },
-            oldOnObservableAssembly));
+    RxJavaPlugins.setScheduleHandler(
+        runnable -> {
+          Context context = Context.current();
+          return () -> {
+            try (Scope ignored = context.makeCurrent()) {
+              runnable.run();
+            }
+          };
+        });
   }
 
   @GuardedBy("TracingAssembly.class")

instrumentation/rxjava/rxjava-2.0/library/src/main/java/io/opentelemetry/instrumentation/rxjava/v2_0/TracingCallableObservable.java

@@ -11,6 +11,7 @@
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
 import io.opentelemetry.instrumentation.testing.util.ThrowingRunnable;
@@ -23,14 +24,18 @@
 import io.reactivex.rxjava3.core.Observable;
 import io.reactivex.rxjava3.core.Scheduler;
 import io.reactivex.rxjava3.core.Single;
+import io.reactivex.rxjava3.disposables.Disposable;
 import io.reactivex.rxjava3.internal.operators.flowable.FlowablePublish;
 import io.reactivex.rxjava3.internal.operators.observable.ObservablePublish;
 import io.reactivex.rxjava3.schedulers.Schedulers;
 import java.util.Comparator;
 import java.util.List;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.stream.Stream;
+import org.assertj.core.api.Assertions;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.TestInstance;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -224,6 +229,37 @@ public void twoOperationsFlowable() {
                             .hasParent(trace.getSpan(0))));
   }
 
+  @Test
+  void observableFromCallableContextPropagation() throws InterruptedException {
+    CountDownLatch latch = new CountDownLatch(1);
+    AtomicReference<String> traceId = new AtomicReference<>();
+    AtomicReference<String> innerObservableTraceId = new AtomicReference<>();
+    AtomicReference<String> endObservableTraceId = new AtomicReference<>();
+
+    createParentSpan(
+        () -> {
+          traceId.set(Span.current().getSpanContext().getTraceId());
+          Disposable unused =
+              Observable.fromCallable(
+                      () -> {
+                        innerObservableTraceId.set(Span.current().getSpanContext().getTraceId());
+                        return "success";
+                      })
+                  .subscribeOn(Schedulers.io())
+                  .observeOn(Schedulers.single())
+                  .subscribe(
+                      data -> {
+                        endObservableTraceId.set(Span.current().getSpanContext().getTraceId());
+                        latch.countDown();
+                      });
+          assertThat(unused).isNotNull();
+        });
+
+    latch.await();
+    Assertions.assertThat(innerObservableTraceId.get()).isEqualTo(traceId.get());
+    Assertions.assertThat(endObservableTraceId.get()).isEqualTo(traceId.get());
+  }
+
   @Test
   public void delayedFlowable() {
     List<Integer> result =

instrumentation/rxjava/rxjava-3-common/testing/src/main/java/io/opentelemetry/instrumentation/rxjava/v3/common/AbstractRxJava3Test.java

@@ -66,6 +66,11 @@ public final class TracingAssembly {
   private static BiFunction<? super Observable, ? super Observer, ? extends Observer>
       oldOnObservableSubscribe;
 
+  @SuppressWarnings("rawtypes")
+  @GuardedBy("TracingAssembly.class")
+  @Nullable
+  private static Function<? super Observable, ? extends Observable> oldOnObservableAssembly;
+
   @SuppressWarnings("rawtypes")
   @GuardedBy("TracingAssembly.class")
   @Nullable
@@ -122,6 +127,8 @@ public void enable() {
 
       enableObservable();
 
+      enableObservableAssembly();
+
       enableCompletable();
 
       enableSingle();
@@ -146,6 +153,8 @@ public void disable() {
 
       disableObservable();
 
+      disableObservableAssembly();
+
       disableCompletable();
 
       disableSingle();
@@ -221,6 +230,19 @@ private static void enableObservable() {
             }));
   }
 
+  @GuardedBy("TracingAssembly.class")
+  private static void enableObservableAssembly() {
+    RxJavaPlugins.setScheduleHandler(
+        runnable -> {
+          Context context = Context.current();
+          return () -> {
+            try (Scope ignored = context.makeCurrent()) {
+              runnable.run();
+            }
+          };
+        });
+  }
+
   @GuardedBy("TracingAssembly.class")
   @SuppressWarnings({"rawtypes", "unchecked"})
   private static void enableSingle() {
@@ -276,6 +298,12 @@ private static void disableObservable() {
     oldOnObservableSubscribe = null;
   }
 
+  @GuardedBy("TracingAssembly.class")
+  private static void disableObservableAssembly() {
+    RxJavaPlugins.setOnObservableAssembly(oldOnObservableAssembly);
+    oldOnObservableAssembly = null;
+  }
+
   @GuardedBy("TracingAssembly.class")
   private static void disableCompletable() {
     RxJavaPlugins.setOnCompletableSubscribe(oldOnCompletableSubscribe);

instrumentation/rxjava/rxjava-3.0/library/src/main/java/io/opentelemetry/instrumentation/rxjava/v3_0/TracingAssembly.java

@@ -66,6 +66,11 @@ public final class TracingAssembly {
   private static BiFunction<? super Observable, ? super Observer, ? extends Observer>
       oldOnObservableSubscribe;
 
+  @SuppressWarnings("rawtypes")
+  @GuardedBy("TracingAssembly.class")
+  @Nullable
+  private static Function<? super Observable, ? extends Observable> oldOnObservableAssembly;
+
   @SuppressWarnings("rawtypes")
   @GuardedBy("TracingAssembly.class")
   @Nullable
@@ -122,6 +127,8 @@ public void enable() {
 
       enableObservable();
 
+      enableObservableAssembly();
+
       enableCompletable();
 
       enableSingle();
@@ -146,6 +153,8 @@ public void disable() {
 
       disableObservable();
 
+      disableObservableAssembly();
+
       disableCompletable();
 
       disableSingle();
@@ -186,6 +195,19 @@ private static void enableCompletable() {
             }));
   }
 
+  @GuardedBy("TracingAssembly.class")
+  private static void enableObservableAssembly() {
+    RxJavaPlugins.setScheduleHandler(
+        runnable -> {
+          Context context = Context.current();
+          return () -> {
+            try (Scope ignored = context.makeCurrent()) {
+              runnable.run();
+            }
+          };
+        });
+  }
+
   @GuardedBy("TracingAssembly.class")
   @SuppressWarnings({"rawtypes", "unchecked"})
   private static void enableFlowable() {
@@ -276,6 +298,12 @@ private static void disableObservable() {
     oldOnObservableSubscribe = null;
   }
 
+  @GuardedBy("TracingAssembly.class")
+  private static void disableObservableAssembly() {
+    RxJavaPlugins.setOnObservableAssembly(oldOnObservableAssembly);
+    oldOnObservableAssembly = null;
+  }
+
   @GuardedBy("TracingAssembly.class")
   private static void disableCompletable() {
     RxJavaPlugins.setOnCompletableSubscribe(oldOnCompletableSubscribe);