Fix flaky spring webflux tests (#3150)

instrumentation/spring/spring-webflux-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/server/AdviceUtils.java

@@ -11,6 +11,7 @@ import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.StatusCode;
 import io.opentelemetry.context.Scope;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Function;
 import org.reactivestreams.Publisher;
 import org.reactivestreams.Subscription;
@@ -79,17 +80,18 @@ public class AdviceUtils {
 
   private static void finishSpanIfPresentInAttributes(
       Map<String, Object> attributes, Throwable throwable) {
-
     io.opentelemetry.context.Context context =
         (io.opentelemetry.context.Context) attributes.remove(CONTEXT_ATTRIBUTE);
     finishSpanIfPresent(context, throwable);
   }
 
-  public static class SpanFinishingSubscriber<T> implements CoreSubscriber<T> {
+  public static class SpanFinishingSubscriber<T> implements CoreSubscriber<T>, Subscription {
 
     private final CoreSubscriber<? super T> subscriber;
     private final io.opentelemetry.context.Context otelContext;
     private final Context context;
+    private final AtomicBoolean completed = new AtomicBoolean();
+    private Subscription subscription;
 
     public SpanFinishingSubscriber(
         CoreSubscriber<? super T> subscriber, io.opentelemetry.context.Context otelContext) {
@@ -99,9 +101,10 @@ public class AdviceUtils {
     }
 
     @Override
-    public void onSubscribe(Subscription s) {
+    public void onSubscribe(Subscription subscription) {
+      this.subscription = subscription;
       try (Scope scope = otelContext.makeCurrent()) {
-        subscriber.onSubscribe(s);
+        subscriber.onSubscribe(this);
       }
     }
 
@@ -114,13 +117,17 @@ public class AdviceUtils {
 
     @Override
     public void onError(Throwable t) {
+      if (completed.compareAndSet(false, true)) {
         finishSpanIfPresent(otelContext, t);
+      }
       subscriber.onError(t);
     }
 
     @Override
     public void onComplete() {
+      if (completed.compareAndSet(false, true)) {
         finishSpanIfPresent(otelContext, null);
+      }
       subscriber.onComplete();
     }
 
@@ -128,5 +135,18 @@ public class AdviceUtils {
     public Context currentContext() {
       return context;
     }
+
+    @Override
+    public void request(long n) {
+      subscription.request(n);
+    }
+
+    @Override
+    public void cancel() {
+      if (completed.compareAndSet(false, true)) {
+        finishSpanIfPresent(otelContext, null);
+      }
+      subscription.cancel();
+    }
   }
 }