Fix rxjava context propagation for ObservableFromCallable

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

@@ -93,6 +93,10 @@ public final class TracingAssembly {
   private static Function<? super ParallelFlowable, ? extends ParallelFlowable>
       oldOnParallelAssembly;
 
+  @GuardedBy("TracingAssembly.class")
+  @Nullable
+  private static Function<? super Runnable, ? extends Runnable> oldScheduleHandler;
+
   @GuardedBy("TracingAssembly.class")
   private static boolean enabled;
 
@@ -118,6 +122,8 @@ public void enable() {
 
       enableObservable();
 
+      enableWrappedScheduleHandler();
+
       enableCompletable();
 
       enableSingle();
@@ -142,6 +148,8 @@ public void disable() {
 
       disableObservable();
 
+      disableObservableAssembly();
+
       disableCompletable();
 
       disableSingle();
@@ -219,6 +227,25 @@ private static void enableObservable() {
     }
   }
 
+  @GuardedBy("TracingAssembly.class")
+  private static void enableWrappedScheduleHandler() {
+    oldScheduleHandler = RxJavaPlugins.getScheduleHandler();
+    RxJavaPlugins.setScheduleHandler(
+        runnable -> {
+          Context context = Context.current();
+          Runnable wrappedRunnable =
+              () -> {
+                try (Scope ignored = context.makeCurrent()) {
+                  runnable.run();
+                }
+              };
+          // If there was a previous handler, apply it to our wrapped runnable
+          return oldScheduleHandler != null
+              ? oldScheduleHandler.apply(wrappedRunnable)
+              : wrappedRunnable;
+        });
+  }
+
   @GuardedBy("TracingAssembly.class")
   @SuppressWarnings({"rawtypes", "unchecked"})
   private static void enableSingle() {
@@ -274,6 +301,12 @@ private static void disableObservable() {
     oldOnObservableSubscribe = null;
   }
 
+  @GuardedBy("TracingAssembly.class")
+  private static void disableObservableAssembly() {
+    RxJavaPlugins.setScheduleHandler(oldScheduleHandler);
+    oldScheduleHandler = null;
+  }
+
   @GuardedBy("TracingAssembly.class")
   private static void disableCompletable() {
     RxJavaPlugins.setOnCompletableSubscribe(oldOnCompletableSubscribe);