content: /*
   * Copyright The OpenTelemetry Authors
   * SPDX-License-Identifier: Apache-2.0
   */
  
  package io.opentelemetry.instrumentation.reactor.v3_1;
  
  import static io.opentelemetry.instrumentation.api.annotation.support.async.AsyncOperationEndSupport.tryToGetResponse;
  
  import io.opentelemetry.api.common.AttributeKey;
  import io.opentelemetry.api.trace.Span;
  import io.opentelemetry.context.Context;
  import io.opentelemetry.instrumentation.api.annotation.support.async.AsyncOperationEndStrategy;
  import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
  import java.util.concurrent.atomic.AtomicBoolean;
  import java.util.function.Consumer;
  import org.reactivestreams.Publisher;
  import reactor.core.publisher.Flux;
  import reactor.core.publisher.Mono;
  
  public final class ReactorAsyncOperationEndStrategy implements AsyncOperationEndStrategy {
    // Visible for testing
    static final AttributeKey<Boolean> CANCELED_ATTRIBUTE_KEY =
        AttributeKey.booleanKey("reactor.canceled");
  
    public static ReactorAsyncOperationEndStrategy create() {
      return builder().build();
    }
  
    public static ReactorAsyncOperationEndStrategyBuilder builder() {
      return new ReactorAsyncOperationEndStrategyBuilder();
    }
  
    private final boolean captureExperimentalSpanAttributes;
  
    ReactorAsyncOperationEndStrategy(boolean captureExperimentalSpanAttributes) {
      this.captureExperimentalSpanAttributes = captureExperimentalSpanAttributes;
    }
  
    @Override
    public boolean supports(Class<?> returnType) {
      return returnType == Publisher.class || returnType == Mono.class || returnType == Flux.class;
    }
  
    @Override
    public <REQUEST, RESPONSE> Object end(
        Instrumenter<REQUEST, RESPONSE> instrumenter,
        Context context,
        REQUEST request,
        Object asyncValue,
        Class<RESPONSE> responseType) {
  
      EndOnFirstNotificationConsumer notificationConsumer =
          new EndOnFirstNotificationConsumer(context) {
            @Override
            protected void end(Object result, Throwable error) {
              instrumenter.end(context, request, tryToGetResponse(responseType, result), error);
            }
          };
  
      if (asyncValue instanceof Mono) {
        Mono<?> mono = (Mono<?>) asyncValue;
  
        return ContextPropagationOperator.runWithContext(mono, context)
            .doOnError(notificationConsumer)
            .doOnSuccess(notificationConsumer::onSuccess)
            .doOnCancel(notificationConsumer::onCancel);
      } else {
        Flux<?> flux = Flux.from((Publisher<?>) asyncValue);
        return ContextPropagationOperator.runWithContext(flux, context)
            .doOnError(notificationConsumer)
            .doOnComplete(notificationConsumer)
            .doOnCancel(notificationConsumer::onCancel);
      }
    }
  
    /**
     * Helper class to ensure that the span is ended exactly once regardless of how many OnComplete or
     * OnError notifications are received. Multiple notifications can happen anytime multiple
     * subscribers subscribe to the same publisher.
     */
    private abstract class EndOnFirstNotificationConsumer extends AtomicBoolean
        implements Runnable, Consumer<Throwable> {
  
      private static final long serialVersionUID = 1L;
  
      private final Context context;
  
      protected EndOnFirstNotificationConsumer(Context context) {
        super(false);
        this.context = context;
      }
  
      public <T> void onSuccess(T result) {
        accept(result, null);
      }
  
      public void onCancel() {
        if (compareAndSet(false, true)) {
          if (captureExperimentalSpanAttributes) {
            Span.fromContext(context).setAttribute(CANCELED_ATTRIBUTE_KEY, true);
          }
          end(null, null);
        }
      }
  
      @Override
      public void run() {
        accept(null, null);
      }
  
      @Override
      public void accept(Throwable exception) {
        end(null, exception);
      }
  
      private void accept(Object result, Throwable error) {
        if (compareAndSet(false, true)) {
          end(result, error);
        }
      }
  
      protected abstract void end(Object result, Throwable error);
    }
  }
  