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"); }