From 8840c95ad3f90a0a5a7f615d9eab117cff31bfaa Mon Sep 17 00:00:00 2001 From: Igor Melnichenko Date: Sat, 5 Sep 2026 15:44:50 +0300 Subject: [PATCH 1/3] 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 | 43 +++++++++- .../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, 190 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 e81268305..5c43490b2 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,11 +34,28 @@ 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 + */ + protected abstract CompletableFuture createNewStream(String debugId); protected abstract void onNext(S stream, R message); + /** + * @param stream the stopped stream, or {@code null} when stream creation itself failed + * @param status status the stream stopped with + */ protected abstract void onRetry(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(S stream, Status status); public void start() { @@ -46,13 +65,33 @@ public void start() { } String streamID = debugId + '.' + streamCount.incrementAndGet(); - S 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); + // there is no stream to report: creation is what failed + onStreamStop(null, errorStatus, retryConfig.getThrowableRetryPolicy(cause)); + return; + } + startStream(stream); + }); + } + + private void startStream(S 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(msg -> onNext(stream, msg)).whenComplete((status, th) -> { if (!realStream.compareAndSet(stream, null)) { return; 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 d99d9801e..ab571000d 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 26735242f..2cb1ce803 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 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 TopicStreamBase 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 @@ -195,6 +208,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"); } From 380415c376c7f0b1247ee66d213f8a0b925eb2cc Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Wed, 23 Sep 2026 12:13:06 +0100 Subject: [PATCH 2/3] Replaced TopicStreamFail by future with Result --- .../ydb/topic/impl/TopicRetryableStream.java | 101 +++++++++----- .../tech/ydb/topic/impl/TopicStreamFail.java | 38 ------ .../tech/ydb/topic/read/impl/ReaderImpl.java | 14 +- .../ydb/topic/write/impl/WriteSession.java | 3 +- .../ydb/topic/write/impl/WriteStream.java | 7 - .../write/impl/WriteStreamDirectFactory.java | 73 ++++------ .../topic/write/impl/WriteStreamFactory.java | 6 +- .../topic/impl/TopicRetryableStreamTest.java | 82 ++++++----- .../impl/WriteStreamDirectFactoryTest.java | 128 +++++++++--------- .../write/impl/WriteStreamFactoryTest.java | 2 +- 10 files changed, 209 insertions(+), 245 deletions(-) delete mode 100644 topic/src/main/java/tech/ydb/topic/impl/TopicStreamFail.java 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 5c43490b2..910809c7a 100644 --- a/topic/src/main/java/tech/ydb/topic/impl/TopicRetryableStream.java +++ b/topic/src/main/java/tech/ydb/topic/impl/TopicRetryableStream.java @@ -6,14 +6,16 @@ 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; -import tech.ydb.core.utils.FutureTools; public abstract class TopicRetryableStream> { protected final String debugId; @@ -21,10 +23,12 @@ public abstract class TopicRetryableStream realStream = new AtomicReference<>(); private final AtomicInteger streamCount = new AtomicInteger(0); private final RetryState state = new RetryState(); + private final AtomicReference realStreamId = new AtomicReference<>(); + private volatile S realStream = null; + private volatile boolean isClosed = false; public TopicRetryableStream(Logger logger, String debugId, RetryConfig config, ScheduledExecutorService scheduler) { @@ -40,9 +44,9 @@ public TopicRetryableStream(Logger logger, String debugId, RetryConfig config, S * transport. * * @param debugId identifier of the new stream for logging - * @return future with the new stream + * @return future with the new stream result */ - protected abstract CompletableFuture createNewStream(String debugId); + protected abstract CompletableFuture> createNewStream(String debugId); protected abstract void onNext(S stream, R message); @@ -50,13 +54,13 @@ public TopicRetryableStream(Logger logger, String debugId, RetryConfig config, S * @param stream the stopped stream, or {@code null} when stream creation itself failed * @param status status the stream stopped with */ - protected abstract void onRetry(S stream, Status status); + 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(S stream, Status status); + protected abstract void onClose(@Nullable S stream, Status status); public void start() { if (isClosed) { @@ -65,35 +69,45 @@ public void start() { } String streamID = debugId + '.' + streamCount.incrementAndGet(); - 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); - // there is no stream to report: creation is what failed - onStreamStop(null, errorStatus, retryConfig.getThrowableRetryPolicy(cause)); - return; - } - - startStream(stream); + createNewStream(streamID).whenComplete((result, th) -> { + tryStartStream(streamID, result, th); }); } - private void startStream(S stream) { - if (!realStream.compareAndSet(null, stream)) { - logger.warn("[{}] double start of stream, skipping", debugId); + private void tryStartStream(String streamID, Result result, Throwable th) { + if (isClosed) { + logger.info("[{}] stream was closed while it was creating, skipping", streamID); 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); + if (!realStreamId.compareAndSet(null, streamID)) { + logger.warn("[{}] double start of stream, skipping", streamID); return; } + if (result != null && result.isSuccess()) { + startStream(streamID, result.getValue()); + return; + } + + if (!realStreamId.compareAndSet(streamID, null)) { + return; + } + + if (result == null) { + logger.warn("[{}] cannot create stream with exeption", streamID, th); + Status wrapped = Status.of(StatusCode.CLIENT_INTERNAL_ERROR, th); + onStreamStop(null, wrapped, retryConfig.getThrowableRetryPolicy(th)); + } else { + logger.warn("[{}] cannot create stream with status {}", streamID, result.getStatus()); + onStreamStop(null, result.getStatus(), retryConfig.getStatusRetryPolicy(result.getStatus())); + } + } + + private void startStream(String streamID, S stream) { + realStream = stream; stream.start(msg -> onNext(stream, msg)).whenComplete((status, th) -> { - if (!realStream.compareAndSet(stream, null)) { + if (!realStreamId.compareAndSet(streamID, null)) { return; } if (status != null) { @@ -104,6 +118,11 @@ private void startStream(S stream) { onStreamStop(stream, wrapped, retryConfig.getThrowableRetryPolicy(th)); } }); + + if (isClosed) { // stream may be closed by other thread + realStream = null; + stream.close(); + } } protected void resetRetries() { @@ -115,32 +134,40 @@ public boolean isClosed() { } public void fail(Status status) { - S closed = realStream.getAndSet(null); - if (closed != null) { - logger.warn("[{}] failed by application-side error {}", debugId, status); - closed.close(); - onStreamStop(closed, status, retryConfig.getStatusRetryPolicy(status)); + String streamId = realStreamId.getAndSet(null); + if (streamId != null) { + logger.warn("[{}] failed by application-side error {}", streamId, status); + S local = realStream; + realStream = null; + if (local != null) { + local.close(); + } + onStreamStop(local, status, retryConfig.getStatusRetryPolicy(status)); } } public void send(W msg) { - S stream = realStream.get(); - if (stream == null) { + S local = realStream; + if (local == null) { logger.warn("[{}] send message before stream is ready", debugId); return; } - stream.send(msg); + local.send(msg); } public boolean close() { isClosed = true; - S stream = realStream.getAndSet(null); - if (stream == null) { + String streamId = realStreamId.getAndSet(null); + if (streamId == null) { return false; } - stream.close(); - onStreamStop(stream, Status.SUCCESS, null); + S local = realStream; + realStream = null; + if (local != null) { + local.close(); + } + onStreamStop(local, Status.SUCCESS, null); return true; } diff --git a/topic/src/main/java/tech/ydb/topic/impl/TopicStreamFail.java b/topic/src/main/java/tech/ydb/topic/impl/TopicStreamFail.java deleted file mode 100644 index 3d1965d7a..000000000 --- a/topic/src/main/java/tech/ydb/topic/impl/TopicStreamFail.java +++ /dev/null @@ -1,38 +0,0 @@ -package tech.ydb.topic.impl; - -import java.util.concurrent.CompletableFuture; -import java.util.function.Consumer; - -import com.google.protobuf.Message; -import org.slf4j.Logger; - -import tech.ydb.core.Status; - -public class TopicStreamFail implements TopicStream { - private final Logger logger; - private final String debugId; - - private final Status status; - - public TopicStreamFail(Logger logger, String debugId, Status status) { - this.logger = logger; - this.debugId = debugId; - this.status = status; - } - - @Override - public CompletableFuture start(Consumer messageHandler) { - return CompletableFuture.completedFuture(status); - } - - @Override - public void send(W req) { - logger.warn("[{}] is failed stream with status {}. Next message with type {} was NOT sent", debugId, status, - req.getDescriptorForType().getName()); - } - - @Override - public void close() { - logger.warn("[{}] is failed stream with status {}. It doesn't need to close", debugId, status); - } -} diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/ReaderImpl.java b/topic/src/main/java/tech/ydb/topic/read/impl/ReaderImpl.java index 02ce252d4..55d7ad8e4 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/ReaderImpl.java +++ b/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> 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); } 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 ab571000d..78d8444a1 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 @@ -7,6 +7,7 @@ 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; @@ -49,7 +50,7 @@ public WriteSession(String debugId, WriteStreamFactory factory, WriterSettings s } @Override - protected CompletableFuture createNewStream(String id) { + protected CompletableFuture> createNewStream(String id) { return streamFactory.createNewStream(id); } diff --git a/topic/src/main/java/tech/ydb/topic/write/impl/WriteStream.java b/topic/src/main/java/tech/ydb/topic/write/impl/WriteStream.java index 3f0b1ae89..d5f300a1b 100644 --- a/topic/src/main/java/tech/ydb/topic/write/impl/WriteStream.java +++ b/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 implements WriteSession.Stream { - public Fail(String id, Status status) { - super(logger, id, status); - } - } } 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 cdc384b2b..0f6034ca6 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,33 +36,23 @@ public WriteStreamDirectFactory(TopicRpc rpc, WriterSettings settings) { } @Override - public CompletableFuture createNewStream(String id) { - CompletableFuture> partitionIdLookup = partitionId == null + public CompletableFuture> createNewStream(String id) { + CompletableFuture> partition = 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())); - } - - long targetPartitionId = lookupResult.getValue(); - return lookupLocation(id, targetPartitionId) - .thenApply(location -> buildDirectStream(id, targetPartitionId, location)); - }); + return partition.thenCompose(Result.compose( + partID -> lookupLocation(id, partID) + .thenApply(r -> r.map(location -> buildDirectStream(id, partID, location))) + )); } - private WriteSession.Stream buildDirectStream(String id, long targetPartitionId, - Result location) { - if (!location.isSuccess()) { - return new WriteStream.Fail(id, location.getStatus()); - } - + private WriteSession.Stream buildDirectStream(String id, long partitionId, YdbTopic.PartitionLocation location) { StreamWriteMessage.InitRequest.Builder req = StreamWriteMessage.InitRequest.newBuilder() .setPath(topicPath) .setPartitionWithGeneration(YdbTopic.PartitionWithGeneration.newBuilder() - .setPartitionId(targetPartitionId) - .setGeneration(location.getValue().getGeneration()) + .setPartitionId(partitionId) + .setGeneration(location.getGeneration()) .build()); if (producerId != null) { @@ -74,7 +64,7 @@ private WriteSession.Stream buildDirectStream(String id, long targetPartitionId, .withTraceId(id) .disableDeadline() .withDirectMode(true) - .withPreferredNodeID(location.getValue().getNodeId()) + .withPreferredNodeID(location.getNodeId()) .build(); return new WriteStream(id, rpc.writeSession(settings), init); @@ -82,28 +72,26 @@ private WriteSession.Stream buildDirectStream(String id, long targetPartitionId, protected CompletableFuture> lookupLocation(String id, long targetPartitionId) { logger.info("[{}] describe topic {} to look up node for partition {}", id, topicPath, targetPartitionId); - return rpc.describeTopic( - YdbTopic.DescribeTopicRequest.newBuilder().setIncludeLocation(true).setPath(topicPath).build(), - GrpcRequestSettings.newBuilder().withDeadline(Duration.ofMinutes(1)).build() - ).thenApply(describeTopic -> parseLocation(id, targetPartitionId, describeTopic)); + YdbTopic.DescribeTopicRequest req = YdbTopic.DescribeTopicRequest.newBuilder() + .setIncludeLocation(true).setPath(topicPath) + .build(); + GrpcRequestSettings settings = GrpcRequestSettings.newBuilder().withDeadline(Duration.ofMinutes(1)).build(); + return rpc.describeTopic(req, settings).thenApply(res -> parseLocation(id, targetPartitionId, res)); } - 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()); + private Result parseLocation(String id, long partitionId, + Result description) { + if (!description.isSuccess()) { + logger.warn("[{}] describe topic {} failed with status {}", id, topicPath, description.getStatus()); + return Result.fail(description.getStatus()); } // lookup for partition location - for (YdbTopic.DescribeTopicResult.PartitionInfo partition : describeTopic.getValue().getPartitionsList()) { - if (partition.getPartitionId() == targetPartitionId) { + for (YdbTopic.DescribeTopicResult.PartitionInfo partition : description.getValue().getPartitionsList()) { + if (partition.getPartitionId() == partitionId) { if (!partition.hasPartitionLocation()) { - logger.warn("[{}] partition {} has no valid location info", id, targetPartitionId); - Issue issue = Issue.of("Partition " + targetPartitionId + " has no location", Issue.Severity.ERROR); + logger.warn("[{}] partition {} has no valid location info", id, partitionId); + Issue issue = Issue.of("Partition " + partitionId + " has no location", Issue.Severity.ERROR); return Result.fail(Status.of(StatusCode.BAD_REQUEST, issue)); } @@ -111,8 +99,8 @@ private Result parseLocation( } } - logger.warn("[{}] topic {} doesn't have partition {}, direct writing failed", id, topicPath, targetPartitionId); - Issue issue = Issue.of("Cannot find partition " + targetPartitionId, Issue.Severity.ERROR); + logger.warn("[{}] topic {} doesn't have partition {}, direct writing failed", id, topicPath, partitionId); + Issue issue = Issue.of("Cannot find partition " + partitionId, Issue.Severity.ERROR); return Result.fail(Status.of(StatusCode.BAD_REQUEST, issue)); } @@ -178,13 +166,8 @@ private CompletableFuture> lookupPartitionId(String id) { .build(); stream.sendNext(init); } catch (Throwable throwable) { - logger.warn( - "[{}] cannot send init request to probe stream of topic {} with producer {}", - id, - topicPath, - producerId, - 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))); } 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 04a98f36b..ad7c6f1ca 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 @@ -2,6 +2,7 @@ import java.util.concurrent.CompletableFuture; +import tech.ydb.core.Result; import tech.ydb.proto.topic.YdbTopic.StreamWriteMessage; import tech.ydb.proto.topic.YdbTopic.StreamWriteMessage.FromClient; import tech.ydb.topic.TopicRpc; @@ -59,8 +60,9 @@ public StreamWriteMessage.InitRequest buildInitRequest() { * @param id identifier of the new stream for logging * @return future with the new stream */ - public CompletableFuture createNewStream(String id) { + public CompletableFuture> createNewStream(String id) { FromClient init = FromClient.newBuilder().setInitRequest(buildInitRequest()).build(); - return CompletableFuture.completedFuture(new WriteStream(id, rpc.writeSession(id), init)); + WriteStream stream = new WriteStream(id, rpc.writeSession(id), init); + return CompletableFuture.completedFuture(Result.success(stream)); } } 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 2cb1ce803..950624297 100644 --- a/topic/src/test/java/tech/ydb/topic/impl/TopicRetryableStreamTest.java +++ b/topic/src/test/java/tech/ydb/topic/impl/TopicRetryableStreamTest.java @@ -2,11 +2,11 @@ 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; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; import com.google.protobuf.Empty; import org.junit.Assert; @@ -17,6 +17,7 @@ import org.slf4j.LoggerFactory; import tech.ydb.common.retry.RetryConfig; +import tech.ydb.core.Result; import tech.ydb.core.Status; import tech.ydb.core.StatusCode; import tech.ydb.core.grpc.GrpcReadWriteStream; @@ -73,33 +74,27 @@ void fail(Throwable th) { } private static class TestStream extends TopicRetryableStream> { - private final List handles; - private int handleIndex = 0; + private final List>> handles = new ArrayList<>(); + private final AtomicInteger handleIndex = new AtomicInteger(0); final List retryStatuses = new ArrayList<>(); 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) { + TestStream(RetryConfig retryConfig, ScheduledExecutorService scheduler, StreamHandle... initList) { super(logger, "test", retryConfig, scheduler); - this.handles = handles; - } - - @Override - protected CompletableFuture> createNewStream(String debugId) { - if (pendingCreation != null) { - CompletableFuture> future = pendingCreation; - pendingCreation = null; - return future; + for (StreamHandle handle: initList) { + handles.add(CompletableFuture.completedFuture(Result.success(handle))); } + } - StreamHandle handle = handles.get(handleIndex); - handleIndex++; + public void addHandle(CompletableFuture> handle) { + handles.add(handle); + } - return CompletableFuture.completedFuture(handle.stream); + @Override + protected CompletableFuture>> createNewStream(String debugId) { + return handles.get(handleIndex.getAndIncrement()).thenApply(r -> r.map(h -> h.stream)); } @Override @@ -125,7 +120,7 @@ private ScheduledExecutorService mockScheduler() { @Test public void simpleStartAndCloseTest() { StreamHandle h = new StreamHandle(); - TestStream retryable = new TestStream(Arrays.asList(h), RetryConfig.noRetries(), mockScheduler()); + TestStream retryable = new TestStream(RetryConfig.noRetries(), mockScheduler(), h); retryable.start(); @@ -149,7 +144,7 @@ public void simpleStartAndCloseTest() { @Test public void failStreamTest() { StreamHandle h = new StreamHandle(); - TestStream retryable = new TestStream(Arrays.asList(h), RetryConfig.noRetries(), mockScheduler()); + TestStream retryable = new TestStream(RetryConfig.noRetries(), mockScheduler(), h); retryable.start(); @@ -176,7 +171,7 @@ public void failStreamTest() { public void doubleStartTest() { StreamHandle h1 = new StreamHandle(); StreamHandle h2 = new StreamHandle(); - TestStream retryable = new TestStream(Arrays.asList(h1, h2), RetryConfig.noRetries(), mockScheduler()); + TestStream retryable = new TestStream(RetryConfig.noRetries(), mockScheduler(), h1, h2); retryable.start(); // sets realStream = h1.topicStream retryable.start(); // compareAndSet fails → h2.topicStream is closed @@ -189,7 +184,7 @@ public void doubleStartTest() { @Test public void doubleCloseTest() { StreamHandle h1 = new StreamHandle(); - TestStream retryable = new TestStream(Arrays.asList(h1), RetryConfig.noRetries(), mockScheduler()); + TestStream retryable = new TestStream(RetryConfig.noRetries(), mockScheduler(), h1); retryable.start(); @@ -203,7 +198,7 @@ public void doubleCloseTest() { @Test public void startAfterCloseTest() { - TestStream retryable = new TestStream(Arrays.asList(), RetryConfig.noRetries(), mockScheduler()); + TestStream retryable = new TestStream(RetryConfig.noRetries(), mockScheduler()); Assert.assertFalse(retryable.close()); retryable.start(); // nothing } @@ -211,10 +206,10 @@ public void startAfterCloseTest() { @Test public void asyncStreamCreationTest() { StreamHandle streamHandle = new StreamHandle(); - TestStream retryable = new TestStream(Collections.emptyList(), RetryConfig.noRetries(), mockScheduler()); + CompletableFuture> creation = new CompletableFuture<>(); - CompletableFuture> creation = new CompletableFuture<>(); - retryable.pendingCreation = creation; + TestStream retryable = new TestStream(RetryConfig.noRetries(), mockScheduler()); + retryable.addHandle(creation); retryable.start(); // must return without waiting for the creation future Mockito.verify(streamHandle.grpc, Mockito.never()).start(Mockito.any()); @@ -222,7 +217,7 @@ public void asyncStreamCreationTest() { retryable.send(EMPTY); // stream is not ready yet, message is skipped Mockito.verify(streamHandle.grpc, Mockito.never()).sendNext(Mockito.any()); - creation.complete(streamHandle.stream); + creation.complete(Result.success(streamHandle)); Mockito.verify(streamHandle.grpc).start(Mockito.any()); retryable.send(EMPTY); @@ -235,15 +230,15 @@ public void asyncStreamCreationTest() { @Test public void closeWhileStreamIsCreatingTest() { StreamHandle streamHandle = new StreamHandle(); - TestStream retryable = new TestStream(Collections.emptyList(), RetryConfig.noRetries(), mockScheduler()); + TestStream retryable = new TestStream(RetryConfig.noRetries(), mockScheduler()); - CompletableFuture> creation = new CompletableFuture<>(); - retryable.pendingCreation = creation; + CompletableFuture> creation = new CompletableFuture<>(); + retryable.addHandle(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 + creation.complete(Result.success(streamHandle)); // the created stream must not be started Mockito.verify(streamHandle.grpc, Mockito.never()).start(Mockito.any()); Assert.assertFalse(retryable.close()); @@ -252,10 +247,10 @@ public void closeWhileStreamIsCreatingTest() { @Test @HideLoggers({TopicRetryableStreamTest.class}) public void streamCreationFailedTest() { - TestStream retryable = new TestStream(Collections.emptyList(), RetryConfig.noRetries(), mockScheduler()); + TestStream retryable = new TestStream(RetryConfig.noRetries(), mockScheduler()); - CompletableFuture> creation = new CompletableFuture<>(); - retryable.pendingCreation = creation; + CompletableFuture> creation = new CompletableFuture<>(); + retryable.addHandle(creation); retryable.start(); creation.completeExceptionally(new RuntimeException("cannot create stream")); @@ -268,7 +263,7 @@ public void streamCreationFailedTest() { @Test public void sendBeforeStartIsIgnoredTest() { StreamHandle h = new StreamHandle(); - TestStream retryable = new TestStream(Arrays.asList(h), RetryConfig.noRetries(), mockScheduler()); + TestStream retryable = new TestStream(RetryConfig.noRetries(), mockScheduler(), h); retryable.send(EMPTY); // just skipping @@ -278,7 +273,7 @@ public void sendBeforeStartIsIgnoredTest() { @Test public void closeBeforeStartIsNoOpTest() { StreamHandle h = new StreamHandle(); - TestStream retryable = new TestStream(Arrays.asList(h), RetryConfig.noRetries(), mockScheduler()); + TestStream retryable = new TestStream(RetryConfig.noRetries(), mockScheduler(), h); Assert.assertFalse(retryable.close()); // no stream yet, should not throw @@ -288,7 +283,7 @@ public void closeBeforeStartIsNoOpTest() { @Test public void noRetriesErrorStatusTest() { StreamHandle h = new StreamHandle(); - TestStream retryable = new TestStream(Arrays.asList(h), RetryConfig.noRetries(), mockScheduler()); + TestStream retryable = new TestStream(RetryConfig.noRetries(), mockScheduler(), h); retryable.start(); h.complete(Status.of(StatusCode.ABORTED)); @@ -301,7 +296,7 @@ public void noRetriesErrorStatusTest() { public void noRetriesExceptionStatusTest() { @SuppressWarnings("unchecked") StreamHandle h = new StreamHandle(Mockito.mock(TopicStreamBase.class)); - TestStream retryable = new TestStream(Arrays.asList(h), RetryConfig.noRetries(), mockScheduler()); + TestStream retryable = new TestStream(RetryConfig.noRetries(), mockScheduler(), h); retryable.start(); RuntimeException ex = new RuntimeException("fail"); @@ -328,7 +323,7 @@ public void immediateRetryTest() { // Policy: immediate retry (0ms) on all attempts, then no more RetryConfig config = status -> (retryCount, elapsed) -> (status.getCode() != StatusCode.BAD_REQUEST) ? 0 : -1; - TestStream retryable = new TestStream(Arrays.asList(h1, h2, h3), config, mockScheduler()); + TestStream retryable = new TestStream(config, mockScheduler(), h1, h2, h3); Assert.assertFalse(retryable.isClosed()); retryable.start(); @@ -372,7 +367,7 @@ public void closeOnWrongSchedulerTest() { long delayMs = 500L; RetryConfig config = status -> (retryCount, elapsed) -> delayMs; - TestStream retryable = new TestStream(Arrays.asList(h), config, null); + TestStream retryable = new TestStream(config, null, h); retryable.start(); Assert.assertFalse(retryable.isClosed()); @@ -391,8 +386,7 @@ public void scheduledRetryWithCorrectDelayTest() { long delayMs = 500L; RetryConfig config = status -> (retryCount, elapsed) -> delayMs; - TestStream retryable = new TestStream( - Arrays.asList(h), config, scheduler); + TestStream retryable = new TestStream(config, scheduler, h); retryable.start(); h.complete(Status.of(StatusCode.UNAVAILABLE)); @@ -411,7 +405,7 @@ public void testResetRetriesAllowsRetryingAgainFromZero() { // Policy: one immediate retry (retryCount 0), then no more RetryConfig config = status -> (retryCount, elapsed) -> retryCount == 0 ? 0 : -1; - TestStream retryable = new TestStream(Arrays.asList(h1, h2, h3), config, mockScheduler()); + TestStream retryable = new TestStream(config, mockScheduler(), h1, h2, h3); Status error = Status.of(StatusCode.UNAVAILABLE); retryable.start(); 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 2f9ac15d7..3a7dbad9e 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 @@ -69,6 +69,10 @@ public void responseWith(Exception ex) { }).when(grpc).sendNext(Mockito.any()); } + public void responseWithException(Exception ex) { + Mockito.doThrow(ex).when(grpc).sendNext(Mockito.any()); + } + public void closeImmediately(Status status) { result.complete(status); } @@ -136,15 +140,15 @@ public void directWriteByPartitionIdTest() { WriteStreamFactory factory = new WriteStreamDirectFactory(rpc, settings); Assert.assertEquals("/local/topic", factory.getTopicPath()); - WriteSession.Stream stream = factory.createNewStream("s1").join(); - Assert.assertTrue(stream instanceof WriteStream); + Result res = factory.createNewStream("s1").join(); + Assert.assertTrue(res.isSuccess()); ArgumentCaptor options = ArgumentCaptor.forClass(GrpcRequestSettings.class); Mockito.verify(rpc).writeSession(options.capture()); Assert.assertTrue(options.getValue().isDirectMode()); Assert.assertEquals(42, options.getValue().getPreferredNodeID().intValue()); - stream.start(null); + res.getValue().start(null); FromClient msg = mocked.verifyNextMsg(); Assert.assertTrue(msg.hasInitRequest()); @@ -166,16 +170,11 @@ public void directWriteByPartitionIdTestDescribeFailTest() { .setDirectWrite(true) .build()); - WriteSession.Stream stream = factory.createNewStream("s1").join(); - + Result res = factory.createNewStream("s1").join(); Mockito.verify(rpc, Mockito.never()).writeSession(Mockito.any(GrpcRequestSettings.class)); - Assert.assertTrue(stream instanceof WriteStream.Fail); - CompletableFuture res = stream.start(null); - Assert.assertTrue(res.isDone()); - Assert.assertEquals(Status.of(StatusCode.UNAVAILABLE), res.join()); - - stream.close(); // no effect + Assert.assertTrue(!res.isSuccess()); + Assert.assertEquals(Status.of(StatusCode.UNAVAILABLE), res.getStatus()); } @Test @@ -190,17 +189,12 @@ public void directWriteByPartitionIdTestPartitionNotFoundTest() { .setDirectWrite(true) .build()); - WriteSession.Stream stream = factory.createNewStream("s1").join(); - + Result res = factory.createNewStream("s1").join(); Mockito.verify(rpc, Mockito.never()).writeSession(Mockito.any(GrpcRequestSettings.class)); - Assert.assertTrue(stream instanceof WriteStream.Fail); - CompletableFuture res = stream.start(null); - Assert.assertTrue(res.isDone()); + Assert.assertTrue(!res.isSuccess()); Status expected = Status.of(StatusCode.BAD_REQUEST, Issue.of("Cannot find partition 3", Issue.Severity.ERROR)); - Assert.assertEquals(expected, res.join()); - - stream.close(); // no effect + Assert.assertEquals(expected, res.getStatus()); } @Test @@ -217,17 +211,12 @@ public void directWriteByPartitionIdTestPartitionHasNoLocationTest() { .setDirectWrite(true) .build()); - WriteSession.Stream stream = factory.createNewStream("s1").join(); - + Result res = factory.createNewStream("s1").join(); Mockito.verify(rpc, Mockito.never()).writeSession(Mockito.any(GrpcRequestSettings.class)); - Assert.assertTrue(stream instanceof WriteStream.Fail); - CompletableFuture res = stream.start(null); - Assert.assertTrue(res.isDone()); + Assert.assertTrue(!res.isSuccess()); Status expected = Status.of(StatusCode.BAD_REQUEST, Issue.of("Partition 3 has no location", Issue.Severity.ERROR)); - Assert.assertEquals(expected, res.join()); - - stream.close(); // no effect + Assert.assertEquals(expected, res.getStatus()); } @Test @@ -258,15 +247,15 @@ public void directWriteByProducerIdTest() { .setDirectWrite(true) .build()); - WriteSession.Stream stream = factory.createNewStream("s1").join(); - Assert.assertTrue(stream instanceof WriteStream); + Result res = factory.createNewStream("s1").join(); + Assert.assertTrue(res.isSuccess()); ArgumentCaptor options = ArgumentCaptor.forClass(GrpcRequestSettings.class); Mockito.verify(rpc, Mockito.times(2)).writeSession(options.capture()); Assert.assertTrue(options.getValue().isDirectMode()); Assert.assertEquals(55, options.getValue().getPreferredNodeID().intValue()); - stream.start(null); + res.getValue().start(null); FromClient msg = actual.verifyNextMsg(); Assert.assertTrue(msg.hasInitRequest()); @@ -292,14 +281,11 @@ public void directWriteByProducerIdProbeFailTest() { .setDirectWrite(true) .build()); - WriteSession.Stream stream = factory.createNewStream("s1").join(); - Assert.assertTrue(stream instanceof WriteStream.Fail); + Result res = factory.createNewStream("s1").join(); Mockito.verify(rpc).writeSession(Mockito.any(GrpcRequestSettings.class)); - CompletableFuture res = stream.start(null); - Assert.assertTrue(res.isDone()); - Assert.assertEquals(Status.of(StatusCode.UNAUTHORIZED), res.join()); - stream.close(); // no effect + Assert.assertFalse(res.isSuccess()); + Assert.assertEquals(Status.of(StatusCode.UNAUTHORIZED), res.getStatus()); } @Test @@ -316,14 +302,11 @@ public void directWriteByProducerIdProbeFailOnSendTest() { .setDirectWrite(true) .build()); - WriteSession.Stream stream = factory.createNewStream("s1").join(); - Assert.assertTrue(stream instanceof WriteStream.Fail); + Result res = factory.createNewStream("s1").join(); Mockito.verify(rpc).writeSession(Mockito.any(GrpcRequestSettings.class)); - CompletableFuture res = stream.start(null); - Assert.assertTrue(res.isDone()); - Assert.assertEquals(Status.of(StatusCode.PRECONDITION_FAILED), res.join()); - stream.close(); // no effect + Assert.assertFalse(res.isSuccess()); + Assert.assertEquals(Status.of(StatusCode.PRECONDITION_FAILED), res.getStatus()); } @Test @@ -340,17 +323,38 @@ public void directWriteByProducerIdProbeExceptionOnSendTest() { .setDirectWrite(true) .build()); - WriteSession.Stream stream = factory.createNewStream("s1").join(); - Assert.assertTrue(stream instanceof WriteStream.Fail); + Result res = factory.createNewStream("s1").join(); Mockito.verify(rpc).writeSession(Mockito.any(GrpcRequestSettings.class)); - CompletableFuture res = stream.start(null); - Assert.assertTrue(res.isDone()); - Status status = res.join(); + Assert.assertFalse(res.isSuccess()); + Status status = res.getStatus(); Assert.assertEquals(StatusCode.CLIENT_INTERNAL_ERROR, status.getCode()); - Assert.assertNotNull(status.getCause()); + Assert.assertTrue(status.getCause() instanceof RuntimeException); Assert.assertEquals("something went wrong", status.getCause().getMessage()); - stream.close(); // no effect + } + + @Test + public void directWriteByProducerIdProbeThrowsExceptionOnSendTest() { + TopicRpc rpc = Mockito.mock(TopicRpc.class); + + MockedStream probe = new MockedStream(); + probe.responseWithException(new IllegalStateException("invalid state")); + Mockito.when(rpc.writeSession(Mockito.any(GrpcRequestSettings.class))).thenReturn(probe.grpc); + + WriteStreamFactory factory = new WriteStreamDirectFactory(rpc, WriterSettings.newBuilder() + .setTopicPath("/test/topic") + .setProducerId("producer-1") + .setDirectWrite(true) + .build()); + + Result res = factory.createNewStream("s1").join(); + Mockito.verify(rpc).writeSession(Mockito.any(GrpcRequestSettings.class)); + + Assert.assertFalse(res.isSuccess()); + Status status = res.getStatus(); + Assert.assertEquals(StatusCode.CLIENT_INTERNAL_ERROR, status.getCode()); + Assert.assertTrue(status.getCause() instanceof IllegalStateException); + Assert.assertEquals("invalid state", status.getCause().getMessage()); } @Test @@ -370,14 +374,11 @@ public void directWriteByProducerIdProbeWrongResponseTest() { .setDirectWrite(true) .build()); - WriteSession.Stream stream = factory.createNewStream("s1").join(); - Assert.assertTrue(stream instanceof WriteStream.Fail); + Result res = factory.createNewStream("s1").join(); Mockito.verify(rpc).writeSession(Mockito.any(GrpcRequestSettings.class)); - CompletableFuture res = stream.start(null); - Assert.assertTrue(res.isDone()); - Assert.assertEquals(Status.of(StatusCode.INTERNAL_ERROR), res.join()); - stream.close(); // no effect + Assert.assertFalse(res.isSuccess()); + Assert.assertEquals(Status.of(StatusCode.INTERNAL_ERROR), res.getStatus()); } @Test @@ -398,15 +399,12 @@ public void directWriteByProducerIdProbeUnexpectedResponseTest() { .setDirectWrite(true) .build()); - WriteSession.Stream stream = factory.createNewStream("s1").join(); - Assert.assertTrue(stream instanceof WriteStream.Fail); + Result res = factory.createNewStream("s1").join(); Mockito.verify(rpc).writeSession(Mockito.any(GrpcRequestSettings.class)); - CompletableFuture res = stream.start(null); - Assert.assertTrue(res.isDone()); + Assert.assertFalse(res.isSuccess()); Issue issue = Issue.of("Unexpected message from stream with producer producer-1", Issue.Severity.ERROR); - Assert.assertEquals(Status.of(StatusCode.BAD_REQUEST, issue), res.join()); - stream.close(); // no effect + Assert.assertEquals(Status.of(StatusCode.BAD_REQUEST, issue), res.getStatus()); } @Test @@ -433,11 +431,9 @@ public void directWriteByProducerIdPartitionNotFoundTest() { .setDirectWrite(true) .build()); - WriteSession.Stream stream = factory.createNewStream("s1").join(); - Assert.assertTrue(stream instanceof WriteStream.Fail); - CompletableFuture res = stream.start(null); - Assert.assertTrue(res.isDone()); + Result res = factory.createNewStream("s1").join(); + Assert.assertFalse(res.isSuccess()); Status expected = Status.of(StatusCode.BAD_REQUEST, Issue.of("Cannot find partition 5", Issue.Severity.ERROR)); - Assert.assertEquals(expected, res.join()); + Assert.assertEquals(expected, res.getStatus()); } } 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 67d89bcec..522afcc94 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").join(); + WriteSession.Stream stream = factory.createNewStream("s1").join().getValue(); Assert.assertTrue(stream instanceof WriteStream); Mockito.verify(rpc).writeSession("s1"); } From 148c03c1fb1e6e080b3e3f10c11ed7a5e1c1ba13 Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Thu, 24 Sep 2026 11:08:33 +0100 Subject: [PATCH 3/3] Revert atomic state for TopicRetryableStream --- .../ydb/topic/impl/TopicRetryableStream.java | 171 ++++++++++-------- .../write/impl/WriteStreamDirectFactory.java | 12 +- .../topic/impl/TopicRetryableStreamTest.java | 47 ++++- .../topic/read/impl/AsyncReaderImplTest.java | 31 ++-- .../topic/read/impl/SyncReaderImplTest.java | 28 +-- 5 files changed, 177 insertions(+), 112 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 910809c7a..c8ac9c174 100644 --- a/topic/src/main/java/tech/ydb/topic/impl/TopicRetryableStream.java +++ b/topic/src/main/java/tech/ydb/topic/impl/TopicRetryableStream.java @@ -23,11 +23,9 @@ public abstract class TopicRetryableStream state = new AtomicReference<>(); private final AtomicInteger streamCount = new AtomicInteger(0); - private final RetryState state = new RetryState(); - - private final AtomicReference realStreamId = new AtomicReference<>(); - private volatile S realStream = null; + private final RetryState retryState = new RetryState(); private volatile boolean isClosed = false; @@ -69,64 +67,23 @@ public void start() { } String streamID = debugId + '.' + streamCount.incrementAndGet(); - createNewStream(streamID).whenComplete((result, th) -> { - tryStartStream(streamID, result, th); - }); - } + State newState = new State(createNewStream(streamID)); - private void tryStartStream(String streamID, Result result, Throwable th) { - if (isClosed) { - logger.info("[{}] stream was closed while it was creating, skipping", streamID); + if (!state.compareAndSet(null, newState)) { + logger.warn("[{}] double start of stream, skipping", debugId); + newState.closeWithoutStart(); return; } - if (!realStreamId.compareAndSet(null, streamID)) { - logger.warn("[{}] double start of stream, skipping", streamID); - return; - } + newState.start(); - if (result != null && result.isSuccess()) { - startStream(streamID, result.getValue()); - return; - } - - if (!realStreamId.compareAndSet(streamID, null)) { - return; - } - - if (result == null) { - logger.warn("[{}] cannot create stream with exeption", streamID, th); - Status wrapped = Status.of(StatusCode.CLIENT_INTERNAL_ERROR, th); - onStreamStop(null, wrapped, retryConfig.getThrowableRetryPolicy(th)); - } else { - logger.warn("[{}] cannot create stream with status {}", streamID, result.getStatus()); - onStreamStop(null, result.getStatus(), retryConfig.getStatusRetryPolicy(result.getStatus())); - } - } - - private void startStream(String streamID, S stream) { - realStream = stream; - stream.start(msg -> onNext(stream, msg)).whenComplete((status, th) -> { - if (!realStreamId.compareAndSet(streamID, 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)); - } - }); - - if (isClosed) { // stream may be closed by other thread - realStream = null; - stream.close(); + if (isClosed) { + close(); } } protected void resetRetries() { - state.reset(); + retryState.reset(); } public boolean isClosed() { @@ -134,40 +91,31 @@ public boolean isClosed() { } public void fail(Status status) { - String streamId = realStreamId.getAndSet(null); - if (streamId != null) { - logger.warn("[{}] failed by application-side error {}", streamId, status); - S local = realStream; - realStream = null; - if (local != null) { - local.close(); - } - onStreamStop(local, status, retryConfig.getStatusRetryPolicy(status)); + State local = state.getAndSet(null); + if (local != null) { + logger.warn("[{}] failed by application-side error {}", debugId, status); + onStreamStop(local.closeStream(), status, retryConfig.getStatusRetryPolicy(status)); } } public void send(W msg) { - S local = realStream; - if (local == null) { + 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; } - local.send(msg); + stream.send(msg); } public boolean close() { isClosed = true; - String streamId = realStreamId.getAndSet(null); - if (streamId == null) { + State local = state.getAndSet(null); + if (local == null) { return false; } - S local = realStream; - realStream = null; - if (local != null) { - local.close(); - } - onStreamStop(local, Status.SUCCESS, null); + onStreamStop(local.closeStream(), Status.SUCCESS, null); return true; } @@ -184,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); @@ -194,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 { @@ -213,6 +161,77 @@ private void onStreamStop(S closed, Status status, RetryPolicy policy) { } } + private class State { + private final CompletableFuture> future; + private final AtomicReference stream = new AtomicReference<>(); + + State(CompletableFuture> 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; 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 0f6034ca6..0aff9b232 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 @@ -79,7 +79,7 @@ protected CompletableFuture> lookupLocation(S return rpc.describeTopic(req, settings).thenApply(res -> parseLocation(id, targetPartitionId, res)); } - private Result parseLocation(String id, long partitionId, + private Result parseLocation(String id, long targetPartitionId, Result description) { if (!description.isSuccess()) { logger.warn("[{}] describe topic {} failed with status {}", id, topicPath, description.getStatus()); @@ -88,10 +88,10 @@ private Result parseLocation(String id, long partiti // lookup for partition location for (YdbTopic.DescribeTopicResult.PartitionInfo partition : description.getValue().getPartitionsList()) { - if (partition.getPartitionId() == partitionId) { + if (partition.getPartitionId() == targetPartitionId) { if (!partition.hasPartitionLocation()) { - logger.warn("[{}] partition {} has no valid location info", id, partitionId); - Issue issue = Issue.of("Partition " + partitionId + " has no location", Issue.Severity.ERROR); + logger.warn("[{}] partition {} has no valid location info", id, targetPartitionId); + Issue issue = Issue.of("Partition " + targetPartitionId + " has no location", Issue.Severity.ERROR); return Result.fail(Status.of(StatusCode.BAD_REQUEST, issue)); } @@ -99,8 +99,8 @@ private Result parseLocation(String id, long partiti } } - logger.warn("[{}] topic {} doesn't have partition {}, direct writing failed", id, topicPath, partitionId); - Issue issue = Issue.of("Cannot find partition " + partitionId, Issue.Severity.ERROR); + logger.warn("[{}] topic {} doesn't have partition {}, direct writing failed", id, topicPath, targetPartitionId); + Issue issue = Issue.of("Cannot find partition " + targetPartitionId, Issue.Severity.ERROR); return Result.fail(Status.of(StatusCode.BAD_REQUEST, issue)); } 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 950624297..4a1159277 100644 --- a/topic/src/test/java/tech/ydb/topic/impl/TopicRetryableStreamTest.java +++ b/topic/src/test/java/tech/ydb/topic/impl/TopicRetryableStreamTest.java @@ -174,11 +174,11 @@ public void doubleStartTest() { TestStream retryable = new TestStream(RetryConfig.noRetries(), mockScheduler(), h1, h2); retryable.start(); // sets realStream = h1.topicStream - retryable.start(); // compareAndSet fails → h2.topicStream is closed + retryable.start(); // compareAndSet fails → h2.topicStream is closed byt not started Mockito.verify(h1.grpc).start(Mockito.any()); Mockito.verify(h2.grpc, Mockito.never()).start(Mockito.any()); // h2 was never started - Mockito.verify(h2.grpc, Mockito.never()).close(); // h2 was never closed + Mockito.verify(h2.grpc).close(); } @Test @@ -228,7 +228,7 @@ public void asyncStreamCreationTest() { } @Test - public void closeWhileStreamIsCreatingTest() { + public void closeWhileAsyncInitializationTest() { StreamHandle streamHandle = new StreamHandle(); TestStream retryable = new TestStream(RetryConfig.noRetries(), mockScheduler()); @@ -236,12 +236,49 @@ public void closeWhileStreamIsCreatingTest() { retryable.addHandle(creation); retryable.start(); - Assert.assertFalse(retryable.close()); // there is no stream to close yet + Assert.assertTrue(retryable.close()); - creation.complete(Result.success(streamHandle)); // the created stream must not be started + creation.complete(Result.success(streamHandle)); Mockito.verify(streamHandle.grpc, Mockito.never()).start(Mockito.any()); + Mockito.verify(streamHandle.grpc).close(); + Assert.assertFalse(retryable.close()); + + Assert.assertTrue(retryable.retryStatuses.isEmpty()); + Assert.assertEquals(Arrays.asList(Status.SUCCESS), retryable.closeStatuses); + } + + @Test + public void closeWhileAsyncInitializationFailedTest() { + TestStream retryable = new TestStream(RetryConfig.noRetries(), mockScheduler()); + + CompletableFuture> creation = new CompletableFuture<>(); + retryable.addHandle(creation); + + retryable.start(); + Assert.assertTrue(retryable.close()); + creation.complete(Result.fail(Status.of(StatusCode.BAD_REQUEST))); // will be lost + + Assert.assertFalse(retryable.close()); + Assert.assertTrue(retryable.retryStatuses.isEmpty()); + Assert.assertEquals(Arrays.asList(Status.SUCCESS), retryable.closeStatuses); + } + + @Test + public void closeWhileAsyncInitializationErrorTest() { + TestStream retryable = new TestStream(RetryConfig.noRetries(), mockScheduler()); + + CompletableFuture> creation = new CompletableFuture<>(); + retryable.addHandle(creation); + + retryable.start(); + Assert.assertTrue(retryable.close()); + creation.completeExceptionally(new IllegalArgumentException("error")); // will be lost + + Assert.assertFalse(retryable.close()); + Assert.assertTrue(retryable.retryStatuses.isEmpty()); + Assert.assertEquals(Arrays.asList(Status.SUCCESS), retryable.closeStatuses); } @Test diff --git a/topic/src/test/java/tech/ydb/topic/read/impl/AsyncReaderImplTest.java b/topic/src/test/java/tech/ydb/topic/read/impl/AsyncReaderImplTest.java index 221905268..4a6d385da 100644 --- a/topic/src/test/java/tech/ydb/topic/read/impl/AsyncReaderImplTest.java +++ b/topic/src/test/java/tech/ydb/topic/read/impl/AsyncReaderImplTest.java @@ -87,34 +87,39 @@ private static void assertMessages(List expected, List messages @Test public void initAndShutdownTest() { - ReadStreamMock mock = new ReadStreamMock(); - AsyncReader reader = reader(settings().setReaderName("test-reader-name").build(), mock); - mock.assertSentMessagesCount(0); + ReadStreamMock m1 = new ReadStreamMock(); + ReadStreamMock m2 = new ReadStreamMock(); + AsyncReader reader = reader(settings().setReaderName("test-reader-name").build(), m1, m2); + m1.assertSentMessagesCount(0); Mockito.verifyNoInteractions(handler); CompletableFuture init = reader.init(); Assert.assertFalse(init.isDone()); - mock.assertSentMessagesCount(1); - mock.assertLastMessage().isInitRequest("consumer", "/test-topic"); + m1.assertSentMessagesCount(1); + m1.assertLastMessage().isInitRequest("consumer", "/test-topic"); - Assert.assertSame(init, reader.init()); // double init is allowed - mock.assertSentMessagesCount(1); - mock.responseInit("read-session-1"); + // double init is allowed, but second stream will be closed without start + Assert.assertSame(init, reader.init()); + m2.assertIsNotStarted(); + m2.assertIsClosed(); + + m1.assertSentMessagesCount(1); + m1.responseInit("read-session-1"); Assert.assertTrue(init.isDone()); Assert.assertFalse(init.isCompletedExceptionally()); ArgumentCaptor started = ArgumentCaptor.forClass(SessionStartedEvent.class); Mockito.verify(handler).onSessionStarted(started.capture()); Assert.assertEquals("read-session-1", started.getValue().getSessionId()); - mock.assertSentMessagesCount(2); - mock.assertLastMessage().isReadRequest(100 * 1024 * 1024); + m1.assertSentMessagesCount(2); + m1.assertLastMessage().isReadRequest(100 * 1024 * 1024); CompletableFuture shutdown = reader.shutdown(); - mock.assertIsClosed(); + m1.assertIsClosed(); Assert.assertTrue(shutdown.isDone()); Assert.assertSame(shutdown, reader.shutdown()); // double shutdown is allowed - mock.assertIsClosed(); - mock.closeStream(Status.SUCCESS); + m1.assertIsClosed(); + m1.closeStream(Status.SUCCESS); Assert.assertTrue(shutdown.isDone()); Assert.assertFalse(shutdown.isCompletedExceptionally()); diff --git a/topic/src/test/java/tech/ydb/topic/read/impl/SyncReaderImplTest.java b/topic/src/test/java/tech/ydb/topic/read/impl/SyncReaderImplTest.java index 6ca412fe2..66edb0746 100644 --- a/topic/src/test/java/tech/ydb/topic/read/impl/SyncReaderImplTest.java +++ b/topic/src/test/java/tech/ydb/topic/read/impl/SyncReaderImplTest.java @@ -53,7 +53,8 @@ private static TopicRpc mockRpc(ReadStreamMock first, ReadStreamMock... rest) { @Test public void initAndShutdownTest() throws InterruptedException { - ReadStreamMock mock = new ReadStreamMock(); + ReadStreamMock m1 = new ReadStreamMock(); + ReadStreamMock m2 = new ReadStreamMock(); ReaderSettings settings = ReaderSettings.newBuilder() .addTopic(TopicReadSettings.newBuilder().setPath("/test-topic").build()) @@ -61,36 +62,39 @@ public void initAndShutdownTest() throws InterruptedException { .setReaderName("test-reader-name") .build(); - SyncReader reader = new SyncReaderImpl(mockRpc(mock), settings, REGISTRY); - mock.assertSentMessagesCount(0); + SyncReader reader = new SyncReaderImpl(mockRpc(m1, m2), settings, REGISTRY); + m1.assertSentMessagesCount(0); // before init there is nothing to read Assert.assertNull(reader.receive(0, TimeUnit.MILLISECONDS)); reader.init(); - mock.assertSentMessagesCount(1); - mock.assertLastMessage().isInitRequest("consumer", "/test-topic"); + m1.assertSentMessagesCount(1); + m1.assertLastMessage().isInitRequest("consumer", "/test-topic"); - reader.init(); // double init is allowed - mock.assertSentMessagesCount(1); + reader.init(); // double init is allowed, but second stream will be closed without start + m2.assertIsNotStarted(); + m2.assertIsClosed(); - mock.responseInit("read-session-1"); + m1.assertSentMessagesCount(1); + + m1.responseInit("read-session-1"); Assert.assertEquals("read-session-1", reader.getSessionId()); - mock.assertSentMessagesCount(2); - mock.assertLastMessage().isReadRequest(100 * 1024 * 1024); + m1.assertSentMessagesCount(2); + m1.assertLastMessage().isReadRequest(100 * 1024 * 1024); Assert.assertNull(reader.receive(0, TimeUnit.MILLISECONDS)); reader.shutdown(); - mock.assertIsClosed(); + m1.assertIsClosed(); reader.shutdown(); // double shutdow is allowed Exception ex = Assert.assertThrows(RuntimeException.class, () -> reader.receive(0, TimeUnit.MILLISECONDS)); Assert.assertEquals("Reader was stopped with Status{code = SUCCESS}", ex.getMessage()); - mock.closeStream(Status.SUCCESS); + m1.closeStream(Status.SUCCESS); } @Test