From 0f004744161b8c3a741c85e25f1b5648e4b7e016 Mon Sep 17 00:00:00 2001 From: Abhishek Pal Date: Sun, 2 Aug 2026 15:40:58 +0530 Subject: [PATCH] RATIS-2632. Apply backpressure on install-snapshot chunk loop --- .../grpc/util/StreamObserverWithTimeout.java | 22 ++++++++ .../util/TestStreamObserverWithTimeout.java | 55 +++++++++++++++++++ 2 files changed, 77 insertions(+) diff --git a/ratis-grpc/src/main/java/org/apache/ratis/grpc/util/StreamObserverWithTimeout.java b/ratis-grpc/src/main/java/org/apache/ratis/grpc/util/StreamObserverWithTimeout.java index 3cc754e565..d49d5676a9 100644 --- a/ratis-grpc/src/main/java/org/apache/ratis/grpc/util/StreamObserverWithTimeout.java +++ b/ratis-grpc/src/main/java/org/apache/ratis/grpc/util/StreamObserverWithTimeout.java @@ -19,6 +19,7 @@ import org.apache.ratis.protocol.exceptions.TimeoutIOException; import org.apache.ratis.thirdparty.io.grpc.ClientInterceptor; +import org.apache.ratis.thirdparty.io.grpc.stub.CallStreamObserver; import org.apache.ratis.thirdparty.io.grpc.stub.StreamObserver; import org.apache.ratis.util.JavaUtils; import org.apache.ratis.util.ResourceSemaphore; @@ -95,11 +96,32 @@ private void acquire(StringSupplier request, TimeDuration timeout) { } } + /** + * Wait while the underlying stream is not ready providing backpressure + * in addition to the outstanding-request semaphore. + */ + private void awaitReady(StringSupplier request) { + if (!(observer instanceof CallStreamObserver)) { + return; + } + final CallStreamObserver callStreamObserver = (CallStreamObserver) observer; + while (!callStreamObserver.isReady() && !isClose.get()) { + try { + TimeDuration.ONE_MILLISECOND.sleep(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IllegalStateException(name + ": Interrupted while waiting for the stream to be ready to send " + + request, e); + } + } + } + @Override public void onNext(T request) { final StringSupplier requestString = StringSupplier.get(() -> requestToStringFunction.apply(request)); final TimeDuration timeout = timeoutSupplier.get(); acquire(requestString, timeout); + awaitReady(requestString); observer.onNext(request); final int id = requestCount.incrementAndGet(); LOG.debug("{}: send {} with timeout={}: {}", name, id, timeout, requestString); diff --git a/ratis-test/src/test/java/org/apache/ratis/grpc/util/TestStreamObserverWithTimeout.java b/ratis-test/src/test/java/org/apache/ratis/grpc/util/TestStreamObserverWithTimeout.java index b279736f39..e019f60d31 100644 --- a/ratis-test/src/test/java/org/apache/ratis/grpc/util/TestStreamObserverWithTimeout.java +++ b/ratis-test/src/test/java/org/apache/ratis/grpc/util/TestStreamObserverWithTimeout.java @@ -20,6 +20,7 @@ import org.apache.ratis.BaseTest; import org.apache.ratis.grpc.util.GrpcTestClient.StreamObserverFactory; import org.apache.ratis.thirdparty.io.grpc.StatusRuntimeException; +import org.apache.ratis.thirdparty.io.grpc.stub.CallStreamObserver; import org.apache.ratis.util.NetUtils; import org.apache.ratis.util.Slf4jUtils; import org.apache.ratis.util.StringUtils; @@ -30,9 +31,11 @@ import org.slf4j.event.Level; import java.util.ArrayList; +import java.util.Collections; import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; import java.util.function.Function; public class TestStreamObserverWithTimeout extends BaseTest { @@ -78,6 +81,58 @@ public void testWithTimeout() throws Exception { runTestTimeout(5, Type.WithTimeout); } + /** A {@link CallStreamObserver} whose readiness and delivered messages can be inspected. */ + private static final class ReadinessControlledObserver extends CallStreamObserver { + private volatile boolean ready; + private final List delivered = Collections.synchronizedList(new ArrayList<>()); + + ReadinessControlledObserver(boolean ready) { + this.ready = ready; + } + + void setReady(boolean isReady) { + this.ready = isReady; + } + + @Override public boolean isReady() { + return ready; + } + @Override public void onNext(String value) { + delivered.add(value); + } + @Override public void onError(Throwable t) { } + @Override public void onCompleted() { } + @Override public void setOnReadyHandler(Runnable onReadyHandler) { } + @Override public void disableAutoInboundFlowControl() { } + @Override public void request(int count) { } + @Override public void setMessageCompression(boolean enable) { } + } + + /** + * When the outstanding-request limit is 0 (semaphore disabled), the only backpressure is + * {@link CallStreamObserver#isReady()}. onNext must stall while the stream is not ready so that + * requests are not buffered. + */ + @Test + public void testOnNextWaitsForReadyWhenUnbounded() throws Exception { + final ReadinessControlledObserver observer = new ReadinessControlledObserver(false); + final StreamObserverWithTimeout withTimeout = StreamObserverWithTimeout.newInstance( + "test", Function.identity(), () -> TimeDuration.valueOf(60, TimeUnit.SECONDS), 0, + interceptor -> observer); + + final CompletableFuture sent = CompletableFuture.runAsync(() -> withTimeout.onNext("m1")); + + // Not ready: the request must be held back rather than delivered. + Thread.sleep(100); + Assertions.assertFalse(sent.isDone(), "onNext should block while the stream is not ready"); + Assertions.assertTrue(observer.delivered.isEmpty(), "request must not be sent while the stream is not ready"); + + // Becomes ready: the request should now be delivered and onNext should return. + observer.setReady(true); + sent.get(5, TimeUnit.SECONDS); + Assertions.assertEquals(Collections.singletonList("m1"), observer.delivered); + } + void runTestTimeout(int slow, Type type) throws Exception { LOG.info("slow = {}, {}", slow, type); final TimeDuration timeout = ONE_SECOND.multiply(0.5);