check to make sure class is ObservableFromCallable

Author: jaydeluca

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

@@ -93,6 +93,11 @@ public final class TracingAssembly {
   private static Function<? super ParallelFlowable, ? extends ParallelFlowable>
       oldOnParallelAssembly;
 
+  @SuppressWarnings("rawtypes")
+  @GuardedBy("TracingAssembly.class")
+  @Nullable
+  private static Function<? super Observable, ? extends Observable> oldOnObservableAssembly;
+
   @GuardedBy("TracingAssembly.class")
   private static boolean enabled;
 
@@ -118,6 +123,8 @@ public void enable() {
 
       enableObservable();
 
+      enableObservableAssembly();
+
       enableCompletable();
 
       enableSingle();
@@ -142,6 +149,8 @@ public void disable() {
 
       disableObservable();
 
+      disableObservableAssembly();
+
       disableCompletable();
 
       disableSingle();
@@ -219,6 +228,24 @@ private static void enableObservable() {
     }
   }
 
+  @GuardedBy("TracingAssembly.class")
+  @SuppressWarnings({"rawtypes", "unchecked"})
+  private static void enableObservableAssembly() {
+    oldOnObservableAssembly = RxJavaPlugins.getOnObservableAssembly();
+    RxJavaPlugins.setOnObservableAssembly(
+        compose(
+            oldOnObservableAssembly,
+            observable -> {
+              if (observable
+                  .getClass()
+                  .getName()
+                  .equals("io.reactivex.internal.operators.observable.ObservableFromCallable")) {
+                return new TracingCallableObservable(observable, Context.current());
+              }
+              return observable;
+            }));
+  }
+
   @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.setOnObservableAssembly(oldOnObservableAssembly);
+    oldOnObservableAssembly = null;
+  }
+
   @GuardedBy("TracingAssembly.class")
   private static void disableCompletable() {
     RxJavaPlugins.setOnCompletableSubscribe(oldOnCompletableSubscribe);

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

@@ -0,0 +1,49 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.rxjava.v2_0;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import io.reactivex.Observable;
+import io.reactivex.Observer;
+import java.lang.reflect.Field;
+import java.util.concurrent.Callable;
+
+/** Wraps an ObservableFromCallable to ensure context propagation for the callable execution. */
+final class TracingCallableObservable<T> extends Observable<T> {
+  private final Observable<T> source;
+  private final Context context;
+
+  TracingCallableObservable(Observable<T> source, Context context) {
+    this.source = source;
+    this.context = context;
+  }
+
+  @Override
+  protected void subscribeActual(Observer<? super T> observer) {
+    try {
+      Field callableField = source.getClass().getDeclaredField("callable");
+      callableField.setAccessible(true);
+      @SuppressWarnings("unchecked")
+      Callable<T> originalCallable = (Callable<T>) callableField.get(source);
+
+      if (originalCallable != null) {
+        Callable<T> wrappedCallable =
+            () -> {
+              try (Scope ignored = context.makeCurrent()) {
+                return originalCallable.call();
+              }
+            };
+
+        callableField.set(source, wrappedCallable);
+      }
+    } catch (Exception e) {
+      // If reflection fails, fall back to original behavior
+    }
+
+    source.subscribe(observer);
+  }
+}

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

@@ -30,6 +30,7 @@
 import io.reactivex.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;
@@ -342,7 +343,8 @@ void basicObservable() {
   }
 
   @Test
-  void basicObservableFromCallable() {
+  void observableFromCallableContextPropagation() throws InterruptedException {
+    CountDownLatch latch = new CountDownLatch(1);
     AtomicReference<String> traceId = new AtomicReference<>();
     AtomicReference<String> innerObservableTraceId = new AtomicReference<>();
     AtomicReference<String> endObservableTraceId = new AtomicReference<>();
@@ -359,11 +361,14 @@ void basicObservableFromCallable() {
                   .subscribeOn(Schedulers.io())
                   .observeOn(Schedulers.single())
                   .subscribe(
-                      data ->
-                          endObservableTraceId.set(Span.current().getSpanContext().getTraceId()));
+                      data -> {
+                        endObservableTraceId.set(Span.current().getSpanContext().getTraceId());
+                        latch.countDown();
+                      });
           assertThat(unused).isNotNull();
         });
 
+    latch.await();
     assertThat(innerObservableTraceId.get()).isEqualTo(traceId.get());
     assertThat(endObservableTraceId.get()).isEqualTo(traceId.get());
   }