Fix flaky spring webflux tests (#3150)

This commit is contained in:
Lauri Tulmin 2021-05-31 18:20:25 +03:00 committed by GitHub
parent c3c5c11acb
commit d2906841bd
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 26 additions and 6 deletions

View File

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