rxjava

Author: jaydeluca

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

@@ -62,6 +62,12 @@ public final class TracingAssembly {
   private static BiFunction<? super Observable, ? super Observer, ? extends Observer>
       oldOnObservableSubscribe;
 
+  @SuppressWarnings("rawtypes")
+  @GuardedBy("TracingAssembly.class")
+  @Nullable
+  static volatile Function<? super Observable, ? extends Observable> oldOnObservableAssembly;
+
+
   @SuppressWarnings("rawtypes")
   @GuardedBy("TracingAssembly.class")
   @Nullable
@@ -118,6 +124,8 @@ public void enable() {
 
       enableObservable();
 
+      enableObservableAssembly();
+
       enableCompletable();
 
       enableSingle();
@@ -142,6 +150,8 @@ public void disable() {
 
       disableObservable();
 
+      disableObservableAssembly();
+
       disableCompletable();
 
       disableSingle();
@@ -219,6 +229,18 @@ private static void enableObservable() {
     }
   }
 
+  @GuardedBy("TracingAssembly.class")
+  @SuppressWarnings({"rawtypes", "unchecked"})
+  private static void enableObservableAssembly() {
+    if (TracingObserver.canEnable()) {
+      oldOnObservableAssembly = RxJavaPlugins.getOnObservableAssembly();
+      RxJavaPlugins.setOnObservableAssembly(
+          compose(
+              oldOnObservableAssembly,
+              observable -> new TracingObservable(observable, Context.current())));
+    }
+  }
+
   @GuardedBy("TracingAssembly.class")
   @SuppressWarnings({"rawtypes", "unchecked"})
   private static void enableSingle() {
@@ -274,6 +296,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/TracingObservable.java

@@ -0,0 +1,41 @@
+/*
+ * 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;
+
+final class TracingObservable<T> extends Observable<T> {
+  private final Observable<T> source;
+  private final Context context;
+
+  TracingObservable(Observable<T> source, Context context) {
+    this.source = source;
+    this.context = context;
+    System.out.println("TracingObservable created with context: " + context);
+  }
+
+  @Override
+  protected void subscribeActual(Observer<? super T> observer) {
+    System.out.println("TracingObservable.subscribeActual called with context: " + context);
+    System.out.println("Current context before making current: " + Context.current());
+
+    try (Scope ignored = context.makeCurrent()) {
+      System.out.println("Current context after making current: " + Context.current());
+
+      // Don't double-wrap if the observer is already a TracingObserver
+      if (observer instanceof TracingObserver) {
+        System.out.println("Observer is already TracingObserver, not wrapping");
+        source.subscribe(observer);
+      } else {
+        System.out.println("Wrapping observer in TracingObserver");
+        source.subscribe(new TracingObserver<>(observer, context));
+      }
+    }
+  }
+}

instrumentation/rxjava/rxjava-2.0/testing/src/main/java/io/opentelemetry/instrumentation/rxjava/v2_0/AbstractRxJava2Test.java

@@ -10,6 +10,7 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
+import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
@@ -23,6 +24,7 @@
 import io.reactivex.Observable;
 import io.reactivex.Scheduler;
 import io.reactivex.Single;
+import io.reactivex.disposables.Disposable;
 import io.reactivex.internal.operators.flowable.FlowablePublish;
 import io.reactivex.internal.operators.observable.ObservablePublish;
 import io.reactivex.schedulers.Schedulers;
@@ -338,6 +340,29 @@ public void basicObservable() {
                             .hasParent(trace.getSpan(0))));
   }
 
+
+  @Test
+  @SuppressWarnings("CanIgnoreReturnValue")
+  public void basicObservableFromCallable() {
+    Observable<String> test = createParentSpan(() -> Observable.fromCallable(() -> "success"));
+    assertThat(test).isNotNull();
+
+    // Actually subscribe to trigger the Observable execution
+    String result = test.blockingFirst();
+    assertThat(result).isEqualTo("success");
+
+    testing()
+        .waitAndAssertTraces(
+            trace ->
+                trace.hasSpansSatisfyingExactly(
+                    span -> span.hasName(PARENT).hasKind(SpanKind.INTERNAL).hasNoParent(),
+                    span ->
+                        span.hasName(ADD_ONE)
+                            .hasKind(SpanKind.INTERNAL)
+                            .hasParent(trace.getSpan(0))));
+  }
+
+
   @Test
   public void connectableFlowable() {
     List<Integer> result =
@@ -779,8 +804,29 @@ public void observableChainHasSubscriptionContext() {
                             .hasParent(trace.getSpan(0))));
   }
 
-  @ParameterizedTest
+  @Test
+  public void basicObservableFromCallableTest() {
+    Disposable disposable = Observable.fromCallable(() -> "success")
+        .onErrorReturnItem("")
+        .subscribeOn(Schedulers.io())
+        .observeOn(Schedulers.single())
+        .subscribe(data -> System.out.print("done"));
+
+    assertThat(disposable).isNotNull();
+
+    testing()
+        .waitAndAssertTraces(
+            trace ->
+                trace.hasSpansSatisfyingExactly(
+                    span -> span.hasName(PARENT).hasKind(SpanKind.INTERNAL).hasNoParent(),
+                    span ->
+                        span.hasName(ADD_ONE)
+                            .hasKind(SpanKind.INTERNAL)
+                            .hasParent(trace.getSpan(0))));
+  }
+
   @MethodSource("schedulers")
+  @ParameterizedTest
   public void flowableMultiResults(Scheduler scheduler) {
     List<Integer> result =
         testing()