TopicRetryableStream async initialization

topic/src/main/java/tech/ydb/topic/impl/TopicRetryableStream.java

@@ -1,15 +1,19 @@
 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;
 import java.util.concurrent.atomic.AtomicReference;
 
+import javax.annotation.Nullable;
+
 import com.google.protobuf.Message;
 import org.slf4j.Logger;
 
 import tech.ydb.common.retry.RetryConfig;
 import tech.ydb.common.retry.RetryPolicy;
+import tech.ydb.core.Result;
 import tech.ydb.core.Status;
 import tech.ydb.core.StatusCode;
 
@@ -19,9 +23,9 @@ public abstract class TopicRetryableStream<R extends Message, W extends Message,
     private final RetryConfig retryConfig;
     private final ScheduledExecutorService scheduler;
 
-    private final AtomicReference<S> realStream = new AtomicReference<>();
+    private final AtomicReference<State> state = new AtomicReference<>();
     private final AtomicInteger streamCount = new AtomicInteger(0);
-    private final RetryState state = new RetryState();
+    private final RetryState retryState = new RetryState();
 
     private volatile boolean isClosed = false;
 
@@ -32,12 +36,29 @@ public TopicRetryableStream(Logger logger, String debugId, RetryConfig config, S
         this.scheduler = scheduler;
     }
 
-    protected abstract S 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 result
+     */
+    protected abstract CompletableFuture<Result<S>> createNewStream(String debugId);
 
     protected abstract void onNext(S stream, R message);
 
