From fdf5ef76b3ba227a71c99be57f2a5288dc74ced0 Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Thu, 10 Sep 2026 14:27:51 +0100 Subject: [PATCH 1/2] Extended TopicRetryableStream: added the current stream processing --- .../ydb/topic/impl/TopicRetryableStream.java | 60 ++++++++++++------- .../ydb/topic/write/impl/WriteSession.java | 8 +-- .../topic/impl/TopicRetryableStreamTest.java | 48 +++++++++++++-- 3 files changed, 87 insertions(+), 29 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 3cdb97368..fcf65b2c0 100644 --- a/topic/src/main/java/tech/ydb/topic/impl/TopicRetryableStream.java +++ b/topic/src/main/java/tech/ydb/topic/impl/TopicRetryableStream.java @@ -13,13 +13,13 @@ import tech.ydb.core.Status; import tech.ydb.core.StatusCode; -public abstract class TopicRetryableStream { +public abstract class TopicRetryableStream> { protected final String debugId; private final Logger logger; private final RetryConfig retryConfig; private final ScheduledExecutorService scheduler; - private final AtomicReference> realStream = new AtomicReference<>(); + private final AtomicReference realStream = new AtomicReference<>(); private final AtomicInteger streamCount = new AtomicInteger(0); private final RetryState state = new RetryState(); @@ -32,34 +32,38 @@ public TopicRetryableStream(Logger logger, String debugId, RetryConfig config, S this.scheduler = scheduler; } - protected abstract TopicStream createNewStream(String debugId); + protected abstract S createNewStream(String debugId); - protected abstract void onNext(R message); + protected abstract void onNext(S stream, R message); - protected abstract void onRetry(Status status); - protected abstract void onClose(Status status); + protected abstract void onRetry(S stream, Status status); + protected abstract void onClose(S stream, Status status); public void start() { if (isClosed) { + logger.warn("[{}] double start of closed stream, ignored", debugId); return; } String streamID = debugId + '.' + streamCount.incrementAndGet(); - TopicStream stream = createNewStream(streamID); + S stream = createNewStream(streamID); if (!realStream.compareAndSet(null, stream)) { logger.warn("[{}] double start of stream, skipping", debugId); return; } - stream.start(this::onNext).whenComplete((status, th) -> { - realStream.compareAndSet(stream, null); + stream.start(msg -> onNext(stream, msg)).whenComplete((status, th) -> { + S closed = realStream.getAndSet(null); + if (closed == null) { + return; + } if (status != null) { - onStreamStop(status, retryConfig.getStatusRetryPolicy(status)); + onStreamStop(closed, status, retryConfig.getStatusRetryPolicy(status)); } if (th != null) { Status wrapped = Status.of(StatusCode.CLIENT_INTERNAL_ERROR, th); - onStreamStop(wrapped, retryConfig.getThrowableRetryPolicy(th)); + onStreamStop(closed, wrapped, retryConfig.getThrowableRetryPolicy(th)); } }); } @@ -68,8 +72,21 @@ protected void resetRetries() { state.reset(); } + public boolean isClosed() { + return 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)); + } + } + public void send(W msg) { - TopicStream stream = realStream.get(); + S stream = realStream.get(); if (stream == null) { logger.warn("[{}] send message before stream is ready", debugId); return; @@ -79,24 +96,26 @@ public void send(W msg) { public boolean close() { isClosed = true; - TopicStream stream = realStream.getAndSet(null); + S stream = realStream.getAndSet(null); if (stream == null) { return false; } stream.close(); + onStreamStop(stream, Status.SUCCESS, null); return true; } - private void onStreamStop(Status status, RetryPolicy policy) { + private void onStreamStop(S closed, Status status, RetryPolicy policy) { if (isClosed) { // stream was already closed (usually with success) - onClose(status); + onClose(closed, status); return; } if (policy == null) { logger.warn("[{}] stopped by non-retryable status {}", debugId, status); - onClose(status); + isClosed = true; + onClose(closed, status); return; } @@ -104,26 +123,27 @@ private void onStreamStop(Status status, RetryPolicy policy) { if (nextRetryMs < 0) { logger.warn("[{}] stopped after retry policy evaluation for status {}", debugId, status); - onClose(status); + isClosed = true; + onClose(closed, status); return; } if (nextRetryMs == 0) { // retry immediately logger.warn("[{}] retry #{}. Retry immediately...", debugId, state.retryNumber()); - onRetry(status); + onRetry(closed, status); start(); return; } // retry scheduling logger.warn("[{}] retry #{}. Scheduling reconnect in {}ms...", debugId, state.retryNumber(), nextRetryMs); - onRetry(status); + onRetry(closed, status); try { scheduler.schedule(this::start, nextRetryMs, TimeUnit.MILLISECONDS); } catch (Exception ex) { logger.error("[{}] cannot schedule reconnect, stopping", debugId, ex); - onClose(status); + onClose(closed, 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 ee716e43b..d99d9801e 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 @@ -19,7 +19,7 @@ /** * @author Nikolay Perfilov */ -public final class WriteSession extends TopicRetryableStream { +public final class WriteSession extends TopicRetryableStream { public interface Stream extends TopicStream { } private static final Logger logger = LoggerFactory.getLogger(WriteSession.class); @@ -108,7 +108,7 @@ WriteAck mapAck(WriteAck.Statistics statistics, YdbTopic.StreamWriteMessage.Writ } @Override - public void onRetry(Status status) { + public void onRetry(Stream stream, Status status) { logger.warn("[{}] Session onRetry with status {} called", debugId, status); listener.onStop(status); if (errorsHandler != null) { @@ -117,7 +117,7 @@ public void onRetry(Status status) { } @Override - public void onClose(Status status) { + public void onClose(Stream stream, Status status) { logger.info("[{}] Session closed with status {}", debugId, status); listener.onClose(status); if (errorsHandler != null && !status.isSuccess()) { @@ -126,7 +126,7 @@ public void onClose(Status status) { } @Override - public void onNext(YdbTopic.StreamWriteMessage.FromServer message) { + public void onNext(Stream stream, YdbTopic.StreamWriteMessage.FromServer message) { if (message.hasInitResponse()) { onInitResponse(message.getInitResponse()); } else if (message.hasWriteResponse()) { diff --git a/topic/src/test/java/tech/ydb/topic/impl/TopicRetryableStreamTest.java b/topic/src/test/java/tech/ydb/topic/impl/TopicRetryableStreamTest.java index f09eab067..4425d6104 100644 --- a/topic/src/test/java/tech/ydb/topic/impl/TopicRetryableStreamTest.java +++ b/topic/src/test/java/tech/ydb/topic/impl/TopicRetryableStreamTest.java @@ -71,7 +71,7 @@ void fail(Throwable th) { } } - private static class TestStream extends TopicRetryableStream { + private static class TestStream extends TopicRetryableStream> { private final List handles; private int handleIndex = 0; @@ -85,22 +85,22 @@ private static class TestStream extends TopicRetryableStream { } @Override - protected TopicStream createNewStream(String debugId) { + protected TopicStreamBase createNewStream(String debugId) { return handles.get(handleIndex++).stream; } @Override - protected void onNext(Empty message) { + protected void onNext(TopicStreamBase stream, Empty message) { receivedMessages.add(message); } @Override - protected void onRetry(Status status) { + protected void onRetry(TopicStreamBase stream, Status status) { retryStatuses.add(status); } @Override - protected void onClose(Status status) { + protected void onClose(TopicStreamBase stream, Status status) { closeStatuses.add(status); } } @@ -115,12 +115,15 @@ public void simpleStartAndCloseTest() { TestStream retryable = new TestStream(Arrays.asList(h), RetryConfig.noRetries(), mockScheduler()); retryable.start(); + retryable.send(EMPTY); Mockito.verify(h.grpc).start(Mockito.any()); Mockito.verify(h.grpc, Mockito.times(2)).sendNext(EMPTY); // init + sent request + Assert.assertFalse(retryable.isClosed()); Assert.assertTrue(retryable.close()); + Assert.assertTrue(retryable.isClosed()); h.complete(Status.SUCCESS); @@ -130,6 +133,32 @@ public void simpleStartAndCloseTest() { Assert.assertEquals(Arrays.asList(Status.SUCCESS), retryable.closeStatuses); } + @Test + public void failStreamTest() { + StreamHandle h = new StreamHandle(); + TestStream retryable = new TestStream(Arrays.asList(h), RetryConfig.noRetries(), mockScheduler()); + + retryable.start(); + + retryable.send(EMPTY); + + Mockito.verify(h.grpc).start(Mockito.any()); + Mockito.verify(h.grpc, Mockito.times(2)).sendNext(EMPTY); // init + sent request + + Assert.assertFalse(retryable.isClosed()); + retryable.fail(Status.of(StatusCode.ABORTED)); + retryable.fail(Status.of(StatusCode.CLIENT_INTERNAL_ERROR)); // will be ignored + + Assert.assertTrue(retryable.isClosed()); + Assert.assertFalse(retryable.close()); + + Mockito.verify(h.grpc).close(); + Mockito.verify(h.grpc, Mockito.never()).cancel(); + + h.complete(Status.SUCCESS); + Assert.assertEquals(Arrays.asList(Status.of(StatusCode.ABORTED)), retryable.closeStatuses); + } + @Test public void doubleStartTest() { StreamHandle h1 = new StreamHandle(); @@ -152,6 +181,7 @@ public void doubleCloseTest() { retryable.start(); Assert.assertTrue(retryable.close()); + Assert.assertFalse(retryable.close()); Mockito.verify(h1.grpc).start(Mockito.any()); @@ -205,10 +235,14 @@ public void noRetriesExceptionStatusTest() { retryable.start(); RuntimeException ex = new RuntimeException("fail"); + Assert.assertFalse(retryable.isClosed()); h.fail(ex); + Assert.assertTrue(retryable.isClosed()); Assert.assertEquals(Arrays.asList(Status.of(StatusCode.CLIENT_INTERNAL_ERROR, ex)), retryable.closeStatuses); Assert.assertTrue(retryable.retryStatuses.isEmpty()); + + Assert.assertFalse(retryable.close()); } @Test @@ -225,6 +259,7 @@ public void immediateRetryTest() { RetryConfig config = status -> (retryCount, elapsed) -> (status.getCode() != StatusCode.BAD_REQUEST) ? 0 : -1; TestStream retryable = new TestStream(Arrays.asList(h1, h2, h3), config, mockScheduler()); + Assert.assertFalse(retryable.isClosed()); retryable.start(); @@ -232,15 +267,18 @@ public void immediateRetryTest() { retryable.send(EMPTY); h1.complete(s1); + Assert.assertFalse(retryable.isClosed()); Mockito.verify(h2.grpc).start(Mockito.any()); // second stream was started retryable.send(EMPTY); retryable.send(EMPTY); h2.complete(s2); + Assert.assertFalse(retryable.isClosed()); Mockito.verify(h3.grpc).start(Mockito.any()); // third stream was started retryable.send(EMPTY); h3.complete(s3); + Assert.assertTrue(retryable.isClosed()); Assert.assertFalse(retryable.close()); // no effect From 84833fe707617d42d9fe6c5a0d90d4497687bbd8 Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Thu, 10 Sep 2026 14:42:54 +0100 Subject: [PATCH 2/2] Added LazyExecutor utility class --- .../ydb/topic/impl/TopicRetryableStream.java | 8 +- .../ydb/topic/read/impl/LazyExecutor.java | 82 +++++++++++++++++++ .../topic/impl/TopicRetryableStreamTest.java | 3 + .../ydb/topic/read/impl/LazyExecutorTest.java | 72 ++++++++++++++++ 4 files changed, 161 insertions(+), 4 deletions(-) create mode 100644 topic/src/main/java/tech/ydb/topic/read/impl/LazyExecutor.java create mode 100644 topic/src/test/java/tech/ydb/topic/read/impl/LazyExecutorTest.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 fcf65b2c0..e81268305 100644 --- a/topic/src/main/java/tech/ydb/topic/impl/TopicRetryableStream.java +++ b/topic/src/main/java/tech/ydb/topic/impl/TopicRetryableStream.java @@ -54,16 +54,15 @@ public void start() { } stream.start(msg -> onNext(stream, msg)).whenComplete((status, th) -> { - S closed = realStream.getAndSet(null); - if (closed == null) { + if (!realStream.compareAndSet(stream, null)) { return; } if (status != null) { - onStreamStop(closed, status, retryConfig.getStatusRetryPolicy(status)); + onStreamStop(stream, status, retryConfig.getStatusRetryPolicy(status)); } if (th != null) { Status wrapped = Status.of(StatusCode.CLIENT_INTERNAL_ERROR, th); - onStreamStop(closed, wrapped, retryConfig.getThrowableRetryPolicy(th)); + onStreamStop(stream, wrapped, retryConfig.getThrowableRetryPolicy(th)); } }); } @@ -143,6 +142,7 @@ private void onStreamStop(S closed, Status status, RetryPolicy policy) { scheduler.schedule(this::start, nextRetryMs, TimeUnit.MILLISECONDS); } catch (Exception ex) { logger.error("[{}] cannot schedule reconnect, stopping", debugId, ex); + isClosed = true; onClose(closed, status); } } diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/LazyExecutor.java b/topic/src/main/java/tech/ydb/topic/read/impl/LazyExecutor.java new file mode 100644 index 000000000..406d981e6 --- /dev/null +++ b/topic/src/main/java/tech/ydb/topic/read/impl/LazyExecutor.java @@ -0,0 +1,82 @@ +package tech.ydb.topic.read.impl; + +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.ThreadFactory; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * + * @author Aleksandr Gorshenin {@literal } + */ +public class LazyExecutor implements Executor, AutoCloseable { + private static final Logger logger = LoggerFactory.getLogger(LazyExecutor.class); + private static final int MAX_EXECUTOR_THREADS_COUNT = 4; + + private final String name; + private final Executor custom; + private final AtomicInteger threadsCount = new AtomicInteger(); + private final AtomicReference service = new AtomicReference<>(); + private volatile boolean isStopped = false; + + public LazyExecutor(String name, Executor custom) { + this.name = name; + this.custom = custom; + } + + @Override + public void execute(Runnable command) { + if (isStopped) { + return; + } + + if (custom != null) { + custom.execute(command); + return; + } + + ExecutorService local = service.get(); + while (!isStopped && local == null) { + ThreadFactory factory = r -> new Thread(r, name + "-" + threadsCount.incrementAndGet()); + ExecutorService pool = Executors.newFixedThreadPool(MAX_EXECUTOR_THREADS_COUNT, factory); + if (!service.compareAndSet(local, pool)) { + pool.shutdown(); + } + local = service.get(); + } + + if (!isStopped) { + local.execute(command); + } else { + shutdown(service.getAndSet(null)); + } + } + + @Override + public void close() { + isStopped = true; + shutdown(service.getAndSet(null)); + } + + private void shutdown(ExecutorService service) { + if (service == null) { + return; + } + + try { + service.shutdown(); + if (!service.awaitTermination(100, TimeUnit.MILLISECONDS)) { + service.shutdownNow(); + } + } catch (InterruptedException e) { + logger.warn("executor {} shutdown interrupted", name, e); + Thread.currentThread().interrupt(); + } + } +} 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 4425d6104..26735242f 100644 --- a/topic/src/test/java/tech/ydb/topic/impl/TopicRetryableStreamTest.java +++ b/topic/src/test/java/tech/ydb/topic/impl/TopicRetryableStreamTest.java @@ -305,8 +305,11 @@ public void closeOnWrongSchedulerTest() { TestStream retryable = new TestStream(Arrays.asList(h), config, null); retryable.start(); + Assert.assertFalse(retryable.isClosed()); + h.complete(Status.of(StatusCode.UNAVAILABLE)); + Assert.assertTrue(retryable.isClosed()); Assert.assertEquals(1, retryable.retryStatuses.size()); Assert.assertEquals(1, retryable.closeStatuses.size()); } diff --git a/topic/src/test/java/tech/ydb/topic/read/impl/LazyExecutorTest.java b/topic/src/test/java/tech/ydb/topic/read/impl/LazyExecutorTest.java new file mode 100644 index 000000000..b74011c93 --- /dev/null +++ b/topic/src/test/java/tech/ydb/topic/read/impl/LazyExecutorTest.java @@ -0,0 +1,72 @@ +package tech.ydb.topic.read.impl; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; + +import org.junit.Assert; +import org.junit.Test; + +/** + * + * @author Aleksandr Gorshenin {@literal } + */ +public class LazyExecutorTest { + + @Test + public void customExecutorTest() { + AtomicInteger counter = new AtomicInteger(); + List queue = new ArrayList<>(); + LazyExecutor custom = new LazyExecutor("test", queue::add); + + custom.execute(counter::incrementAndGet); + custom.execute(counter::incrementAndGet); + custom.execute(counter::incrementAndGet); + + Assert.assertEquals(0, counter.get()); + Assert.assertEquals(3, queue.size()); + queue.forEach(Runnable::run); + Assert.assertEquals(3, counter.get()); + + custom.close(); + + custom.execute(counter::incrementAndGet); + custom.execute(counter::incrementAndGet); + + Assert.assertEquals(3, counter.get()); + Assert.assertEquals(3, queue.size()); + + custom.close(); // no effect + } + + @Test + public void lazyExecutorTest() throws InterruptedException { + AtomicInteger counter = new AtomicInteger(); + + LazyExecutor lazy = new LazyExecutor("lazy", null); + CountDownLatch latch = new CountDownLatch(600); + + ExecutorService producer = Executors.newFixedThreadPool(6); + for (int i = 0; i < 6; i++) { + producer.execute(() -> { + for (int j = 0; j < 100; j++) { + lazy.execute(counter::incrementAndGet); + latch.countDown(); + } + }); + } + + Assert.assertTrue(latch.await(5, TimeUnit.SECONDS)); + producer.shutdown(); + Assert.assertTrue(producer.awaitTermination(5, TimeUnit.SECONDS)); + + lazy.close(); + Assert.assertEquals(600, counter.get()); + lazy.close(); + Assert.assertEquals(600, counter.get()); + } +}