Fix flaky spring webflux tests
@@ -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();
+ }
}
}
