diff --git a/rsocket-test/src/main/java/io/rsocket/test/TransportTest.java b/rsocket-test/src/main/java/io/rsocket/test/TransportTest.java index 725d621f5..570a7de2f 100644 --- a/rsocket-test/src/main/java/io/rsocket/test/TransportTest.java +++ b/rsocket-test/src/main/java/io/rsocket/test/TransportTest.java @@ -57,6 +57,7 @@ import org.reactivestreams.Subscription; import reactor.core.CoreSubscriber; import reactor.core.Disposable; +import reactor.core.Disposables; import reactor.core.Exceptions; import reactor.core.Fuseable; import reactor.core.publisher.Flux; @@ -710,6 +711,7 @@ private static class DisconnectingDuplexConnection implements DuplexConnection { private final String tag; final DuplexConnection source; final Duration delay; + final Disposable.Swap disposables = Disposables.swap(); DisconnectingDuplexConnection(String tag, DuplexConnection source, Duration delay) { this.tag = tag; @@ -719,6 +721,7 @@ private static class DisconnectingDuplexConnection implements DuplexConnection { @Override public void dispose() { + disposables.dispose(); source.dispose(); } @@ -747,14 +750,15 @@ public Flux receive() { bb -> { if (!receivedFirst) { receivedFirst = true; - Mono.delay(delay) - .takeUntilOther(source.onClose()) - .subscribe( - __ -> { - logger.warn( - "Tag {}. Disposing Connection[{}]", tag, source.hashCode()); - source.dispose(); - }); + disposables.replace( + Mono.delay(delay) + .takeUntilOther(source.onClose()) + .subscribe( + __ -> { + logger.warn( + "Tag {}. Disposing Connection[{}]", tag, source.hashCode()); + source.dispose(); + })); } }); }