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..c8ac9c174 100644 --- a/topic/src/main/java/tech/ydb/topic/impl/TopicRetryableStream.java +++ b/topic/src/main/java/tech/ydb/topic/impl/TopicRetryableStream.java @@ -1,15 +1,19 @@ package tech.ydb.topic.impl; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; +import javax.annotation.Nullable; + import com.google.protobuf.Message; import org.slf4j.Logger; import tech.ydb.common.retry.RetryConfig; import tech.ydb.common.retry.RetryPolicy; +import tech.ydb.core.Result; import tech.ydb.core.Status; import tech.ydb.core.StatusCode; @@ -19,9 +23,9 @@ public abstract class TopicRetryableStream realStream = new AtomicReference<>(); + private final AtomicReference state = new AtomicReference<>(); private final AtomicInteger streamCount = new AtomicInteger(0); - private final RetryState state = new RetryState(); + private final RetryState retryState = new RetryState(); private volatile boolean isClosed = false; @@ -32,12 +36,29 @@ public TopicRetryableStream(Logger logger, String debugId, RetryConfig config, S this.scheduler = scheduler; } - protected abstract S createNewStream(String debugId); + /** + * Creates a new stream. Implementations must not block the calling thread: this method is invoked from the shared + * scheduler on every reconnect, and blocking there stalls discovery, session pools and timeouts of the whole + * transport. + * + * @param debugId identifier of the new stream for logging + * @return future with the new stream result + */ + protected abstract CompletableFuture> createNewStream(String debugId); protected abstract void onNext(S stream, R message); - protected abstract void onRetry(S stream, Status status); - protected abstract void onClose(S stream, Status status); + /** + * @param stream the stopped stream, or {@code null} when stream creation itself failed + * @param status status the stream stopped with + */ + protected abstract void onRetry(@Nullable S stream, Status status); + + /** + * @param stream the closed stream, or {@code null} when stream creation itself failed + * @param status status the stream was closed with + */ + protected abstract void onClose(@Nullable S stream, Status status); public void start() { if (isClosed) { @@ -46,29 +67,23 @@ public void start() { } String streamID = debugId + '.' + streamCount.incrementAndGet(); - S stream = createNewStream(streamID); + State newState = new State(createNewStream(streamID)); - if (!realStream.compareAndSet(null, stream)) { + if (!state.compareAndSet(null, newState)) { logger.warn("[{}] double start of stream, skipping", debugId); + newState.closeWithoutStart(); return; } - stream.start(msg -> onNext(stream, msg)).whenComplete((status, th) -> { - if (!realStream.compareAndSet(stream, null)) { - return; - } - if (status != null) { - onStreamStop(stream, status, retryConfig.getStatusRetryPolicy(status)); - } - if (th != null) { - Status wrapped = Status.of(StatusCode.CLIENT_INTERNAL_ERROR, th); - onStreamStop(stream, wrapped, retryConfig.getThrowableRetryPolicy(th)); - } - }); + newState.start(); + + if (isClosed) { + close(); + } } protected void resetRetries() { - state.reset(); + retryState.reset(); } public boolean isClosed() { @@ -76,16 +91,16 @@ public boolean isClosed() { } public void fail(Status status) { - S closed = realStream.getAndSet(null); - if (closed != null) { + State local = state.getAndSet(null); + if (local != null) { logger.warn("[{}] failed by application-side error {}", debugId, status); - closed.close(); - onStreamStop(closed, status, retryConfig.getStatusRetryPolicy(status)); + onStreamStop(local.closeStream(), status, retryConfig.getStatusRetryPolicy(status)); } } public void send(W msg) { - S stream = realStream.get(); + State local = state.get(); + S stream = local != null ? local.stream.get() : null; if (stream == null) { logger.warn("[{}] send message before stream is ready", debugId); return; @@ -95,13 +110,12 @@ public void send(W msg) { public boolean close() { isClosed = true; - S stream = realStream.getAndSet(null); - if (stream == null) { + State local = state.getAndSet(null); + if (local == null) { return false; } - stream.close(); - onStreamStop(stream, Status.SUCCESS, null); + onStreamStop(local.closeStream(), Status.SUCCESS, null); return true; } @@ -118,7 +132,7 @@ private void onStreamStop(S closed, Status status, RetryPolicy policy) { return; } - long nextRetryMs = state.nextRetryMs(policy); + long nextRetryMs = retryState.nextRetryMs(policy); if (nextRetryMs < 0) { logger.warn("[{}] stopped after retry policy evaluation for status {}", debugId, status); @@ -128,14 +142,14 @@ private void onStreamStop(S closed, Status status, RetryPolicy policy) { } if (nextRetryMs == 0) { // retry immediately - logger.warn("[{}] retry #{}. Retry immediately...", debugId, state.retryNumber()); + logger.warn("[{}] retry #{}. Retry immediately...", debugId, retryState.retryNumber()); onRetry(closed, status); start(); return; } // retry scheduling - logger.warn("[{}] retry #{}. Scheduling reconnect in {}ms...", debugId, state.retryNumber(), nextRetryMs); + logger.warn("[{}] retry #{}. Scheduling reconnect in {}ms...", debugId, retryState.retryNumber(), nextRetryMs); onRetry(closed, status); try { @@ -147,6 +161,77 @@ private void onStreamStop(S closed, Status status, RetryPolicy policy) { } } + private class State { + private final CompletableFuture> 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/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 d99d9801e..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 @@ -1,11 +1,13 @@ package tech.ydb.topic.write.impl; import java.util.List; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.ScheduledExecutorService; import java.util.function.BiConsumer; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import tech.ydb.core.Result; import tech.ydb.core.Status; import tech.ydb.core.utils.ProtobufUtils; import tech.ydb.proto.topic.YdbTopic; @@ -48,7 +50,7 @@ public WriteSession(String debugId, WriteStreamFactory factory, WriterSettings s } @Override - protected Stream createNewStream(String id) { + protected CompletableFuture> 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 c3cf7e425..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 @@ -36,26 +36,23 @@ 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()); - } - targetPartitionId = pid.getValue(); - } - - Result location = lookupLocation(id, targetPartitionId); - if (!location.isSuccess()) { - return new WriteStream.Fail(id, location.getStatus()); - } + public CompletableFuture> createNewStream(String id) { + CompletableFuture> partition = partitionId == null + ? lookupPartitionId(id) + : CompletableFuture.completedFuture(Result.success(partitionId)); + + 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 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) { @@ -67,26 +64,30 @@ public WriteSession.Stream createNewStream(String id) { .withTraceId(id) .disableDeadline() .withDirectMode(true) - .withPreferredNodeID(location.getValue().getNodeId()) + .withPreferredNodeID(location.getNodeId()) .build(); 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( - YdbTopic.DescribeTopicRequest.newBuilder().setIncludeLocation(true).setPath(topicPath).build(), - GrpcRequestSettings.newBuilder().withDeadline(Duration.ofMinutes(1)).build() - ).join(); - - if (!describeTopic.isSuccess()) { - logger.warn("[{}] describe topic {} failed with status {}", id, topicPath, describeTopic.getStatus()); - return Result.fail(describeTopic.getStatus()); + 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 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()) { + for (YdbTopic.DescribeTopicResult.PartitionInfo partition : description.getValue().getPartitionsList()) { if (partition.getPartitionId() == targetPartitionId) { if (!partition.hasPartitionLocation()) { logger.warn("[{}] partition {} has no valid location info", id, targetPartitionId); @@ -103,7 +104,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 +142,35 @@ 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..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 @@ -1,7 +1,8 @@ package tech.ydb.topic.write.impl; +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; @@ -52,8 +53,16 @@ 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); + 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 26735242f..4a1159277 100644 --- a/topic/src/test/java/tech/ydb/topic/impl/TopicRetryableStreamTest.java +++ b/topic/src/test/java/tech/ydb/topic/impl/TopicRetryableStreamTest.java @@ -6,6 +6,7 @@ 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; @@ -16,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; @@ -72,21 +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<>(); - TestStream(List handles, RetryConfig retryConfig, ScheduledExecutorService scheduler) { + TestStream(RetryConfig retryConfig, ScheduledExecutorService scheduler, StreamHandle... initList) { super(logger, "test", retryConfig, scheduler); - this.handles = handles; + for (StreamHandle handle: initList) { + handles.add(CompletableFuture.completedFuture(Result.success(handle))); + } + } + + public void addHandle(CompletableFuture> handle) { + handles.add(handle); } @Override - protected TopicStreamBase createNewStream(String debugId) { - return handles.get(handleIndex++).stream; + protected CompletableFuture>> createNewStream(String debugId) { + return handles.get(handleIndex.getAndIncrement()).thenApply(r -> r.map(h -> h.stream)); } @Override @@ -112,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(); @@ -136,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(); @@ -163,20 +171,20 @@ 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 + 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 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(); @@ -190,15 +198,109 @@ 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 } + @Test + public void asyncStreamCreationTest() { + StreamHandle streamHandle = new StreamHandle(); + CompletableFuture> creation = new CompletableFuture<>(); + + 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()); + + retryable.send(EMPTY); // stream is not ready yet, message is skipped + Mockito.verify(streamHandle.grpc, Mockito.never()).sendNext(Mockito.any()); + + creation.complete(Result.success(streamHandle)); + + 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 closeWhileAsyncInitializationTest() { + StreamHandle streamHandle = new StreamHandle(); + TestStream retryable = new TestStream(RetryConfig.noRetries(), mockScheduler()); + + CompletableFuture> creation = new CompletableFuture<>(); + retryable.addHandle(creation); + + retryable.start(); + Assert.assertTrue(retryable.close()); + + 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 + @HideLoggers({TopicRetryableStreamTest.class}) + public void streamCreationFailedTest() { + TestStream retryable = new TestStream(RetryConfig.noRetries(), mockScheduler()); + + CompletableFuture> creation = new CompletableFuture<>(); + retryable.addHandle(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(); - TestStream retryable = new TestStream(Arrays.asList(h), RetryConfig.noRetries(), mockScheduler()); + TestStream retryable = new TestStream(RetryConfig.noRetries(), mockScheduler(), h); retryable.send(EMPTY); // just skipping @@ -208,7 +310,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 @@ -218,7 +320,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)); @@ -231,7 +333,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"); @@ -258,7 +360,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(); @@ -302,7 +404,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()); @@ -321,8 +423,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)); @@ -341,7 +442,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/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 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..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"); - 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"); - + 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"); - + 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"); - + 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"); - 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"); - 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"); - 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"); - 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"); - 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"); - 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"); - 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 267fdae10..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"); + WriteSession.Stream stream = factory.createNewStream("s1").join().getValue(); Assert.assertTrue(stream instanceof WriteStream); Mockito.verify(rpc).writeSession("s1"); }