content: /*
   * Copyright The OpenTelemetry Authors
   * SPDX-License-Identifier: Apache-2.0
   */
  
  package io.opentelemetry.instrumentation.rxjava.v1_0;
  
  import io.opentelemetry.context.Context;
  import io.opentelemetry.context.Scope;
  import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
  import java.util.concurrent.atomic.AtomicReference;
  import rx.Subscriber;
  
  final class TracedSubscriber<T, REQUEST> extends Subscriber<T> {
  
    private final Subscriber<T> delegate;
    private final Instrumenter<REQUEST, ?> instrumenter;
    private final AtomicReference<Context> contextRef;
    private final REQUEST request;
  
    TracedSubscriber(
        Subscriber<T> delegate,
        Instrumenter<REQUEST, ?> instrumenter,
        Context context,
        REQUEST request) {
      this.delegate = delegate;
      this.instrumenter = instrumenter;
      this.contextRef = new AtomicReference<>(context);
      this.request = request;
  
      delegate.add(new SpanFinishingSubscription<>(instrumenter, contextRef, request));
    }
  
    @Override
    public void onStart() {
      Context context = contextRef.get();
      if (context != null) {
        try (Scope ignored = context.makeCurrent()) {
          delegate.onStart();
        }
      } else {
        delegate.onStart();
      }
    }
  
    @Override
    public void onNext(T value) {
      Context context = contextRef.get();
      if (context != null) {
        try (Scope ignored = context.makeCurrent()) {
          delegate.onNext(value);
        }
      } else {
        delegate.onNext(value);
      }
    }
  
    @Override
    public void onCompleted() {
      Context context = contextRef.getAndSet(null);
      if (context != null) {
        try (Scope ignored = context.makeCurrent()) {
          delegate.onCompleted();
        } catch (Throwable t) {
          instrumenter.end(context, request, null, t);
          throw t;
        }
        instrumenter.end(context, request, null, null);
      } else {
        delegate.onCompleted();
      }
    }
  
    @Override
    public void onError(Throwable e) {
      Context context = contextRef.getAndSet(null);
      if (context != null) {
        instrumenter.end(context, request, null, e);
      }
      // TODO (trask) should this be wrapped in parent of context(?)
      delegate.onError(e);
    }
  }
  