-    protected abstract void onRetry(S stream, Status status);
-    protected abstract void onClose(S stream, Status status);
+    /**
+     * @param stream the stopped stream, or {@code null} when stream creation itself failed
+     * @param status status the stream stopped with
+     */
+    protected abstract void onRetry(@Nullable S stream, Status status);
+
+    /**
+     * @param stream the closed stream, or {@code null} when stream creation itself failed
+     * @param status status the stream was closed with
+     */
+    protected abstract void onClose(@Nullable S stream, Status status);
 
     public void start() {
         if (isClosed) {
@@ -46,46 +67,40 @@ public void start() {
         }
 
         String streamID = debugId + '.' + streamCount.incrementAndGet();
-        S stream = createNewStream(streamID);
+        State newState = new State(createNewStream(streamID));
 
-        if (!realStream.compareAndSet(null, stream)) {
+        if (!state.compareAndSet(null, newState)) {
             logger.warn("[{}] double start of stream, skipping", debugId);
+            newState.closeWithoutStart();
             return;
         }
 
-        stream.start(msg -> onNext(stream, msg)).whenComplete((status, th) -> {
-            if (!realStream.compareAndSet(stream, null)) {
-                return;
-            }
-            if (status != null) {
-                onStreamStop(stream, status, retryConfig.getStatusRetryPolicy(status));
-            }
-            if (th != null) {
-                Status wrapped = Status.of(StatusCode.CLIENT_INTERNAL_ERROR, th);
-                onStreamStop(stream, wrapped, retryConfig.getThrowableRetryPolicy(th));
-            }
-        });
+        newState.start();
+
+        if (isClosed) {
+            close();
+        }
     }
 
     protected void resetRetries() {
-        state.reset();
+        retryState.reset();
     }
 
     public boolean isClosed() {
         return isClosed;
     }
 
     public void fail(Status status) {
-        S closed = realStream.getAndSet(null);
-        if (closed != null) {
+        State local = state.getAndSet(null);
+        if (local != null) {
             logger.warn("[{}] failed by application-side error {}", debugId, status);
-            closed.close();
-            onStreamStop(closed, status, retryConfig.getStatusRetryPolicy(status));
+            onStreamStop(local.closeStream(), status, retryConfig.getStatusRetryPolicy(status));
         }
     }
 
     public void send(W msg) {
-        S stream = realStream.get();
+        State local = state.get();
+        S stream = local != null ? local.stream.get() : null;
         if (stream == null) {
             logger.warn("[{}] send message before stream is ready", debugId);
             return;
@@ -95,13 +110,12 @@ public void send(W msg) {
 
     public boolean close() {
         isClosed = true;
-        S stream = realStream.getAndSet(null);
-        if (stream == null) {
+        State local = state.getAndSet(null);
+        if (local == null) {
             return false;
         }
 
-        stream.close();
-        onStreamStop(stream, Status.SUCCESS, null);
+        onStreamStop(local.closeStream(), Status.SUCCESS, null);
         return true;
     }
 
@@ -118,7 +132,7 @@ private void onStreamStop(S closed, Status status, RetryPolicy policy) {
             return;
         }
 
-        long nextRetryMs = state.nextRetryMs(policy);
+        long nextRetryMs = retryState.nextRetryMs(policy);
 
         if (nextRetryMs < 0) {
             logger.warn("[{}] stopped after retry policy evaluation for status {}", debugId, status);
@@ -128,14 +142,14 @@ private void onStreamStop(S closed, Status status, RetryPolicy policy) {
         }
 
         if (nextRetryMs == 0) { // retry immediately
-            logger.warn("[{}] retry #{}. Retry immediately...", debugId, state.retryNumber());
+            logger.warn("[{}] retry #{}. Retry immediately...", debugId, retryState.retryNumber());
             onRetry(closed, status);
             start();
             return;
         }
 
         // retry scheduling
-        logger.warn("[{}] retry #{}. Scheduling reconnect in {}ms...", debugId, state.retryNumber(), nextRetryMs);
+        logger.warn("[{}] retry #{}. Scheduling reconnect in {}ms...", debugId, retryState.retryNumber(), nextRetryMs);
         onRetry(closed, status);
 
         try {
@@ -147,6 +161,77 @@ private void onStreamStop(S closed, Status status, RetryPolicy policy) {
         }
     }
 
+    private class State {
+        private final CompletableFuture<Result<S>> future;
+        private final AtomicReference<S> stream = new AtomicReference<>();
+
+        State(CompletableFuture<Result<S>> future) {
+            this.future = future;
+        }
+
+        public void start() {
+            this.future.whenComplete((res, th) -> {
+                if (res == null) {
+                    if (state.compareAndSet(this, null)) {
+                        Status wrapped = Status.of(StatusCode.CLIENT_INTERNAL_ERROR, th);
+                        onStreamStop(null, wrapped, retryConfig.getThrowableRetryPolicy(th));
+                    }
+                    return;
+                }
+
+                if (!res.isSuccess()) {
+                    if (state.compareAndSet(this, null)) {
+                        onStreamStop(null, res.getStatus(), retryConfig.getStatusRetryPolicy(res.getStatus()));
+                    }
+                    return;
+                }
+
+                startStream(res.getValue());
+            });
+        }
+
+        public void closeWithoutStart() {
+            this.future.whenComplete((res, th) -> {
+                if (res != null && res.isSuccess()) {
+                    res.getValue().close();
+                }
+            });
+        }
+
+        private void startStream(S local) {
+            if (state.get() != this) {
+                local.close();
+                return;
+            }
+
+            stream.set(local);
+            local.start(msg -> onNext(local, msg)).whenComplete((status, th) -> {
+                if (!state.compareAndSet(this, null)) {
+                    return;
+                }
+
+                if (status != null) {
+                    onStreamStop(local, status, retryConfig.getStatusRetryPolicy(status));
+                } else {
+                    Status wrapped = Status.of(StatusCode.CLIENT_INTERNAL_ERROR, th);
+                    onStreamStop(local, wrapped, retryConfig.getThrowableRetryPolicy(th));
+                }
+            });
+
+            if (state.get() != this) {
+                closeStream();
+            }
+        }
+
+        public S closeStream() {
+            S local = stream.getAndSet(null);
+            if (local != null) {
+                local.close();
+            }
+            return local;
+        }
+    }
+
     private static class RetryState {
         private final AtomicInteger count = new AtomicInteger();
         private volatile long startedAt = 0;

topic/src/main/java/tech/ydb/topic/read/impl/ReaderImpl.java

@@ -14,6 +14,7 @@
 
 import tech.ydb.common.transaction.YdbTransaction;
 import tech.ydb.core.Issue;
+import tech.ydb.core.Result;
 import tech.ydb.core.Status;
 import tech.ydb.core.StatusCode;
 import tech.ydb.core.grpc.GrpcRequestSettings;
@@ -76,8 +77,9 @@ public ReaderImpl(TopicRpc rpc, String id, ReaderSettings settings, ReadConfig c
     }
 
     @Override
-    protected ReadSession createNewStream(String id) {
-        return new ReadSession(id, rpc.readSession(id), initRequest, handler::handleDataReceivedEvent, config);
+    protected CompletableFuture<Result<ReadSession>> createNewStream(String id) {
+        ReadSession s = new ReadSession(id, rpc.readSession(id), initRequest, handler::handleDataReceivedEvent, config);
+        return CompletableFuture.completedFuture(Result.success(s));
     }
 
     @Override
@@ -91,7 +93,9 @@ protected void onRetry(ReadSession stream, Status status) {
                 logger.error("[{}] errorHandler onRetry processing throws exception", debugId, ex);
             }
         }
-        stream.closeAll().forEach(ps -> handler.handleClosePartitionSession(ps));
+        if (stream != null) {
+            stream.closeAll().forEach(ps -> handler.handleClosePartitionSession(ps));
+        }
     }
 
     @Override
@@ -109,7 +113,9 @@ protected void onClose(ReadSession stream, Status status) {
                 logger.error("[{}] errorHandler onClose processing throws exception", debugId, ex);
             }
         }
-        stream.closeAll().forEach(ps -> handler.handleClosePartitionSession(ps));
+        if (stream != null) {
+            stream.closeAll().forEach(ps -> handler.handleClosePartitionSession(ps));
+        }
         handler.handleReaderClosed(status);
     }
 

topic/src/main/java/tech/ydb/topic/write/impl/WriteSession.java

@@ -1,11 +1,13 @@
 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;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import tech.ydb.core.Result;
 import tech.ydb.core.Status;
 import tech.ydb.core.utils.ProtobufUtils;
 import tech.ydb.proto.topic.YdbTopic;
@@ -48,7 +50,7 @@ public WriteSession(String debugId, WriteStreamFactory factory, WriterSettings s
     }
 
     @Override
-    protected Stream createNewStream(String id) {
+    protected CompletableFuture<Result<Stream>> createNewStream(String id) {
         return streamFactory.createNewStream(id);
     }
 

topic/src/main/java/tech/ydb/topic/write/impl/WriteStream.java

@@ -11,7 +11,6 @@
 import tech.ydb.proto.topic.YdbTopic.StreamWriteMessage.FromServer;
 import tech.ydb.proto.topic.YdbTopic.UpdateTokenRequest;
 import tech.ydb.topic.impl.TopicStreamBase;
-import tech.ydb.topic.impl.TopicStreamFail;
 
 /**
  *
@@ -35,10 +34,4 @@ protected FromClient updateTokenMessage(String token) {
     protected Status parseMessageStatus(FromServer message) {
         return Status.of(StatusCode.fromProto(message.getStatus()), Issue.fromPb(message.getIssuesList()));
     }
-
-    public static class Fail extends TopicStreamFail<FromServer, FromClient> implements WriteSession.Stream {
-        public Fail(String id, Status status) {
-            super(logger, id, status);
-        }
-    }
 }