From ecb00ef257b2783070f35ca4e92f3f196dde5f8a Mon Sep 17 00:00:00 2001 From: Igor Melnichenko Date: Sat, 5 Sep 2026 15:44:50 +0300 Subject: [PATCH] Make topic write stream creation non-blocking TopicRetryableStream.start() is called from the shared transport scheduler on every reconnect. With directWrite enabled, WriteStreamDirectFactory resolved the target partition and its location synchronously inside createNewStream: lookupPartitionId() joined a probe stream future (1 min deadline) and lookupLocation() joined describeTopic() (1 min deadline). Each reconnect of an unresponsive destination could therefore occupy a scheduler thread for up to two minutes. The shared scheduler is sized max(cores / 2, 2) and is also used by discovery, session pools, retry contexts and operation tray, so a handful of stalled writers could stall the whole transport: session acquire timeouts stop firing and discovery ticks stop running. Make createNewStream() return CompletableFuture and compose the partition and location lookups instead of joining them, so no shared scheduler thread is held while a stream is being created. Since stream creation is now asynchronous, close() may happen while it is in progress. TopicRetryableStream handles that by re-checking isClosed after publishing the new stream: close() sets the volatile flag before clearing the stream reference, so a creation that wins the race always observes the flag and drops the stream without starting it. Co-Authored-By: Claude Opus 5 --- .../ydb/topic/impl/TopicRetryableStream.java | 33 +++++++- .../ydb/topic/write/impl/WriteSession.java | 3 +- .../write/impl/WriteStreamDirectFactory.java | 79 +++++++++++++------ .../topic/write/impl/WriteStreamFactory.java | 13 ++- .../topic/impl/TopicRetryableStreamTest.java | 74 ++++++++++++++++- .../impl/WriteStreamDirectFactoryTest.java | 22 +++--- .../write/impl/WriteStreamFactoryTest.java | 2 +- 7 files changed, 180 insertions(+), 46 deletions(-) diff --git a/topic/src/main/java/tech/ydb/topic/impl/TopicRetryableStream.java b/topic/src/main/java/tech/ydb/topic/impl/TopicRetryableStream.java index 3cdb97368..b030ed5ed 100644 --- a/topic/src/main/java/tech/ydb/topic/impl/TopicRetryableStream.java +++ b/topic/src/main/java/tech/ydb/topic/impl/TopicRetryableStream.java @@ -1,5 +1,6 @@ package tech.ydb.topic.impl; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; @@ -12,6 +13,7 @@ import tech.ydb.common.retry.RetryPolicy; import tech.ydb.core.Status; import tech.ydb.core.StatusCode; +import tech.ydb.core.utils.FutureTools; public abstract class TopicRetryableStream { protected final String debugId; @@ -32,7 +34,15 @@ public TopicRetryableStream(Logger logger, String debugId, RetryConfig config, S this.scheduler = scheduler; } - protected abstract TopicStream createNewStream(String debugId); + /** + * Creates a new stream. Implementations must not block the calling thread: this method is invoked from the shared + * scheduler on every reconnect, and blocking there stalls discovery, session pools and timeouts of the whole + * transport. + * + * @param debugId identifier of the new stream for logging + * @return future with the new stream + */ + protected abstract CompletableFuture> createNewStream(String debugId); protected abstract void onNext(R message); @@ -45,13 +55,32 @@ public void start() { } String streamID = debugId + '.' + streamCount.incrementAndGet(); - TopicStream stream = createNewStream(streamID); + createNewStream(streamID).whenComplete((stream, throwable) -> { + if (throwable != null) { + // creation may be composed of several futures, so the error comes wrapped in a CompletionException + Throwable cause = FutureTools.unwrapCompletionException(throwable); + logger.warn("[{}] cannot create stream", debugId, cause); + Status errorStatus = Status.of(StatusCode.CLIENT_INTERNAL_ERROR, cause); + onStreamStop(errorStatus, retryConfig.getThrowableRetryPolicy(cause)); + return; + } + + startStream(stream); + }); + } + private void startStream(TopicStream stream) { if (!realStream.compareAndSet(null, stream)) { logger.warn("[{}] double start of stream, skipping", debugId); return; } + // stream creation is asynchronous, so close() may have happened while it was in progress + if (isClosed && realStream.compareAndSet(stream, null)) { + logger.info("[{}] stream was closed while it was creating, skipping", debugId); + return; + } + stream.start(this::onNext).whenComplete((status, th) -> { realStream.compareAndSet(stream, null); if (status != null) { diff --git a/topic/src/main/java/tech/ydb/topic/write/impl/WriteSession.java b/topic/src/main/java/tech/ydb/topic/write/impl/WriteSession.java index ee716e43b..49af8b070 100644 --- a/topic/src/main/java/tech/ydb/topic/write/impl/WriteSession.java +++ b/topic/src/main/java/tech/ydb/topic/write/impl/WriteSession.java @@ -1,5 +1,6 @@ package tech.ydb.topic.write.impl; import java.util.List; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.ScheduledExecutorService; import java.util.function.BiConsumer; @@ -48,7 +49,7 @@ public WriteSession(String debugId, WriteStreamFactory factory, WriterSettings s } @Override - protected Stream createNewStream(String id) { + protected CompletableFuture createNewStream(String id) { return streamFactory.createNewStream(id); } diff --git a/topic/src/main/java/tech/ydb/topic/write/impl/WriteStreamDirectFactory.java b/topic/src/main/java/tech/ydb/topic/write/impl/WriteStreamDirectFactory.java index c3cf7e425..cdc384b2b 100644 --- a/topic/src/main/java/tech/ydb/topic/write/impl/WriteStreamDirectFactory.java +++ b/topic/src/main/java/tech/ydb/topic/write/impl/WriteStreamDirectFactory.java @@ -36,17 +36,24 @@ public WriteStreamDirectFactory(TopicRpc rpc, WriterSettings settings) { } @Override - public WriteSession.Stream createNewStream(String id) { - Long targetPartitionId = partitionId; - if (targetPartitionId == null) { - Result pid = lookupPartitionId(id); - if (!pid.isSuccess()) { - return new WriteStream.Fail(id, pid.getStatus()); + public CompletableFuture createNewStream(String id) { + CompletableFuture> partitionIdLookup = partitionId == null + ? lookupPartitionId(id) + : CompletableFuture.completedFuture(Result.success(partitionId)); + + return partitionIdLookup.thenCompose(lookupResult -> { + if (!lookupResult.isSuccess()) { + return CompletableFuture.completedFuture(new WriteStream.Fail(id, lookupResult.getStatus())); } - targetPartitionId = pid.getValue(); - } - Result location = lookupLocation(id, targetPartitionId); + long targetPartitionId = lookupResult.getValue(); + return lookupLocation(id, targetPartitionId) + .thenApply(location -> buildDirectStream(id, targetPartitionId, location)); + }); + } + + private WriteSession.Stream buildDirectStream(String id, long targetPartitionId, + Result location) { if (!location.isSuccess()) { return new WriteStream.Fail(id, location.getStatus()); } @@ -73,13 +80,19 @@ public WriteSession.Stream createNewStream(String id) { return new WriteStream(id, rpc.writeSession(settings), init); } - protected Result lookupLocation(String id, long targetPartitionId) { + protected CompletableFuture> lookupLocation(String id, long targetPartitionId) { logger.info("[{}] describe topic {} to look up node for partition {}", id, topicPath, targetPartitionId); - Result describeTopic = rpc.describeTopic( + return rpc.describeTopic( YdbTopic.DescribeTopicRequest.newBuilder().setIncludeLocation(true).setPath(topicPath).build(), GrpcRequestSettings.newBuilder().withDeadline(Duration.ofMinutes(1)).build() - ).join(); + ).thenApply(describeTopic -> parseLocation(id, targetPartitionId, describeTopic)); + } + private Result parseLocation( + String id, + long targetPartitionId, + Result describeTopic + ) { if (!describeTopic.isSuccess()) { logger.warn("[{}] describe topic {} failed with status {}", id, topicPath, describeTopic.getStatus()); return Result.fail(describeTopic.getStatus()); @@ -103,7 +116,7 @@ protected Result lookupLocation(String id, long targ return Result.fail(Status.of(StatusCode.BAD_REQUEST, issue)); } - private Result lookupPartitionId(String id) { + private CompletableFuture> lookupPartitionId(String id) { CompletableFuture> pidFuture = new CompletableFuture<>(); // create one-shot stream to detect partitionID for this producer @@ -141,26 +154,40 @@ private Result lookupPartitionId(String id) { if (streamFuture.isDone()) { logger.warn("[{}] probe stream to topic {} with producer {} failed with status {}", id, topicPath, producerId, streamFuture.join()); - return Result.fail(streamFuture.join()); + return CompletableFuture.completedFuture(Result.fail(streamFuture.join())); } - try { - streamFuture.whenComplete((st, th) -> { - Status status = st != null ? st : Status.of(StatusCode.CLIENT_INTERNAL_ERROR, th); - if (pidFuture.complete(Result.fail(status))) { - logger.warn("[{}] probe stream to topic {} with producer {} failed with status {}", id, topicPath, + streamFuture.whenComplete((st, th) -> { + Status status = st != null ? st : Status.of(StatusCode.CLIENT_INTERNAL_ERROR, th); + if (pidFuture.complete(Result.fail(status))) { + logger.warn("[{}] probe stream to topic {} with producer {} failed with status {}", id, topicPath, producerId, status); - } - }); + } + }); + + // the probe stream is closed as soon as the partition is known, whichever thread completes the future + CompletableFuture> result = pidFuture.whenComplete((__, ___) -> { + if (!streamFuture.isDone()) { + stream.close(); + } + }); + + try { YdbTopic.StreamWriteMessage.FromClient init = YdbTopic.StreamWriteMessage.FromClient.newBuilder() .setInitRequest(buildInitRequest()) .build(); stream.sendNext(init); - return pidFuture.join(); - } finally { - if (!streamFuture.isDone()) { - stream.close(); - } + } catch (Throwable throwable) { + logger.warn( + "[{}] cannot send init request to probe stream of topic {} with producer {}", + id, + topicPath, + producerId, + throwable + ); + pidFuture.complete(Result.fail(Status.of(StatusCode.CLIENT_INTERNAL_ERROR, throwable))); } + + return result; } } diff --git a/topic/src/main/java/tech/ydb/topic/write/impl/WriteStreamFactory.java b/topic/src/main/java/tech/ydb/topic/write/impl/WriteStreamFactory.java index 2cc0459d5..04a98f36b 100644 --- a/topic/src/main/java/tech/ydb/topic/write/impl/WriteStreamFactory.java +++ b/topic/src/main/java/tech/ydb/topic/write/impl/WriteStreamFactory.java @@ -1,6 +1,6 @@ package tech.ydb.topic.write.impl; - +import java.util.concurrent.CompletableFuture; import tech.ydb.proto.topic.YdbTopic.StreamWriteMessage; import tech.ydb.proto.topic.YdbTopic.StreamWriteMessage.FromClient; @@ -52,8 +52,15 @@ public StreamWriteMessage.InitRequest buildInitRequest() { return req.build(); } - public WriteSession.Stream createNewStream(String id) { + /** + * Creates a new write stream. The returned future may be completed asynchronously, the method itself never blocks + * the caller. + * + * @param id identifier of the new stream for logging + * @return future with the new stream + */ + public CompletableFuture createNewStream(String id) { FromClient init = FromClient.newBuilder().setInitRequest(buildInitRequest()).build(); - return new WriteStream(id, rpc.writeSession(id), init); + return CompletableFuture.completedFuture(new WriteStream(id, rpc.writeSession(id), init)); } } diff --git a/topic/src/test/java/tech/ydb/topic/impl/TopicRetryableStreamTest.java b/topic/src/test/java/tech/ydb/topic/impl/TopicRetryableStreamTest.java index f09eab067..9c2db1375 100644 --- a/topic/src/test/java/tech/ydb/topic/impl/TopicRetryableStreamTest.java +++ b/topic/src/test/java/tech/ydb/topic/impl/TopicRetryableStreamTest.java @@ -2,6 +2,7 @@ import java.util.ArrayList; import java.util.Arrays; +import java.util.Collections; import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ScheduledExecutorService; @@ -79,14 +80,26 @@ private static class TestStream extends TopicRetryableStream { final List closeStatuses = new ArrayList<>(); final List receivedMessages = new ArrayList<>(); + /** If set, the next createNewStream() returns this future instead of the next handle */ + CompletableFuture> pendingCreation = null; + TestStream(List handles, RetryConfig retryConfig, ScheduledExecutorService scheduler) { super(logger, "test", retryConfig, scheduler); this.handles = handles; } @Override - protected TopicStream createNewStream(String debugId) { - return handles.get(handleIndex++).stream; + protected CompletableFuture> createNewStream(String debugId) { + if (pendingCreation != null) { + CompletableFuture> future = pendingCreation; + pendingCreation = null; + return future; + } + + StreamHandle handle = handles.get(handleIndex); + handleIndex++; + + return CompletableFuture.completedFuture(handle.stream); } @Override @@ -165,6 +178,63 @@ public void startAfterCloseTest() { retryable.start(); // nothing } + @Test + public void asyncStreamCreationTest() { + StreamHandle streamHandle = new StreamHandle(); + TestStream retryable = new TestStream(Collections.emptyList(), RetryConfig.noRetries(), mockScheduler()); + + CompletableFuture> creation = new CompletableFuture<>(); + retryable.pendingCreation = creation; + + retryable.start(); // must return without waiting for the creation future + Mockito.verify(streamHandle.grpc, Mockito.never()).start(Mockito.any()); + + retryable.send(EMPTY); // stream is not ready yet, message is skipped + Mockito.verify(streamHandle.grpc, Mockito.never()).sendNext(Mockito.any()); + + creation.complete(streamHandle.stream); + + Mockito.verify(streamHandle.grpc).start(Mockito.any()); + retryable.send(EMPTY); + Mockito.verify(streamHandle.grpc, Mockito.times(2)).sendNext(EMPTY); // init + sent request + + Assert.assertTrue(retryable.close()); + Mockito.verify(streamHandle.grpc).close(); + } + + @Test + public void closeWhileStreamIsCreatingTest() { + StreamHandle streamHandle = new StreamHandle(); + TestStream retryable = new TestStream(Collections.emptyList(), RetryConfig.noRetries(), mockScheduler()); + + CompletableFuture> creation = new CompletableFuture<>(); + retryable.pendingCreation = creation; + + retryable.start(); + Assert.assertFalse(retryable.close()); // there is no stream to close yet + + creation.complete(streamHandle.stream); // the created stream must not be started + + Mockito.verify(streamHandle.grpc, Mockito.never()).start(Mockito.any()); + Assert.assertFalse(retryable.close()); + } + + @Test + @HideLoggers({TopicRetryableStreamTest.class}) + public void streamCreationFailedTest() { + TestStream retryable = new TestStream(Collections.emptyList(), RetryConfig.noRetries(), mockScheduler()); + + CompletableFuture> creation = new CompletableFuture<>(); + retryable.pendingCreation = creation; + + retryable.start(); + creation.completeExceptionally(new RuntimeException("cannot create stream")); + + Assert.assertEquals(1, retryable.closeStatuses.size()); + Assert.assertEquals(StatusCode.CLIENT_INTERNAL_ERROR, retryable.closeStatuses.get(0).getCode()); + Assert.assertTrue(retryable.retryStatuses.isEmpty()); + } + @Test public void sendBeforeStartIsIgnoredTest() { StreamHandle h = new StreamHandle(); diff --git a/topic/src/test/java/tech/ydb/topic/write/impl/WriteStreamDirectFactoryTest.java b/topic/src/test/java/tech/ydb/topic/write/impl/WriteStreamDirectFactoryTest.java index 9dbc2ed90..2f9ac15d7 100644 --- a/topic/src/test/java/tech/ydb/topic/write/impl/WriteStreamDirectFactoryTest.java +++ b/topic/src/test/java/tech/ydb/topic/write/impl/WriteStreamDirectFactoryTest.java @@ -136,7 +136,7 @@ public void directWriteByPartitionIdTest() { WriteStreamFactory factory = new WriteStreamDirectFactory(rpc, settings); Assert.assertEquals("/local/topic", factory.getTopicPath()); - WriteSession.Stream stream = factory.createNewStream("s1"); + WriteSession.Stream stream = factory.createNewStream("s1").join(); Assert.assertTrue(stream instanceof WriteStream); ArgumentCaptor options = ArgumentCaptor.forClass(GrpcRequestSettings.class); @@ -166,7 +166,7 @@ public void directWriteByPartitionIdTestDescribeFailTest() { .setDirectWrite(true) .build()); - WriteSession.Stream stream = factory.createNewStream("s1"); + WriteSession.Stream stream = factory.createNewStream("s1").join(); Mockito.verify(rpc, Mockito.never()).writeSession(Mockito.any(GrpcRequestSettings.class)); @@ -190,7 +190,7 @@ public void directWriteByPartitionIdTestPartitionNotFoundTest() { .setDirectWrite(true) .build()); - WriteSession.Stream stream = factory.createNewStream("s1"); + WriteSession.Stream stream = factory.createNewStream("s1").join(); Mockito.verify(rpc, Mockito.never()).writeSession(Mockito.any(GrpcRequestSettings.class)); @@ -217,7 +217,7 @@ public void directWriteByPartitionIdTestPartitionHasNoLocationTest() { .setDirectWrite(true) .build()); - WriteSession.Stream stream = factory.createNewStream("s1"); + WriteSession.Stream stream = factory.createNewStream("s1").join(); Mockito.verify(rpc, Mockito.never()).writeSession(Mockito.any(GrpcRequestSettings.class)); @@ -258,7 +258,7 @@ public void directWriteByProducerIdTest() { .setDirectWrite(true) .build()); - WriteSession.Stream stream = factory.createNewStream("s1"); + WriteSession.Stream stream = factory.createNewStream("s1").join(); Assert.assertTrue(stream instanceof WriteStream); ArgumentCaptor options = ArgumentCaptor.forClass(GrpcRequestSettings.class); @@ -292,7 +292,7 @@ public void directWriteByProducerIdProbeFailTest() { .setDirectWrite(true) .build()); - WriteSession.Stream stream = factory.createNewStream("s1"); + WriteSession.Stream stream = factory.createNewStream("s1").join(); Assert.assertTrue(stream instanceof WriteStream.Fail); Mockito.verify(rpc).writeSession(Mockito.any(GrpcRequestSettings.class)); @@ -316,7 +316,7 @@ public void directWriteByProducerIdProbeFailOnSendTest() { .setDirectWrite(true) .build()); - WriteSession.Stream stream = factory.createNewStream("s1"); + WriteSession.Stream stream = factory.createNewStream("s1").join(); Assert.assertTrue(stream instanceof WriteStream.Fail); Mockito.verify(rpc).writeSession(Mockito.any(GrpcRequestSettings.class)); @@ -340,7 +340,7 @@ public void directWriteByProducerIdProbeExceptionOnSendTest() { .setDirectWrite(true) .build()); - WriteSession.Stream stream = factory.createNewStream("s1"); + WriteSession.Stream stream = factory.createNewStream("s1").join(); Assert.assertTrue(stream instanceof WriteStream.Fail); Mockito.verify(rpc).writeSession(Mockito.any(GrpcRequestSettings.class)); @@ -370,7 +370,7 @@ public void directWriteByProducerIdProbeWrongResponseTest() { .setDirectWrite(true) .build()); - WriteSession.Stream stream = factory.createNewStream("s1"); + WriteSession.Stream stream = factory.createNewStream("s1").join(); Assert.assertTrue(stream instanceof WriteStream.Fail); Mockito.verify(rpc).writeSession(Mockito.any(GrpcRequestSettings.class)); @@ -398,7 +398,7 @@ public void directWriteByProducerIdProbeUnexpectedResponseTest() { .setDirectWrite(true) .build()); - WriteSession.Stream stream = factory.createNewStream("s1"); + WriteSession.Stream stream = factory.createNewStream("s1").join(); Assert.assertTrue(stream instanceof WriteStream.Fail); Mockito.verify(rpc).writeSession(Mockito.any(GrpcRequestSettings.class)); @@ -433,7 +433,7 @@ public void directWriteByProducerIdPartitionNotFoundTest() { .setDirectWrite(true) .build()); - WriteSession.Stream stream = factory.createNewStream("s1"); + WriteSession.Stream stream = factory.createNewStream("s1").join(); Assert.assertTrue(stream instanceof WriteStream.Fail); CompletableFuture res = stream.start(null); Assert.assertTrue(res.isDone()); diff --git a/topic/src/test/java/tech/ydb/topic/write/impl/WriteStreamFactoryTest.java b/topic/src/test/java/tech/ydb/topic/write/impl/WriteStreamFactoryTest.java index 267fdae10..67d89bcec 100644 --- a/topic/src/test/java/tech/ydb/topic/write/impl/WriteStreamFactoryTest.java +++ b/topic/src/test/java/tech/ydb/topic/write/impl/WriteStreamFactoryTest.java @@ -31,7 +31,7 @@ public void regularWriteTest() { WriteStreamFactory factory = new WriteStreamFactory(rpc, settings); Assert.assertEquals("/local/topic", factory.getTopicPath()); - WriteSession.Stream stream = factory.createNewStream("s1"); + WriteSession.Stream stream = factory.createNewStream("s1").join(); Assert.assertTrue(stream instanceof WriteStream); Mockito.verify(rpc).writeSession("s1"); }