Making topic write stream creation non-blocking

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<R extends Message, W extends Message> {
     protected final String debugId;
@@ -32,7 +34,15 @@ public TopicRetryableStream(Logger logger, String debugId, RetryConfig config, S
         this.scheduler = scheduler;
     }
 
-    protected abstract TopicStream<R, W> 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<? extends TopicStream<R, W>> createNewStream(String debugId);
 
     protected abstract void onNext(R message);
 
@@ -45,13 +55,32 @@ public void start() {
         }
 
         String streamID = debugId + '.' + streamCount.incrementAndGet();
-        TopicStream<R, W> 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<R, W> 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) {

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<Stream> createNewStream(String id) {
         return streamFactory.createNewStream(id);
     }
 

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<Long> pid = lookupPartitionId(id);
-            if (!pid.isSuccess()) {
-                return new WriteStream.Fail(id, pid.getStatus());
+    public CompletableFuture<WriteSession.Stream> createNewStream(String id) {
+        CompletableFuture<Result<Long>> 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<YdbTopic.PartitionLocation> 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<YdbTopic.PartitionLocation> 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<YdbTopic.PartitionLocation> lookupLocation(String id, long targetPartitionId) {
+    protected CompletableFuture<Result<YdbTopic.PartitionLocation>> lookupLocation(String id, long targetPartitionId) {
         logger.info("[{}] describe topic {} to look up node for partition {}", id, topicPath, targetPartitionId);
-        Result<YdbTopic.DescribeTopicResult> 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<YdbTopic.PartitionLocation> parseLocation(
+            String id,
+            long targetPartitionId,
+            Result<YdbTopic.DescribeTopicResult> 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<YdbTopic.PartitionLocation> lookupLocation(String id, long targ
         return Result.fail(Status.of(StatusCode.BAD_REQUEST, issue));
     }
 
-    private Result<Long> lookupPartitionId(String id) {
+    private CompletableFuture<Result<Long>> lookupPartitionId(String id) {
         CompletableFuture<Result<Long>> pidFuture = new CompletableFuture<>();
 
         // create one-shot stream to detect partitionID for this producer
@@ -141,26 +154,40 @@ private Result<Long> 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<Long>> 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;
     }
 }

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<WriteSession.Stream> 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));
     }
 }

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<Empty, Empty> {
         final List<Status> closeStatuses = new ArrayList<>();
         final List<Empty> receivedMessages = new ArrayList<>();
 
+        /** If set, the next createNewStream() returns this future instead of the next handle */
+        CompletableFuture<TopicStream<Empty, Empty>> pendingCreation = null;
+
         TestStream(List<StreamHandle> handles, RetryConfig retryConfig, ScheduledExecutorService scheduler) {
             super(logger, "test", retryConfig, scheduler);
             this.handles = handles;
         }
 
         @Override
-        protected TopicStream<Empty, Empty> createNewStream(String debugId) {
-            return handles.get(handleIndex++).stream;
+        protected CompletableFuture<TopicStream<Empty, Empty>> createNewStream(String debugId) {
+            if (pendingCreation != null) {
+                CompletableFuture<TopicStream<Empty, Empty>> 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<TopicStream<Empty, Empty>> 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<TopicStream<Empty, Empty>> 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<TopicStream<Empty, Empty>> 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();