Fix flaky spring webflux tests (#3150)
This commit is contained in:
parent
c3c5c11acb
commit
d2906841bd
|
@ -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();
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in New Issue