From 0741d00e337740bda196974d86cd113a007b7b8d Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Wed, 14 Oct 2026 15:14:07 +0100 Subject: [PATCH 1/2] Added graceful shutdown for topic streams --- .../ydb/topic/impl/TopicRetryableStream.java | 96 ++++++++++--------- .../tech/ydb/topic/impl/TopicStreamBase.java | 51 ++++++---- .../ydb/topic/read/impl/AsyncReaderImpl.java | 1 + .../tech/ydb/topic/read/impl/ReadSession.java | 4 +- .../tech/ydb/topic/read/impl/ReaderImpl.java | 8 +- .../topic/impl/TopicRetryableStreamTest.java | 52 ++++++---- .../tech/ydb/topic/impl/TopicStreamTest.java | 19 ++++ .../topic/read/impl/AsyncReaderImplTest.java | 10 +- .../ydb/topic/read/impl/ReadStreamMock.java | 12 +++ .../ydb/topic/read/impl/ReaderImplTest.java | 19 +++- .../topic/read/impl/ReaderMetricsTest.java | 18 ++-- .../topic/read/impl/SyncReaderImplTest.java | 16 ++-- 12 files changed, 194 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 c8ac9c174..7517e431c 100644 --- a/topic/src/main/java/tech/ydb/topic/impl/TopicRetryableStream.java +++ b/topic/src/main/java/tech/ydb/topic/impl/TopicRetryableStream.java @@ -19,11 +19,12 @@ public abstract class TopicRetryableStream> { protected final String debugId; + private final Logger logger; private final RetryConfig retryConfig; private final ScheduledExecutorService scheduler; - private final AtomicReference state = new AtomicReference<>(); + private final AtomicReference state = new AtomicReference<>(null); private final AtomicInteger streamCount = new AtomicInteger(0); private final RetryState retryState = new RetryState(); @@ -71,14 +72,13 @@ public void start() { if (!state.compareAndSet(null, newState)) { logger.warn("[{}] double start of stream, skipping", debugId); - newState.closeWithoutStart(); + newState.close(); return; } newState.start(); - if (isClosed) { - close(); + newState.close(); } } @@ -94,29 +94,32 @@ public void fail(Status 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)); + local.close(); + onStreamStop(local.stream, status, retryConfig.getStatusRetryPolicy(status)); } } public void send(W msg) { State local = state.get(); - S stream = local != null ? local.stream.get() : null; - if (stream == null) { + if (isClosed || local == null) { + logger.warn("[{}] send message before starting or after closing", debugId); + return; + } + if (local.stream == null) { logger.warn("[{}] send message before stream is ready", debugId); return; } - stream.send(msg); + + local.stream.send(msg); } public boolean close() { isClosed = true; - State local = state.getAndSet(null); - if (local == null) { - return false; + State local = state.get(); + if (local != null) { + local.close(); } - - onStreamStop(local.closeStream(), Status.SUCCESS, null); - return true; + return local != null; } private void onStreamStop(S closed, Status status, RetryPolicy policy) { @@ -163,72 +166,73 @@ private void onStreamStop(S closed, Status status, RetryPolicy policy) { private class State { private final CompletableFuture> future; - private final AtomicReference stream = new AtomicReference<>(); + private final S stream; State(CompletableFuture> future) { + this(future, null); + } + + State(CompletableFuture> future, S stream) { this.future = future; + this.stream = stream; } 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)); + if (res == null || !res.isSuccess()) { + if (!state.compareAndSet(this, null)) { + return; } - return; - } - if (!res.isSuccess()) { - if (state.compareAndSet(this, null)) { + if (res != null) { onStreamStop(null, res.getStatus(), retryConfig.getStatusRetryPolicy(res.getStatus())); + } else { + Status wrapped = Status.of(StatusCode.CLIENT_INTERNAL_ERROR, th); + onStreamStop(null, wrapped, retryConfig.getThrowableRetryPolicy(th)); } return; } - startStream(res.getValue()); + State next = new State(future, res.getValue()); + if (state.compareAndSet(this, next)) { + next.startStream(); + } }); } - public void closeWithoutStart() { + public boolean close() { + if (stream != null) { + stream.close(); + return true; + } + 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; + if (!state.compareAndSet(this, null)) { + return false; } - stream.set(local); - local.start(msg -> onNext(local, msg)).whenComplete((status, th) -> { + onStreamStop(null, Status.SUCCESS, retryConfig.getStatusRetryPolicy(Status.SUCCESS)); + return true; + } + + private void startStream() { + stream.start(msg -> onNext(stream, msg)).whenComplete((status, th) -> { if (!state.compareAndSet(this, null)) { return; } if (status != null) { - onStreamStop(local, status, retryConfig.getStatusRetryPolicy(status)); + onStreamStop(stream, status, retryConfig.getStatusRetryPolicy(status)); } else { Status wrapped = Status.of(StatusCode.CLIENT_INTERNAL_ERROR, th); - onStreamStop(local, wrapped, retryConfig.getThrowableRetryPolicy(th)); + onStreamStop(stream, wrapped, retryConfig.getThrowableRetryPolicy(th)); } }); - - if (state.get() != this) { - closeStream(); - } - } - - public S closeStream() { - S local = stream.getAndSet(null); - if (local != null) { - local.close(); - } - return local; } } diff --git a/topic/src/main/java/tech/ydb/topic/impl/TopicStreamBase.java b/topic/src/main/java/tech/ydb/topic/impl/TopicStreamBase.java index 522143003..530f63f6b 100644 --- a/topic/src/main/java/tech/ydb/topic/impl/TopicStreamBase.java +++ b/topic/src/main/java/tech/ydb/topic/impl/TopicStreamBase.java @@ -2,6 +2,7 @@ import java.util.Objects; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; import com.google.protobuf.Message; @@ -14,16 +15,15 @@ public abstract class TopicStreamBase implements TopicStream { private final Logger logger; private final String debugId; - private final GrpcReadWriteStream stream; private final W initRequest; private final CompletableFuture streamStatus = new CompletableFuture<>(); + private final AtomicReference> ref; private volatile String token; - private volatile boolean isStopped = false; public TopicStreamBase(Logger logger, String debugId, GrpcReadWriteStream stream, W initRequest) { this.logger = logger; this.debugId = debugId; - this.stream = stream; + this.ref = new AtomicReference<>(stream); this.initRequest = initRequest; this.token = stream.authToken(); } @@ -34,50 +34,65 @@ public TopicStreamBase(Logger logger, String debugId, GrpcReadWriteStream @Override public CompletableFuture start(Consumer messageHandler) { this.logger.debug("[{}] is about to start", debugId); - this.stream.start((R msg) -> { + GrpcReadWriteStream stream = ref.get(); + if (stream == null) { + return CompletableFuture.completedFuture(Status.of(StatusCode.CLIENT_INTERNAL_ERROR)); + } + + stream.start((R msg) -> { Status messageStatus = parseMessageStatus(msg); if (messageStatus.isSuccess()) { messageHandler.accept(msg); } else { logger.warn("[{}] stopped by getting status {}", debugId, messageStatus); - if (streamStatus.complete(messageStatus)) { - stream.close(); - } + preDestroy(); + streamStatus.complete(messageStatus); + closeStream(); } }).whenComplete((st, th) -> { - isStopped = true; + ref.set(null); Status status = st != null ? st : Status.of(StatusCode.CLIENT_INTERNAL_ERROR, th); logger.debug("[{}] finished with status {}", debugId, status); + preDestroy(); streamStatus.complete(status); }); - if (!streamStatus.isDone()) { - stream.sendNext(initRequest); + GrpcReadWriteStream started = ref.get(); + if (started != null) { + started.sendNext(initRequest); } - return streamStatus; } @Override public void close() { logger.debug("[{}] closed by app", debugId); - if (!streamStatus.isDone()) { - isStopped = true; + preDestroy(); + closeStream(); + } + + protected void preDestroy() { + // nothign + } + + private void closeStream() { + GrpcReadWriteStream stream = ref.getAndSet(null); + if (stream != null) { stream.close(); } } protected void fail(Status status) { logger.warn("[{}] stopped by fail {}", debugId, status); - if (streamStatus.complete(status)) { - isStopped = true; - stream.close(); - } + preDestroy(); + streamStatus.complete(status); + closeStream(); } @Override public void send(W req) { - if (isStopped) { + GrpcReadWriteStream stream = ref.get(); + if (stream == null) { logger.warn("[{}] is already closed. Next message with type {} was NOT sent", debugId, req.getDescriptorForType().getName()); return; diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/AsyncReaderImpl.java b/topic/src/main/java/tech/ydb/topic/read/impl/AsyncReaderImpl.java index 352847cfa..5da01583a 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/AsyncReaderImpl.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/AsyncReaderImpl.java @@ -91,6 +91,7 @@ public CompletableFuture shutdown() { // implicit closing because stream will never call onClose close(Status.SUCCESS.withIssues(Issue.of("Closed by client", Issue.Severity.INFO))); } + initFuture.completeExceptionally(new RuntimeException("Reader closed by user")); return shutdownFuture; } diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/ReadSession.java b/topic/src/main/java/tech/ydb/topic/read/impl/ReadSession.java index 6b26f3aa1..3b6969ae0 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/ReadSession.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/ReadSession.java @@ -101,13 +101,13 @@ Handler getHandler() { return handler; } - public void closeAll() { + @Override + protected void preDestroy() { isClosed = true; decoder.stop(); locks.clear(); partitions.values().forEach(ReadPartition::close); partitions.clear(); - config.getMetrics().unregister(); } public boolean commitOffsets(PartitionSession session, List rangesToCommit) { 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 e3548c093..4fcf00823 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 @@ -69,6 +69,7 @@ protected CompletableFuture> createNewStream(String id) { protected void onRetry(ReadSession stream, Status status) { logger.warn("[{}] stopped by status {}", debugId, status); currentSessionId = null; + config.getMetrics().unregister(); if (errorHandler != null) { try { errorHandler.accept(status, null); @@ -76,14 +77,12 @@ protected void onRetry(ReadSession stream, Status status) { logger.error("[{}] errorHandler onRetry processing throws exception", debugId, ex); } } - if (stream != null) { - stream.closeAll(); - } } @Override protected void onClose(ReadSession stream, Status status) { currentSessionId = null; + config.getMetrics().unregister(); if (!status.isSuccess()) { logger.warn("[{}] closed by status {}", debugId, status); } else { @@ -96,9 +95,6 @@ protected void onClose(ReadSession stream, Status status) { logger.error("[{}] errorHandler onClose processing throws exception", debugId, ex); } } - if (stream != null) { - stream.closeAll(); - } onReaderClosed(status); } 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 4a1159277..0f6ce3d77 100644 --- a/topic/src/test/java/tech/ydb/topic/impl/TopicRetryableStreamTest.java +++ b/topic/src/test/java/tech/ydb/topic/impl/TopicRetryableStreamTest.java @@ -133,11 +133,10 @@ public void simpleStartAndCloseTest() { Assert.assertTrue(retryable.close()); Assert.assertTrue(retryable.isClosed()); - h.complete(Status.SUCCESS); - Mockito.verify(h.grpc).close(); Mockito.verify(h.grpc, Mockito.never()).cancel(); + h.complete(Status.SUCCESS); Assert.assertEquals(Arrays.asList(Status.SUCCESS), retryable.closeStatuses); } @@ -165,6 +164,7 @@ public void failStreamTest() { h.complete(Status.SUCCESS); Assert.assertEquals(Arrays.asList(Status.of(StatusCode.ABORTED)), retryable.closeStatuses); + Assert.assertTrue(retryable.retryStatuses.isEmpty()); } @Test @@ -186,21 +186,47 @@ public void doubleCloseTest() { StreamHandle h1 = new StreamHandle(); TestStream retryable = new TestStream(RetryConfig.noRetries(), mockScheduler(), h1); + Assert.assertFalse(retryable.isClosed()); retryable.start(); + Mockito.verify(h1.grpc).start(Mockito.any()); + + Assert.assertFalse(retryable.isClosed()); Assert.assertTrue(retryable.close()); + Assert.assertTrue(retryable.isClosed()); - Assert.assertFalse(retryable.close()); + Mockito.verify(h1.grpc).close(); - Mockito.verify(h1.grpc).start(Mockito.any()); + Assert.assertTrue(retryable.close()); // close is idempotent + Mockito.verify(h1.grpc).close(); + h1.complete(Status.SUCCESS); + + Assert.assertFalse(retryable.close()); // no effect Mockito.verify(h1.grpc).close(); + h1.complete(Status.SUCCESS); } @Test public void startAfterCloseTest() { - TestStream retryable = new TestStream(RetryConfig.noRetries(), mockScheduler()); + StreamHandle h = new StreamHandle(); + TestStream retryable = new TestStream(RetryConfig.noRetries(), mockScheduler(), h); + + Assert.assertFalse(retryable.isClosed()); Assert.assertFalse(retryable.close()); + Assert.assertTrue(retryable.isClosed()); + + Mockito.verify(h.grpc, Mockito.never()).start(Mockito.any()); + Mockito.verify(h.grpc, Mockito.never()).close(); + retryable.start(); // nothing + Assert.assertFalse(retryable.close()); // nothing + + retryable.send(EMPTY); // skipped + + Mockito.verify(h.grpc, Mockito.never()).start(Mockito.any()); + Mockito.verify(h.grpc, Mockito.never()).close(); + + Assert.assertTrue(retryable.isClosed()); } @Test @@ -307,16 +333,6 @@ public void sendBeforeStartIsIgnoredTest() { Mockito.verify(h.grpc, Mockito.never()).sendNext(Mockito.any()); } - @Test - public void closeBeforeStartIsNoOpTest() { - StreamHandle h = new StreamHandle(); - TestStream retryable = new TestStream(RetryConfig.noRetries(), mockScheduler(), h); - - Assert.assertFalse(retryable.close()); // no stream yet, should not throw - - Mockito.verify(h.grpc, Mockito.never()).close(); - } - @Test public void noRetriesErrorStatusTest() { StreamHandle h = new StreamHandle(); @@ -385,13 +401,13 @@ public void immediateRetryTest() { Assert.assertFalse(retryable.close()); // no effect Mockito.verify(h1.grpc, Mockito.times(2)).sendNext(EMPTY); // init req + send - Mockito.verify(h1.grpc, Mockito.never()).close(); // stream was closed by error + Mockito.verify(h1.grpc, Mockito.never()).close(); // stream was already closed with error Mockito.verify(h2.grpc, Mockito.times(3)).sendNext(EMPTY); // init req + 2 * send - Mockito.verify(h2.grpc, Mockito.never()).close(); // stream was closed by error + Mockito.verify(h2.grpc, Mockito.never()).close(); // stream was already closed with error Mockito.verify(h3.grpc, Mockito.times(2)).sendNext(EMPTY); // init req + send - Mockito.verify(h3.grpc, Mockito.never()).close(); // stream was closed by error + Mockito.verify(h3.grpc, Mockito.never()).close(); // stream was already closed with error Assert.assertEquals(Arrays.asList(s1, s2), retryable.retryStatuses); Assert.assertEquals(Arrays.asList(s3), retryable.closeStatuses); diff --git a/topic/src/test/java/tech/ydb/topic/impl/TopicStreamTest.java b/topic/src/test/java/tech/ydb/topic/impl/TopicStreamTest.java index c63b77e65..900ab22db 100644 --- a/topic/src/test/java/tech/ydb/topic/impl/TopicStreamTest.java +++ b/topic/src/test/java/tech/ydb/topic/impl/TopicStreamTest.java @@ -96,6 +96,25 @@ public void baseTest() { stream.close(); // no effect } + @Test + public void closedStreamTest() { + CompletableFuture status = new CompletableFuture<>(); + MockedStream mock = buildMockedStream("token", status); + + TestStream stream = new TestStream(mock); + stream.close(); + + stream.send(msg("1")); + CompletableFuture start = stream.start(msg -> {}); + Assert.assertTrue(start.isDone()); + Assert.assertEquals(Status.of(StatusCode.CLIENT_INTERNAL_ERROR), start.join()); + + Mockito.verify(mock, Mockito.never()).sendNext(Mockito.any()); + Mockito.verify(mock, Mockito.never()).start(Mockito.any()); + Mockito.verify(mock, Mockito.never()).cancel(); + Mockito.verify(mock).close(); + } + @Test public void startStreamAndImmediatelyFinishTest() { // Stream may be closed immediately after start, for example if token is invalid 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 6b2dbc181..efb9be165 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 @@ -116,13 +116,14 @@ public void initAndShutdownTest() { CompletableFuture shutdown = reader.shutdown(); m1.assertIsClosed(); - Assert.assertTrue(shutdown.isDone()); + Assert.assertFalse(shutdown.isDone()); // stream is not closed yet Assert.assertSame(shutdown, reader.shutdown()); // double shutdown is allowed m1.assertIsClosed(); - m1.closeStream(Status.SUCCESS); + m1.closeStream(Status.SUCCESS); Assert.assertTrue(shutdown.isDone()); Assert.assertFalse(shutdown.isCompletedExceptionally()); + Mockito.verify(handler).onReaderClosed(Mockito.any(ReaderClosedEvent.class)); Mockito.verify(handler, Mockito.never()).onMessages(Mockito.any(DataReceivedEvent.class)); } @@ -137,14 +138,15 @@ public void shutdownBeforeInitTest() { CompletableFuture shutdown = reader.shutdown(); mock.assertIsClosed(); + Assert.assertTrue(init.isCompletedExceptionally()); - Assert.assertTrue(shutdown.isDone()); + Assert.assertFalse(shutdown.isDone()); mock.closeStream(Status.SUCCESS); Assert.assertTrue(shutdown.isDone()); Assert.assertFalse(shutdown.isCompletedExceptionally()); CompletionException ex = Assert.assertThrows(CompletionException.class, init::join); - Assert.assertEquals("Reader closed with Status{code = SUCCESS}", ex.getCause().getMessage()); + Assert.assertEquals("Reader closed by user", ex.getCause().getMessage()); Assert.assertSame(init, reader.init()); Mockito.verify(handler).onReaderClosed(Mockito.any(ReaderClosedEvent.class)); Mockito.verify(handler, Mockito.never()).onSessionStarted(Mockito.any(SessionStartedEvent.class)); diff --git a/topic/src/test/java/tech/ydb/topic/read/impl/ReadStreamMock.java b/topic/src/test/java/tech/ydb/topic/read/impl/ReadStreamMock.java index 99f7d8109..e61dff391 100644 --- a/topic/src/test/java/tech/ydb/topic/read/impl/ReadStreamMock.java +++ b/topic/src/test/java/tech/ydb/topic/read/impl/ReadStreamMock.java @@ -39,6 +39,15 @@ public class ReadStreamMock implements GrpcReadWriteStream observer = null; private final AtomicInteger isClosed = new AtomicInteger(); private final AtomicInteger isCanceled = new AtomicInteger(); + private final boolean autoClose; + + public ReadStreamMock() { + this(false); + } + + public ReadStreamMock(boolean autoClose) { + this.autoClose = autoClose; + } @Override public String authToken() { @@ -53,6 +62,9 @@ public void sendNext(FromClient message) { @Override public void close() { isClosed.incrementAndGet(); + if (autoClose) { + future.complete(Status.SUCCESS); + } } @Override diff --git a/topic/src/test/java/tech/ydb/topic/read/impl/ReaderImplTest.java b/topic/src/test/java/tech/ydb/topic/read/impl/ReaderImplTest.java index f152f49eb..15f3417a2 100644 --- a/topic/src/test/java/tech/ydb/topic/read/impl/ReaderImplTest.java +++ b/topic/src/test/java/tech/ydb/topic/read/impl/ReaderImplTest.java @@ -153,8 +153,11 @@ public void partitionStoppedBeforeStartEventTest() { reader.verifyHandler(Mockito.never()).onPartitionClosed(Mockito.any()); reader.close(); - reader.assertReaderIsStopped(Status.SUCCESS); mock.assertIsClosed(); + reader.assertReaderIsNotStopped(); + + mock.closeStream(Status.SUCCESS); + reader.assertReaderIsStopped(Status.SUCCESS); } @Test @@ -209,8 +212,11 @@ public void partitionStoppedBeforeConfirmTest() { reader.verifyHandler(Mockito.times(2)).onPartitionClosed(Mockito.any()); reader.close(); - reader.assertReaderIsStopped(Status.SUCCESS); mock.assertIsClosed(); + reader.assertReaderIsNotStopped(); + + mock.closeStream(Status.SUCCESS); + reader.assertReaderIsStopped(Status.SUCCESS); } @Test @@ -352,8 +358,11 @@ public void onPartitionSessionStatusTest() { mock.responseEmpty(); // nothing reader.close(); - reader.assertReaderIsStopped(Status.SUCCESS); mock.assertIsClosed(); + reader.assertReaderIsNotStopped(); + + mock.closeStream(Status.SUCCESS); + reader.assertReaderIsStopped(Status.SUCCESS); } @Test @@ -647,6 +656,10 @@ public void resetHandler() { Mockito.reset(handler); } + public void assertReaderIsNotStopped() { + Assert.assertTrue(stops.isEmpty()); + } + public void assertReaderIsStopped(Status expected) { Assert.assertEquals(1, stops.size()); Assert.assertEquals(expected, stops.get(0)); diff --git a/topic/src/test/java/tech/ydb/topic/read/impl/ReaderMetricsTest.java b/topic/src/test/java/tech/ydb/topic/read/impl/ReaderMetricsTest.java index c300e764d..6c656a8d3 100644 --- a/topic/src/test/java/tech/ydb/topic/read/impl/ReaderMetricsTest.java +++ b/topic/src/test/java/tech/ydb/topic/read/impl/ReaderMetricsTest.java @@ -31,13 +31,17 @@ public class ReaderMetricsTest { + private static final byte[] MSG1 = new byte[] { 1 }; + private static final byte[] MSG2 = new byte[] { 2 }; + private static final byte[] MSG3 = new byte[] { 3 }; + @Test public void readerCountersIncrementOnReceive() throws InterruptedException { String delivered = "ydb.topic.reader.delivered.messages"; String receivedMessages = "ydb.topic.reader.received.messages"; String receivedBytes = "ydb.topic.reader.received.bytes"; RecordingMeter meter = new RecordingMeter(); - ReadStreamMock stream = new ReadStreamMock(); + ReadStreamMock stream = new ReadStreamMock(true); TopicRpc rpc = Mockito.mock(TopicRpc.class); Mockito.when(rpc.getScheduler()).thenReturn(Mockito.mock(ScheduledExecutorService.class)); Mockito.when(rpc.readSession(Mockito.any(String.class))).thenReturn(stream); @@ -54,16 +58,14 @@ public void readerCountersIncrementOnReceive() throws InterruptedException { stream.assertLastMessage().isInitRequest("consumer", "/topic").hasReaderName("reader"); stream.responseInit("read-session"); stream.responseStartPartition("/topic", 42, 0); - byte[] msg1 = new byte[] { 1 }; - byte[] msg2 = new byte[] { 2 }; stream.responseData(25).partition(1, 0) - .batch(Codec.RAW, msg1, msg2) + .batch(Codec.RAW, MSG1, MSG2) .and().send(); Assert.assertEquals(2, meter.value(receivedMessages)); Assert.assertEquals(25, meter.value(receivedBytes)); stream.responseData(30).partition(1, 2) - .batch(Codec.RAW, msg1, msg2, msg1) + .batch(Codec.RAW, MSG1, MSG2, MSG3) .and().send(); Assert.assertEquals(5, meter.value(receivedMessages)); Assert.assertEquals(55, meter.value(receivedBytes)); @@ -86,7 +88,7 @@ public void gaugesObservePartitionSessionsAndProtocolCredit() throws Interrupted String partitions = "ydb.topic.reader.partition_session.count"; String credit = "ydb.topic.reader.credit_balance_bytes"; RecordingMeter meter = new RecordingMeter(); - ReadStreamMock stream = new ReadStreamMock(); + ReadStreamMock stream = new ReadStreamMock(true); TopicRpc rpc = Mockito.mock(TopicRpc.class); Mockito.when(rpc.getScheduler()).thenReturn(Mockito.mock(ScheduledExecutorService.class)); Mockito.when(rpc.readSession(Mockito.anyString())).thenReturn(stream); @@ -104,7 +106,9 @@ public void gaugesObservePartitionSessionsAndProtocolCredit() throws Interrupted stream.responseStartPartition("/topic", 42, 0); Assert.assertEquals(1, meter.collect(partitions)); Assert.assertEquals(100, meter.collect(credit)); - stream.responseData(20).partition(1, 0).batch(Codec.RAW, new byte[]{1}).and().send(); + stream.responseData(20).partition(1, 0).batch(Codec.RAW, MSG1, MSG2).and().send(); + Assert.assertEquals(80, meter.collect(credit)); + Assert.assertNotNull(reader.receive(1, TimeUnit.SECONDS)); Assert.assertEquals(80, meter.collect(credit)); Assert.assertNotNull(reader.receive(1, TimeUnit.SECONDS)); Assert.assertEquals(100, meter.collect(credit)); 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 99f6c2b8d..00ccaa270 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,8 +53,8 @@ private static TopicRpc mockRpc(ReadStreamMock first, ReadStreamMock... rest) { @Test public void initAndShutdownTest() throws InterruptedException { - ReadStreamMock m1 = new ReadStreamMock(); - ReadStreamMock m2 = new ReadStreamMock(); + ReadStreamMock m1 = new ReadStreamMock(true); + ReadStreamMock m2 = new ReadStreamMock(true); ReaderSettings settings = ReaderSettings.newBuilder() .addTopic(TopicReadSettings.newBuilder().setPath("/test-topic").build()) @@ -99,7 +99,7 @@ public void initAndShutdownTest() throws InterruptedException { @Test public void shutdownBeforeInitTest() throws InterruptedException { - ReadStreamMock mock = new ReadStreamMock(); + ReadStreamMock mock = new ReadStreamMock(true); ReaderSettings settings = ReaderSettings.newBuilder() .addTopic(TopicReadSettings.newBuilder().setPath("/test-topic").build()) @@ -125,7 +125,7 @@ public void shutdownBeforeInitTest() throws InterruptedException { @Test public void shutdownWithoutInitTest() throws InterruptedException { - ReadStreamMock mock = new ReadStreamMock(); + ReadStreamMock mock = new ReadStreamMock(true); ReaderSettings settings = ReaderSettings.newBuilder() .addTopic(TopicReadSettings.newBuilder().setPath("/test-topic").build()) @@ -148,7 +148,7 @@ public void shutdownWithoutInitTest() throws InterruptedException { @Test public void readClosedPartitionTest() throws InterruptedException { - ReadStreamMock mock = new ReadStreamMock(); + ReadStreamMock mock = new ReadStreamMock(true); ReaderSettings settings = ReaderSettings.newBuilder() .addTopic(TopicReadSettings.newBuilder().setPath("/test-topic").build()) @@ -199,7 +199,7 @@ public void readClosedPartitionTest() throws InterruptedException { @Test @HideLoggers({ BufferManager.class, ReadPartition.class }) public void invalidBatchesTest() throws InterruptedException { - ReadStreamMock mock = new ReadStreamMock(); + ReadStreamMock mock = new ReadStreamMock(true); ReaderSettings settings = ReaderSettings.newBuilder() .addTopic(TopicReadSettings.newBuilder().setPath("/test-topic").build()) @@ -236,8 +236,8 @@ public void invalidBatchesTest() throws InterruptedException { public void retrySkipsReadMessagesTest() throws InterruptedException { ErrorsHandler errorsHandler = new ErrorsHandler(); - ReadStreamMock m1 = new ReadStreamMock(); - ReadStreamMock m2 = new ReadStreamMock(); + ReadStreamMock m1 = new ReadStreamMock(true); + ReadStreamMock m2 = new ReadStreamMock(true); ReaderSettings settings = ReaderSettings.newBuilder() .addTopic(TopicReadSettings.newBuilder().setPath("/test-topic").build()) From d22442e220984330a41c8a6ae242450a510447be Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Fri, 16 Oct 2026 12:18:43 +0100 Subject: [PATCH 2/2] Small fixes --- .../java/tech/ydb/topic/impl/TopicRetryableStream.java | 9 +++++---- .../main/java/tech/ydb/topic/impl/TopicStreamBase.java | 2 +- .../main/java/tech/ydb/topic/read/impl/ReadSession.java | 5 +++++ .../main/java/tech/ydb/topic/read/impl/ReaderImpl.java | 4 ++-- 4 files changed, 13 insertions(+), 7 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 7517e431c..319204f81 100644 --- a/topic/src/main/java/tech/ydb/topic/impl/TopicRetryableStream.java +++ b/topic/src/main/java/tech/ydb/topic/impl/TopicRetryableStream.java @@ -206,18 +206,19 @@ public boolean close() { return true; } + boolean effective = state.compareAndSet(this, null); this.future.whenComplete((res, th) -> { if (res != null && res.isSuccess()) { res.getValue().close(); } }); - if (!state.compareAndSet(this, null)) { - return false; + if (effective) { + onStreamStop(null, Status.SUCCESS, retryConfig.getStatusRetryPolicy(Status.SUCCESS)); + return true; } - onStreamStop(null, Status.SUCCESS, retryConfig.getStatusRetryPolicy(Status.SUCCESS)); - return true; + return false; } private void startStream() { diff --git a/topic/src/main/java/tech/ydb/topic/impl/TopicStreamBase.java b/topic/src/main/java/tech/ydb/topic/impl/TopicStreamBase.java index 530f63f6b..9a36605cd 100644 --- a/topic/src/main/java/tech/ydb/topic/impl/TopicStreamBase.java +++ b/topic/src/main/java/tech/ydb/topic/impl/TopicStreamBase.java @@ -72,7 +72,7 @@ public void close() { } protected void preDestroy() { - // nothign + // nothing } private void closeStream() { diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/ReadSession.java b/topic/src/main/java/tech/ydb/topic/read/impl/ReadSession.java index 3b6969ae0..a8d48ad51 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/ReadSession.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/ReadSession.java @@ -213,6 +213,11 @@ public void onStartPartition(YdbTopic.StreamReadMessage.StartPartitionSessionReq return; } + if (isClosed) { + preDestroy(); + return; + } + partition.start(committed, offsets); TopicPartition lockKey = new TopicPartition(ps); if (locks.computeIfPresent(lockKey, (key, lock) -> lock.bind(partition)) != null) { 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 4fcf00823..346152185 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 @@ -69,7 +69,6 @@ protected CompletableFuture> createNewStream(String id) { protected void onRetry(ReadSession stream, Status status) { logger.warn("[{}] stopped by status {}", debugId, status); currentSessionId = null; - config.getMetrics().unregister(); if (errorHandler != null) { try { errorHandler.accept(status, null); @@ -77,12 +76,12 @@ protected void onRetry(ReadSession stream, Status status) { logger.error("[{}] errorHandler onRetry processing throws exception", debugId, ex); } } + config.getMetrics().unregister(); } @Override protected void onClose(ReadSession stream, Status status) { currentSessionId = null; - config.getMetrics().unregister(); if (!status.isSuccess()) { logger.warn("[{}] closed by status {}", debugId, status); } else { @@ -95,6 +94,7 @@ protected void onClose(ReadSession stream, Status status) { logger.error("[{}] errorHandler onClose processing throws exception", debugId, ex); } } + config.getMetrics().unregister(); onReaderClosed(status); }