content: /*
   * Copyright The OpenTelemetry Authors
   * SPDX-License-Identifier: Apache-2.0
   */
  
  // Includes work from:
  /*
   * Copyright 2018 LINE Corporation
   *
   * LINE Corporation licenses this file to you under the Apache License,
   * version 2.0 (the "License"); you may not use this file except in compliance
   * with the License. You may obtain a copy of the License at:
   *
   *   https://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
   * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
   * License for the specific language governing permissions and limitations
   * under the License.
   */
  
  package io.opentelemetry.instrumentation.rxjava.v3_1_1;
  
  import io.opentelemetry.context.Context;
  import io.opentelemetry.context.Scope;
  import io.opentelemetry.instrumentation.api.annotation.support.async.AsyncOperationEndStrategies;
  import io.opentelemetry.instrumentation.api.internal.GuardedBy;
  import io.opentelemetry.instrumentation.rxjava.v3.common.RxJava3AsyncOperationEndStrategy;
  import io.opentelemetry.instrumentation.rxjava.v3.common.TracingCompletableObserver;
  import io.opentelemetry.instrumentation.rxjava.v3.common.TracingMaybeObserver;
  import io.opentelemetry.instrumentation.rxjava.v3.common.TracingSingleObserver;
  import io.reactivex.rxjava3.core.Completable;
  import io.reactivex.rxjava3.core.CompletableObserver;
  import io.reactivex.rxjava3.core.Flowable;
  import io.reactivex.rxjava3.core.Maybe;
  import io.reactivex.rxjava3.core.MaybeObserver;
  import io.reactivex.rxjava3.core.Observable;
  import io.reactivex.rxjava3.core.Observer;
  import io.reactivex.rxjava3.core.Single;
  import io.reactivex.rxjava3.core.SingleObserver;
  import io.reactivex.rxjava3.functions.BiFunction;
  import io.reactivex.rxjava3.functions.Function;
  import io.reactivex.rxjava3.operators.ConditionalSubscriber;
  import io.reactivex.rxjava3.parallel.ParallelFlowable;
  import io.reactivex.rxjava3.plugins.RxJavaPlugins;
  import javax.annotation.Nullable;
  import org.reactivestreams.Subscriber;
  
  /**
   * RxJava3 library instrumentation.
   *
   * <p>In order to enable RxJava3 instrumentation one has to call the {@link
   * TracingAssembly#enable()} method.
   *
   * <p>Instrumentation uses <code>on*Assembly</code> and <code>on*Subscribe</code> RxJavaPlugin hooks
   * to wrap RxJava3 classes in their tracing equivalents.
   *
   * <p>Instrumentation can be disabled by calling the {@link TracingAssembly#disable()} method.
   */
  public final class TracingAssembly {
  
    @SuppressWarnings("rawtypes")
    @GuardedBy("TracingAssembly.class")
    @Nullable
    private static BiFunction<? super Observable, ? super Observer, ? extends Observer>
        oldOnObservableSubscribe;
  
    @SuppressWarnings("rawtypes")
    @GuardedBy("TracingAssembly.class")
    @Nullable
    private static BiFunction<
            ? super Completable, ? super CompletableObserver, ? extends CompletableObserver>
        oldOnCompletableSubscribe;
  
    @SuppressWarnings("rawtypes")
    @GuardedBy("TracingAssembly.class")
    @Nullable
    private static BiFunction<? super Single, ? super SingleObserver, ? extends SingleObserver>
        oldOnSingleSubscribe;
  
    @SuppressWarnings("rawtypes")
    @GuardedBy("TracingAssembly.class")
    @Nullable
    private static BiFunction<? super Maybe, ? super MaybeObserver, ? extends MaybeObserver>
        oldOnMaybeSubscribe;
  
    @SuppressWarnings("rawtypes")
    @GuardedBy("TracingAssembly.class")
    @Nullable
    private static BiFunction<? super Flowable, ? super Subscriber, ? extends Subscriber>
        oldOnFlowableSubscribe;
  
    @SuppressWarnings("rawtypes")
    @GuardedBy("TracingAssembly.class")
    @Nullable
    private static Function<? super ParallelFlowable, ? extends ParallelFlowable>
        oldOnParallelAssembly;
  
    @GuardedBy("TracingAssembly.class")
    private static boolean enabled;
  
    public static TracingAssembly create() {
      return builder().build();
    }
  
    public static TracingAssemblyBuilder builder() {
      return new TracingAssemblyBuilder();
    }
  
    private final boolean captureExperimentalSpanAttributes;
  
    TracingAssembly(boolean captureExperimentalSpanAttributes) {
      this.captureExperimentalSpanAttributes = captureExperimentalSpanAttributes;
    }
  
    public void enable() {
      synchronized (TracingAssembly.class) {
        if (enabled) {
          return;
        }
  
        enableObservable();
  
        enableCompletable();
  
        enableSingle();
  
        enableMaybe();
  
        enableFlowable();
  
        enableParallel();
  
        enableWithSpanStrategy(captureExperimentalSpanAttributes);
  
        enabled = true;
      }
    }
  
    public void disable() {
      synchronized (TracingAssembly.class) {
        if (!enabled) {
          return;
        }
  
        disableObservable();
  
        disableCompletable();
  
        disableSingle();
  
        disableMaybe();
  
        disableFlowable();
  
        disableParallel();
  
        disableWithSpanStrategy();
  
        enabled = false;
      }
    }
  
    @GuardedBy("TracingAssembly.class")
    @SuppressWarnings({"rawtypes", "unchecked"})
    private static void enableParallel() {
      oldOnParallelAssembly = RxJavaPlugins.getOnParallelAssembly();
      RxJavaPlugins.setOnParallelAssembly(
          compose(
              oldOnParallelAssembly,
              parallelFlowable -> new TracingParallelFlowable(parallelFlowable, Context.current())));
    }
  
    @GuardedBy("TracingAssembly.class")
    private static void enableCompletable() {
      oldOnCompletableSubscribe = RxJavaPlugins.getOnCompletableSubscribe();
      RxJavaPlugins.setOnCompletableSubscribe(
          biCompose(
              oldOnCompletableSubscribe,
              (completable, observer) -> {
                Context context = Context.current();
                try (Scope ignored = context.makeCurrent()) {
                  return new TracingCompletableObserver(observer, context);
                }
              }));
    }
  
    @GuardedBy("TracingAssembly.class")
    @SuppressWarnings({"rawtypes", "unchecked"})
    private static void enableFlowable() {
      oldOnFlowableSubscribe = RxJavaPlugins.getOnFlowableSubscribe();
      RxJavaPlugins.setOnFlowableSubscribe(
          biCompose(
              oldOnFlowableSubscribe,
              (flowable, subscriber) -> {
                Context context = Context.current();
                try (Scope ignored = context.makeCurrent()) {
                  if (subscriber instanceof ConditionalSubscriber) {
                    return new TracingConditionalSubscriber<>(
                        (ConditionalSubscriber) subscriber, context);
                  } else {
                    return new TracingSubscriber<>(subscriber, context);
                  }
                }
              }));
    }
  
    @GuardedBy("TracingAssembly.class")
    @SuppressWarnings({"rawtypes", "unchecked"})
    private static void enableObservable() {
      oldOnObservableSubscribe = RxJavaPlugins.getOnObservableSubscribe();
      RxJavaPlugins.setOnObservableSubscribe(
          biCompose(
              oldOnObservableSubscribe,
              (observable, observer) -> {
                Context context = Context.current();
                try (Scope ignored = context.makeCurrent()) {
                  return new TracingObserver(observer, context);
                }
              }));
    }
  
    @GuardedBy("TracingAssembly.class")
    @SuppressWarnings({"rawtypes", "unchecked"})
    private static void enableSingle() {
      oldOnSingleSubscribe = RxJavaPlugins.getOnSingleSubscribe();
      RxJavaPlugins.setOnSingleSubscribe(
          biCompose(
              oldOnSingleSubscribe,
              (single, singleObserver) -> {
                Context context = Context.current();
                try (Scope ignored = context.makeCurrent()) {
                  return new TracingSingleObserver(singleObserver, context);
                }
              }));
    }
  
    @GuardedBy("TracingAssembly.class")
    @SuppressWarnings({"rawtypes", "unchecked"})
    private static void enableMaybe() {
      oldOnMaybeSubscribe = RxJavaPlugins.getOnMaybeSubscribe();
      RxJavaPlugins.setOnMaybeSubscribe(
          (BiFunction<? super Maybe, MaybeObserver, ? extends MaybeObserver>)
              biCompose(
                  oldOnMaybeSubscribe,
                  (BiFunction<Maybe, MaybeObserver, MaybeObserver>)
                      (maybe, maybeObserver) -> {
                        Context context = Context.current();
                        try (Scope ignored = context.makeCurrent()) {
                          return new TracingMaybeObserver(maybeObserver, context);
                        }
                      }));
    }
  
    private static RxJava3AsyncOperationEndStrategy asyncOperationEndStrategy;
  
    private static void enableWithSpanStrategy(boolean captureExperimentalSpanAttributes) {
      asyncOperationEndStrategy =
          RxJava3AsyncOperationEndStrategy.builder()
              .setCaptureExperimentalSpanAttributes(captureExperimentalSpanAttributes)
              .build();
  
      AsyncOperationEndStrategies.instance().registerStrategy(asyncOperationEndStrategy);
    }
  
    @GuardedBy("TracingAssembly.class")
    private static void disableParallel() {
      RxJavaPlugins.setOnParallelAssembly(oldOnParallelAssembly);
      oldOnParallelAssembly = null;
    }
  
    @GuardedBy("TracingAssembly.class")
    private static void disableObservable() {
      RxJavaPlugins.setOnObservableSubscribe(oldOnObservableSubscribe);
      oldOnObservableSubscribe = null;
    }
  
    @GuardedBy("TracingAssembly.class")
    private static void disableCompletable() {
      RxJavaPlugins.setOnCompletableSubscribe(oldOnCompletableSubscribe);
      oldOnCompletableSubscribe = null;
    }
  
    @GuardedBy("TracingAssembly.class")
    private static void disableFlowable() {
      RxJavaPlugins.setOnFlowableSubscribe(oldOnFlowableSubscribe);
      oldOnFlowableSubscribe = null;
    }
  
    @GuardedBy("TracingAssembly.class")
    private static void disableSingle() {
      RxJavaPlugins.setOnSingleSubscribe(oldOnSingleSubscribe);
      oldOnSingleSubscribe = null;
    }
  
    @GuardedBy("TracingAssembly.class")
    @SuppressWarnings({"rawtypes", "unchecked"})
    private static void disableMaybe() {
      RxJavaPlugins.setOnMaybeSubscribe(
          (BiFunction<? super Maybe, MaybeObserver, ? extends MaybeObserver>) oldOnMaybeSubscribe);
      oldOnMaybeSubscribe = null;
    }
  
    private static void disableWithSpanStrategy() {
      if (asyncOperationEndStrategy != null) {
        AsyncOperationEndStrategies.instance().unregisterStrategy(asyncOperationEndStrategy);
        asyncOperationEndStrategy = null;
      }
    }
  
    private static <T> Function<? super T, ? extends T> compose(
        Function<? super T, ? extends T> before, Function<? super T, ? extends T> after) {
      if (before == null) {
        return after;
      }
      return (T v) -> after.apply(before.apply(v));
    }
  
    private static <T, U> BiFunction<? super T, ? super U, ? extends U> biCompose(
        BiFunction<? super T, ? super U, ? extends U> before,
        BiFunction<? super T, ? super U, ? extends U> after) {
      if (before == null) {
        return after;
      }
      return (T v, U u) -> after.apply(v, before.apply(v, u));
    }
  }
  