From 7bac8e41394d7689ab0cf3d6db56aa576e61dc40 Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Thu, 10 Sep 2026 14:42:54 +0100 Subject: [PATCH 01/12] Added support of RetryPolicy to topic readers --- .../ydb/topic/read/impl/AsyncReaderImpl.java | 228 +++---- .../topic/read/impl/MessageCommitterImpl.java | 43 +- .../ydb/topic/read/impl/MessageDecoder.java | 6 +- .../tech/ydb/topic/read/impl/ReadConfig.java | 52 ++ .../topic/read/impl/ReadPartitionSession.java | 59 +- .../tech/ydb/topic/read/impl/ReadSession.java | 558 ++++++------------ .../tech/ydb/topic/read/impl/ReaderImpl.java | 302 +++++++--- .../ydb/topic/read/impl/SyncReaderImpl.java | 250 ++++---- .../ydb/topic/settings/ReaderSettings.java | 35 ++ 9 files changed, 796 insertions(+), 737 deletions(-) create mode 100644 topic/src/main/java/tech/ydb/topic/read/impl/ReadConfig.java 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 bf5d4a10a..ed81e1be1 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 @@ -3,9 +3,7 @@ import java.util.List; import java.util.Map; import java.util.concurrent.CompletableFuture; -import java.util.concurrent.Executor; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; +import java.util.stream.Collectors; import javax.annotation.Nonnull; @@ -13,9 +11,12 @@ import org.slf4j.LoggerFactory; import tech.ydb.common.transaction.YdbTransaction; +import tech.ydb.core.Issue; import tech.ydb.core.Status; +import tech.ydb.core.StatusCode; import tech.ydb.topic.TopicRpc; import tech.ydb.topic.description.CodecRegistry; +import tech.ydb.topic.impl.DebugTools; import tech.ydb.topic.impl.SerialExecutor; import tech.ydb.topic.read.AsyncReader; import tech.ydb.topic.read.PartitionOffsets; @@ -25,6 +26,7 @@ import tech.ydb.topic.read.events.ReaderClosedEvent; import tech.ydb.topic.read.events.StartPartitionSessionEvent; import tech.ydb.topic.read.events.StopPartitionSessionEvent; +import tech.ydb.topic.read.impl.ReaderImpl.Releaser; import tech.ydb.topic.read.impl.events.CommitOffsetAcknowledgementEventImpl; import tech.ydb.topic.read.impl.events.PartitionSessionClosedEventImpl; import tech.ydb.topic.read.impl.events.SessionStartedEvent; @@ -35,159 +37,171 @@ /** * @author Nikolay Perfilov */ -public class AsyncReaderImpl extends ReaderImpl implements AsyncReader { +public class AsyncReaderImpl implements AsyncReader { private static final Logger logger = LoggerFactory.getLogger(AsyncReaderImpl.class); - private static final int DEFAULT_HANDLER_THREAD_COUNT = 4; - private final Executor handlerExecutor; - private final ExecutorService defaultHandlerExecutorService; + private final String debugId; + private final LazyExecutor processor; + private final LazyExecutor decompressor; private final ReadEventHandler eventHandler; private final SerialExecutor controlEventsExecutor; + private final ReadConfig config; + private final ReaderImpl impl; + + private final CompletableFuture initFuture = new CompletableFuture<>(); + private final CompletableFuture shutdownFuture = new CompletableFuture<>(); public AsyncReaderImpl(TopicRpc topicRpc, ReaderSettings settings, ReadEventHandlersSettings handlersSettings, @Nonnull CodecRegistry codecRegistry) { - super(topicRpc, settings, codecRegistry); + this.debugId = DebugTools.createDebugId(settings.getLogPrefix()); this.eventHandler = handlersSettings.getEventHandler(); - - if (handlersSettings.getExecutor() != null) { - logger.debug("Using handler executor provided by user"); - this.defaultHandlerExecutorService = null; - this.handlerExecutor = handlersSettings.getExecutor(); - } else { - logger.debug("Using default handler executor"); - this.defaultHandlerExecutorService = Executors.newFixedThreadPool(DEFAULT_HANDLER_THREAD_COUNT); - this.handlerExecutor = defaultHandlerExecutorService; - } - - this.controlEventsExecutor = new SerialExecutor(handlerExecutor); + this.processor = new LazyExecutor("reader[" + debugId + "]-handler", handlersSettings.getExecutor()); + this.decompressor = new LazyExecutor("reader[" + debugId + "]-decoder", settings.getDecompressionExecutor()); + this.controlEventsExecutor = new SerialExecutor(processor); + + this.config = new ReadConfig(codecRegistry, processor, decompressor, settings); + this.impl = new ReaderImpl(topicRpc, debugId, settings, config, new AsyncHandler()); + + String readerName = settings.getReaderName(); + String consumerName = settings.getConsumerName(); + logger.info("Reader{} (generated id {}) created for topic(s) {} and {}", + readerName != null ? (" '" + readerName + "'") : "", + debugId, + settings.getTopics().stream().map(t -> "\"" + t.getPath() + "\"").collect(Collectors.joining(", ")), + consumerName != null ? (" consumer \"" + consumerName + "\"") : "without a consumer" + ); } @Override public CompletableFuture init() { - return initImpl(); + impl.start(); + return initFuture; } @Override public CompletableFuture updateOffsetsInTransaction(YdbTransaction transaction, Map> offsets, UpdateOffsetsInTransactionSettings settings) { - return super.updateOffsetsInTransaction(transaction, offsets, settings); - } - - @Override - Executor getDataHandlerExecutor() { - return handlerExecutor; + return impl.updateOffsetsInTransaction(transaction, offsets, settings); } - @Override - protected void handleSessionStarted(String sessionId) { - controlEventsExecutor.execute(() -> { + protected CompletableFuture handleReaderClosed() { + return CompletableFuture.runAsync(() -> { try { - eventHandler.onSessionStarted(new SessionStartedEvent(sessionId)); + eventHandler.onReaderClosed(new ReaderClosedEvent()); } catch (Throwable th) { - logUserThrowableAndStopWorking(th, "onSessionStarted"); + failSession(th, "onReaderClosed"); throw th; } - }); + }, controlEventsExecutor); } @Override - protected void handleDataReceivedEvent(ReadPartitionSession session, DataReceivedEvent event) { - try { - int messagesCount = event.getMessages().size(); - long offsetStart = event.getMessages().get(0).getOffset(); - long offsetEnd = event.getMessages().get(event.getMessages().size() - 1).getOffset(); - logger.debug("{} DataReceivedEvent callback with {} message(s) (offsets {}-{}) is about " - + "to be called...", session, messagesCount, offsetStart, offsetEnd); - eventHandler.onMessages(event); - logger.debug("{} DataReceivedEvent callback with {} message(s) (offsets {}-{}) " - + "successfully finished", session, messagesCount, offsetStart, offsetEnd); - } catch (Throwable th) { - logUserThrowableAndStopWorking(th, "onMessages"); - throw th; - } finally { - session.releaseRange(event.getRangeToCommit()); - } + public CompletableFuture shutdown() { + impl.close(); + return shutdownFuture; } - @Override - protected void handleCommitResponse(long committedOffset, PartitionSession partition) { - handlerExecutor.execute(() -> { - try { - eventHandler.onCommitResponse(new CommitOffsetAcknowledgementEventImpl(partition, committedOffset)); - } catch (Throwable th) { - logUserThrowableAndStopWorking(th, "onCommitResponse"); - throw th; - } - }); + private void close() { + decompressor.close(); + processor.close(); + shutdownFuture.complete(null); } - @Override - protected void handleStartPartitionSessionRequest(StartPartitionSessionEvent event) { - controlEventsExecutor.execute(() -> { - try { - eventHandler.onStartPartitionSession(event); - } catch (Throwable th) { - logUserThrowableAndStopWorking(th, "onStartPartitionSession"); - throw th; - } - }); + private void failSession(Throwable th, String callbackName) { + String errorMessage = "Unhandled throwable in " + callbackName + " user callback: " + th.getMessage(); + logger.error(errorMessage, th); + impl.fail(Status.of(StatusCode.CLIENT_INTERNAL_ERROR, th, Issue.of(errorMessage, Issue.Severity.ERROR))); } - @Override - protected void handleStopPartitionSession(StopPartitionSessionEvent event) { - controlEventsExecutor.execute(() -> { + private class AsyncHandler implements ReaderImpl.Handler { + @Override + public void handleSessionStarted(String sessionId) { + initFuture.complete(null); try { - eventHandler.onStopPartitionSession(event); + eventHandler.onSessionStarted(new SessionStartedEvent(sessionId)); } catch (Throwable th) { - logUserThrowableAndStopWorking(th, "onStopPartitionSession"); - throw th; + failSession(th, "onSessionStarted"); } - }); - } + } - @Override - protected void handleClosePartitionSession(PartitionSession partition) { - controlEventsExecutor.execute(() -> { + @Override + public void handleReaderClosed(Status status) { try { - eventHandler.onPartitionSessionClosed(new PartitionSessionClosedEventImpl(partition)); + eventHandler.onReaderClosed(new ReaderClosedEvent()); } catch (Throwable th) { - logUserThrowableAndStopWorking(th, "onPartitionSessionClosed"); - throw th; + failSession(th, "onReaderClosed"); + } finally { + close(); } - }); - } + } - protected CompletableFuture handleReaderClosed() { - return CompletableFuture.runAsync(() -> { + @Override + public void handleDataReceivedEvent(Releaser releaser, DataReceivedEvent event) { try { - eventHandler.onReaderClosed(new ReaderClosedEvent()); + int messagesCount = event.getMessages().size(); + long offsetStart = event.getMessages().get(0).getOffset(); + long offsetEnd = event.getMessages().get(event.getMessages().size() - 1).getOffset(); + logger.debug("[{}] DataReceivedEvent callback with {} message(s) (offsets {}-{}) is about " + + "to be called...", debugId, messagesCount, offsetStart, offsetEnd); + eventHandler.onMessages(event); + logger.debug("[{}] DataReceivedEvent callback with {} message(s) (offsets {}-{}) " + + "successfully finished", debugId, messagesCount, offsetStart, offsetEnd); } catch (Throwable th) { - logUserThrowableAndStopWorking(th, "onReaderClosed"); + failSession(th, "onMessages"); throw th; + } finally { + releaser.releaseRange(event.getPartitionSession(), event.getRangeToCommit()); } - }, controlEventsExecutor); - } + } - @Override - protected void onShutdown(String reason) { - super.onShutdown(reason); - handleReaderClosed().join(); - if (defaultHandlerExecutorService != null) { - logger.debug("Shutting down default handler executor"); - defaultHandlerExecutorService.shutdown(); + @Override + public void handleCommitResponse(long committedOffset, PartitionSession partition) { + processor.execute(() -> { + try { + eventHandler.onCommitResponse(new CommitOffsetAcknowledgementEventImpl(partition, committedOffset)); + } catch (Throwable th) { + failSession(th, "onCommitResponse"); + throw th; + } + }); } - } - @Override - public CompletableFuture shutdown() { - return shutdownImpl(); - } + @Override + public void handleStartPartitionSessionRequest(StartPartitionSessionEvent event) { + controlEventsExecutor.execute(() -> { + try { + eventHandler.onStartPartitionSession(event); + } catch (Throwable th) { + failSession(th, "onStartPartitionSession"); + throw th; + } + }); + } - private void logUserThrowableAndStopWorking(Throwable th, String callbackName) { - String errorMessage = "Unhandled throwable in " + callbackName + " user callback: " + th; - logger.error(errorMessage); - shutdownImpl(errorMessage); + @Override + public void handleStopPartitionSession(StopPartitionSessionEvent event) { + controlEventsExecutor.execute(() -> { + try { + eventHandler.onStopPartitionSession(event); + } catch (Throwable th) { + failSession(th, "onStopPartitionSession"); + throw th; + } + }); + } + + @Override + public void handleClosePartitionSession(PartitionSession partition) { + controlEventsExecutor.execute(() -> { + try { + eventHandler.onPartitionSessionClosed(new PartitionSessionClosedEventImpl(partition)); + } catch (Throwable th) { + failSession(th, "onPartitionSessionClosed"); + throw th; + } + }); + } } } diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/MessageCommitterImpl.java b/topic/src/main/java/tech/ydb/topic/read/impl/MessageCommitterImpl.java index 09dce0dfb..f66957c11 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/MessageCommitterImpl.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/MessageCommitterImpl.java @@ -13,6 +13,7 @@ import tech.ydb.topic.description.OffsetsRange; import tech.ydb.topic.read.MessageCommitter; +import tech.ydb.topic.read.PartitionSession; /** * @@ -21,26 +22,30 @@ class MessageCommitterImpl implements MessageCommitter { private static final Logger logger = LoggerFactory.getLogger(ReaderImpl.class); - private final ReadPartitionSession session; + private final String debugId; + private final ReadSession stream; + private final PartitionSession partition; private final NavigableMap> commitFutures = new TreeMap<>(); private final ReentrantLock commitFuturesLock = new ReentrantLock(); private volatile long lastCommittedOffset; - MessageCommitterImpl(ReadPartitionSession session, long lastCommittedOffset) { - this.session = session; + MessageCommitterImpl(String debugId, ReadSession stream, PartitionSession partition, long lastCommittedOffset) { + this.debugId = debugId; + this.stream = stream; + this.partition = partition; this.lastCommittedOffset = lastCommittedOffset; } private RuntimeException partitionIsClosedException() { - return new RuntimeException("" + session.getPartition() + " is already stopped"); + return new RuntimeException("" + partition + " is already stopped"); } public void confirmCommit(long committedOffset) { if (committedOffset <= lastCommittedOffset) { // never happens - logger.error("{} received commit response. Committed offset: {} which is less than previous " + - "committed offset: {}.", session, committedOffset, lastCommittedOffset); + logger.error("[{}] received commit response. Committed offset: {} which is less than previous " + + "committed offset: {}.", debugId, committedOffset, lastCommittedOffset); return; } @@ -48,8 +53,8 @@ public void confirmCommit(long committedOffset) { try { Map> confirmed = commitFutures.headMap(committedOffset, true); - logger.debug("{} received commit response. Committed offset: {}. " - + "Previous committed offset: {} (diff is {} message(s)). Completing {} commit futures", session, + logger.debug("[{}] received commit response. Committed offset: {}. " + + "Previous committed offset: {} (diff is {} message(s)). Completing {} commit futures", debugId, committedOffset, lastCommittedOffset, committedOffset - lastCommittedOffset, confirmed.size()); lastCommittedOffset = committedOffset; @@ -62,8 +67,10 @@ public void confirmCommit(long committedOffset) { @Override public CompletableFuture commit(OffsetsRange range) { - logger.debug("{} Offset range {} is requested to be committed. Last committed offset is {} (commit lag is {})", - session, range, lastCommittedOffset, range.getStart() - lastCommittedOffset); + logger.debug( + "[{}] Offset range {} is requested to be committed. Last committed offset is {} (commit lag is {})", + debugId, range, lastCommittedOffset, range.getStart() - lastCommittedOffset + ); CompletableFuture future; commitFuturesLock.lock(); @@ -77,9 +84,9 @@ public CompletableFuture commit(OffsetsRange range) { commitFuturesLock.unlock(); } - if (!session.commitOffsets(Collections.singletonList(range))) { - logger.info("{} Offset range {} is requested to be committed, but partition session is already stopped", - session, range); + if (!stream.commitOffsets(partition, Collections.singletonList(range))) { + logger.info("[{}] Offset range {} is requested to be committed, but partition session is already stopped", + debugId, range); future.completeExceptionally(partitionIsClosedException()); commitFuturesLock.lock(); @@ -95,14 +102,18 @@ public CompletableFuture commit(OffsetsRange range) { @Override public void commitRanges(List ranges) { - session.commitOffsets(ranges); + stream.commitOffsets(partition, ranges); } public void failPendingCommits() { commitFuturesLock.lock(); try { - logger.info("{} for {} is stopping. Failing {} commit futures...", session, - session.getPartition().getPath(), commitFutures.size()); + if (commitFutures.isEmpty()) { + return; + } + + logger.info("[{}] for {} is stopping. Failing {} commit futures...", debugId, partition.getPath(), + commitFutures.size()); commitFutures.values().forEach(f -> f.completeExceptionally(partitionIsClosedException())); commitFutures.clear(); } finally { diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/MessageDecoder.java b/topic/src/main/java/tech/ydb/topic/read/impl/MessageDecoder.java index 8357af332..4cc15f526 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/MessageDecoder.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/MessageDecoder.java @@ -25,7 +25,11 @@ public class MessageDecoder { private final SerialRunnable decodeNext = new SerialRunnable(new DecodeNext()); private volatile boolean isStopped = false; - public MessageDecoder(long maxBufferSize, Executor decompressionExecutor, CodecRegistry codecRegistry) { + public MessageDecoder(ReadConfig config) { + this(config.getMaxMemoryUsageBytes(), config.getDecompressor(), config.getCodecRegistry()); + } + + MessageDecoder(long maxBufferSize, Executor decompressionExecutor, CodecRegistry codecRegistry) { this.totalAvailable = new AtomicLong(maxBufferSize); this.decompressionExecutor = decompressionExecutor; this.codecRegistry = codecRegistry; diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/ReadConfig.java b/topic/src/main/java/tech/ydb/topic/read/impl/ReadConfig.java new file mode 100644 index 000000000..353fcd181 --- /dev/null +++ b/topic/src/main/java/tech/ydb/topic/read/impl/ReadConfig.java @@ -0,0 +1,52 @@ +package tech.ydb.topic.read.impl; + +import java.util.concurrent.Executor; + +import tech.ydb.topic.description.CodecRegistry; +import tech.ydb.topic.settings.ReaderSettings; + +/** + * + * @author Aleksandr Gorshenin {@literal } + */ +public class ReadConfig { + private final CodecRegistry codecRegistry; + private final Executor processor; + private final Executor decompressor; + private final String consumerName; + private final long maxMemoryUsageBytes; + private final int maxBatchSize; + + public ReadConfig(CodecRegistry codecRegistry, Executor processor, Executor decompressor, ReaderSettings settings) { + this.codecRegistry = codecRegistry; + this.processor = processor; + this.decompressor = decompressor; + this.consumerName = settings.getConsumerName(); + this.maxMemoryUsageBytes = settings.getMaxMemoryUsageBytes(); + this.maxBatchSize = settings.getMaxBatchSize(); + } + + public CodecRegistry getCodecRegistry() { + return codecRegistry; + } + + public Executor getDecompressor() { + return decompressor; + } + + public Executor getProcessor() { + return processor; + } + + public long getMaxMemoryUsageBytes() { + return maxMemoryUsageBytes; + } + + public String getConsumerName() { + return consumerName; + } + + public int getMaxBatchSize() { + return maxBatchSize; + } +} diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionSession.java b/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionSession.java index b20764953..eeabf1d29 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionSession.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionSession.java @@ -5,8 +5,7 @@ import java.util.List; import java.util.Queue; import java.util.concurrent.ConcurrentLinkedQueue; -import java.util.concurrent.Executor; -import java.util.stream.Collectors; +import java.util.function.Consumer; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -23,35 +22,33 @@ /** * @author Nikolay Perfilov */ -public abstract class ReadPartitionSession { - +public class ReadPartitionSession { private static final Logger logger = LoggerFactory.getLogger(ReaderImpl.class); private final String traceID; - private final ReadSession session; private final PartitionSession partition; - private final int maxBatchSize; - private final BufferManager bufferManager; - private final SerialExecutor executor; private final MessageCommitterImpl committer; private final ReadPartitionDecoder decoder; + private final Consumer eventConsumer; + + private final int maxBatchSize; + private final SerialExecutor executor; private volatile long lastReadOffset; private volatile boolean isStopped = false; private final Queue readingQueue = new ConcurrentLinkedQueue<>(); - ReadPartitionSession(String traceID, ReadSession session, PartitionSession partition, Executor executor, - long lastCommittedOffset) { + ReadPartitionSession(String traceID, ReadConfig config, PartitionSession partition, MessageCommitterImpl committer, + MessageDecoder decoder, Consumer eventConsumer, long lastCommittedOffset) { this.traceID = traceID; - this.session = session; this.partition = partition; - this.maxBatchSize = session.getMaxBatchSize(); - this.bufferManager = session.getBufferManager(); - this.executor = new SerialExecutor(executor); - this.committer = new MessageCommitterImpl(this, lastCommittedOffset); - this.decoder = new ReadPartitionDecoder(traceID, session.getMessageDecoder(), partition, committer, - this::sendDataToReaders); + this.committer = committer; + this.decoder = new ReadPartitionDecoder(traceID, decoder, partition, committer, this::sendDataToReaders); + + this.maxBatchSize = config.getMaxBatchSize(); + this.executor = new SerialExecutor(config.getProcessor()); + this.eventConsumer = eventConsumer; this.lastReadOffset = lastCommittedOffset; } @@ -59,38 +56,21 @@ public PartitionSession getPartition() { return partition; } - @Override - public String toString() { - return "[" + traceID + "]"; - } - public boolean isStopped() { return isStopped; } - boolean commitOffsets(List ranges) { - if (isStopped) { - logger.info("[{}] Offset ranges {} are requested to be committed, but partition session is already closed", - traceID, ranges.stream().map(OffsetsRange::toString).collect(Collectors.joining(","))); - return false; - } - session.sendCommitOffsetRequest(partition, ranges); - return true; - } - - void confirmCommit(long committedOffset) { + public void confirmCommittedOffset(long committedOffset) { committer.confirmCommit(committedOffset); } public void stop() { isStopped = true; - committer.failPendingCommits(); decoder.close(); + committer.failPendingCommits(); logger.info("[{}] stopped", traceID); } - public abstract void handleDataReceivedEvent(DataReceivedEvent event); - public boolean addBatches(List batchList) { if (isStopped) { return false; @@ -135,11 +115,10 @@ public boolean addBatches(List ba public void releaseRange(OffsetsRange range) { decoder.releaseRange(range); - bufferManager.releaseRange(partition.getId(), range); sendDataToReaders(); } - private void sendDataToReaders() { + public void sendDataToReaders() { executor.execute(() -> { while (!isStopped) { Iterator it = readingQueue.iterator(); @@ -159,9 +138,7 @@ private void sendDataToReaders() { next = it.hasNext() ? it.next() : null; } - // Should be called maximum in 1 thread at a time - DataReceivedEvent event = new DataReceivedEventImpl(partition, committer, messagesToRead); - handleDataReceivedEvent(event); + eventConsumer.accept(new DataReceivedEventImpl(partition, committer, messagesToRead)); } }); } 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 1d108fb29..018ac52fc 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 @@ -1,168 +1,126 @@ package tech.ydb.topic.read.impl; -import java.time.Duration; -import java.time.Instant; +import java.util.HashSet; import java.util.List; import java.util.Map; -import java.util.UUID; -import java.util.concurrent.CompletableFuture; +import java.util.Set; import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.Executor; +import java.util.function.BiConsumer; +import java.util.function.Consumer; import java.util.stream.Collectors; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import tech.ydb.common.transaction.YdbTransaction; import tech.ydb.core.Issue; import tech.ydb.core.Status; import tech.ydb.core.StatusCode; -import tech.ydb.core.grpc.GrpcRequestSettings; -import tech.ydb.core.utils.ProtobufUtils; -import tech.ydb.proto.StatusCodesProtos; +import tech.ydb.core.grpc.GrpcReadWriteStream; import tech.ydb.proto.topic.YdbTopic; -import tech.ydb.topic.TopicRpc; +import tech.ydb.proto.topic.YdbTopic.StreamReadMessage.CommitOffsetRequest; +import tech.ydb.proto.topic.YdbTopic.StreamReadMessage.CommitOffsetResponse; +import tech.ydb.proto.topic.YdbTopic.StreamReadMessage.FromClient; +import tech.ydb.proto.topic.YdbTopic.StreamReadMessage.FromServer; +import tech.ydb.proto.topic.YdbTopic.StreamReadMessage.StartPartitionSessionResponse; import tech.ydb.topic.description.OffsetsRange; -import tech.ydb.topic.impl.SessionBase; -import tech.ydb.topic.read.PartitionOffsets; +import tech.ydb.topic.impl.TopicStreamBase; import tech.ydb.topic.read.PartitionSession; import tech.ydb.topic.read.events.DataReceivedEvent; +import tech.ydb.topic.read.events.StartPartitionSessionEvent; +import tech.ydb.topic.read.events.StopPartitionSessionEvent; import tech.ydb.topic.read.impl.events.StartPartitionSessionEventImpl; import tech.ydb.topic.read.impl.events.StopPartitionSessionEventImpl; -import tech.ydb.topic.settings.ReaderSettings; import tech.ydb.topic.settings.StartPartitionSessionSettings; -import tech.ydb.topic.settings.TopicReadSettings; -import tech.ydb.topic.settings.UpdateOffsetsInTransactionSettings; /** - * @author Nikolay Perfilov + * + * @author Aleksandr Gorshenin {@literal } */ -public final class ReadSession extends SessionBase { - private static final Logger logger = LoggerFactory.getLogger(ReaderImpl.class); +public class ReadSession extends TopicStreamBase implements ReaderImpl.Releaser { + private static final Logger logger = LoggerFactory.getLogger(ReadSession.class); - private final TopicRpc rpc; - private final ReaderImpl reader; - - private final String consumerName; - private final YdbTopic.StreamReadMessage.InitRequest initRequest; - - private final int maxBatchSize; + private final String debugId; + private final ReadConfig config; private final MessageDecoder decoder; private final BufferManager bufferManager; + private final BiConsumer eventConsumer; private final Map partitions = new ConcurrentHashMap<>(); - private final Map partSessions = new ConcurrentHashMap<>(); - - public ReadSession(TopicRpc rpc, ReaderImpl reader, MessageDecoder decoder, String id, ReaderSettings settings) { - super(rpc.readSession(id), id); - this.reader = reader; - this.rpc = rpc; - this.decoder = decoder; - this.bufferManager = new BufferManager(id, settings.getMaxMemoryUsageBytes(), this::sendReadRequest); - - this.consumerName = settings.getConsumerName(); - this.maxBatchSize = settings.getMaxBatchSize(); - this.initRequest = buildInitRequest(settings); - } - - @Override - protected Logger getLogger() { - return logger; - } - - int getMaxBatchSize() { - return maxBatchSize; - } - - MessageDecoder getMessageDecoder() { - return decoder; - } - - BufferManager getBufferManager() { - return bufferManager; + private final Map readQueues = new ConcurrentHashMap<>(); + private volatile boolean isClosed = false; + + public ReadSession(String id, GrpcReadWriteStream stream, FromClient initReq, + BiConsumer eventConsumer, ReadConfig config) { + super(logger, id, stream, initReq); + this.debugId = id; + this.config = config; + this.decoder = new MessageDecoder(config); + this.bufferManager = new BufferManager(id, config.getMaxMemoryUsageBytes(), new ReadRequest()); + this.eventConsumer = eventConsumer; } @Override - protected void sendUpdateTokenRequest(String token) { - streamConnection.sendNext(YdbTopic.StreamReadMessage.FromClient.newBuilder() - .setUpdateTokenRequest(YdbTopic.UpdateTokenRequest.newBuilder() - .setToken(token) - .build()) - .build() - ); + protected FromClient updateTokenMessage(String token) { + YdbTopic.UpdateTokenRequest req = YdbTopic.UpdateTokenRequest.newBuilder().setToken(token).build(); + return FromClient.newBuilder().setUpdateTokenRequest(req).build(); } @Override - public void startAndInitialize() { - logger.debug("[{}] Session startAndInitialize called", streamId); - start(this::processMessage).whenComplete(this::closeDueToError); - - send(YdbTopic.StreamReadMessage.FromClient.newBuilder().setInitRequest(initRequest).build()); + protected Status parseMessageStatus(FromServer message) { + return Status.of(StatusCode.fromProto(message.getStatus()), Issue.fromPb(message.getIssuesList())); } - @Override - protected void onStop() { - logger.debug("[{}] Session onStop called", streamId); - + public Set closeAll() { decoder.stop(); - partSessions.values().forEach(ReadPartitionSession::stop); - partSessions.clear(); - - partitions.values().forEach(reader::handleClosePartitionSession); + Set closed = new HashSet<>(partitions.values()); partitions.clear(); - } - protected void closeDueToError(Status status, Throwable th) { - logger.info("[{}] Session closeDueToError called", streamId); - if (shutdown()) { - // Signal reader to retry - reader.onSessionClosed(status, th); - } - } + readQueues.values().forEach(ReadPartitionSession::stop); + readQueues.clear(); - private void sendReadRequest(long sizeToRequest) { - logger.debug("[{}] Sending DataRequest with {} bytes", streamId, sizeToRequest); + return closed; + } - send(YdbTopic.StreamReadMessage.FromClient.newBuilder() - .setReadRequest(YdbTopic.StreamReadMessage.ReadRequest.newBuilder() - .setBytesSize(sizeToRequest) - .build()) - .build()); + @Override + public void releaseRange(PartitionSession partition, OffsetsRange range) { + bufferManager.releaseRange(partition.getId(), range); + ReadPartitionSession queue = readQueues.get(partition.getId()); + if (queue != null) { + queue.releaseRange(range); + } } - void sendCommitOffsetRequest(PartitionSession session, List rangesToCommit) { - if (isStopped()) { + public boolean commitOffsets(PartitionSession session, List rangesToCommit) { + if (isClosed) { logger.atInfo() .setMessage("[{}] Need to send CommitRequest for {} with offset ranges {}, " + "but reading session is already closed") - .addArgument(streamId) + .addArgument(debugId) .addArgument(session) .addArgument(() -> rangesToCommit.stream().map(Object::toString).collect(Collectors.joining(", "))) .log(); - return; + return false; } - send(YdbTopic.StreamReadMessage.FromClient.newBuilder() - .setCommitOffsetRequest(YdbTopic.StreamReadMessage.CommitOffsetRequest.newBuilder() - .addCommitOffsets(YdbTopic.StreamReadMessage.CommitOffsetRequest.PartitionCommitOffset - .newBuilder() - .setPartitionSessionId(session.getId()) - .addAllOffsets(rangesToCommit.stream() - .map(ReadSession::buildOffsetRange) - .collect(Collectors.toList())) - .build()) + CommitOffsetRequest req = CommitOffsetRequest.newBuilder() + .addCommitOffsets(CommitOffsetRequest.PartitionCommitOffset.newBuilder() + .setPartitionSessionId(session.getId()) + .addAllOffsets(rangesToCommit.stream() + .map(ReaderImpl::buildOffsetRange) + .collect(Collectors.toList())) .build()) - .build()); + .build(); + + send(FromClient.newBuilder().setCommitOffsetRequest(req).build()); + return true; } - private void onInitResponse(YdbTopic.StreamReadMessage.InitResponse response) { - reader.onSessionStarted(response.getSessionId()); + public void onInit(YdbTopic.StreamReadMessage.InitResponse response) { bufferManager.init(response.getSessionId()); } - private void onStartPartitionSessionRequest(YdbTopic.StreamReadMessage.StartPartitionSessionRequest req) { + public StartPartitionSessionEvent onStartPartition(YdbTopic.StreamReadMessage.StartPartitionSessionRequest req) { long psid = req.getPartitionSession().getPartitionSessionId(); long pid = req.getPartitionSession().getPartitionId(); long committed = req.getCommittedOffset(); @@ -173,135 +131,56 @@ private void onStartPartitionSessionRequest(YdbTopic.StreamReadMessage.StartPart req.getPartitionOffsets().getEnd() ); - String traceID = streamId + '/' + psid + "-p" + pid; + String traceID = debugId + '/' + psid + "-p" + pid; logger.info("[{}] Received StartPartitionSessionRequest for {} and consumer \"{}\" with committedOffset {}" - + " and partitionOffsets {}", traceID, partition, consumerName, committed, offsets); + + " and partitionOffsets {}", traceID, partition, config.getConsumerName(), committed, offsets); partitions.put(psid, partition); - - reader.handleStartPartitionSessionRequest(new StartPartitionSessionEventImpl(partition, committed, offsets) { - @Override - public void confirm(StartPartitionSessionSettings options) { - if (isStopped()) { - logger.info("[{}] Need to send StartPartitionSessionResponse, but reading session is " - + "already closed", traceID); - return; - } - - PartitionSession partition = partitions.get(psid); - if (partition == null) { - logger.info("[{}] Need to send StartPartitionSessionResponse, but have no such active partition " - + "session anymore", traceID); - return; - } - - long readFrom = committed; - long commitTo = committed; - - YdbTopic.StreamReadMessage.StartPartitionSessionResponse.Builder resp = YdbTopic.StreamReadMessage - .StartPartitionSessionResponse.newBuilder() - .setPartitionSessionId(psid); - - if (options != null) { - if (options.getReadOffset() != null) { - readFrom = options.getReadOffset(); - resp.setReadOffset(readFrom); - } - if (options.getCommitOffset() != null) { - commitTo = options.getCommitOffset(); - resp.setCommitOffset(commitTo); - } - } - - ReadSession self = ReadSession.this; - Executor executor = reader.getDataHandlerExecutor(); - partSessions.put(psid, new ReadPartitionSession(traceID, self, partition, executor, commitTo) { - @Override - public void handleDataReceivedEvent(DataReceivedEvent event) { - reader.handleDataReceivedEvent(this, event); - } - }); - - logger.info("[{}] Sending StartPartitionSessionResponse for {} and consumer \"{}\" with readOffset " - + "{} and commitOffset {}", traceID, partition, consumerName, readFrom, commitTo); - send(YdbTopic.StreamReadMessage.FromClient.newBuilder() - .setStartPartitionSessionResponse(resp.build()) - .build()); - } - }); + return new StartPartitionRequest(traceID, partition, committed, offsets); } - protected void onStopPartitionSessionRequest(YdbTopic.StreamReadMessage.StopPartitionSessionRequest request) { - if (!request.getGraceful()) { - long psid = request.getPartitionSessionId(); - PartitionSession partition = partitions.remove(psid); - if (partition == null) { - logger.warn("[{}] Received force StopPartitionSessionRequest for partition session {}, " + - "but have no such partition session running", streamId, request.getPartitionSessionId()); - return; - } - - ReadPartitionSession rps = partSessions.remove(psid); - if (rps != null) { - logger.info("[{}] Received force StopPartitionSessionRequest for {} ", streamId, rps.getPartition()); - rps.stop(); - bufferManager.releasePartition(psid); - } + public PartitionSession onClosePartition(long partitionSessionId) { + PartitionSession partition = partitions.remove(partitionSessionId); + if (partition == null) { + logger.warn("[{}] Received force StopPartitionSessionRequest for partition session {}, " + + "but have no such partition session running", debugId, partitionSessionId); + return null; + } - reader.handleClosePartitionSession(partition); - return; + ReadPartitionSession queue = readQueues.remove(partitionSessionId); + if (queue != null) { + logger.info("[{}] Received force StopPartitionSessionRequest for {} ", debugId, queue.getPartition()); + queue.stop(); + bufferManager.releasePartition(partitionSessionId); } + return partition; + } + + public StopPartitionSessionEvent onStopPartition(YdbTopic.StreamReadMessage.StopPartitionSessionRequest request) { long committedOffset = request.getCommittedOffset(); long psid = request.getPartitionSessionId(); PartitionSession partition = partitions.get(psid); if (partition == null) { logger.error("[{}] Received graceful StopPartitionSessionRequest for partition session {}, " + - "but have no such partition session active", streamId, psid); - closeDueToError(null, new RuntimeException("Restarting read session due to receiving " - + "StopPartitionSessionRequest with PartitionSessionId " + psid + " that SDK knows nothing about")); - return; + "but have no such partition session active", debugId, psid); + return null; } - logger.info("[{}] Received graceful StopPartitionSessionRequest for {}", streamId, partition); - reader.handleStopPartitionSession(new StopPartitionSessionEventImpl(partition, committedOffset) { - @Override - public void confirm() { - if (isStopped()) { - logger.info("[{}] Need to send StopPartitionSessionResponse for {}, " + - "but reading session is already closed", streamId, partition); - return; - } - - if (partitions.remove(psid, partition)) { - logger.info("[{}] Sending StopPartitionSessionResponse for {}", streamId, partition); - send(YdbTopic.StreamReadMessage.FromClient.newBuilder().setStopPartitionSessionResponse( - YdbTopic.StreamReadMessage.StopPartitionSessionResponse.newBuilder() - .setPartitionSessionId(psid) - .build()) - .build()); - - ReadPartitionSession session = partSessions.remove(psid); - if (session != null) { - session.stop(); - } - } - - bufferManager.releasePartition(psid); - } - }); + logger.info("[{}] Received graceful StopPartitionSessionRequest for {}", debugId, partition); + return new StopPartitionRequest(partition, committedOffset); } - private void onReadResponse(YdbTopic.StreamReadMessage.ReadResponse response) { - logger.debug("[{}] Received ReadResponse of {} bytes", streamId, response.getBytesSize()); + public void onRead(YdbTopic.StreamReadMessage.ReadResponse response) { + logger.debug("[{}] Received ReadResponse of {} bytes", debugId, response.getBytesSize()); bufferManager.allocate(response.getBytesSize(), response.getPartitionDataList()); for (YdbTopic.StreamReadMessage.ReadResponse.PartitionData data: response.getPartitionDataList()) { long psid = data.getPartitionSessionId(); - ReadPartitionSession session = partSessions.get(psid); - if (session == null || !session.addBatches(data.getBatchesList())) { + ReadPartitionSession queue = readQueues.get(psid); + if (queue == null || !queue.addBatches(data.getBatchesList())) { logger.warn("[{}] Received PartitionData for unknown(most likely already closed) PartitionSessionId={}", - streamId, psid); + debugId, psid); bufferManager.releasePartition(psid); } } @@ -309,27 +188,29 @@ private void onReadResponse(YdbTopic.StreamReadMessage.ReadResponse response) { decoder.decodeNext(); } - protected void onCommitOffsetResponse(YdbTopic.StreamReadMessage.CommitOffsetResponse response) { - logger.trace("[{}] Received CommitOffsetResponse", streamId); - response.getPartitionsCommittedOffsetsList().forEach(offset -> { - ReadPartitionSession session = partSessions.get(offset.getPartitionSessionId()); - if (session == null) { + public void onCommitOffset(YdbTopic.StreamReadMessage.CommitOffsetResponse response, + BiConsumer callback) { + logger.trace("[{}] Received CommitOffsetResponse", debugId); + + for (CommitOffsetResponse.PartitionCommittedOffset offset: response.getPartitionsCommittedOffsetsList()) { + ReadPartitionSession queue = readQueues.get(offset.getPartitionSessionId()); + if (queue == null) { logger.info("[{}] Received CommitOffsetResponse for unknown (most likely already closed) " + - "e session with id={}", streamId, offset.getPartitionSessionId()); + "partition session with id={}", debugId, offset.getPartitionSessionId()); return; } // Handling CompletableFuture completions for single commits - session.confirmCommit(offset.getCommittedOffset()); + queue.confirmCommittedOffset(offset.getCommittedOffset()); // Handling onCommitResponse callback - reader.handleCommitResponse(offset.getCommittedOffset(), session.getPartition()); - }); + callback.accept(offset.getCommittedOffset(), queue.getPartition()); + } } - protected void onPartitionSessionStatusResponse(YdbTopic.StreamReadMessage.PartitionSessionStatusResponse resp) { + public void onPartitionSessionStatus(YdbTopic.StreamReadMessage.PartitionSessionStatusResponse resp) { PartitionSession partition = partitions.get(resp.getPartitionSessionId()); logger.info("[{}] Received PartitionSessionStatusResponse: partition session {} (partition {})." + - " Partition offsets: [{}, {}). Committed offset: {}", streamId, + " Partition offsets: [{}, {}). Committed offset: {}", debugId, resp.getPartitionSessionId(), partition == null ? "unknown" : partition.getPartitionId(), resp.getPartitionOffsets().getStart(), @@ -337,174 +218,103 @@ protected void onPartitionSessionStatusResponse(YdbTopic.StreamReadMessage.Parti resp.getCommittedOffset()); } - private void processMessage(YdbTopic.StreamReadMessage.FromServer message) { - if (isStopped()) { - logger.debug("[{}] processMessage called, but read session is already closed", streamId); - return; - } - logger.trace("[{}] processMessage called", streamId); - if (message.getStatus() != StatusCodesProtos.StatusIds.StatusCode.SUCCESS) { - Status status = Status.of(StatusCode.fromProto(message.getStatus()), - Issue.fromPb(message.getIssuesList())); - logger.warn("[{}] Got non-success status in processMessage method: {}", streamId, status); - closeDueToError(status, null); - return; - } - - if (message.hasInitResponse()) { - onInitResponse(message.getInitResponse()); - } else if (message.hasStartPartitionSessionRequest()) { - onStartPartitionSessionRequest(message.getStartPartitionSessionRequest()); - } else if (message.hasStopPartitionSessionRequest()) { - onStopPartitionSessionRequest(message.getStopPartitionSessionRequest()); - } else if (message.hasReadResponse()) { - onReadResponse(message.getReadResponse()); - } else if (message.hasCommitOffsetResponse()) { - onCommitOffsetResponse(message.getCommitOffsetResponse()); - } else if (message.hasPartitionSessionStatusResponse()) { - onPartitionSessionStatusResponse(message.getPartitionSessionStatusResponse()); - } else if (message.hasUpdateTokenResponse()) { - logger.debug("[{}] Received UpdateTokenResponse", streamId); - } else { - logger.error("[{}] Unhandled message from server: {}", streamId, message); + private class ReadRequest implements Consumer { + @Override + public void accept(Long sizeToRequest) { + logger.debug("[{}] Sending DataRequest with {} bytes", debugId, sizeToRequest); + send(YdbTopic.StreamReadMessage.FromClient.newBuilder() + .setReadRequest(YdbTopic.StreamReadMessage.ReadRequest.newBuilder() + .setBytesSize(sizeToRequest) + .build()) + .build()); } } - public CompletableFuture sendUpdateOffsetsInTransaction(YdbTransaction transaction, - Map> offsets, - UpdateOffsetsInTransactionSettings settings) { - if (offsets.isEmpty()) { - throw new IllegalArgumentException("Empty topic list to update in transaction"); - } - for (List offset: offsets.values()) { - if (offset.isEmpty()) { - throw new IllegalArgumentException("Empty offsets range to update in transaction"); - } - } + private class StartPartitionRequest extends StartPartitionSessionEventImpl { + private final String traceID; - if (logger.isDebugEnabled()) { - StringBuilder str = new StringBuilder("Updating "); - boolean first = true; - for (Map.Entry> topicOffsets : offsets.entrySet()) { - for (PartitionOffsets partitionOffsets : topicOffsets.getValue()) { - if (!first) { - str.append(", "); - } else { - first = false; - } - str.append("offsets [").append(partitionOffsets.getOffsets().get(0).getStart()).append("..") - .append(partitionOffsets.getOffsets().get(partitionOffsets.getOffsets().size() - 1) - .getEnd()).append(") for partition ") - .append(partitionOffsets.getPartitionSession().getPartitionId()) - .append(" [topic ").append(topicOffsets.getKey()).append("]"); - } - } - logger.debug(str.toString()); + StartPartitionRequest(String traceID, PartitionSession ps, long committed, OffsetsRange offsets) { + super(ps, committed, offsets); + this.traceID = traceID; } - transaction.getStatusFuture().whenComplete((status, error) -> { - if (error != null) { - closeDueToError(null, - new RuntimeException("Restarting read session due to transaction " + transaction.getId() + - " with partition offsets from read session " + getStreamId() + - " was not committed with reason: " + error)); - } else if (!status.isSuccess()) { - closeDueToError(null, - new RuntimeException("Restarting read session due to transaction " + transaction.getId() + - " with partition offsets from read session " + getStreamId() + - " was not committed with status: " + status)); + @Override + public void confirm(StartPartitionSessionSettings options) { + if (isClosed) { + logger.info("[{}] Need to send StartPartitionSessionResponse, but reading session is " + + "already closed", traceID); + return; } - }); - - YdbTopic.UpdateOffsetsInTransactionRequest req = YdbTopic.UpdateOffsetsInTransactionRequest.newBuilder() - .setTx(YdbTopic.TransactionIdentity.newBuilder() - .setId(transaction.getId()) - .setSession(transaction.getSessionId()) - .build()) - .setConsumer(consumerName) - .addAllTopics(offsets.entrySet().stream() - .map(entry -> buildTopicOffsets(entry.getKey(), entry.getValue())) - .collect(Collectors.toList())) - .build(); - - String traceId = settings.getTraceId() == null ? UUID.randomUUID().toString() : settings.getTraceId(); - final GrpcRequestSettings grpcRequestSettings = GrpcRequestSettings.newBuilder() - .withDeadline(settings.getRequestTimeout()) - .withTraceId(traceId) - .build(); - - return rpc.updateOffsetsInTransaction(req, grpcRequestSettings); - } - - private static YdbTopic.UpdateOffsetsInTransactionRequest.TopicOffsets.PartitionOffsets buildPartitionOffsets( - PartitionOffsets partitionOffsets) { - return YdbTopic.UpdateOffsetsInTransactionRequest.TopicOffsets.PartitionOffsets.newBuilder() - .setPartitionId(partitionOffsets.getPartitionSession().getPartitionId()) - .addAllPartitionOffsets(partitionOffsets.getOffsets().stream() - .map(ReadSession::buildOffsetRange) - .collect(Collectors.toList())) - .build(); - } - private static YdbTopic.UpdateOffsetsInTransactionRequest.TopicOffsets buildTopicOffsets(String topicPath, - List partitions) { + long psid = getPartitionSession().getId(); + long readFrom = getCommittedOffset(); + long commitTo = getCommittedOffset(); - return YdbTopic.UpdateOffsetsInTransactionRequest.TopicOffsets.newBuilder() - .setPath(topicPath) - .addAllPartitions(partitions.stream() - .map(ReadSession::buildPartitionOffsets) - .collect(Collectors.toList())) - .build(); - } + PartitionSession partition = partitions.get(psid); + if (partition == null) { + logger.info("[{}] Need to send StartPartitionSessionResponse, but have no such active partition " + + "session anymore", traceID); + return; + } - private static YdbTopic.OffsetsRange buildOffsetRange(OffsetsRange range) { - return YdbTopic.OffsetsRange.newBuilder() - .setStart(range.getStart()) - .setEnd(range.getEnd()) - .build(); - } + StartPartitionSessionResponse.Builder resp = StartPartitionSessionResponse.newBuilder() + .setPartitionSessionId(psid); - private static YdbTopic.StreamReadMessage.InitRequest.TopicReadSettings buildTopicSettings(TopicReadSettings trs) { - String topicPath = trs.getPath(); - List partitions = trs.getPartitionIds(); - Instant readFrom = trs.getReadFrom(); - Duration maxLag = trs.getMaxLag(); + if (options != null) { + if (options.getReadOffset() != null) { + readFrom = options.getReadOffset(); + resp.setReadOffset(readFrom); + } + if (options.getCommitOffset() != null) { + commitTo = options.getCommitOffset(); + resp.setCommitOffset(commitTo); + } + } - YdbTopic.StreamReadMessage.InitRequest.TopicReadSettings.Builder builder = YdbTopic.StreamReadMessage - .InitRequest.TopicReadSettings.newBuilder(); + MessageCommitterImpl committer = new MessageCommitterImpl(traceID, ReadSession.this, partition, commitTo); + ReadPartitionSession queue = new ReadPartitionSession(traceID, config, partition, committer, decoder, + event -> eventConsumer.accept(ReadSession.this, event), commitTo); + if (readQueues.putIfAbsent(psid, queue) != null) { + logger.warn("[{}] partition {} is already started", traceID, partition); + return; + } - builder.setPath(topicPath); - if (partitions != null && !partitions.isEmpty()) { - builder.addAllPartitionIds(partitions); + logger.info("[{}] Sending StartPartitionSessionResponse for {} and consumer \"{}\" with readOffset " + + "{} and commitOffset {}", traceID, partition, config.getConsumerName(), readFrom, commitTo); + send(FromClient.newBuilder().setStartPartitionSessionResponse(resp.build()).build()); } - if (readFrom != null) { - builder.setReadFrom(ProtobufUtils.instantToProto(readFrom)); - } - if (maxLag != null) { - builder.setMaxLag(ProtobufUtils.durationToProto(maxLag)); + }; + + private class StopPartitionRequest extends StopPartitionSessionEventImpl { + StopPartitionRequest(PartitionSession partition, long committedOffset) { + super(partition, committedOffset); } - return builder.build(); - } + @Override + public void confirm() { + PartitionSession partition = getPartitionSession(); + long psid = getPartitionSessionId(); + if (isClosed) { + logger.info("[{}] Need to send StopPartitionSessionResponse for {}, " + + "but reading session is already closed", debugId, partition); + return; + } - private static YdbTopic.StreamReadMessage.InitRequest buildInitRequest(ReaderSettings settings) { - String consumerName = settings.getConsumerName(); - String readerName = settings.getReaderName(); - List topics = settings.getTopics(); + if (partitions.remove(psid, partition)) { + logger.info("[{}] Sending StopPartitionSessionResponse for {}", debugId, partition); + send(YdbTopic.StreamReadMessage.FromClient.newBuilder().setStopPartitionSessionResponse( + YdbTopic.StreamReadMessage.StopPartitionSessionResponse.newBuilder() + .setPartitionSessionId(psid) + .build()) + .build()); - YdbTopic.StreamReadMessage.InitRequest.Builder builder = YdbTopic.StreamReadMessage.InitRequest.newBuilder(); + ReadPartitionSession session = readQueues.remove(psid); + if (session != null) { + session.stop(); + } + } - builder.setPartitionMaxInFlightBytes(settings.getPartitionMaxInFlightBytes()); - if (consumerName != null && !consumerName.isEmpty()) { - builder.setConsumer(consumerName); - } - if (readerName != null && !readerName.isEmpty()) { - builder.setReaderName(readerName); + bufferManager.releasePartition(psid); } - for (TopicReadSettings trs: topics) { - builder.addTopicsReadSettings(buildTopicSettings(trs)); - } - - return builder.build(); } } 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 f27236991..93f8d56ec 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 @@ -1,148 +1,272 @@ package tech.ydb.topic.read.impl; +import java.time.Duration; +import java.time.Instant; import java.util.List; import java.util.Map; +import java.util.UUID; import java.util.concurrent.CompletableFuture; -import java.util.concurrent.Executor; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.atomic.AtomicLong; import java.util.stream.Collectors; -import javax.annotation.Nonnull; - import org.slf4j.Logger; import org.slf4j.LoggerFactory; import tech.ydb.common.transaction.YdbTransaction; +import tech.ydb.core.Issue; import tech.ydb.core.Status; +import tech.ydb.core.StatusCode; +import tech.ydb.core.grpc.GrpcRequestSettings; +import tech.ydb.core.utils.ProtobufUtils; +import tech.ydb.proto.topic.YdbTopic; +import tech.ydb.proto.topic.YdbTopic.StreamReadMessage.FromClient; +import tech.ydb.proto.topic.YdbTopic.StreamReadMessage.FromServer; import tech.ydb.topic.TopicRpc; -import tech.ydb.topic.description.CodecRegistry; -import tech.ydb.topic.impl.GrpcStreamRetrier; +import tech.ydb.topic.description.OffsetsRange; +import tech.ydb.topic.impl.TopicRetryableStream; import tech.ydb.topic.read.PartitionOffsets; import tech.ydb.topic.read.PartitionSession; import tech.ydb.topic.read.events.DataReceivedEvent; import tech.ydb.topic.read.events.StartPartitionSessionEvent; import tech.ydb.topic.read.events.StopPartitionSessionEvent; import tech.ydb.topic.settings.ReaderSettings; +import tech.ydb.topic.settings.TopicReadSettings; import tech.ydb.topic.settings.UpdateOffsetsInTransactionSettings; /** * @author Nikolay Perfilov */ -public abstract class ReaderImpl extends GrpcStreamRetrier { - private static final Logger logger = LoggerFactory.getLogger(ReaderImpl.class); - - private static final int DEFAULT_DECOMPRESSION_THREAD_COUNT = 4; - private final ExecutorService defaultDecompressionExecutorService; - private final ReadSessionFactory sessionFactory; +public class ReaderImpl extends TopicRetryableStream { + public interface Releaser { + void releaseRange(PartitionSession partition, OffsetsRange range); + } + public interface Handler { + void handleSessionStarted(String sessionId); - private final CompletableFuture sessionReady = new CompletableFuture<>(); - private volatile ReadSession session = null; + void handleStartPartitionSessionRequest(StartPartitionSessionEvent event); + void handleStopPartitionSession(StopPartitionSessionEvent event); + void handleClosePartitionSession(PartitionSession partition); - public ReaderImpl(TopicRpc topicRpc, ReaderSettings settings, @Nonnull CodecRegistry codecRegistry) { - super(settings.getLogPrefix(), topicRpc.getScheduler(), settings.getErrorsHandler()); + void handleDataReceivedEvent(Releaser releaser, DataReceivedEvent event); + void handleCommitResponse(long committedOffset, PartitionSession partition); - Executor decompressionExecutor = settings.getDecompressionExecutor(); - if (decompressionExecutor != null) { - this.defaultDecompressionExecutorService = null; - } else { - this.defaultDecompressionExecutorService = Executors.newFixedThreadPool(DEFAULT_DECOMPRESSION_THREAD_COUNT); - decompressionExecutor = defaultDecompressionExecutorService; - } - this.sessionFactory = new ReadSessionFactory(topicRpc, settings, decompressionExecutor, codecRegistry); + void handleReaderClosed(Status status); + } - String consumerName = settings.getConsumerName(); - String readerName = settings.getReaderName(); + private static final Logger logger = LoggerFactory.getLogger(ReaderImpl.class); - logger.info("Reader{} (generated id {}) created for topic(s) {} and {}", - readerName != null ? (" '" + readerName + "'") : "", - id, - settings.getTopics().stream().map(t -> "\"" + t.getPath() + "\"").collect(Collectors.joining(", ")), - consumerName != null ? (" consumer \"" + consumerName + "\"") : "without a consumer" - ); - } + private final TopicRpc rpc; + private final ReadConfig config; + private final Handler handler; - abstract Executor getDataHandlerExecutor(); - protected abstract void handleDataReceivedEvent(ReadPartitionSession session, DataReceivedEvent event); - protected abstract void handleSessionStarted(String sessionId); - protected abstract void handleCommitResponse(long committedOffset, PartitionSession partition); - protected abstract void handleStartPartitionSessionRequest(StartPartitionSessionEvent event); - protected abstract void handleStopPartitionSession(StopPartitionSessionEvent event); - protected abstract void handleClosePartitionSession(PartitionSession partition); + private final FromClient initRequest; - @Override - protected Logger getLogger() { - return logger; + public ReaderImpl(TopicRpc rpc, String id, ReaderSettings settings, ReadConfig config, Handler handler) { + super(logger, id, settings.getRetryConfig(), rpc.getScheduler()); + this.rpc = rpc; + this.initRequest = FromClient.newBuilder().setInitRequest(buildInitRequest(settings)).build(); + this.config = config; + this.handler = handler; } @Override - protected String getStreamName() { - return "Reader"; + protected ReadSession createNewStream(String id) { + return new ReadSession(id, rpc.readSession(id), initRequest, handler::handleDataReceivedEvent, config); } @Override - protected void onStreamReconnect() { - session = sessionFactory.createNextSession(); - session.startAndInitialize(); + protected void onRetry(ReadSession stream, Status status) { + logger.warn("[{}] paused by status {}", debugId, status); + stream.closeAll().forEach(ps -> handler.handleClosePartitionSession(ps)); } - protected CompletableFuture initImpl() { - logger.info("[{}] initImpl called", id); - if (session == null) { - onStreamReconnect(); + @Override + protected void onClose(ReadSession stream, Status status) { + if (!status.isSuccess()) { + logger.warn("[{}] closed by status {}", debugId, status); } else { - logger.warn("[{}] Init is called on this reader more than once. Nothing is done", id); + logger.info("[{}] closed by status {}", debugId, status); } - - return sessionReady; + stream.closeAll().forEach(ps -> handler.handleClosePartitionSession(ps)); + handler.handleReaderClosed(status); } - void onSessionStarted(String sessionId) { - sessionReady.complete(null); - reconnectCounter.set(0); - handleSessionStarted(sessionId); + @Override + protected void onNext(ReadSession stream, FromServer message) { + logger.trace("[{}] processMessage called", debugId); + + if (message.hasInitResponse()) { + resetRetries(); + handler.handleSessionStarted(message.getInitResponse().getSessionId()); + stream.onInit(message.getInitResponse()); + } else if (message.hasStartPartitionSessionRequest()) { + StartPartitionSessionEvent event = stream.onStartPartition(message.getStartPartitionSessionRequest()); + handler.handleStartPartitionSessionRequest(event); + } else if (message.hasStopPartitionSessionRequest()) { + YdbTopic.StreamReadMessage.StopPartitionSessionRequest req = message.getStopPartitionSessionRequest(); + if (req.getGraceful()) { + StopPartitionSessionEvent event = stream.onStopPartition(req); + if (event != null) { + handler.handleStopPartitionSession(event); + } + } else { + PartitionSession closed = stream.onClosePartition(req.getPartitionSessionId()); + if (closed != null) { + handler.handleClosePartitionSession(closed); + } + } + } else if (message.hasReadResponse()) { + stream.onRead(message.getReadResponse()); + } else if (message.hasCommitOffsetResponse()) { + stream.onCommitOffset(message.getCommitOffsetResponse(), handler::handleCommitResponse); + } else if (message.hasPartitionSessionStatusResponse()) { + stream.onPartitionSessionStatus(message.getPartitionSessionStatusResponse()); + } else if (message.hasUpdateTokenResponse()) { + logger.debug("[{}] Received UpdateTokenResponse", debugId); + } else { + logger.error("[{}] Unhandled message from server: {}", debugId, message); + } } - protected CompletableFuture updateOffsetsInTransaction(YdbTransaction transaction, - Map> offsets, - UpdateOffsetsInTransactionSettings settings) { + public CompletableFuture updateOffsetsInTransaction(YdbTransaction transaction, + Map> offsets, + UpdateOffsetsInTransactionSettings settings) { if (!transaction.isActive()) { throw new IllegalArgumentException("Transaction is not active. " + "Can only read topic messages in already running transactions from other services"); } - return session.sendUpdateOffsetsInTransaction(transaction, offsets, settings); + if (offsets.isEmpty()) { + throw new IllegalArgumentException("Empty topic list to update in transaction"); + } + for (List offset: offsets.values()) { + if (offset.isEmpty()) { + throw new IllegalArgumentException("Empty offsets range to update in transaction"); + } + } + + if (logger.isDebugEnabled()) { + StringBuilder str = new StringBuilder("Updating "); + boolean first = true; + for (Map.Entry> topicOffsets : offsets.entrySet()) { + for (PartitionOffsets partitionOffsets : topicOffsets.getValue()) { + if (!first) { + str.append(", "); + } else { + first = false; + } + str.append("offsets [").append(partitionOffsets.getOffsets().get(0).getStart()).append("..") + .append(partitionOffsets.getOffsets().get(partitionOffsets.getOffsets().size() - 1) + .getEnd()).append(") for partition ") + .append(partitionOffsets.getPartitionSession().getPartitionId()) + .append(" [topic ").append(topicOffsets.getKey()).append("]"); + } + } + logger.debug(str.toString()); + } + + transaction.getStatusFuture().whenComplete((status, error) -> { + if (status != null && !status.isSuccess()) { + String msg = "Restarting read session due to transaction " + transaction.getId() + + " with partition offsets from read session " + debugId + + " was not committed with status: " + status; + fail(Status.of(StatusCode.CLIENT_INTERNAL_ERROR, Issue.of(msg, Issue.Severity.ERROR))); + } + if (error != null) { + String msg = "Restarting read session due to transaction " + transaction.getId() + + " with partition offsets from read session " + debugId + + " was not committed with reason: " + error.getMessage(); + fail(Status.of(StatusCode.CLIENT_INTERNAL_ERROR, error, Issue.of(msg, Issue.Severity.ERROR))); + } + }); + + YdbTopic.UpdateOffsetsInTransactionRequest req = YdbTopic.UpdateOffsetsInTransactionRequest.newBuilder() + .setTx(YdbTopic.TransactionIdentity.newBuilder() + .setId(transaction.getId()) + .setSession(transaction.getSessionId()) + .build()) + .setConsumer(config.getConsumerName()) + .addAllTopics(offsets.entrySet().stream() + .map(entry -> buildTopicOffsets(entry.getKey(), entry.getValue())) + .collect(Collectors.toList())) + .build(); + + String traceId = settings.getTraceId() == null ? UUID.randomUUID().toString() : settings.getTraceId(); + final GrpcRequestSettings grpcRequestSettings = GrpcRequestSettings.newBuilder() + .withDeadline(settings.getRequestTimeout()) + .withTraceId(traceId) + .build(); + + return rpc.updateOffsetsInTransaction(req, grpcRequestSettings); } - @Override - protected void onShutdown(String reason) { - if (session != null) { - session.shutdown(); + private static YdbTopic.UpdateOffsetsInTransactionRequest.TopicOffsets.PartitionOffsets buildPartitionOffsets( + PartitionOffsets partitionOffsets) { + return YdbTopic.UpdateOffsetsInTransactionRequest.TopicOffsets.PartitionOffsets.newBuilder() + .setPartitionId(partitionOffsets.getPartitionSession().getPartitionId()) + .addAllPartitionOffsets(partitionOffsets.getOffsets().stream() + .map(ReaderImpl::buildOffsetRange) + .collect(Collectors.toList())) + .build(); + } + + private static YdbTopic.UpdateOffsetsInTransactionRequest.TopicOffsets buildTopicOffsets(String topicPath, + List partitions) { + + return YdbTopic.UpdateOffsetsInTransactionRequest.TopicOffsets.newBuilder() + .setPath(topicPath) + .addAllPartitions(partitions.stream() + .map(ReaderImpl::buildPartitionOffsets) + .collect(Collectors.toList())) + .build(); + } + + public static YdbTopic.OffsetsRange buildOffsetRange(OffsetsRange range) { + return YdbTopic.OffsetsRange.newBuilder() + .setStart(range.getStart()) + .setEnd(range.getEnd()) + .build(); + } + + private static YdbTopic.StreamReadMessage.InitRequest.TopicReadSettings buildTopicSettings(TopicReadSettings trs) { + String topicPath = trs.getPath(); + List partitions = trs.getPartitionIds(); + Instant readFrom = trs.getReadFrom(); + Duration maxLag = trs.getMaxLag(); + + YdbTopic.StreamReadMessage.InitRequest.TopicReadSettings.Builder builder = YdbTopic.StreamReadMessage + .InitRequest.TopicReadSettings.newBuilder(); + + builder.setPath(topicPath); + if (partitions != null && !partitions.isEmpty()) { + builder.addAllPartitionIds(partitions); + } + if (readFrom != null) { + builder.setReadFrom(ProtobufUtils.instantToProto(readFrom)); } - sessionReady.completeExceptionally(new RuntimeException(reason)); - if (defaultDecompressionExecutorService != null) { - defaultDecompressionExecutorService.shutdown(); + if (maxLag != null) { + builder.setMaxLag(ProtobufUtils.durationToProto(maxLag)); } + + return builder.build(); } - private class ReadSessionFactory { - private final TopicRpc rpc; - private final ReaderSettings settings; - private final Executor decompressor; - private final CodecRegistry codecRegistry; - private final AtomicLong sessionCounter = new AtomicLong(0); - - ReadSessionFactory(TopicRpc rpc, ReaderSettings settings, Executor decompressor, CodecRegistry codecRegistry) { - this.rpc = rpc; - this.settings = settings; - this.decompressor = decompressor; - this.codecRegistry = codecRegistry; - } + private static YdbTopic.StreamReadMessage.InitRequest buildInitRequest(ReaderSettings settings) { + String consumerName = settings.getConsumerName(); + String readerName = settings.getReaderName(); + List topics = settings.getTopics(); + + YdbTopic.StreamReadMessage.InitRequest.Builder builder = YdbTopic.StreamReadMessage.InitRequest.newBuilder(); - public ReadSession createNextSession() { - String streamID = id + '.' + sessionCounter.incrementAndGet(); - MessageDecoder decoder = new MessageDecoder(settings.getMaxMemoryUsageBytes(), decompressor, codecRegistry); - return new ReadSession(rpc, ReaderImpl.this, decoder, streamID, settings); + builder.setPartitionMaxInFlightBytes(settings.getPartitionMaxInFlightBytes()); + if (consumerName != null && !consumerName.isEmpty()) { + builder.setConsumer(consumerName); + } + if (readerName != null && !readerName.isEmpty()) { + builder.setReaderName(readerName); } + for (TopicReadSettings trs: topics) { + builder.addTopicsReadSettings(buildTopicSettings(trs)); + } + + return builder.build(); } } diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java b/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java index 3914e8db6..86e314887 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java @@ -6,11 +6,12 @@ import java.util.List; import java.util.Queue; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentLinkedQueue; -import java.util.concurrent.Executor; import java.util.concurrent.TimeUnit; import java.util.concurrent.locks.Condition; import java.util.concurrent.locks.ReentrantLock; +import java.util.stream.Collectors; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -21,7 +22,7 @@ import tech.ydb.core.Status; import tech.ydb.topic.TopicRpc; import tech.ydb.topic.description.CodecRegistry; -import tech.ydb.topic.description.OffsetsRange; +import tech.ydb.topic.impl.DebugTools; import tech.ydb.topic.read.Message; import tech.ydb.topic.read.PartitionOffsets; import tech.ydb.topic.read.PartitionSession; @@ -36,9 +37,20 @@ /** * @author Nikolay Perfilov */ -public class SyncReaderImpl extends ReaderImpl implements SyncReader { +public class SyncReaderImpl implements SyncReader { private static final Logger logger = LoggerFactory.getLogger(SyncReaderImpl.class); + private static final int POLL_INTERVAL_SECONDS = 5; + + private final String debugId; + private final LazyExecutor decompressor; + private final ReadConfig config; + private final ReaderImpl impl; + + private final CompletableFuture initFuture = new CompletableFuture<>(); + private final CompletableFuture shutdownFuture = new CompletableFuture<>(); + + private final ConcurrentHashMap activePartitions = new ConcurrentHashMap<>(); private final Queue queue = new ConcurrentLinkedQueue<>(); private final ReentrantLock waitingLock = new ReentrantLock(); private final Condition waitingCondition = waitingLock.newCondition(); @@ -46,33 +58,20 @@ public class SyncReaderImpl extends ReaderImpl implements SyncReader { private volatile String sessionId = null; public SyncReaderImpl(TopicRpc topicRpc, ReaderSettings settings, @Nonnull CodecRegistry codecRegistry) { - super(topicRpc, settings, codecRegistry); - } - - private static class MessageWrapper { - private final Message msg; - private final ReadPartitionSession session; - private final OffsetsRange rangeToRelease; - - private MessageWrapper(Message msg, ReadPartitionSession session, OffsetsRange rangeToRelease) { - this.msg = msg; - this.session = session; - this.rangeToRelease = rangeToRelease; - } - - boolean isActive() { - return !session.isStopped(); - } - - Message getMessage() { - return msg; - } - - void release() { - if (rangeToRelease != null) { - session.releaseRange(rangeToRelease); - } - } + this.debugId = DebugTools.createDebugId(settings.getLogPrefix()); + this.decompressor = new LazyExecutor("reader[" + debugId + "]-decoder", settings.getDecompressionExecutor()); + + this.config = new ReadConfig(codecRegistry, Runnable::run, decompressor, settings); + this.impl = new ReaderImpl(topicRpc, debugId, settings, config, new SyncHandler()); + + String readerName = settings.getReaderName(); + String consumerName = settings.getConsumerName(); + logger.info("Reader{} (generated id {}) created for topic(s) {} and {}", + readerName != null ? (" '" + readerName + "'") : "", + debugId, + settings.getTopics().stream().map(t -> "\"" + t.getPath() + "\"").collect(Collectors.joining(", ")), + consumerName != null ? (" consumer \"" + consumerName + "\"") : "without a consumer" + ); } @Override @@ -82,12 +81,42 @@ public String getSessionId() { @Override public void init() { - initImpl(); + impl.start(); } @Override public void initAndWait() { - initImpl().join(); + impl.start(); + initFuture.join(); + } + + + @Override + public void shutdown() { + impl.close(); + + waitingLock.lock(); + try { + waitingCondition.signalAll(); + } finally { + waitingLock.unlock(); + } + + shutdownFuture.join(); + } + + @Override + public Message receive(ReceiveSettings receiveSettings) throws InterruptedException { + if (receiveSettings.getTimeout() != null) { + return receiveInternal(receiveSettings, receiveSettings.getTimeout(), receiveSettings.getTimeoutTimeUnit()); + } + + Message result; + // Poll to prevent infinite wait in case if reader was stopped + do { + result = receiveInternal(receiveSettings, POLL_INTERVAL_SECONDS, TimeUnit.SECONDS); + } while (result == null); + return result; } private MessageWrapper waitReadyMessage(long timeout, TimeUnit unit) throws InterruptedException { @@ -106,7 +135,7 @@ private MessageWrapper waitReadyMessage(long timeout, TimeUnit unit) throws Inte logger.trace("No messages in queue. Waiting for {} ms...", millisToWait); waitingCondition.await(millisToWait, TimeUnit.MILLISECONDS); - if (isStopped.get()) { + if (impl.isClosed()) { throw new RuntimeException("Reader was stopped"); } next = queue.poll(); @@ -120,7 +149,7 @@ private MessageWrapper waitReadyMessage(long timeout, TimeUnit unit) throws Inte @Nullable public Message receiveInternal(ReceiveSettings receiveSettings, long timeout, TimeUnit unit) throws InterruptedException { - if (isStopped.get()) { + if (impl.isClosed()) { throw new RuntimeException("Reader was stopped"); } @@ -133,8 +162,7 @@ public Message receiveInternal(ReceiveSettings receiveSettings, long timeout, Ti } } - if (!next.isActive()) { - next.release(); + if (!activePartitions.containsKey(next.getPartition())) { continue; } @@ -145,7 +173,7 @@ public Message receiveInternal(ReceiveSettings receiveSettings, long timeout, Ti result.getPartitionSession(), Collections.singletonList(result.getRangeToCommit()) )); - Status updateStatus = updateOffsetsInTransaction( + Status updateStatus = impl.updateOffsetsInTransaction( receiveSettings.getTransaction(), Collections.singletonMap(result.getPartitionSession().getPath(), offsets), UpdateOffsetsInTransactionSettings.newBuilder().build() @@ -156,100 +184,104 @@ public Message receiveInternal(ReceiveSettings receiveSettings, long timeout, Ti } } - next.release(); + next.confirm(); return result; } } - @Override - public Message receive(ReceiveSettings receiveSettings) throws InterruptedException { - if (receiveSettings.getTimeout() != null) { - return receiveInternal(receiveSettings, receiveSettings.getTimeout(), receiveSettings.getTimeoutTimeUnit()); + private class SyncHandler implements ReaderImpl.Handler { + @Override + public void handleSessionStarted(String sessionId) { + SyncReaderImpl.this.sessionId = sessionId; + initFuture.complete(null); } - Message result; - // Poll to prevent infinite wait in case if reader was stopped - do { - result = receiveInternal(receiveSettings, POLL_INTERVAL_SECONDS, TimeUnit.SECONDS); - } while (result == null); - return result; - } + @Override + public void handleReaderClosed(Status status) { + shutdownFuture.complete(null); + } - @Override - Executor getDataHandlerExecutor() { - return Runnable::run; - } + @Override + public void handleDataReceivedEvent(ReaderImpl.Releaser releaser, DataReceivedEvent event) { + if (impl.isClosed()) { + return; + } + if (event.getMessages().isEmpty()) { + releaser.releaseRange(event.getPartitionSession(), event.getRangeToCommit()); + return; + } - @Override - protected void handleDataReceivedEvent(ReadPartitionSession session, DataReceivedEvent event) { - if (isStopped.get() || event.getMessages().isEmpty()) { - session.releaseRange(event.getRangeToCommit()); - return; - } - - int messagesCount = event.getMessages().size(); - long offsetStart = event.getMessages().get(0).getOffset(); - long offsetEnd = event.getMessages().get(event.getMessages().size() - 1).getOffset(); - logger.debug("{} Putting a batch into queueData with {} message(s) (offsets {}-{})", - session, messagesCount, offsetStart, offsetEnd); - - for (Message msg: event.getMessages()) { - if (msg.getRangeToCommit().getEnd() == event.getRangeToCommit().getEnd()) { // last message in batch - queue.offer(new MessageWrapper(msg, session, event.getRangeToCommit())); - } else { - queue.offer(new MessageWrapper(msg, session, null)); + PartitionSession ps = event.getPartitionSession(); + int messagesCount = event.getMessages().size(); + long offsetStart = event.getMessages().get(0).getOffset(); + long offsetEnd = event.getMessages().get(event.getMessages().size() - 1).getOffset(); + logger.debug("{} Putting a batch into queueData with {} message(s) (offsets {}-{}) from {}", + debugId, messagesCount, offsetStart, offsetEnd, ps); + + Runnable confirm = () -> releaser.releaseRange(ps, event.getRangeToCommit()); + for (Message msg: event.getMessages()) { + if (msg.getRangeToCommit().getEnd() == event.getRangeToCommit().getEnd()) { // last message in batch + queue.offer(new MessageWrapper(ps, msg, confirm)); + } else { + queue.offer(new MessageWrapper(ps, msg, null)); + } + } + + waitingLock.lock(); + try { + waitingCondition.signalAll(); + } finally { + waitingLock.unlock(); } } - waitingLock.lock(); - try { - waitingCondition.signalAll(); - } finally { - waitingLock.unlock(); + @Override + public void handleCommitResponse(long committedOffset, PartitionSession partitionSession) { + logger.debug("CommitResponse received for{} with committedOffset {}", partitionSession, committedOffset); } - } - @Override - protected void handleSessionStarted(String sessionId) { - this.sessionId = sessionId; - } + @Override + public void handleStartPartitionSessionRequest(StartPartitionSessionEvent event) { + activePartitions.put(event.getPartitionSession(), event.getPartitionSession()); + event.confirm(); + } - @Override - protected void handleCommitResponse(long committedOffset, PartitionSession partitionSession) { - if (logger.isDebugEnabled()) { - logger.debug("CommitResponse received for partition session {} (partition {}) with committedOffset {}", - partitionSession.getId(), partitionSession.getPartitionId(), committedOffset); + @Override + public void handleStopPartitionSession(StopPartitionSessionEvent event) { + activePartitions.remove(event.getPartitionSession()); + // TODO: wait for all commits + event.confirm(); } - } - @Override - protected void handleStartPartitionSessionRequest(StartPartitionSessionEvent event) { - event.confirm(); + @Override + public void handleClosePartitionSession(PartitionSession partition) { + activePartitions.remove(partition); + } } - @Override - protected void handleStopPartitionSession(StopPartitionSessionEvent event) { - // TODO: wait for all commits - event.confirm(); - } + private static class MessageWrapper { + private final PartitionSession partition; + private final Message msg; + private final Runnable confirm; - @Override - protected void handleClosePartitionSession(PartitionSession partition) { - // TODO: clean reading queue - logger.debug("ClosePartitionSession event received. Ignoring."); - } + private MessageWrapper(PartitionSession partition, Message msg, Runnable confirm) { + this.partition = partition; + this.msg = msg; + this.confirm = confirm; + } - @Override - public void shutdown() { - CompletableFuture impl = shutdownImpl(); + Message getMessage() { + return msg; + } - waitingLock.lock(); - try { - waitingCondition.signalAll(); - } finally { - waitingLock.unlock(); + PartitionSession getPartition() { + return partition; } - impl.join(); + void confirm() { + if (confirm != null) { + confirm.run(); + } + } } } diff --git a/topic/src/main/java/tech/ydb/topic/settings/ReaderSettings.java b/topic/src/main/java/tech/ydb/topic/settings/ReaderSettings.java index 175ef1667..4b46d3535 100644 --- a/topic/src/main/java/tech/ydb/topic/settings/ReaderSettings.java +++ b/topic/src/main/java/tech/ydb/topic/settings/ReaderSettings.java @@ -9,6 +9,7 @@ import com.google.common.collect.ImmutableList; +import tech.ydb.common.retry.RetryConfig; import tech.ydb.core.Status; import tech.ydb.topic.read.events.DataReceivedEvent; @@ -26,6 +27,7 @@ public class ReaderSettings { private final int maxBatchSize; private final long partitionMaxInFlightBytes; private final Executor decompressionExecutor; + private final RetryConfig retryConfig; private final BiConsumer errorsHandler; private ReaderSettings(Builder builder) { @@ -37,6 +39,7 @@ private ReaderSettings(Builder builder) { this.maxBatchSize = builder.maxBatchSize; this.partitionMaxInFlightBytes = builder.partitionMaxInFlightBytes; this.decompressionExecutor = builder.decompressionExecutor; + this.retryConfig = builder.retryConfig; this.errorsHandler = builder.errorsHandler; } @@ -61,6 +64,10 @@ public BiConsumer getErrorsHandler() { return errorsHandler; } + public RetryConfig getRetryConfig() { + return retryConfig; + } + public long getMaxMemoryUsageBytes() { return maxMemoryUsageBytes; } @@ -94,6 +101,7 @@ public static class Builder { private long partitionMaxInFlightBytes = 0; private int maxBatchSize = 0; private Executor decompressionExecutor = null; + private RetryConfig retryConfig = TopicRetryConfig.FOREVER; private BiConsumer errorsHandler = null; /** @@ -178,6 +186,33 @@ public Builder setErrorsHandler(BiConsumer handler) { return this; } + /** + * Set retry configuration for the reader's underlying stream connection. + * Controls how the reader reconnects when the stream is interrupted. + *

+ * The default value is {@link TopicRetryConfig#FOREVER}, which retries any disconnection + * indefinitely with exponential backoff (up to ~65 seconds between attempts). + *

+ * Use {@link TopicRetryConfig#NEVER} to disable retries and surface errors immediately + * via the errors handler set by {@link #setErrorsHandler}. + * Use {@link TopicRetryConfig#STANDARD} to retry only transient errors and treat + * permanent status codes (e.g. {@code UNAUTHORIZED}, {@code BAD_REQUEST}) as terminal. + * + * @param config retry configuration, must not be {@code null} + * @return this builder + * @throws NullPointerException if {@code config} is {@code null} + * @see TopicRetryConfig#FOREVER + * @see TopicRetryConfig#NEVER + * @see TopicRetryConfig#STANDARD + */ + public Builder setRetryConfig(RetryConfig config) { + if (config == null) { + throw new NullPointerException("RetryConfig must not be null"); + } + this.retryConfig = config; + return this; + } + /** * Set executor for decompression tasks. * If not set, default executor will be used. From d4edff18c5bdf4f4000be91dd06bac9f6ba9ff56 Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Thu, 10 Sep 2026 16:03:14 +0100 Subject: [PATCH 02/12] Removed old retrier implementation --- .../ydb/topic/impl/GrpcStreamRetrier.java | 161 ------------------ .../java/tech/ydb/topic/impl/Session.java | 9 - .../java/tech/ydb/topic/impl/SessionBase.java | 117 ------------- 3 files changed, 287 deletions(-) delete mode 100644 topic/src/main/java/tech/ydb/topic/impl/GrpcStreamRetrier.java delete mode 100644 topic/src/main/java/tech/ydb/topic/impl/Session.java delete mode 100644 topic/src/main/java/tech/ydb/topic/impl/SessionBase.java diff --git a/topic/src/main/java/tech/ydb/topic/impl/GrpcStreamRetrier.java b/topic/src/main/java/tech/ydb/topic/impl/GrpcStreamRetrier.java deleted file mode 100644 index cab650fdc..000000000 --- a/topic/src/main/java/tech/ydb/topic/impl/GrpcStreamRetrier.java +++ /dev/null @@ -1,161 +0,0 @@ -package tech.ydb.topic.impl; - -import java.util.Random; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.RejectedExecutionException; -import java.util.concurrent.ScheduledExecutorService; -import java.util.concurrent.ThreadLocalRandom; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.atomic.AtomicInteger; -import java.util.function.BiConsumer; - -import org.slf4j.Logger; - -import tech.ydb.core.Status; - -/** - * @author Nikolay Perfilov - */ -public abstract class GrpcStreamRetrier { - // TODO: add retry policy - private static final int MAX_RECONNECT_COUNT = 0; // Inf - private static final int EXP_BACKOFF_BASE_MS = 256; - private static final int EXP_BACKOFF_CEILING_MS = 40000; // 40 sec (max delays would be 40-80 sec) - private static final int EXP_BACKOFF_MAX_POWER = 7; - private static final int ID_LENGTH = 6; - private static final char[] ID_ALPHABET = "abcdefghijklmnopqrstuvwxyzABSDEFGHIJKLMNOPQRSTUVWXYZ1234567890" - .toCharArray(); - - protected final String id; - protected final AtomicBoolean isReconnecting = new AtomicBoolean(false); - protected final AtomicBoolean isStopped = new AtomicBoolean(false); - protected final AtomicInteger reconnectCounter = new AtomicInteger(0); - - private final ScheduledExecutorService scheduler; - private final BiConsumer errorsHandler; - - protected GrpcStreamRetrier( - String id, - ScheduledExecutorService scheduler, - BiConsumer errorsHandler - ) { - this.scheduler = scheduler; - this.id = id == null ? generateRandomId(ID_LENGTH) : id; - this.errorsHandler = errorsHandler; - } - - protected abstract Logger getLogger(); - protected abstract String getStreamName(); - protected abstract void onStreamReconnect(); - protected abstract void onShutdown(String reason); - - protected static String generateRandomId(int length) { - return new Random().ints(0, ID_ALPHABET.length) - .limit(length) - .map(charId -> ID_ALPHABET[charId]) - .collect(StringBuilder::new, StringBuilder::appendCodePoint, StringBuilder::append) - .toString(); - } - - private void tryScheduleReconnect() { - int currentReconnectCounter = reconnectCounter.get() + 1; - if (MAX_RECONNECT_COUNT > 0 && currentReconnectCounter > MAX_RECONNECT_COUNT) { - if (isStopped.compareAndSet(false, true)) { - String errorMessage = "[" + id + "] Maximum retry count (" + MAX_RECONNECT_COUNT - + ") exceeded. Shutting down " + getStreamName(); - getLogger().error(errorMessage); - shutdownImpl(errorMessage); - return; - } else { - getLogger().info("[{}] Maximum retry count ({}}) exceeded. Need to shutdown {} but it's already " + - "shut down.", id, MAX_RECONNECT_COUNT, getStreamName()); - } - } - if (isReconnecting.compareAndSet(false, true)) { - reconnectCounter.set(currentReconnectCounter); - int delayMs = currentReconnectCounter <= EXP_BACKOFF_MAX_POWER - ? EXP_BACKOFF_BASE_MS * (1 << currentReconnectCounter) - : EXP_BACKOFF_CEILING_MS; - // Add jitter - delayMs = delayMs + ThreadLocalRandom.current().nextInt(delayMs); - getLogger().warn("[{}] Retry #{}. Scheduling {} reconnect in {}ms...", id, currentReconnectCounter, - getStreamName(), delayMs); - try { - scheduler.schedule(this::reconnect, delayMs, TimeUnit.MILLISECONDS); - } catch (RejectedExecutionException exception) { - String errorMessage = "[" + id + "] Couldn't schedule reconnect: scheduler is already shut down. " + - "Shutting down " + getStreamName(); - getLogger().error(errorMessage); - shutdownImpl(errorMessage); - } - } else { - getLogger().info("[{}] should reconnect {} stream, but reconnect is already in progress", id, - getStreamName()); - } - } - - void reconnect() { - if (isStopped.get()) { - getLogger().info("[{}] {} is already stopped, no need to reconnect", id, getStreamName()); - return; - } - - getLogger().info("[{}] {} reconnect #{} started", id, getStreamName(), reconnectCounter.get()); - if (!isReconnecting.compareAndSet(true, false)) { - getLogger().warn("Couldn't reset reconnect flag. Shouldn't happen"); - } - onStreamReconnect(); - } - - protected CompletableFuture shutdownImpl() { - return shutdownImpl(""); - } - - protected CompletableFuture shutdownImpl(String reason) { - getLogger().info( - "[{}] Shutting down {}{}", - id, - getStreamName(), - reason == null || reason.isEmpty() ? "" : " with reason: " + reason - ); - isStopped.set(true); - return CompletableFuture.runAsync(() -> { - onShutdown(reason); - }); - } - - public void onSessionClosed(Status status, Throwable th) { - getLogger().info("[{}] onSessionClosed called", id); - - if (status != null) { - if (status.isSuccess()) { - if (isStopped.get()) { - getLogger().info("[{}] {} stream session closed successfully", id, getStreamName()); - return; - } else { - getLogger().warn("[{}] {} stream session was closed on working {}", id, getStreamName(), - getStreamName()); - } - } else { - getLogger().warn("[{}] Error in {} stream session: {}", id, getStreamName(), status); - } - } else { - getLogger().error("[{}] Exception in {} stream session: ", id, getStreamName(), th); - } - - if (errorsHandler != null) { - try { - errorsHandler.accept(status, th); - } catch (Exception ex) { - getLogger().error("[{}] error handler throws exception", id, ex); - } - } - - if (!isStopped.get()) { - tryScheduleReconnect(); - } else { - getLogger().info("[{}] {} is already stopped, no need to schedule reconnect", id, getStreamName()); - } - } -} diff --git a/topic/src/main/java/tech/ydb/topic/impl/Session.java b/topic/src/main/java/tech/ydb/topic/impl/Session.java deleted file mode 100644 index f9745d173..000000000 --- a/topic/src/main/java/tech/ydb/topic/impl/Session.java +++ /dev/null @@ -1,9 +0,0 @@ -package tech.ydb.topic.impl; - -/** - * @author Nikolay Perfilov - */ -public interface Session { - void startAndInitialize(); - boolean shutdown(); -} diff --git a/topic/src/main/java/tech/ydb/topic/impl/SessionBase.java b/topic/src/main/java/tech/ydb/topic/impl/SessionBase.java deleted file mode 100644 index 5e2c7198a..000000000 --- a/topic/src/main/java/tech/ydb/topic/impl/SessionBase.java +++ /dev/null @@ -1,117 +0,0 @@ -package tech.ydb.topic.impl; - -import java.util.Objects; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.locks.ReentrantLock; - -import org.slf4j.Logger; - -import tech.ydb.core.Status; -import tech.ydb.core.grpc.GrpcReadStream; -import tech.ydb.core.grpc.GrpcReadWriteStream; - -/** - * @author Nikolay Perfilov - * @param type of message received from the server - * @param type of message to be sent to the server - */ -public abstract class SessionBase implements Session { - - protected final GrpcReadWriteStream streamConnection; - protected final String streamId; - private final AtomicBoolean isWorking = new AtomicBoolean(true); - private final ReentrantLock lock = new ReentrantLock(); - private String token; - - public SessionBase(GrpcReadWriteStream streamConnection, String streamId) { - this.streamConnection = streamConnection; - this.streamId = streamId; - this.token = streamConnection.authToken(); - } - - public String getStreamId() { - return streamId; - } - - public boolean isStopped() { - return !isWorking.get(); - } - - protected abstract Logger getLogger(); - - protected abstract void sendUpdateTokenRequest(String token); - - protected abstract void onStop(); - - protected CompletableFuture start(GrpcReadStream.Observer streamObserver) { - lock.lock(); - - try { - getLogger().info("[{}] Session start", streamId); - return streamConnection.start(message -> { - if (getLogger().isTraceEnabled()) { - getLogger().trace("[{}] Message received:\n{}", streamId, message); - } - - if (isWorking.get()) { - streamObserver.onNext(message); - } - }); - } finally { - lock.unlock(); - } - } - - public void send(W request) { - lock.lock(); - - try { - if (!isWorking.get()) { - if (getLogger().isTraceEnabled()) { - getLogger().trace( - "[{}] Session is already closed. This message is NOT sent:\n{}", - streamId, - request - ); - } - return; - } - String currentToken = streamConnection.authToken(); - if (!Objects.equals(token, currentToken)) { - token = currentToken; - getLogger().info("[{}] Sending new token", streamId); - sendUpdateTokenRequest(token); - } - - if (getLogger().isTraceEnabled()) { - getLogger().trace("[{}] Sending request:\n{}", streamId, request); - } - streamConnection.sendNext(request); - } finally { - lock.unlock(); - } - } - - private boolean stop() { - getLogger().info("[{}] Session stop", streamId); - return isWorking.compareAndSet(true, false); - } - - @Override - public boolean shutdown() { - lock.lock(); - - try { - getLogger().info("[{}] Session shutdown", streamId); - if (stop()) { - onStop(); - streamConnection.close(); - return true; - } - return false; - } finally { - lock.unlock(); - } - } -} From 1e2af0cc41749cf28b33bbd8ae68222e48f1efaf Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Wed, 16 Sep 2026 09:55:05 +0100 Subject: [PATCH 03/12] Fixed small errors --- .../ydb/topic/read/impl/AsyncReaderImpl.java | 6 +++++- .../tech/ydb/topic/read/impl/ReadSession.java | 8 +++++++- .../tech/ydb/topic/read/impl/ReaderImpl.java | 10 ++++++++++ .../ydb/topic/read/impl/SyncReaderImpl.java | 18 ++++++++++++++---- 4 files changed, 36 insertions(+), 6 deletions(-) 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 ed81e1be1..451d1f461 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 @@ -99,13 +99,17 @@ protected CompletableFuture handleReaderClosed() { @Override public CompletableFuture shutdown() { - impl.close(); + if (!impl.close()) { + // implicit closing because stream will never call onClose + close(); + } return shutdownFuture; } private void close() { decompressor.close(); processor.close(); + initFuture.complete(null); shutdownFuture.complete(null); } 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 018ac52fc..627f05730 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 @@ -71,6 +71,7 @@ protected Status parseMessageStatus(FromServer message) { } public Set closeAll() { + isClosed = true; decoder.stop(); Set closed = new HashSet<>(partitions.values()); @@ -164,6 +165,11 @@ public StopPartitionSessionEvent onStopPartition(YdbTopic.StreamReadMessage.Stop if (partition == null) { logger.error("[{}] Received graceful StopPartitionSessionRequest for partition session {}, " + "but have no such partition session active", debugId, psid); + send(YdbTopic.StreamReadMessage.FromClient.newBuilder().setStopPartitionSessionResponse( + YdbTopic.StreamReadMessage.StopPartitionSessionResponse.newBuilder() + .setPartitionSessionId(psid) + .build()) + .build()); return null; } @@ -197,7 +203,7 @@ public void onCommitOffset(YdbTopic.StreamReadMessage.CommitOffsetResponse respo if (queue == null) { logger.info("[{}] Received CommitOffsetResponse for unknown (most likely already closed) " + "partition session with id={}", debugId, offset.getPartitionSessionId()); - return; + continue; } // Handling CompletableFuture completions for single commits 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 93f8d56ec..0230a211b 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 @@ -6,6 +6,7 @@ import java.util.Map; import java.util.UUID; import java.util.concurrent.CompletableFuture; +import java.util.function.BiConsumer; import java.util.stream.Collectors; import org.slf4j.Logger; @@ -36,6 +37,7 @@ * @author Nikolay Perfilov */ public class ReaderImpl extends TopicRetryableStream { + public interface Releaser { void releaseRange(PartitionSession partition, OffsetsRange range); } @@ -57,6 +59,7 @@ public interface Handler { private final TopicRpc rpc; private final ReadConfig config; private final Handler handler; + private final BiConsumer errorHandler; private final FromClient initRequest; @@ -66,6 +69,7 @@ public ReaderImpl(TopicRpc rpc, String id, ReaderSettings settings, ReadConfig c this.initRequest = FromClient.newBuilder().setInitRequest(buildInitRequest(settings)).build(); this.config = config; this.handler = handler; + this.errorHandler = settings.getErrorsHandler(); } @Override @@ -76,6 +80,9 @@ protected ReadSession createNewStream(String id) { @Override protected void onRetry(ReadSession stream, Status status) { logger.warn("[{}] paused by status {}", debugId, status); + if (errorHandler != null) { + errorHandler.accept(status, null); + } stream.closeAll().forEach(ps -> handler.handleClosePartitionSession(ps)); } @@ -86,6 +93,9 @@ protected void onClose(ReadSession stream, Status status) { } else { logger.info("[{}] closed by status {}", debugId, status); } + if (errorHandler != null) { + errorHandler.accept(status, null); + } stream.closeAll().forEach(ps -> handler.handleClosePartitionSession(ps)); handler.handleReaderClosed(status); } diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java b/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java index 86e314887..5967cf28b 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java @@ -93,7 +93,19 @@ public void initAndWait() { @Override public void shutdown() { - impl.close(); + if (!impl.close()) { + // implicit closing because stream will never call onClose + close(); + } + + shutdownFuture.join(); + } + + private void close() { + initFuture.complete(null); + shutdownFuture.complete(null); + + decompressor.close(); waitingLock.lock(); try { @@ -101,8 +113,6 @@ public void shutdown() { } finally { waitingLock.unlock(); } - - shutdownFuture.join(); } @Override @@ -198,7 +208,7 @@ public void handleSessionStarted(String sessionId) { @Override public void handleReaderClosed(Status status) { - shutdownFuture.complete(null); + close(); } @Override From dcf5b0ce8c4fd6d7e6f47797f83f7a52d0fa6b21 Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Wed, 16 Sep 2026 10:35:02 +0100 Subject: [PATCH 04/12] Added partition session contol --- .../ydb/topic/read/impl/AsyncReaderImpl.java | 5 +- .../topic/read/impl/ReadPartitionSession.java | 44 ++++++++++------- .../tech/ydb/topic/read/impl/ReadSession.java | 49 ++++++++++--------- .../tech/ydb/topic/read/impl/ReaderImpl.java | 9 ++-- .../ydb/topic/read/impl/SyncReaderImpl.java | 36 ++++++-------- 5 files changed, 74 insertions(+), 69 deletions(-) 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 451d1f461..76683f2f7 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 @@ -26,7 +26,6 @@ import tech.ydb.topic.read.events.ReaderClosedEvent; import tech.ydb.topic.read.events.StartPartitionSessionEvent; import tech.ydb.topic.read.events.StopPartitionSessionEvent; -import tech.ydb.topic.read.impl.ReaderImpl.Releaser; import tech.ydb.topic.read.impl.events.CommitOffsetAcknowledgementEventImpl; import tech.ydb.topic.read.impl.events.PartitionSessionClosedEventImpl; import tech.ydb.topic.read.impl.events.SessionStartedEvent; @@ -142,7 +141,7 @@ public void handleReaderClosed(Status status) { } @Override - public void handleDataReceivedEvent(Releaser releaser, DataReceivedEvent event) { + public void handleDataReceivedEvent(ReaderImpl.PartitionControl control, DataReceivedEvent event) { try { int messagesCount = event.getMessages().size(); long offsetStart = event.getMessages().get(0).getOffset(); @@ -156,7 +155,7 @@ public void handleDataReceivedEvent(Releaser releaser, DataReceivedEvent event) failSession(th, "onMessages"); throw th; } finally { - releaser.releaseRange(event.getPartitionSession(), event.getRangeToCommit()); + control.confirmRangeProcessed(event.getRangeToCommit()); } } diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionSession.java b/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionSession.java index eeabf1d29..8d15f3c11 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionSession.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionSession.java @@ -5,7 +5,7 @@ import java.util.List; import java.util.Queue; import java.util.concurrent.ConcurrentLinkedQueue; -import java.util.function.Consumer; +import java.util.function.BiConsumer; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -22,14 +22,15 @@ /** * @author Nikolay Perfilov */ -public class ReadPartitionSession { +public class ReadPartitionSession implements ReaderImpl.PartitionControl { private static final Logger logger = LoggerFactory.getLogger(ReaderImpl.class); private final String traceID; private final PartitionSession partition; private final MessageCommitterImpl committer; private final ReadPartitionDecoder decoder; - private final Consumer eventConsumer; + private final BufferManager bufferManager; + private final BiConsumer eventConsumer; private final int maxBatchSize; private final SerialExecutor executor; @@ -39,25 +40,35 @@ public class ReadPartitionSession { private final Queue readingQueue = new ConcurrentLinkedQueue<>(); - ReadPartitionSession(String traceID, ReadConfig config, PartitionSession partition, MessageCommitterImpl committer, - MessageDecoder decoder, Consumer eventConsumer, long lastCommittedOffset) { + ReadPartitionSession(String traceID, ReadSession session, PartitionSession partition, + MessageCommitterImpl committer, long lastCommittedOffset) { this.traceID = traceID; this.partition = partition; this.committer = committer; - this.decoder = new ReadPartitionDecoder(traceID, decoder, partition, committer, this::sendDataToReaders); + this.decoder = new ReadPartitionDecoder(traceID, session.getDecoder(), partition, committer, + this::sendDataToReaders); - this.maxBatchSize = config.getMaxBatchSize(); - this.executor = new SerialExecutor(config.getProcessor()); - this.eventConsumer = eventConsumer; + this.maxBatchSize = session.getConfig().getMaxBatchSize(); + this.executor = new SerialExecutor(session.getConfig().getProcessor()); + this.bufferManager = session.getBufferManager(); + this.eventConsumer = session.getEventConsumer(); this.lastReadOffset = lastCommittedOffset; } - public PartitionSession getPartition() { - return partition; + @Override + public boolean isActive() { + return !isStopped; } - public boolean isStopped() { - return isStopped; + @Override + public void confirmRangeProcessed(OffsetsRange range) { + bufferManager.releaseRange(partition.getId(), range); + decoder.releaseRange(range); + sendDataToReaders(); + } + + public PartitionSession getPartition() { + return partition; } public void confirmCommittedOffset(long committedOffset) { @@ -113,11 +124,6 @@ public boolean addBatches(List ba return !isStopped; } - public void releaseRange(OffsetsRange range) { - decoder.releaseRange(range); - sendDataToReaders(); - } - public void sendDataToReaders() { executor.execute(() -> { while (!isStopped) { @@ -138,7 +144,7 @@ public void sendDataToReaders() { next = it.hasNext() ? it.next() : null; } - eventConsumer.accept(new DataReceivedEventImpl(partition, committer, messagesToRead)); + eventConsumer.accept(this, new DataReceivedEventImpl(partition, committer, messagesToRead)); } }); } 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 627f05730..3fe0dd4d0 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 @@ -36,21 +36,21 @@ * * @author Aleksandr Gorshenin {@literal } */ -public class ReadSession extends TopicStreamBase implements ReaderImpl.Releaser { +public class ReadSession extends TopicStreamBase { private static final Logger logger = LoggerFactory.getLogger(ReadSession.class); private final String debugId; private final ReadConfig config; private final MessageDecoder decoder; private final BufferManager bufferManager; - private final BiConsumer eventConsumer; + private final BiConsumer eventConsumer; private final Map partitions = new ConcurrentHashMap<>(); private final Map readQueues = new ConcurrentHashMap<>(); private volatile boolean isClosed = false; public ReadSession(String id, GrpcReadWriteStream stream, FromClient initReq, - BiConsumer eventConsumer, ReadConfig config) { + BiConsumer eventConsumer, ReadConfig config) { super(logger, id, stream, initReq); this.debugId = id; this.config = config; @@ -70,6 +70,22 @@ protected Status parseMessageStatus(FromServer message) { return Status.of(StatusCode.fromProto(message.getStatus()), Issue.fromPb(message.getIssuesList())); } + BufferManager getBufferManager() { + return bufferManager; + } + + ReadConfig getConfig() { + return config; + } + + MessageDecoder getDecoder() { + return decoder; + } + + BiConsumer getEventConsumer() { + return eventConsumer; + } + public Set closeAll() { isClosed = true; decoder.stop(); @@ -83,24 +99,12 @@ public Set closeAll() { return closed; } - @Override - public void releaseRange(PartitionSession partition, OffsetsRange range) { - bufferManager.releaseRange(partition.getId(), range); - ReadPartitionSession queue = readQueues.get(partition.getId()); - if (queue != null) { - queue.releaseRange(range); - } - } - public boolean commitOffsets(PartitionSession session, List rangesToCommit) { - if (isClosed) { - logger.atInfo() - .setMessage("[{}] Need to send CommitRequest for {} with offset ranges {}, " - + "but reading session is already closed") - .addArgument(debugId) - .addArgument(session) - .addArgument(() -> rangesToCommit.stream().map(Object::toString).collect(Collectors.joining(", "))) - .log(); + ReadPartitionSession partition = readQueues.get(session.getId()); + if (isClosed || partition == null || !partition.isActive()) { + logger.info("[{}] Need to send CommitRequest for {} with offset ranges {}, " + + "but reading partition session is already closed", debugId, session, + rangesToCommit.stream().map(Object::toString).collect(Collectors.joining(", "))); return false; } @@ -277,9 +281,8 @@ public void confirm(StartPartitionSessionSettings options) { } } - MessageCommitterImpl committer = new MessageCommitterImpl(traceID, ReadSession.this, partition, commitTo); - ReadPartitionSession queue = new ReadPartitionSession(traceID, config, partition, committer, decoder, - event -> eventConsumer.accept(ReadSession.this, event), commitTo); + MessageCommitterImpl comm = new MessageCommitterImpl(traceID, ReadSession.this, partition, commitTo); + ReadPartitionSession queue = new ReadPartitionSession(traceID, ReadSession.this, partition, comm, commitTo); if (readQueues.putIfAbsent(psid, queue) != null) { logger.warn("[{}] partition {} is already started", traceID, partition); return; 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 0230a211b..77a65cf5a 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 @@ -37,10 +37,11 @@ * @author Nikolay Perfilov */ public class ReaderImpl extends TopicRetryableStream { - - public interface Releaser { - void releaseRange(PartitionSession partition, OffsetsRange range); + public interface PartitionControl { + boolean isActive(); + void confirmRangeProcessed(OffsetsRange range); } + public interface Handler { void handleSessionStarted(String sessionId); @@ -48,7 +49,7 @@ public interface Handler { void handleStopPartitionSession(StopPartitionSessionEvent event); void handleClosePartitionSession(PartitionSession partition); - void handleDataReceivedEvent(Releaser releaser, DataReceivedEvent event); + void handleDataReceivedEvent(PartitionControl control, DataReceivedEvent event); void handleCommitResponse(long committedOffset, PartitionSession partition); void handleReaderClosed(Status status); diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java b/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java index 5967cf28b..6beb44bc0 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java @@ -6,7 +6,6 @@ import java.util.List; import java.util.Queue; import java.util.concurrent.CompletableFuture; -import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.TimeUnit; import java.util.concurrent.locks.Condition; @@ -22,6 +21,7 @@ import tech.ydb.core.Status; import tech.ydb.topic.TopicRpc; import tech.ydb.topic.description.CodecRegistry; +import tech.ydb.topic.description.OffsetsRange; import tech.ydb.topic.impl.DebugTools; import tech.ydb.topic.read.Message; import tech.ydb.topic.read.PartitionOffsets; @@ -50,7 +50,6 @@ public class SyncReaderImpl implements SyncReader { private final CompletableFuture initFuture = new CompletableFuture<>(); private final CompletableFuture shutdownFuture = new CompletableFuture<>(); - private final ConcurrentHashMap activePartitions = new ConcurrentHashMap<>(); private final Queue queue = new ConcurrentLinkedQueue<>(); private final ReentrantLock waitingLock = new ReentrantLock(); private final Condition waitingCondition = waitingLock.newCondition(); @@ -172,7 +171,7 @@ public Message receiveInternal(ReceiveSettings receiveSettings, long timeout, Ti } } - if (!activePartitions.containsKey(next.getPartition())) { + if (!next.isActive()) { continue; } @@ -212,12 +211,12 @@ public void handleReaderClosed(Status status) { } @Override - public void handleDataReceivedEvent(ReaderImpl.Releaser releaser, DataReceivedEvent event) { + public void handleDataReceivedEvent(ReaderImpl.PartitionControl control, DataReceivedEvent event) { if (impl.isClosed()) { return; } if (event.getMessages().isEmpty()) { - releaser.releaseRange(event.getPartitionSession(), event.getRangeToCommit()); + control.confirmRangeProcessed(event.getRangeToCommit()); return; } @@ -228,12 +227,11 @@ public void handleDataReceivedEvent(ReaderImpl.Releaser releaser, DataReceivedEv logger.debug("{} Putting a batch into queueData with {} message(s) (offsets {}-{}) from {}", debugId, messagesCount, offsetStart, offsetEnd, ps); - Runnable confirm = () -> releaser.releaseRange(ps, event.getRangeToCommit()); for (Message msg: event.getMessages()) { if (msg.getRangeToCommit().getEnd() == event.getRangeToCommit().getEnd()) { // last message in batch - queue.offer(new MessageWrapper(ps, msg, confirm)); + queue.offer(new MessageWrapper(control, msg, event.getRangeToCommit())); } else { - queue.offer(new MessageWrapper(ps, msg, null)); + queue.offer(new MessageWrapper(control, msg, null)); } } @@ -252,45 +250,43 @@ public void handleCommitResponse(long committedOffset, PartitionSession partitio @Override public void handleStartPartitionSessionRequest(StartPartitionSessionEvent event) { - activePartitions.put(event.getPartitionSession(), event.getPartitionSession()); event.confirm(); } @Override public void handleStopPartitionSession(StopPartitionSessionEvent event) { - activePartitions.remove(event.getPartitionSession()); // TODO: wait for all commits event.confirm(); } @Override public void handleClosePartitionSession(PartitionSession partition) { - activePartitions.remove(partition); + // Nothing } } private static class MessageWrapper { - private final PartitionSession partition; + private final ReaderImpl.PartitionControl control; private final Message msg; - private final Runnable confirm; + private final OffsetsRange rangeToConfirm; - private MessageWrapper(PartitionSession partition, Message msg, Runnable confirm) { - this.partition = partition; + private MessageWrapper(ReaderImpl.PartitionControl control, Message msg, OffsetsRange rangeToConfirm) { + this.control = control; this.msg = msg; - this.confirm = confirm; + this.rangeToConfirm = rangeToConfirm; } Message getMessage() { return msg; } - PartitionSession getPartition() { - return partition; + boolean isActive() { + return control.isActive(); } void confirm() { - if (confirm != null) { - confirm.run(); + if (rangeToConfirm != null) { + control.confirmRangeProcessed(rangeToConfirm); } } } From af21cbaa3de49d73935dde1e81156cb9b3d80552 Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Wed, 16 Sep 2026 12:20:18 +0100 Subject: [PATCH 05/12] Updated YdbTopicsIntegrationTest --- .../ydb/topic/YdbTopicsIntegrationTest.java | 98 +++++++++++-------- 1 file changed, 56 insertions(+), 42 deletions(-) diff --git a/topic/src/test/java/tech/ydb/topic/YdbTopicsIntegrationTest.java b/topic/src/test/java/tech/ydb/topic/YdbTopicsIntegrationTest.java index f4ef9d738..7d48cb693 100644 --- a/topic/src/test/java/tech/ydb/topic/YdbTopicsIntegrationTest.java +++ b/topic/src/test/java/tech/ydb/topic/YdbTopicsIntegrationTest.java @@ -13,11 +13,9 @@ import org.junit.Assert; import org.junit.BeforeClass; import org.junit.ClassRule; -import org.junit.FixMethodOrder; import org.junit.Rule; import org.junit.Test; import org.junit.rules.Timeout; -import org.junit.runners.MethodSorters; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -53,7 +51,6 @@ * * @author Aleksandr Gorshenin */ -@FixMethodOrder(MethodSorters.NAME_ASCENDING) public class YdbTopicsIntegrationTest { private final static Logger logger = LoggerFactory.getLogger(YdbTopicsIntegrationTest.class); @@ -100,7 +97,16 @@ public static void dropTopic() { } @Test - public void step01_writeWithoutDeduplication() throws InterruptedException, ExecutionException, TimeoutException { + public void writeAndReadTest() throws Exception { + step01_writeWithoutDeduplication(); + step02_readHalfWithoutCommit(); + step03_readHalfWithCommit(); + step04_readNextHalfWithoutCommit(); + step05_readNextHalfWithCommit(); + step06_readAllByAsyncReader(); + } + + private void step01_writeWithoutDeduplication() throws InterruptedException, ExecutionException, TimeoutException { WriterSettings settings = WriterSettings.newBuilder() .setTopicPath(TEST_TOPIC) .build(); @@ -119,8 +125,7 @@ public void step01_writeWithoutDeduplication() throws InterruptedException, Exec writer.shutdown(1, TimeUnit.MINUTES); } - @Test - public void step02_readHalfWithoutCommit() throws InterruptedException { + private void step02_readHalfWithoutCommit() throws InterruptedException { ReaderSettings settings = ReaderSettings.newBuilder() .addTopic(TopicReadSettings.newBuilder().setPath(TEST_TOPIC).build()) .setConsumerName(TEST_CONSUMER1) @@ -137,8 +142,7 @@ public void step02_readHalfWithoutCommit() throws InterruptedException { reader.shutdown(); } - @Test - public void step03_readHalfWithCommit() throws InterruptedException { + private void step03_readHalfWithCommit() throws InterruptedException { ReaderSettings settings = ReaderSettings.newBuilder() .addTopic(TopicReadSettings.newBuilder().setPath(TEST_TOPIC).build()) .setConsumerName(TEST_CONSUMER1) @@ -156,8 +160,7 @@ public void step03_readHalfWithCommit() throws InterruptedException { reader.shutdown(); } - @Test - public void step03_readNextHalfWithoutCommit() throws InterruptedException { + private void step04_readNextHalfWithoutCommit() throws InterruptedException { ReaderSettings settings = ReaderSettings.newBuilder() .addTopic(TopicReadSettings.newBuilder().setPath(TEST_TOPIC).build()) .setConsumerName(TEST_CONSUMER1) @@ -178,8 +181,7 @@ public void step03_readNextHalfWithoutCommit() throws InterruptedException { reader.shutdown(); } - @Test - public void step04_readNextHalfWithCommit() throws InterruptedException { + private void step05_readNextHalfWithCommit() throws InterruptedException { ReaderSettings settings = ReaderSettings.newBuilder() .addTopic(TopicReadSettings.newBuilder().setPath(TEST_TOPIC).build()) .setConsumerName(TEST_CONSUMER1) @@ -203,20 +205,7 @@ public void step04_readNextHalfWithCommit() throws InterruptedException { reader.shutdown(); } - @Test - public void step05_describeTopic() { - TopicDescription description = client.describeTopic(TEST_TOPIC).join().getValue(); - - Assert.assertNull(description.getTopicStats()); - List consumers = description.getConsumers(); - Assert.assertEquals(2, consumers.size()); - - Assert.assertEquals(TEST_CONSUMER1, consumers.get(0).getName()); - Assert.assertEquals(TEST_CONSUMER2, consumers.get(1).getName()); - } - - @Test - public void step06_readAllByAsyncReader() throws InterruptedException { + private void step06_readAllByAsyncReader() throws InterruptedException { ReaderSettings settings = ReaderSettings.newBuilder() .addTopic(TopicReadSettings.newBuilder().setPath(TEST_TOPIC).build()) .setConsumerName(TEST_CONSUMER2) @@ -256,7 +245,19 @@ public void onMessages(DataReceivedEvent dre) { } @Test - public void step07_alterTopicWithAutoPartitioning() { + public void describeTopic() { + TopicDescription description = client.describeTopic(TEST_TOPIC).join().getValue(); + + Assert.assertNull(description.getTopicStats()); + List consumers = description.getConsumers(); + Assert.assertEquals(2, consumers.size()); + + Assert.assertEquals(TEST_CONSUMER1, consumers.get(0).getName()); + Assert.assertEquals(TEST_CONSUMER2, consumers.get(1).getName()); + } + + @Test + public void alterTopicWithAutoPartitioning() { client.alterTopic(TEST_TOPIC, AlterTopicSettings.newBuilder() .setAlterPartitioningSettings(AlterPartitioningSettings.newBuilder() .setAutoPartitioningStrategy(AutoPartitioningStrategy.SCALE_UP) @@ -287,7 +288,7 @@ public void step07_alterTopicWithAutoPartitioning() { } @Test - public void step08_createTopicWithAutoPartitioning() { + public void createTopicWithAutoPartitioning() { PartitioningSettings expectedPartitioningSettings = PartitioningSettings.newBuilder() .setMaxActivePartitions(8) .setMinActivePartitions(4) @@ -299,9 +300,9 @@ public void step08_createTopicWithAutoPartitioning() { .build()) .build(); - CompletableFuture secondaryTopicCreated = client.createTopic(TEST_OTHER_TOPIC, CreateTopicSettings.newBuilder() - .setPartitioningSettings(expectedPartitioningSettings) - .build()); + CompletableFuture secondaryTopicCreated = client.createTopic(TEST_OTHER_TOPIC, + CreateTopicSettings.newBuilder().setPartitioningSettings(expectedPartitioningSettings).build() + ); secondaryTopicCreated.join().expectSuccess("can't create the topic"); @@ -311,7 +312,7 @@ public void step08_createTopicWithAutoPartitioning() { } @Test - public void step09_describeTopicStats() { + public void describeTopicStats() { DescribeTopicSettings on = DescribeTopicSettings.newBuilder().withIncludeStats(true).build(); DescribeTopicSettings off = DescribeTopicSettings.newBuilder().withIncludeStats(false).build(); @@ -323,23 +324,21 @@ public void step09_describeTopicStats() { for (Consumer consumer: withoutStats.getConsumers()) { Assert.assertNull(consumer.getStats()); - Assert.assertNull(consumer.getAvailabilityPeriod()); - } - for (Consumer consumer: withStats.getConsumers()) { - Assert.assertNotNull(consumer.getStats()); - Assert.assertNull(consumer.getAvailabilityPeriod()); } - for (PartitionInfo partition: withoutStats.getPartitions()) { Assert.assertNull(partition.getPartitionStats()); } + + for (Consumer consumer: withStats.getConsumers()) { + Assert.assertNotNull(consumer.getStats()); + } for (PartitionInfo partition: withStats.getPartitions()) { Assert.assertNotNull(partition.getPartitionStats()); } } @Test - public void step10_invalidAddConsumerTest() { + public void invalidAddConsumerTest() { AlterTopicSettings settings = AlterTopicSettings.newBuilder() .addAddConsumer(Consumer.newBuilder() .setName("WRONG_CONSUMER") @@ -355,7 +354,7 @@ public void step10_invalidAddConsumerTest() { } @Test - public void step11_invalidAlterConsumerTest() { + public void invalidAlterConsumerTest() { AlterTopicSettings settings = AlterTopicSettings.newBuilder() .addAlterConsumer(AlterConsumerSettings.newBuilder() .setName(TEST_CONSUMER2) @@ -371,7 +370,7 @@ public void step11_invalidAlterConsumerTest() { } @Test - public void step12_alterConsumerTest() { + public void alterConsumerTest() { AlterTopicSettings settings = AlterTopicSettings.newBuilder() .addAlterConsumer(AlterConsumerSettings.newBuilder() .setName(TEST_CONSUMER2) @@ -388,5 +387,20 @@ public void step12_alterConsumerTest() { Assert.assertEquals(TEST_CONSUMER2, description.getConsumer().getName()); Assert.assertEquals(Instant.EPOCH.plusSeconds(10), description.getConsumer().getReadFrom()); Assert.assertEquals(Duration.ofMinutes(5), description.getConsumer().getAvailabilityPeriod()); - } + + TopicDescription topicDesc = client.describeTopic(TEST_TOPIC).join().getValue(); + + Assert.assertNull(topicDesc.getTopicStats()); + + for (Consumer consumer: topicDesc.getConsumers()) { + Assert.assertNull(consumer.getStats()); + if (TEST_CONSUMER2.equals(consumer.getName())) { + Assert.assertEquals(Duration.ofMinutes(5), consumer.getAvailabilityPeriod()); + Assert.assertEquals(Instant.EPOCH.plusSeconds(10), consumer.getReadFrom()); + } else { + Assert.assertNull(consumer.getAvailabilityPeriod()); + Assert.assertEquals(Instant.EPOCH, consumer.getReadFrom()); + } + } + } } From 2df36d3348552c31c215ab041bf08fd1e7dba11a Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Wed, 16 Sep 2026 13:34:28 +0100 Subject: [PATCH 06/12] Added tests for SyncReaderImpl --- .../tech/ydb/topic/read/impl/ReaderImpl.java | 11 +- .../ydb/topic/read/impl/SyncReaderImpl.java | 14 +- .../topic/TopicWritersIntegrationTest.java | 35 +--- .../ydb/topic/read/impl/ReadStreamMock.java | 84 +++++++- .../topic/read/impl/SyncReaderImplTest.java | 191 +++++++++++++++++- .../tech/ydb/topic/utils/ErrorsHandler.java | 42 ++++ 6 files changed, 327 insertions(+), 50 deletions(-) create mode 100644 topic/src/test/java/tech/ydb/topic/utils/ErrorsHandler.java 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 77a65cf5a..c2347126a 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 @@ -82,7 +82,11 @@ protected ReadSession createNewStream(String id) { protected void onRetry(ReadSession stream, Status status) { logger.warn("[{}] paused by status {}", debugId, status); if (errorHandler != null) { - errorHandler.accept(status, null); + try { + errorHandler.accept(status, null); + } catch (RuntimeException ex) { + logger.error("[{}] errorHandler onRetry processing throws exception", debugId, ex); + } } stream.closeAll().forEach(ps -> handler.handleClosePartitionSession(ps)); } @@ -95,6 +99,11 @@ protected void onClose(ReadSession stream, Status status) { logger.info("[{}] closed by status {}", debugId, status); } if (errorHandler != null) { + try { + errorHandler.accept(status, null); + } catch (RuntimeException ex) { + logger.error("[{}] errorHandler onClose processing throws exception", debugId, ex); + } errorHandler.accept(status, null); } stream.closeAll().forEach(ps -> handler.handleClosePartitionSession(ps)); diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java b/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java index 6beb44bc0..47f0ce2c3 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java @@ -105,7 +105,10 @@ private void close() { shutdownFuture.complete(null); decompressor.close(); + wakeUp(); + } + private void wakeUp() { waitingLock.lock(); try { waitingCondition.signalAll(); @@ -212,10 +215,10 @@ public void handleReaderClosed(Status status) { @Override public void handleDataReceivedEvent(ReaderImpl.PartitionControl control, DataReceivedEvent event) { - if (impl.isClosed()) { + if (impl.isClosed()) { // never happens return; } - if (event.getMessages().isEmpty()) { + if (event.getMessages().isEmpty()) { // never happens control.confirmRangeProcessed(event.getRangeToCommit()); return; } @@ -235,12 +238,7 @@ public void handleDataReceivedEvent(ReaderImpl.PartitionControl control, DataRec } } - waitingLock.lock(); - try { - waitingCondition.signalAll(); - } finally { - waitingLock.unlock(); - } + wakeUp(); } @Override diff --git a/topic/src/test/java/tech/ydb/topic/TopicWritersIntegrationTest.java b/topic/src/test/java/tech/ydb/topic/TopicWritersIntegrationTest.java index 4b86c69f0..10a09de87 100644 --- a/topic/src/test/java/tech/ydb/topic/TopicWritersIntegrationTest.java +++ b/topic/src/test/java/tech/ydb/topic/TopicWritersIntegrationTest.java @@ -11,7 +11,6 @@ import java.util.concurrent.CountDownLatch; import java.util.concurrent.ThreadLocalRandom; import java.util.concurrent.TimeUnit; -import java.util.function.BiConsumer; import org.junit.After; import org.junit.AfterClass; @@ -44,6 +43,7 @@ import tech.ydb.topic.settings.TopicReadSettings; import tech.ydb.topic.settings.TopicRetryConfig; import tech.ydb.topic.settings.WriterSettings; +import tech.ydb.topic.utils.ErrorsHandler; import tech.ydb.topic.write.AsyncWriter; import tech.ydb.topic.write.InitResult; import tech.ydb.topic.write.Message; @@ -283,7 +283,7 @@ public void defaultRetryPolicyWriter() throws Exception { StatusCode.TRANSPORT_UNAVAILABLE }; - ErrorsHolder errorsHolder = new ErrorsHolder(); + ErrorsHandler errorsHolder = new ErrorsHandler(); WriterSettings settings = WriterSettings.newBuilder() .setTopicPath(ONE_PART_TOPIC) .setProducerId(TEST_PRODUCER) @@ -625,7 +625,7 @@ public void txWriteWithSplitsTest() throws Exception { public void invalidTxWriteTest() throws Exception { createTopicWithOnePartition(); - ErrorsHolder errorsHolder = new ErrorsHolder(); + ErrorsHandler errorsHolder = new ErrorsHandler(); WriterSettings settings = WriterSettings.newBuilder() .setTopicPath(ONE_PART_TOPIC) .setProducerId(TEST_PRODUCER) @@ -679,7 +679,7 @@ public void txRetryWriteTest() throws Exception { PROXY.unavailableOnAckWithSeqNo(2); - ErrorsHolder errorsHolder = new ErrorsHolder(); + ErrorsHandler errorsHolder = new ErrorsHandler(); WriterSettings settings = WriterSettings.newBuilder() .setTopicPath(ONE_PART_TOPIC) .setProducerId(TEST_PRODUCER) @@ -715,31 +715,4 @@ public void txRetryWriteTest() throws Exception { assertTopicContent(ONE_PART_TOPIC, Arrays.asList(msg1, msg2)); } - - private class ErrorsHolder implements BiConsumer { - private final List problems = new ArrayList<>(); - - @Override - public void accept(Status st, Throwable th) { - if (st != null) { - problems.add(st.getCode()); - } - if (th != null) { - problems.add(StatusCode.CLIENT_INTERNAL_ERROR); - } - } - - public void assertEmpty() { - Assert.assertTrue("No reties was expected", problems.isEmpty()); - } - - public void assertCodes(StatusCode... codes) { - Iterator it = problems.iterator(); - for (StatusCode code: codes) { - Assert.assertTrue("Expected " + code + ", but has nothing", it.hasNext()); - Assert.assertEquals(code, it.next()); - } - Assert.assertFalse("Unexpected error code", it.hasNext()); - } - } } 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 cf0bb0166..e4496e717 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 @@ -4,6 +4,7 @@ import java.io.OutputStream; import java.util.ArrayDeque; import java.util.Deque; +import java.util.List; import java.util.Set; import java.util.concurrent.CompletableFuture; import java.util.concurrent.atomic.AtomicInteger; @@ -21,6 +22,7 @@ import tech.ydb.proto.topic.YdbTopic.StreamReadMessage.FromServer; import tech.ydb.topic.description.Codec; import tech.ydb.topic.description.CodecRegistry; +import tech.ydb.topic.description.OffsetsRange; /** * @@ -33,8 +35,8 @@ public class ReadStreamMock implements GrpcReadWriteStream messages = new ArrayDeque<>(); private final AtomicInteger partCounter = new AtomicInteger(); private Observer observer = null; - private boolean isClosed = false; - private boolean isCanceled = false; + private final AtomicInteger isClosed = new AtomicInteger(); + private final AtomicInteger isCanceled = new AtomicInteger(); @Override public String authToken() { @@ -48,7 +50,7 @@ public void sendNext(FromClient message) { @Override public void close() { - this.isClosed = true; + isClosed.incrementAndGet(); } @Override @@ -59,7 +61,7 @@ public CompletableFuture start(Observer observer) { @Override public void cancel() { - this.isCanceled = true; + isCanceled.incrementAndGet(); } public void closeStream(Status status) { @@ -76,7 +78,7 @@ public void responseInit(String sessionId) { observer.onNext(msg); } - public void responseStartPartition(String topicPath, long partitionID) { + public void responseStartPartition(String topicPath, long partitionID, long committedOffset) { FromServer msg = FromServer.newBuilder() .setStatus(StatusCodesProtos.StatusIds.StatusCode.SUCCESS) .setStartPartitionSessionRequest(YdbTopic.StreamReadMessage.StartPartitionSessionRequest.newBuilder() @@ -85,6 +87,7 @@ public void responseStartPartition(String topicPath, long partitionID) { .setPartitionId(partitionID) .setPartitionSessionId(partCounter.incrementAndGet()) .build()) + .setCommittedOffset(committedOffset) .build()) .build(); observer.onNext(msg); @@ -101,6 +104,10 @@ public void responseStopPartition(long psid, boolean graceful) { observer.onNext(msg); } + public CommitAckResponse responseCommitAck() { + return new CommitAckResponse(); + } + public DataResponse responseData(long bytesSize) { return new DataResponse(bytesSize); } @@ -110,11 +117,11 @@ public void assertSentMessagesCount(int expectedCount) { } public void assertIsClosed() { - Assert.assertTrue("Read stream is closed", isClosed); + Assert.assertEquals("Read stream is closed", 1, isClosed.get()); } public void assertIsCancelled() { - Assert.assertTrue("Read stream is cancelled", isCanceled); + Assert.assertEquals("Read stream is cancelled", 1, isCanceled.get()); } public MessageAssert assertLastMessage() { @@ -158,7 +165,7 @@ public Partition batch(int codec, byte[]... messages) { batch.addMessageData(YdbTopic.StreamReadMessage.ReadResponse.MessageData.newBuilder() .setUncompressedSize(msg.length) .setData(encode(codec, msg)) - .setOffset(offset.incrementAndGet()) + .setOffset(offset.getAndIncrement()) .build()); } part.addBatches(batch.build()); @@ -172,6 +179,32 @@ public DataResponse and() { } } + public class CommitAckResponse { + private final YdbTopic.StreamReadMessage.CommitOffsetResponse.Builder resp; + + public CommitAckResponse() { + this.resp = YdbTopic.StreamReadMessage.CommitOffsetResponse.newBuilder(); + } + + public CommitAckResponse partition(long psid, long offset) { + resp.addPartitionsCommittedOffsets( + YdbTopic.StreamReadMessage.CommitOffsetResponse.PartitionCommittedOffset.newBuilder() + .setPartitionSessionId(psid) + .setCommittedOffset(offset) + .build() + ); + return this; + } + + public void send() { + FromServer msg = FromServer.newBuilder() + .setStatus(StatusCodesProtos.StatusIds.StatusCode.SUCCESS) + .setCommitOffsetResponse(resp.build()) + .build(); + observer.onNext(msg); + } + } + public static class MessageAssert { private final FromClient msg; @@ -212,6 +245,41 @@ public MessageAssert isStopPartition(long psid) { Assert.assertEquals("Stop partition has incorrect id", psid, resp.getPartitionSessionId()); return this; } + + public CommitAssert isCommit(long count) { + Assert.assertTrue("Msg is not commit offset request", msg.hasCommitOffsetRequest()); + YdbTopic.StreamReadMessage.CommitOffsetRequest resp = msg.getCommitOffsetRequest(); + Assert.assertEquals("Commit offset request has incorrect size", count, resp.getCommitOffsetsCount()); + return new CommitAssert(resp); + } + + public class CommitAssert { + private final YdbTopic.StreamReadMessage.CommitOffsetRequest resp; + + public CommitAssert(YdbTopic.StreamReadMessage.CommitOffsetRequest resp) { + this.resp = resp; + } + + public CommitAssert hasPartitionOffset(long psid, OffsetsRange... expected) { + Assert.assertEquals("Commit offset request has no partition " + psid, 1, + resp.getCommitOffsetsList().stream().filter(co -> co.getPartitionSessionId() == psid).count()); + + List committed = resp.getCommitOffsetsList().stream() + .filter(co -> co.getPartitionSessionId() == psid).findFirst().get().getOffsetsList(); + + Assert.assertEquals("Commit offset request for " + psid + " has incorrect count", + expected.length, committed.size()); + + for (int idx = 0; idx < expected.length; idx++) { + YdbTopic.OffsetsRange c = committed.get(idx); + OffsetsRange range = OffsetsRange.of(c.getStart(), c.getEnd()); + Assert.assertEquals("Incorrect commit offset for " + psid, expected[idx], range); + } + + return this; + } + + } } private static ByteString encode(int code, byte[] data) { 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 4eca94bf9..d8ad7e36b 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 @@ -1,19 +1,29 @@ package tech.ydb.topic.read.impl; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import org.junit.Assert; +import org.junit.Rule; import org.junit.Test; import org.mockito.Mockito; +import tech.ydb.common.retry.RetryConfig; import tech.ydb.core.Status; +import tech.ydb.core.StatusCode; import tech.ydb.topic.TopicRpc; import tech.ydb.topic.description.Codec; import tech.ydb.topic.description.CodecRegistry; +import tech.ydb.topic.description.OffsetsRange; +import tech.ydb.topic.read.DeferredCommitter; +import tech.ydb.topic.read.Message; import tech.ydb.topic.read.SyncReader; import tech.ydb.topic.settings.ReaderSettings; import tech.ydb.topic.settings.TopicReadSettings; +import tech.ydb.topic.utils.ErrorsHandler; +import tech.ydb.topic.utils.HideLoggers; +import tech.ydb.topic.utils.HideLoggersRule; /** * @@ -29,6 +39,9 @@ public class SyncReaderImplTest { (byte) 0x89, (byte) 0xAB, (byte) 0xCD,(byte) 0xEF }; private static final byte[] MSG5 = "utf8 encoded message".getBytes(); + @Rule + public final HideLoggersRule hideLogger = new HideLoggersRule(); + private static TopicRpc mockRpc(ReadStreamMock first, ReadStreamMock... rest) { TopicRpc rpc = Mockito.mock(TopicRpc.class); Mockito.when(rpc.getScheduler()).thenReturn(Mockito.mock(ScheduledExecutorService.class)); @@ -69,12 +82,37 @@ public void initAndShutdownTest() throws InterruptedException { reader.shutdown(); mock.assertIsClosed(); + reader.shutdown(); // double shutdow is allowed + Exception ex = Assert.assertThrows(RuntimeException.class, () -> reader.receive(0, TimeUnit.MILLISECONDS)); Assert.assertEquals("Reader was stopped", ex.getMessage()); mock.closeStream(Status.SUCCESS); } + @Test + public void shutdownBeforeInitTest() throws InterruptedException { + ReadStreamMock mock = new ReadStreamMock(); + + ReaderSettings settings = ReaderSettings.newBuilder() + .addTopic(TopicReadSettings.newBuilder().setPath("/test-topic").build()) + .setConsumerName("consumer") + .build(); + + SyncReader reader = new SyncReaderImpl(mockRpc(mock), settings, REGISTRY); + + reader.init(); + mock.assertSentMessagesCount(1); + mock.assertLastMessage().isInitRequest("consumer", "/test-topic"); + + reader.shutdown(); // shutdown before successful init + mock.assertIsClosed(); + mock.closeStream(Status.SUCCESS); + + Exception ex = Assert.assertThrows(RuntimeException.class, () -> reader.receive(0, TimeUnit.MILLISECONDS)); + Assert.assertEquals("Reader was stopped", ex.getMessage()); + } + @Test public void readClosedPartitionTest() throws InterruptedException { ReadStreamMock mock = new ReadStreamMock(); @@ -93,12 +131,12 @@ public void readClosedPartitionTest() throws InterruptedException { mock.assertLastMessage().isReadRequest(200000); // partition start is auto confirmed - mock.responseStartPartition("/test-topic", 123); + mock.responseStartPartition("/test-topic", 123, 0); mock.assertSentMessagesCount(3); mock.assertLastMessage().isStartPartition(1); // partition session id != partition id // start second partition - mock.responseStartPartition("/test-topic", 345); + mock.responseStartPartition("/test-topic", 345, 0); mock.assertSentMessagesCount(4); mock.assertLastMessage().isStartPartition(2); @@ -124,4 +162,153 @@ public void readClosedPartitionTest() throws InterruptedException { mock.assertSentMessagesCount(6); mock.assertLastMessage().isReadRequest(20000); } + + @Test + @HideLoggers({ BufferManager.class, ReaderImpl.class }) + public void invalidBatchesTest() throws InterruptedException { + ReadStreamMock mock = new ReadStreamMock(); + + ReaderSettings settings = ReaderSettings.newBuilder() + .addTopic(TopicReadSettings.newBuilder().setPath("/test-topic").build()) + .setMaxMemoryUsageBytes(2000) + .setConsumerName("consumer") + .build(); + + SyncReader reader = new SyncReaderImpl(mockRpc(mock), settings, REGISTRY); + reader.init(); + + mock.responseInit("read-session-1"); + mock.assertSentMessagesCount(2); + mock.assertLastMessage().isReadRequest(2000); + + // partition start is auto confirmed + mock.responseStartPartition("/test-topic", 123, 0); + mock.assertSentMessagesCount(3); + mock.assertLastMessage().isStartPartition(1); // partition session id != partition id + + // batch without messages is just auto released + mock.responseData(1000).partition(1, 0).batch(Codec.RAW).and().send(); + mock.assertSentMessagesCount(4); + mock.assertLastMessage().isReadRequest(1000); + + reader.shutdown(); + // batch after shutdown is just skipped + mock.responseData(1200).partition(1, 1000).batch(Codec.RAW, MSG1, MSG2, MSG3, MSG4, MSG5).and().send(); + mock.assertSentMessagesCount(5); + mock.assertLastMessage().isReadRequest(1200); + + mock.closeStream(Status.of(StatusCode.INTERNAL_ERROR)); + } + + @Test + public void retrySkipsReadMessagesTest() throws InterruptedException { + ErrorsHandler errorsHandler = new ErrorsHandler(); + // Policy: immediate retry (0ms) on all attempts, then no more + RetryConfig config = status -> (retryCount, elapsed) -> (status.getCode() != StatusCode.BAD_REQUEST) ? 0 : -1; + + ReadStreamMock m1 = new ReadStreamMock(); + ReadStreamMock m2 = new ReadStreamMock(); + + ReaderSettings settings = ReaderSettings.newBuilder() + .addTopic(TopicReadSettings.newBuilder().setPath("/test-topic").build()) + .setMaxMemoryUsageBytes(2000) + .setConsumerName("consumer") + .setRetryConfig(config) + .setErrorsHandler(errorsHandler) + .build(); + + SyncReader reader = new SyncReaderImpl(mockRpc(m1, m2), settings, REGISTRY); + reader.init(); + + m1.assertSentMessagesCount(1); + m1.assertLastMessage().isInitRequest("consumer", "/test-topic"); + + m1.responseInit("read-session-1"); + m1.assertSentMessagesCount(2); + m1.assertLastMessage().isReadRequest(2000); + + // start partition read and send 4 messages + m1.responseStartPartition("/test-topic", 123, 0); + m1.assertSentMessagesCount(3); + m1.assertLastMessage().isStartPartition(1); // partition session id != partition id + m1.responseData(1000).partition(1, 0).batch(Codec.RAW, MSG1, MSG2, MSG3, MSG4).and().send(); + + // read 3 messages + Message msg1 = reader.receive(1, TimeUnit.SECONDS); + Assert.assertNotNull(msg1); + Assert.assertArrayEquals(MSG1, msg1.getData()); + Message msg2 = reader.receive(1, TimeUnit.SECONDS); + Assert.assertNotNull(msg2); + Assert.assertArrayEquals(MSG2, msg2.getData()); + Message msg3 = reader.receive(1, TimeUnit.SECONDS); + Assert.assertNotNull(msg3); + Assert.assertArrayEquals(MSG3, msg3.getData()); + + // commit 2 messages + CompletableFuture c1 = msg1.commit(); + CompletableFuture c2 = msg2.commit(); + + // reader doesn't merge commits + m1.assertSentMessagesCount(5); + m1.assertLastMessage().isCommit(1).hasPartitionOffset(1, OffsetsRange.of(1, 2)); + + Assert.assertFalse(c1.isDone()); + Assert.assertFalse(c2.isDone()); + + // comfirm commit for 1 message + m1.responseCommitAck().partition(1, 1).send(); + + Assert.assertTrue(c1.isDone()); + Assert.assertFalse(c2.isDone()); + + // get a stream error + errorsHandler.assertEmpty(); + m1.closeStream(Status.of(StatusCode.TRANSPORT_UNAVAILABLE)); + errorsHandler.assertCodes(StatusCode.TRANSPORT_UNAVAILABLE); + + // commit for 2 message is failed + Assert.assertTrue(c2.isCompletedExceptionally()); + + // init second stream + m2.assertSentMessagesCount(1); + m2.assertLastMessage().isInitRequest("consumer", "/test-topic"); + + m2.responseInit("read-session-2"); + m2.assertSentMessagesCount(2); + m2.assertLastMessage().isReadRequest(2000); + + // partition start is auto confirmed + m2.responseStartPartition("/test-topic", 123, 1); + m2.assertSentMessagesCount(3); + m2.assertLastMessage().isStartPartition(1); // partition session id != partition id + + // no messages in reader queue + Assert.assertNull(reader.receive(0, TimeUnit.SECONDS)); + + m2.responseData(1000).partition(1, 1).batch(Codec.RAW, MSG2, MSG3, MSG4, MSG5).and().send(); + + // commit for the lost stream is failed + Assert.assertTrue(msg3.commit().isCompletedExceptionally()); + + // read 3 messages, message from the lost message will be deleted + msg2 = reader.receive(1, TimeUnit.SECONDS); + Assert.assertNotNull(msg2); + Assert.assertArrayEquals(MSG2, msg2.getData()); + msg3 = reader.receive(1, TimeUnit.SECONDS); + Assert.assertNotNull(msg3); + Assert.assertArrayEquals(MSG3, msg3.getData()); + Message msg4 = reader.receive(1, TimeUnit.SECONDS); + Assert.assertNotNull(msg4); + Assert.assertArrayEquals(MSG4, msg4.getData()); + + DeferredCommitter committer = DeferredCommitter.newInstance(); + committer.add(msg2); + committer.add(msg4); + committer.commit(); + + m2.assertLastMessage().isCommit(1).hasPartitionOffset(1, OffsetsRange.of(1, 2), OffsetsRange.of(3, 4)); + reader.shutdown(); + + m2.assertIsClosed(); + } } diff --git a/topic/src/test/java/tech/ydb/topic/utils/ErrorsHandler.java b/topic/src/test/java/tech/ydb/topic/utils/ErrorsHandler.java new file mode 100644 index 000000000..bb137bd43 --- /dev/null +++ b/topic/src/test/java/tech/ydb/topic/utils/ErrorsHandler.java @@ -0,0 +1,42 @@ +package tech.ydb.topic.utils; + +import java.util.ArrayList; +import java.util.Iterator; +import java.util.List; +import java.util.function.BiConsumer; + +import org.junit.Assert; + +import tech.ydb.core.Status; +import tech.ydb.core.StatusCode; + +/** + * + * @author Aleksandr Gorshenin {@literal } + */ +public class ErrorsHandler implements BiConsumer { + private final List problems = new ArrayList<>(); + + @Override + public void accept(Status st, Throwable th) { + if (st != null) { + problems.add(st.getCode()); + } + if (th != null) { + problems.add(StatusCode.CLIENT_INTERNAL_ERROR); + } + } + + public void assertEmpty() { + Assert.assertTrue("No reties was expected", problems.isEmpty()); + } + + public void assertCodes(StatusCode... codes) { + Iterator it = problems.iterator(); + for (StatusCode code: codes) { + Assert.assertTrue("Expected " + code + ", but has nothing", it.hasNext()); + Assert.assertEquals(code, it.next()); + } + Assert.assertFalse("Unexpected error code", it.hasNext()); + } +} From a63763f98262e2404f727440f635acf4aee7d693 Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Wed, 16 Sep 2026 16:24:44 +0100 Subject: [PATCH 07/12] Fixed problems and updated tests --- .../ydb/topic/read/impl/AsyncReaderImpl.java | 47 ++-- .../ydb/topic/read/impl/LazyExecutor.java | 3 +- .../tech/ydb/topic/read/impl/ReaderImpl.java | 3 +- .../ydb/topic/read/impl/SyncReaderImpl.java | 10 +- .../topic/TopicReadersIntegrationTest.java | 52 +++++ .../ydb/topic/YdbTopicsIntegrationTest.java | 38 ++-- .../ydb/topic/read/impl/ReadStreamMock.java | 5 + .../ydb/topic/read/impl/ReaderImplTest.java | 202 ++++++++++++++++++ .../topic/read/impl/SyncReaderImplTest.java | 7 +- 9 files changed, 317 insertions(+), 50 deletions(-) create mode 100644 topic/src/test/java/tech/ydb/topic/read/impl/ReaderImplTest.java 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 76683f2f7..b24d275d2 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 @@ -85,30 +85,19 @@ public CompletableFuture updateOffsetsInTransaction(YdbTransaction trans return impl.updateOffsetsInTransaction(transaction, offsets, settings); } - protected CompletableFuture handleReaderClosed() { - return CompletableFuture.runAsync(() -> { - try { - eventHandler.onReaderClosed(new ReaderClosedEvent()); - } catch (Throwable th) { - failSession(th, "onReaderClosed"); - throw th; - } - }, controlEventsExecutor); - } - @Override public CompletableFuture shutdown() { if (!impl.close()) { // implicit closing because stream will never call onClose - close(); + close(Status.SUCCESS.withIssues(Issue.of("Closed by client", Issue.Severity.INFO))); } return shutdownFuture; } - private void close() { + private void close(Status status) { decompressor.close(); processor.close(); - initFuture.complete(null); + initFuture.completeExceptionally(new RuntimeException("Reader closed with " + status)); shutdownFuture.complete(null); } @@ -122,22 +111,28 @@ private class AsyncHandler implements ReaderImpl.Handler { @Override public void handleSessionStarted(String sessionId) { initFuture.complete(null); - try { - eventHandler.onSessionStarted(new SessionStartedEvent(sessionId)); - } catch (Throwable th) { - failSession(th, "onSessionStarted"); - } + processor.execute(() -> { + try { + eventHandler.onSessionStarted(new SessionStartedEvent(sessionId)); + } catch (Throwable th) { + failSession(th, "onSessionStarted"); + throw th; + } + }); } @Override public void handleReaderClosed(Status status) { - try { - eventHandler.onReaderClosed(new ReaderClosedEvent()); - } catch (Throwable th) { - failSession(th, "onReaderClosed"); - } finally { - close(); - } + processor.execute(() -> { + try { + eventHandler.onReaderClosed(new ReaderClosedEvent()); + } catch (Throwable th) { + failSession(th, "onSessionStarted"); + throw th; + } + }); + + close(status); // wait while processer finished all tasks } @Override 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 index 406d981e6..c5d575d26 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/LazyExecutor.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/LazyExecutor.java @@ -71,7 +71,8 @@ private void shutdown(ExecutorService service) { try { service.shutdown(); - if (!service.awaitTermination(100, TimeUnit.MILLISECONDS)) { + if (!service.awaitTermination(5, TimeUnit.SECONDS)) { + logger.warn("executor {} shutdown timeout exceeded, interrupt all tasks", name); service.shutdownNow(); } } catch (InterruptedException e) { 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 c2347126a..99fb68f07 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 @@ -98,13 +98,12 @@ protected void onClose(ReadSession stream, Status status) { } else { logger.info("[{}] closed by status {}", debugId, status); } - if (errorHandler != null) { + if (errorHandler != null && !status.isSuccess()) { try { errorHandler.accept(status, null); } catch (RuntimeException ex) { logger.error("[{}] errorHandler onClose processing throws exception", debugId, ex); } - errorHandler.accept(status, null); } stream.closeAll().forEach(ps -> handler.handleClosePartitionSession(ps)); handler.handleReaderClosed(status); diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java b/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java index 47f0ce2c3..3db79f532 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java @@ -18,6 +18,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import tech.ydb.core.Issue; import tech.ydb.core.Status; import tech.ydb.topic.TopicRpc; import tech.ydb.topic.description.CodecRegistry; @@ -94,14 +95,15 @@ public void initAndWait() { public void shutdown() { if (!impl.close()) { // implicit closing because stream will never call onClose - close(); + // implicit closing because stream will never call onClose + close(Status.SUCCESS.withIssues(Issue.of("Closed by client", Issue.Severity.INFO))); } shutdownFuture.join(); } - private void close() { - initFuture.complete(null); + private void close(Status status) { + initFuture.completeExceptionally(new RuntimeException("Reader closed with " + status)); shutdownFuture.complete(null); decompressor.close(); @@ -210,7 +212,7 @@ public void handleSessionStarted(String sessionId) { @Override public void handleReaderClosed(Status status) { - close(); + close(status); } @Override diff --git a/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java b/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java index 254283984..607cbeb7e 100644 --- a/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java +++ b/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java @@ -13,6 +13,7 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; import org.junit.AfterClass; import org.junit.Assert; @@ -25,7 +26,11 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import tech.ydb.common.transaction.TxMode; import tech.ydb.core.Status; +import tech.ydb.table.SessionRetryContext; +import tech.ydb.table.TableClient; +import tech.ydb.table.transaction.TableTransaction; import tech.ydb.test.junit4.GrpcTransportRule; import tech.ydb.topic.description.Consumer; import tech.ydb.topic.description.ConsumerDescription; @@ -48,6 +53,7 @@ import tech.ydb.topic.settings.ReaderSettings; import tech.ydb.topic.settings.StartPartitionSessionSettings; import tech.ydb.topic.settings.TopicReadSettings; +import tech.ydb.topic.settings.UpdateOffsetsInTransactionSettings; import tech.ydb.topic.settings.WriterSettings; import tech.ydb.topic.utils.HideLoggers; import tech.ydb.topic.utils.HideLoggersRule; @@ -583,4 +589,50 @@ public void directDecompressorTest() throws InterruptedException { reader.shutdown().join(); } } + + @Test + public void readAllInTxTest() throws InterruptedException { + ReaderSettings readerSettings = ReaderSettings.newBuilder() + .addTopic(TopicReadSettings.newBuilder().setPath(TEST_TOPIC).build()) + .setConsumerName(TEST_CONSUMER1) + .build(); + + AtomicLong[] offsets = new AtomicLong[] { new AtomicLong(), new AtomicLong(), new AtomicLong() }; + CountDownLatch read = new CountDownLatch(3600); + + try (TableClient tableClient = TableClient.newClient(ydbTransport).build()) { + SessionRetryContext retryCtx = SessionRetryContext.create(tableClient).idempotent(true).build(); + UpdateOffsetsInTransactionSettings settings = UpdateOffsetsInTransactionSettings.newBuilder().build(); + + AtomicReference ref = new AtomicReference<>(); + @SuppressWarnings("deprecation") + AsyncReader reader = client.createAsyncReader(readerSettings, ReadEventHandlersSettings.newBuilder() + .setEventHandler((DataReceivedEvent event) -> { + AtomicLong offset = offsets[(int) event.getPartitionSession().getPartitionId()]; + for (Message msg : event.getMessages()) { + Assert.assertEquals(offset.getAndIncrement(), msg.getOffset()); + } + + retryCtx.supplyStatus(session -> { + TableTransaction tx = session.beginTransaction(TxMode.SERIALIZABLE_RW).join().getValue(); + ref.get().updateOffsetsInTransaction(tx, event.getPartitionOffsets(), settings).join(); + return tx.commit(); + }).join().expectSuccess(); + + event.getMessages().forEach(msg -> read.countDown()); + }).build()); + + ref.set(reader); + reader.init().join(); + try { + Assert.assertTrue(read.await(30, TimeUnit.SECONDS)); + Assert.assertEquals(1000, offsets[0].get()); + Assert.assertEquals(500, offsets[1].get()); + Assert.assertEquals(2100, offsets[2].get()); + } finally { + reader.shutdown().join(); + } + } + + } } diff --git a/topic/src/test/java/tech/ydb/topic/YdbTopicsIntegrationTest.java b/topic/src/test/java/tech/ydb/topic/YdbTopicsIntegrationTest.java index 7d48cb693..ba71da853 100644 --- a/topic/src/test/java/tech/ydb/topic/YdbTopicsIntegrationTest.java +++ b/topic/src/test/java/tech/ydb/topic/YdbTopicsIntegrationTest.java @@ -62,6 +62,7 @@ public class YdbTopicsIntegrationTest { private final static String TEST_TOPIC = "integration_test_topic"; private final static String TEST_OTHER_TOPIC = "integration_test_other_topic"; + private final static String TEST_TMP_TOPIC = "integration_test_tmp_topic"; private final static String TEST_CONSUMER1 = "consumer"; private final static String TEST_CONSUMER2 = "other_consumer"; @@ -86,14 +87,21 @@ public static void initTopic() { .addConsumer(Consumer.newBuilder().setName(TEST_CONSUMER2).build()) .build() ).join().expectSuccess("can't create a new topic"); + + client.createTopic(TEST_OTHER_TOPIC, CreateTopicSettings.newBuilder() + .addConsumer(Consumer.newBuilder().setName(TEST_CONSUMER1).build()) + .addConsumer(Consumer.newBuilder().setName(TEST_CONSUMER2).build()) + .build() + ).join().expectSuccess("can't create a new topic"); } @AfterClass public static void dropTopic() { logger.info("Drop test topic {} ...", TEST_TOPIC); - Status dropStatus = client.dropTopic(TEST_TOPIC).join(); + client.dropTopic(TEST_TOPIC).join(); + logger.info("Drop test topic {} ...", TEST_OTHER_TOPIC); + client.dropTopic(TEST_OTHER_TOPIC).join(); client.close(); - dropStatus.expectSuccess("can't drop test topic"); } @Test @@ -246,7 +254,7 @@ public void onMessages(DataReceivedEvent dre) { @Test public void describeTopic() { - TopicDescription description = client.describeTopic(TEST_TOPIC).join().getValue(); + TopicDescription description = client.describeTopic(TEST_OTHER_TOPIC).join().getValue(); Assert.assertNull(description.getTopicStats()); List consumers = description.getConsumers(); @@ -258,7 +266,7 @@ public void describeTopic() { @Test public void alterTopicWithAutoPartitioning() { - client.alterTopic(TEST_TOPIC, AlterTopicSettings.newBuilder() + client.alterTopic(TEST_OTHER_TOPIC, AlterTopicSettings.newBuilder() .setAlterPartitioningSettings(AlterPartitioningSettings.newBuilder() .setAutoPartitioningStrategy(AutoPartitioningStrategy.SCALE_UP) .setMaxActivePartitions(10) @@ -270,7 +278,7 @@ public void alterTopicWithAutoPartitioning() { .build()) .build()).join().expectSuccess("can't alter the topic"); - TopicDescription description = client.describeTopic(TEST_TOPIC).join().getValue(); + TopicDescription description = client.describeTopic(TEST_OTHER_TOPIC).join().getValue(); PartitioningSettings actualPartitioningSettings = description.getPartitioningSettings(); PartitioningSettings expectedPartitioningSettings = PartitioningSettings.newBuilder() @@ -300,15 +308,17 @@ public void createTopicWithAutoPartitioning() { .build()) .build(); - CompletableFuture secondaryTopicCreated = client.createTopic(TEST_OTHER_TOPIC, + CompletableFuture secondaryTopicCreated = client.createTopic(TEST_TMP_TOPIC, CreateTopicSettings.newBuilder().setPartitioningSettings(expectedPartitioningSettings).build() ); secondaryTopicCreated.join().expectSuccess("can't create the topic"); - TopicDescription description = client.describeTopic(TEST_OTHER_TOPIC).join().getValue(); + TopicDescription description = client.describeTopic(TEST_TMP_TOPIC).join().getValue(); Assert.assertEquals(expectedPartitioningSettings, description.getPartitioningSettings()); + + client.dropTopic(TEST_TMP_TOPIC).join().expectSuccess("can't drop the test topic"); } @Test @@ -316,8 +326,8 @@ public void describeTopicStats() { DescribeTopicSettings on = DescribeTopicSettings.newBuilder().withIncludeStats(true).build(); DescribeTopicSettings off = DescribeTopicSettings.newBuilder().withIncludeStats(false).build(); - TopicDescription withStats = client.describeTopic(TEST_TOPIC, on).join().getValue(); - TopicDescription withoutStats = client.describeTopic(TEST_TOPIC, off).join().getValue(); + TopicDescription withStats = client.describeTopic(TEST_OTHER_TOPIC, on).join().getValue(); + TopicDescription withoutStats = client.describeTopic(TEST_OTHER_TOPIC, off).join().getValue(); Assert.assertNull(withoutStats.getTopicStats()); Assert.assertNotNull(withStats.getTopicStats()); @@ -348,7 +358,7 @@ public void invalidAddConsumerTest() { .build() ).build(); - Status status = client.alterTopic(TEST_TOPIC, settings).join(); + Status status = client.alterTopic(TEST_OTHER_TOPIC, settings).join(); Assert.assertFalse("Alter must fail, but get status " + status, status.isSuccess()); Assert.assertEquals("Alter must fail, but get status " + status, StatusCode.BAD_REQUEST, status.getCode()); } @@ -364,7 +374,7 @@ public void invalidAlterConsumerTest() { .build() ).build(); - Status status = client.alterTopic(TEST_TOPIC, settings).join(); + Status status = client.alterTopic(TEST_OTHER_TOPIC, settings).join(); Assert.assertFalse("Alter must fail, but get status " + status, status.isSuccess()); Assert.assertEquals("Alter must fail, but get status " + status, StatusCode.BAD_REQUEST, status.getCode()); } @@ -379,16 +389,16 @@ public void alterConsumerTest() { .build() ).build(); - Status status = client.alterTopic(TEST_TOPIC, settings).join(); + Status status = client.alterTopic(TEST_OTHER_TOPIC, settings).join(); Assert.assertTrue("Alter must be OK, but got status " + status, status.isSuccess()); - ConsumerDescription description = client.describeConsumer(TEST_TOPIC, TEST_CONSUMER2).join().getValue(); + ConsumerDescription description = client.describeConsumer(TEST_OTHER_TOPIC, TEST_CONSUMER2).join().getValue(); Assert.assertEquals(TEST_CONSUMER2, description.getConsumer().getName()); Assert.assertEquals(Instant.EPOCH.plusSeconds(10), description.getConsumer().getReadFrom()); Assert.assertEquals(Duration.ofMinutes(5), description.getConsumer().getAvailabilityPeriod()); - TopicDescription topicDesc = client.describeTopic(TEST_TOPIC).join().getValue(); + TopicDescription topicDesc = client.describeTopic(TEST_OTHER_TOPIC).join().getValue(); Assert.assertNull(topicDesc.getTopicStats()); 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 e4496e717..1c0228438 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 @@ -116,6 +116,11 @@ public void assertSentMessagesCount(int expectedCount) { Assert.assertEquals("Read stream sent messages count", expectedCount, messages.size()); } + public void assertIsActive() { + Assert.assertEquals("Read stream is active", 0, isClosed.get()); + Assert.assertEquals("Read stream is cancelled", 0, isCanceled.get()); + } + public void assertIsClosed() { Assert.assertEquals("Read stream is closed", 1, isClosed.get()); } 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 new file mode 100644 index 000000000..dafbc27e9 --- /dev/null +++ b/topic/src/test/java/tech/ydb/topic/read/impl/ReaderImplTest.java @@ -0,0 +1,202 @@ +package tech.ydb.topic.read.impl; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ScheduledExecutorService; + +import org.junit.Assert; +import org.junit.Test; +import org.junit.function.ThrowingRunnable; +import org.mockito.Mockito; + +import tech.ydb.common.transaction.TxMode; +import tech.ydb.common.transaction.YdbTransaction; +import tech.ydb.core.Status; +import tech.ydb.core.StatusCode; +import tech.ydb.topic.TopicRpc; +import tech.ydb.topic.description.CodecRegistry; +import tech.ydb.topic.description.OffsetsRange; +import tech.ydb.topic.read.PartitionOffsets; +import tech.ydb.topic.read.PartitionSession; +import tech.ydb.topic.settings.ReaderSettings; +import tech.ydb.topic.settings.TopicReadSettings; +import tech.ydb.topic.settings.TopicRetryConfig; +import tech.ydb.topic.settings.UpdateOffsetsInTransactionSettings; + +public class ReaderImplTest { + private static final CodecRegistry REGISTRY = new CodecRegistry(); + + private static final String TOPIC = "/test/topic"; + + private static TopicRpc mockRpc(ReadStreamMock first, ReadStreamMock... rest) { + TopicRpc rpc = Mockito.mock(TopicRpc.class); + Mockito.when(rpc.getScheduler()).thenReturn(Mockito.mock(ScheduledExecutorService.class)); + Mockito.when(rpc.readSession(Mockito.any(String.class))).thenReturn(first, rest); + Mockito.when(rpc.updateOffsetsInTransaction(Mockito.any(), Mockito.any())) + .thenReturn(CompletableFuture.completedFuture(Status.SUCCESS)); + return rpc; + } + + private static void assertIllegalArgument(String msg, ThrowingRunnable runnable) { + IllegalArgumentException ex = Assert.assertThrows("Must be thrown IllegalArgumentException", + IllegalArgumentException.class, runnable); + Assert.assertEquals(msg, ex.getMessage()); + } + + @Test + public void updateOffsetsInTxValidationTest() { + ReadStreamMock mock = new ReadStreamMock(); + + ReaderSettings settings = ReaderSettings.newBuilder() + .addTopic(TopicReadSettings.newBuilder().setPath(TOPIC).build()) + .setConsumerName("consumer") + .setRetryConfig(TopicRetryConfig.NEVER) + .setMaxMemoryUsageBytes(1000) + .build(); + + ReadConfig config = new ReadConfig(REGISTRY, Runnable::run, Runnable::run, settings); + ReaderImpl.Handler handler = Mockito.mock(ReaderImpl.Handler.class); + + ReaderImpl reader = new ReaderImpl(mockRpc(mock), "test-reader", settings, config, handler); + reader.start(); + + UpdateOffsetsInTransactionSettings updateSettings = UpdateOffsetsInTransactionSettings.newBuilder() + .withTraceId("test-trace").build(); + + TxMock finished = new TxMock(CompletableFuture.completedFuture(Status.SUCCESS)); + assertIllegalArgument("Transaction is not active. " + + "Can only read topic messages in already running transactions from other services", + () -> reader.updateOffsetsInTransaction(finished, Collections.emptyMap(), updateSettings) + ); + + CompletableFuture txStatus = new CompletableFuture<>(); + TxMock active = new TxMock(txStatus); + assertIllegalArgument("Empty topic list to update in transaction", + () -> reader.updateOffsetsInTransaction(active, Collections.emptyMap(), updateSettings) + ); + + assertIllegalArgument("Empty offsets range to update in transaction", + () -> reader.updateOffsetsInTransaction( + active, Collections.singletonMap(TOPIC, new ArrayList()), updateSettings + ) + ); + + List offsets = Arrays.asList( + new PartitionOffsets(new PartitionSession(1, 1, TOPIC), Arrays.asList(OffsetsRange.of(0, 10))), + new PartitionOffsets(new PartitionSession(2, 2, TOPIC), + Arrays.asList(OffsetsRange.of(0, 1), OffsetsRange.of(2, 3))) + ); + reader.updateOffsetsInTransaction(active, Collections.singletonMap(TOPIC, offsets), updateSettings); + + + txStatus.complete(Status.SUCCESS); + mock.assertIsActive(); + reader.close(); + + mock.assertIsClosed(); + } + + @Test + public void updateOffsetsInTxFailTest() { + ReadStreamMock mock = new ReadStreamMock(); + + ReaderSettings settings = ReaderSettings.newBuilder() + .addTopic(TopicReadSettings.newBuilder().setPath(TOPIC).build()) + .setConsumerName("consumer") + .setRetryConfig(TopicRetryConfig.NEVER) + .setMaxMemoryUsageBytes(1000) + .build(); + + ReadConfig config = new ReadConfig(REGISTRY, Runnable::run, Runnable::run, settings); + ReaderImpl.Handler handler = Mockito.mock(ReaderImpl.Handler.class); + + ReaderImpl reader = new ReaderImpl(mockRpc(mock), "test-reader", settings, config, handler); + reader.start(); + + UpdateOffsetsInTransactionSettings updateSettings = UpdateOffsetsInTransactionSettings.newBuilder().build(); + + CompletableFuture txStatus = new CompletableFuture<>(); + TxMock active = new TxMock(txStatus); + List offsets = Arrays.asList( + new PartitionOffsets(new PartitionSession(1, 1, TOPIC), Arrays.asList(OffsetsRange.of(0, 10))), + new PartitionOffsets(new PartitionSession(2, 2, TOPIC), + Arrays.asList(OffsetsRange.of(0, 1), OffsetsRange.of(2, 3))) + ); + reader.updateOffsetsInTransaction(active, Collections.singletonMap(TOPIC, offsets), updateSettings); + + txStatus.complete(Status.of(StatusCode.ABORTED)); + mock.assertIsClosed(); + reader.close(); + mock.assertIsClosed(); + } + + @Test + public void updateOffsetsInTxErrorTest() { + ReadStreamMock mock = new ReadStreamMock(); + + ReaderSettings settings = ReaderSettings.newBuilder() + .addTopic(TopicReadSettings.newBuilder().setPath(TOPIC).build()) + .setConsumerName("consumer") + .setRetryConfig(TopicRetryConfig.NEVER) + .setMaxMemoryUsageBytes(1000) + .build(); + + ReadConfig config = new ReadConfig(REGISTRY, Runnable::run, Runnable::run, settings); + ReaderImpl.Handler handler = Mockito.mock(ReaderImpl.Handler.class); + + ReaderImpl reader = new ReaderImpl(mockRpc(mock), "test-reader", settings, config, handler); + reader.start(); + + UpdateOffsetsInTransactionSettings updateSettings = UpdateOffsetsInTransactionSettings.newBuilder().build(); + + CompletableFuture txStatus = new CompletableFuture<>(); + TxMock active = new TxMock(txStatus); + List offsets = Arrays.asList( + new PartitionOffsets(new PartitionSession(1, 1, TOPIC), Arrays.asList(OffsetsRange.of(0, 10))), + new PartitionOffsets(new PartitionSession(2, 2, TOPIC), + Arrays.asList(OffsetsRange.of(0, 1), OffsetsRange.of(2, 3))) + ); + reader.updateOffsetsInTransaction(active, Collections.singletonMap(TOPIC, offsets), updateSettings); + + txStatus.completeExceptionally(new RuntimeException("tx problem")); + mock.assertIsClosed(); + reader.close(); + mock.assertIsClosed(); + } + + private class TxMock implements YdbTransaction { + private final CompletableFuture status; + + public TxMock(CompletableFuture status) { + this.status = status; + } + + @Override + public boolean isActive() { + return !status.isDone(); + } + + @Override + public String getId() { + return "tx-id"; + } + + @Override + public TxMode getTxMode() { + return TxMode.NONE; + } + + @Override + public String getSessionId() { + return "session-tx-id"; + } + + @Override + public CompletableFuture getStatusFuture() { + return status; + } + } +} 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 d8ad7e36b..ec07402e9 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 @@ -31,6 +31,8 @@ */ public class SyncReaderImplTest { private static final CodecRegistry REGISTRY = new CodecRegistry(); + private static final RetryConfig IMMEDIATE_RETRY = status -> (number, elapsed) -> 0; + private static final byte[] MSG1 = new byte[] { 0x00 }; private static final byte[] MSG2 = new byte[] { }; private static final byte[] MSG3 = new byte[] { 0x01, 0x01, 0x01, 0x01, 0x02, 0x02, 0x02, 0x02 }; @@ -56,6 +58,7 @@ public void initAndShutdownTest() throws InterruptedException { ReaderSettings settings = ReaderSettings.newBuilder() .addTopic(TopicReadSettings.newBuilder().setPath("/test-topic").build()) .setConsumerName("consumer") + .setReaderName("test-reader-name") .build(); SyncReader reader = new SyncReaderImpl(mockRpc(mock), settings, REGISTRY); @@ -203,8 +206,6 @@ public void invalidBatchesTest() throws InterruptedException { @Test public void retrySkipsReadMessagesTest() throws InterruptedException { ErrorsHandler errorsHandler = new ErrorsHandler(); - // Policy: immediate retry (0ms) on all attempts, then no more - RetryConfig config = status -> (retryCount, elapsed) -> (status.getCode() != StatusCode.BAD_REQUEST) ? 0 : -1; ReadStreamMock m1 = new ReadStreamMock(); ReadStreamMock m2 = new ReadStreamMock(); @@ -213,7 +214,7 @@ public void retrySkipsReadMessagesTest() throws InterruptedException { .addTopic(TopicReadSettings.newBuilder().setPath("/test-topic").build()) .setMaxMemoryUsageBytes(2000) .setConsumerName("consumer") - .setRetryConfig(config) + .setRetryConfig(IMMEDIATE_RETRY) .setErrorsHandler(errorsHandler) .build(); From 088e983a3a3e9c3b166cfe8a14ca0a41452ebbf8 Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Thu, 17 Sep 2026 10:30:36 +0100 Subject: [PATCH 08/12] Update tests --- .../ydb/topic/read/impl/SyncReaderImpl.java | 21 +- .../ydb/topic/settings/ReaderSettings.java | 5 + .../topic/read/impl/AsyncReaderImplTest.java | 323 ++++++++++++++++++ .../ydb/topic/read/impl/ReadStreamMock.java | 46 ++- .../ydb/topic/read/impl/ReaderImplTest.java | 71 +++- .../topic/read/impl/SyncReaderImplTest.java | 32 +- .../topic/settings/ReaderSettingsTest.java | 51 +++ 7 files changed, 518 insertions(+), 31 deletions(-) create mode 100644 topic/src/test/java/tech/ydb/topic/read/impl/AsyncReaderImplTest.java create mode 100644 topic/src/test/java/tech/ydb/topic/settings/ReaderSettingsTest.java diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java b/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java index 3db79f532..ccf968ef1 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java @@ -6,6 +6,7 @@ import java.util.List; import java.util.Queue; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.TimeUnit; import java.util.concurrent.locks.Condition; @@ -49,7 +50,7 @@ public class SyncReaderImpl implements SyncReader { private final ReaderImpl impl; private final CompletableFuture initFuture = new CompletableFuture<>(); - private final CompletableFuture shutdownFuture = new CompletableFuture<>(); + private final CompletableFuture shutdownFuture = new CompletableFuture<>(); private final Queue queue = new ConcurrentLinkedQueue<>(); private final ReentrantLock waitingLock = new ReentrantLock(); @@ -87,14 +88,20 @@ public void init() { @Override public void initAndWait() { impl.start(); - initFuture.join(); + try { + initFuture.join(); + } catch (CompletionException ex) { + if (ex.getCause() instanceof RuntimeException) { + throw (RuntimeException) ex.getCause(); + } + throw ex; + } } @Override public void shutdown() { if (!impl.close()) { - // implicit closing because stream will never call onClose // implicit closing because stream will never call onClose close(Status.SUCCESS.withIssues(Issue.of("Closed by client", Issue.Severity.INFO))); } @@ -103,8 +110,8 @@ public void shutdown() { } private void close(Status status) { - initFuture.completeExceptionally(new RuntimeException("Reader closed with " + status)); - shutdownFuture.complete(null); + initFuture.completeExceptionally(new RuntimeException("Reader was closed with " + status)); + shutdownFuture.complete(status); decompressor.close(); wakeUp(); @@ -150,7 +157,7 @@ private MessageWrapper waitReadyMessage(long timeout, TimeUnit unit) throws Inte logger.trace("No messages in queue. Waiting for {} ms...", millisToWait); waitingCondition.await(millisToWait, TimeUnit.MILLISECONDS); if (impl.isClosed()) { - throw new RuntimeException("Reader was stopped"); + throw new RuntimeException("Reader was stopped with " + shutdownFuture.join()); } next = queue.poll(); } @@ -164,7 +171,7 @@ private MessageWrapper waitReadyMessage(long timeout, TimeUnit unit) throws Inte public Message receiveInternal(ReceiveSettings receiveSettings, long timeout, TimeUnit unit) throws InterruptedException { if (impl.isClosed()) { - throw new RuntimeException("Reader was stopped"); + throw new RuntimeException("Reader was stopped with " + shutdownFuture.join()); } while (true) { diff --git a/topic/src/main/java/tech/ydb/topic/settings/ReaderSettings.java b/topic/src/main/java/tech/ydb/topic/settings/ReaderSettings.java index 4b46d3535..eaf570291 100644 --- a/topic/src/main/java/tech/ydb/topic/settings/ReaderSettings.java +++ b/topic/src/main/java/tech/ydb/topic/settings/ReaderSettings.java @@ -147,6 +147,11 @@ public Builder addTopic(TopicReadSettings topic) { return this; } + public Builder addTopic(String topicPath) { + topics.add(TopicReadSettings.newBuilder().setPath(topicPath).build()); + return this; + } + public Builder setTopics(List topics) { this.topics = topics; return this; 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 new file mode 100644 index 000000000..e474f1d12 --- /dev/null +++ b/topic/src/test/java/tech/ydb/topic/read/impl/AsyncReaderImplTest.java @@ -0,0 +1,323 @@ +package tech.ydb.topic.read.impl; + +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Queue; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; +import java.util.concurrent.ScheduledExecutorService; + +import org.junit.Assert; +import org.junit.Rule; +import org.junit.Test; +import org.mockito.ArgumentCaptor; +import org.mockito.Mockito; + +import tech.ydb.common.retry.RetryConfig; +import tech.ydb.core.Status; +import tech.ydb.core.StatusCode; +import tech.ydb.topic.TopicRpc; +import tech.ydb.topic.description.Codec; +import tech.ydb.topic.description.CodecRegistry; +import tech.ydb.topic.description.OffsetsRange; +import tech.ydb.topic.read.AsyncReader; +import tech.ydb.topic.read.DeferredCommitter; +import tech.ydb.topic.read.Message; +import tech.ydb.topic.read.events.CommitOffsetAcknowledgementEvent; +import tech.ydb.topic.read.events.DataReceivedEvent; +import tech.ydb.topic.read.events.PartitionSessionClosedEvent; +import tech.ydb.topic.read.events.ReadEventHandler; +import tech.ydb.topic.read.events.ReaderClosedEvent; +import tech.ydb.topic.read.events.StartPartitionSessionEvent; +import tech.ydb.topic.read.events.StopPartitionSessionEvent; +import tech.ydb.topic.read.impl.events.SessionStartedEvent; +import tech.ydb.topic.settings.ReadEventHandlersSettings; +import tech.ydb.topic.settings.ReaderSettings; +import tech.ydb.topic.utils.ErrorsHandler; +import tech.ydb.topic.utils.HideLoggers; +import tech.ydb.topic.utils.HideLoggersRule; + +public class AsyncReaderImplTest { + private static final CodecRegistry REGISTRY = new CodecRegistry(); + private static final RetryConfig IMMEDIATE_RETRY = status -> (number, elapsed) -> 0; + + private static final byte[] MSG1 = new byte[] { 0x00 }; + private static final byte[] MSG2 = new byte[] { }; + private static final byte[] MSG3 = new byte[] { 0x01, 0x01, 0x01, 0x01, 0x02, 0x02, 0x02, 0x02 }; + private static final byte[] MSG4 = new byte[] { + 0x01, 0x23, 0x34, 0x45, 0x67, (byte) 0x89, (byte) 0xAB, (byte) 0xCD, (byte) 0xEF }; + private static final byte[] MSG5 = "utf8 encoded message".getBytes(); + + @Rule + public final HideLoggersRule hideLogger = new HideLoggersRule(); + + private final ReadEventHandler handler = Mockito.mock(ReadEventHandler.class, Mockito.CALLS_REAL_METHODS); + + private static TopicRpc mockRpc(ReadStreamMock first, ReadStreamMock... rest) { + TopicRpc rpc = Mockito.mock(TopicRpc.class); + Mockito.when(rpc.getScheduler()).thenReturn(Mockito.mock(ScheduledExecutorService.class)); + Mockito.when(rpc.readSession(Mockito.any(String.class))).thenReturn(first, rest); + return rpc; + } + + private static ReaderSettings.Builder settings() { + return ReaderSettings.newBuilder() + .addTopic("/test-topic") + .setConsumerName("consumer") + .setDecompressionExecutor(Runnable::run); + } + + private AsyncReader reader(ReaderSettings settings, ReadStreamMock first, ReadStreamMock... rest) { + ReadEventHandlersSettings handlers = ReadEventHandlersSettings.newBuilder() + .setEventHandler(handler) + .setExecutor(Runnable::run) + .build(); + return new AsyncReaderImpl(mockRpc(first, rest), settings, handlers, REGISTRY); + } + + private static void assertMessages(List expected, List messages, long offset) { + Assert.assertEquals(expected.size(), messages.size()); + for (int idx = 0; idx < expected.size(); idx++) { + Assert.assertEquals(offset + idx, messages.get(idx).getOffset()); + Assert.assertArrayEquals(expected.get(idx), messages.get(idx).getData()); + } + } + + @Test + public void initAndShutdownTest() { + ReadStreamMock mock = new ReadStreamMock(); + AsyncReader reader = reader(settings().setReaderName("test-reader-name").build(), mock); + mock.assertSentMessagesCount(0); + Mockito.verifyNoInteractions(handler); + + CompletableFuture init = reader.init(); + Assert.assertFalse(init.isDone()); + mock.assertSentMessagesCount(1); + mock.assertLastMessage().isInitRequest("consumer", "/test-topic"); + + Assert.assertSame(init, reader.init()); // double init is allowed + mock.assertSentMessagesCount(1); + mock.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); + + CompletableFuture shutdown = reader.shutdown(); + mock.assertIsClosed(); + Assert.assertTrue(shutdown.isDone()); + Assert.assertSame(shutdown, reader.shutdown()); // double shutdown is allowed + mock.assertIsClosed(); + mock.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)); + } + + @Test + public void shutdownBeforeInitTest() { + ReadStreamMock mock = new ReadStreamMock(); + AsyncReader reader = reader(settings().build(), mock); + CompletableFuture init = reader.init(); + mock.assertSentMessagesCount(1); + mock.assertLastMessage().isInitRequest("consumer", "/test-topic"); + + CompletableFuture shutdown = reader.shutdown(); + mock.assertIsClosed(); + Assert.assertTrue(init.isCompletedExceptionally()); + Assert.assertTrue(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.assertSame(init, reader.init()); + Mockito.verify(handler).onReaderClosed(Mockito.any(ReaderClosedEvent.class)); + Mockito.verify(handler, Mockito.never()).onSessionStarted(Mockito.any(SessionStartedEvent.class)); + } + + @Test + public void shutdownWithoutInitTest() { + ReadStreamMock mock = new ReadStreamMock(); + AsyncReader reader = reader(settings().build(), mock); + CompletableFuture shutdown = reader.shutdown(); + Assert.assertTrue(shutdown.isDone()); + Assert.assertFalse(shutdown.isCompletedExceptionally()); + Assert.assertSame(shutdown, reader.shutdown()); + + CompletionException ex = Assert.assertThrows(CompletionException.class, () -> reader.init().join()); + Assert.assertEquals("Reader closed with Status{code = SUCCESS, issues = [Closed by client (S_INFO)]}", + ex.getCause().getMessage()); + mock.assertIsNotStarted(); + mock.assertSentMessagesCount(0); + Mockito.verifyNoInteractions(handler); + } + + @Test + public void readClosedPartitionTest() { + ReadStreamMock mock = new ReadStreamMock(); + Queue decoding = new ArrayDeque<>(); + AsyncReader reader = reader(settings().setMaxMemoryUsageBytes(200000) + .setDecompressionExecutor(decoding::add).build(), mock); + reader.init(); + mock.responseInit("read-session-1"); + mock.assertSentMessagesCount(2); + mock.assertLastMessage().isReadRequest(200000); + + mock.responseStartPartition("/test-topic", 123, 0); + mock.assertLastMessage().isStartPartition(1); + mock.responseStartPartition("/test-topic", 345, 0); + mock.assertSentMessagesCount(4); + mock.assertLastMessage().isStartPartition(2); + Mockito.verify(handler, Mockito.times(2)).onStartPartitionSession(Mockito.any(StartPartitionSessionEvent.class)); + + mock.responseData(10000).partition(1, 0).batch(Codec.GZIP, MSG1, MSG2, MSG3, MSG4, MSG5).and().send(); + mock.responseData(10000).partition(2, 0).batch(Codec.GZIP, MSG5, MSG4, MSG3, MSG2, MSG1).and().send(); + Mockito.verify(handler, Mockito.never()).onMessages(Mockito.any(DataReceivedEvent.class)); + + // Stop the first partition before its queued messages are decoded. + mock.responseStopPartition(1, true); + mock.assertSentMessagesCount(5); + mock.assertLastMessage().isStopPartition(1); + ArgumentCaptor stopped = ArgumentCaptor.forClass(StopPartitionSessionEvent.class); + Mockito.verify(handler).onStopPartitionSession(stopped.capture()); + Assert.assertEquals(1, stopped.getValue().getPartitionSessionId()); + + while (!decoding.isEmpty()) { + decoding.remove().run(); + } + ArgumentCaptor data = ArgumentCaptor.forClass(DataReceivedEvent.class); + Mockito.verify(handler, Mockito.times(5)).onMessages(data.capture()); + + List messages = new ArrayList<>(); + data.getAllValues().forEach(ev -> messages.addAll(ev.getMessages())); + assertMessages(Arrays.asList(MSG5, MSG4, MSG3, MSG2, MSG1), messages, 0); + + Assert.assertEquals(345, data.getValue().getMessages().get(0).getPartitionSession().getPartitionId()); + mock.assertSentMessagesCount(6); + mock.assertLastMessage().isReadRequest(20000); + + reader.shutdown(); + mock.closeStream(Status.SUCCESS); + } + + @Test + @HideLoggers({ BufferManager.class, ReaderImpl.class }) + public void invalidBatchesTest() { + ReadStreamMock mock = new ReadStreamMock(); + AsyncReader reader = reader(settings().setMaxMemoryUsageBytes(2000).build(), mock); + reader.init(); + mock.responseInit("read-session-1"); + mock.assertSentMessagesCount(2); + mock.assertLastMessage().isReadRequest(2000); + mock.responseStartPartition("/test-topic", 123, 0); + mock.assertSentMessagesCount(3); + mock.assertLastMessage().isStartPartition(1); + + // Empty batches release their memory without invoking the message handler. + mock.responseData(1000).partition(1, 0).batch(Codec.RAW).and().send(); + mock.assertSentMessagesCount(4); + mock.assertLastMessage().isReadRequest(1000); + + reader.shutdown(); + mock.responseData(1200).partition(1, 1000).batch(Codec.RAW, MSG1, MSG2, MSG3, MSG4, MSG5).and().send(); + mock.assertSentMessagesCount(5); + mock.assertLastMessage().isReadRequest(1200); + Mockito.verify(handler, Mockito.never()).onMessages(Mockito.any(DataReceivedEvent.class)); + mock.closeStream(Status.of(StatusCode.INTERNAL_ERROR)); + } + + @Test + public void retrySkipsQueuedMessagesTest() { + ErrorsHandler errorsHandler = new ErrorsHandler(); + ReadStreamMock m1 = new ReadStreamMock(); + ReadStreamMock m2 = new ReadStreamMock(); + Queue decoding = new ArrayDeque<>(); + AsyncReader reader = reader(settings().setMaxMemoryUsageBytes(2000) + .setDecompressionExecutor(decoding::add) + .setRetryConfig(IMMEDIATE_RETRY) + .setErrorsHandler(errorsHandler) + .build(), m1, m2); + reader.init(); + m1.assertLastMessage().isInitRequest("consumer", "/test-topic"); + m1.responseInit("read-session-1"); + m1.assertLastMessage().isReadRequest(2000); + m1.responseStartPartition("/test-topic", 123, 0); + m1.assertLastMessage().isStartPartition(1); + m1.responseData(1000).partition(1, 0).batch(Codec.RAW, MSG1, MSG2, MSG3).and().send(); + while (!decoding.isEmpty()) { + decoding.remove().run(); + } + + ArgumentCaptor data = ArgumentCaptor.forClass(DataReceivedEvent.class); + Mockito.verify(handler).onMessages(data.capture()); + assertMessages(Arrays.asList(MSG1, MSG2, MSG3), data.getValue().getMessages(), 0); + List messages = data.getValue().getMessages(); + CompletableFuture c1 = messages.get(0).commit(); + CompletableFuture c2 = messages.get(1).commit(); + m1.assertLastMessage().isCommit(1).hasPartitionOffset(1, OffsetsRange.of(1, 2)); + Assert.assertFalse(c1.isDone()); + Assert.assertFalse(c2.isDone()); + + m1.responseCommitAck().partition(1, 1).send(); + Assert.assertTrue(c1.isDone()); + Assert.assertFalse(c1.isCompletedExceptionally()); + Assert.assertFalse(c2.isDone()); + ArgumentCaptor ack = + ArgumentCaptor.forClass(CommitOffsetAcknowledgementEvent.class); + Mockito.verify(handler).onCommitResponse(ack.capture()); + Assert.assertEquals(1, ack.getValue().getCommittedOffset()); + Assert.assertEquals(123, ack.getValue().getPartitionSession().getPartitionId()); + + // Leave a batch queued on the failed stream to check that it is discarded. + m1.responseData(1000).partition(1, 3).batch(Codec.GZIP, MSG4, MSG2).and().send(); + errorsHandler.assertEmpty(); + m1.closeStream(Status.of(StatusCode.TRANSPORT_UNAVAILABLE)); + errorsHandler.assertCodes(StatusCode.TRANSPORT_UNAVAILABLE); + Assert.assertTrue(c2.isCompletedExceptionally()); + Assert.assertTrue(messages.get(2).commit().isCompletedExceptionally()); + ArgumentCaptor closed = ArgumentCaptor.forClass(PartitionSessionClosedEvent.class); + Mockito.verify(handler).onPartitionSessionClosed(closed.capture()); + Assert.assertEquals(123, closed.getValue().getPartitionSession().getPartitionId()); + Mockito.verify(handler, Mockito.never()).onReaderClosed(Mockito.any(ReaderClosedEvent.class)); + + m2.assertSentMessagesCount(1); + m2.assertLastMessage().isInitRequest("consumer", "/test-topic"); + m2.responseInit("read-session-2"); + m2.assertLastMessage().isReadRequest(2000); + m2.responseStartPartition("/test-topic", 123, 1); + m2.assertLastMessage().isStartPartition(1); + while (!decoding.isEmpty()) { + decoding.remove().run(); + } + Mockito.verify(handler).onMessages(Mockito.any(DataReceivedEvent.class)); + + m2.responseData(1000).partition(1, 1).batch(Codec.RAW, MSG2, MSG3, MSG4, MSG5).and().send(); + while (!decoding.isEmpty()) { + decoding.remove().run(); + } + Mockito.verify(handler, Mockito.times(2)).onMessages(data.capture()); + assertMessages(Arrays.asList(MSG2, MSG3, MSG4, MSG5), data.getValue().getMessages(), 1); + ArgumentCaptor started = ArgumentCaptor.forClass(SessionStartedEvent.class); + Mockito.verify(handler, Mockito.times(2)).onSessionStarted(started.capture()); + Assert.assertEquals("read-session-2", started.getValue().getSessionId()); + + DeferredCommitter committer = DeferredCommitter.newInstance(); + committer.add(data.getValue().getMessages().get(0)); + committer.add(data.getValue().getMessages().get(2)); + committer.commit(); + m2.assertLastMessage().isCommit(1).hasPartitionOffset(1, OffsetsRange.of(1, 2), OffsetsRange.of(3, 4)); + reader.shutdown(); + m2.assertIsClosed(); + m2.closeStream(Status.SUCCESS); + } +} 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 1c0228438..cc2b4fec5 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 @@ -2,13 +2,13 @@ import java.io.IOException; import java.io.OutputStream; -import java.util.ArrayDeque; -import java.util.Deque; +import java.util.ArrayList; import java.util.List; import java.util.Set; import java.util.concurrent.CompletableFuture; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; import java.util.stream.Collectors; import com.google.protobuf.ByteString; @@ -31,8 +31,9 @@ public class ReadStreamMock implements GrpcReadWriteStream { private static final CodecRegistry REGISTRY = new CodecRegistry(); + private final AtomicReference token = new AtomicReference<>("token-value"); private final CompletableFuture future = new CompletableFuture<>(); - private final Deque messages = new ArrayDeque<>(); + private final List messages = new ArrayList<>(); private final AtomicInteger partCounter = new AtomicInteger(); private Observer observer = null; private final AtomicInteger isClosed = new AtomicInteger(); @@ -40,7 +41,7 @@ public class ReadStreamMock implements GrpcReadWriteStream topics = msg.getInitRequest().getTopicsReadSettingsList().stream() .map(YdbTopic.StreamReadMessage.InitRequest.TopicReadSettings::getPath) @@ -258,6 +284,13 @@ public CommitAssert isCommit(long count) { return new CommitAssert(resp); } + public MessageAssert isUpdateToken(String tokenValue) { + Assert.assertTrue("Msg is not update token request", msg.hasUpdateTokenRequest()); + YdbTopic.UpdateTokenRequest resp = msg.getUpdateTokenRequest(); + Assert.assertEquals("Update token request has incorrect value", tokenValue, resp.getToken()); + return this; + } + public class CommitAssert { private final YdbTopic.StreamReadMessage.CommitOffsetRequest resp; @@ -283,7 +316,6 @@ public CommitAssert hasPartitionOffset(long psid, OffsetsRange... expected) { return this; } - } } 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 dafbc27e9..11a47f20f 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 @@ -1,5 +1,7 @@ package tech.ydb.topic.read.impl; +import java.time.Duration; +import java.time.Instant; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; @@ -29,7 +31,8 @@ public class ReaderImplTest { private static final CodecRegistry REGISTRY = new CodecRegistry(); - private static final String TOPIC = "/test/topic"; + private static final String TOPIC1 = "/test/topic"; + private static final String TOPIC2 = "/test/topic2"; private static TopicRpc mockRpc(ReadStreamMock first, ReadStreamMock... rest) { TopicRpc rpc = Mockito.mock(TopicRpc.class); @@ -46,12 +49,52 @@ private static void assertIllegalArgument(String msg, ThrowingRunnable runnable) Assert.assertEquals(msg, ex.getMessage()); } + @Test + public void updateTokenTest() { + ReadStreamMock mock = new ReadStreamMock(); + + ReaderSettings settings = ReaderSettings.newBuilder() + .addTopic(TopicReadSettings.newBuilder() + .setPath(TOPIC1) + .setMaxLag(Duration.ofDays(1)) + .setReadFrom(Instant.EPOCH.plusSeconds(1000000)) + .build()) + .addTopic(TopicReadSettings.newBuilder().setPath(TOPIC2).build()) + .setRetryConfig(TopicRetryConfig.NEVER) + .setMaxMemoryUsageBytes(1000) + .withoutConsumer() + .build(); + + ReadConfig config = new ReadConfig(REGISTRY, Runnable::run, Runnable::run, settings); + ReaderImpl.Handler handler = Mockito.mock(ReaderImpl.Handler.class); + + ReaderImpl reader = new ReaderImpl(mockRpc(mock), "test-reader", settings, config, handler); + reader.start(); + + mock.assertSentMessagesCount(1); + mock.assertLastMessage().isInitRequest(null, TOPIC1, TOPIC2); + + mock.updateTokenValue("new-token-value"); + + mock.assertSentMessagesCount(1); + mock.responseInit("read-session-1"); + + mock.assertSentMessagesCount(3); // update token + read request + mock.assertPreLastMessage().isUpdateToken("new-token-value"); + mock.assertLastMessage().isReadRequest(1000); + + mock.responseUpdateToken(); + + reader.close(); + mock.assertIsClosed(); + } + @Test public void updateOffsetsInTxValidationTest() { ReadStreamMock mock = new ReadStreamMock(); ReaderSettings settings = ReaderSettings.newBuilder() - .addTopic(TopicReadSettings.newBuilder().setPath(TOPIC).build()) + .addTopic(TopicReadSettings.newBuilder().setPath(TOPIC1).build()) .setConsumerName("consumer") .setRetryConfig(TopicRetryConfig.NEVER) .setMaxMemoryUsageBytes(1000) @@ -80,16 +123,16 @@ public void updateOffsetsInTxValidationTest() { assertIllegalArgument("Empty offsets range to update in transaction", () -> reader.updateOffsetsInTransaction( - active, Collections.singletonMap(TOPIC, new ArrayList()), updateSettings + active, Collections.singletonMap(TOPIC1, new ArrayList()), updateSettings ) ); List offsets = Arrays.asList( - new PartitionOffsets(new PartitionSession(1, 1, TOPIC), Arrays.asList(OffsetsRange.of(0, 10))), - new PartitionOffsets(new PartitionSession(2, 2, TOPIC), + new PartitionOffsets(new PartitionSession(1, 1, TOPIC1), Arrays.asList(OffsetsRange.of(0, 10))), + new PartitionOffsets(new PartitionSession(2, 2, TOPIC1), Arrays.asList(OffsetsRange.of(0, 1), OffsetsRange.of(2, 3))) ); - reader.updateOffsetsInTransaction(active, Collections.singletonMap(TOPIC, offsets), updateSettings); + reader.updateOffsetsInTransaction(active, Collections.singletonMap(TOPIC1, offsets), updateSettings); txStatus.complete(Status.SUCCESS); @@ -104,7 +147,7 @@ public void updateOffsetsInTxFailTest() { ReadStreamMock mock = new ReadStreamMock(); ReaderSettings settings = ReaderSettings.newBuilder() - .addTopic(TopicReadSettings.newBuilder().setPath(TOPIC).build()) + .addTopic(TopicReadSettings.newBuilder().setPath(TOPIC1).build()) .setConsumerName("consumer") .setRetryConfig(TopicRetryConfig.NEVER) .setMaxMemoryUsageBytes(1000) @@ -121,11 +164,11 @@ public void updateOffsetsInTxFailTest() { CompletableFuture txStatus = new CompletableFuture<>(); TxMock active = new TxMock(txStatus); List offsets = Arrays.asList( - new PartitionOffsets(new PartitionSession(1, 1, TOPIC), Arrays.asList(OffsetsRange.of(0, 10))), - new PartitionOffsets(new PartitionSession(2, 2, TOPIC), + new PartitionOffsets(new PartitionSession(1, 1, TOPIC1), Arrays.asList(OffsetsRange.of(0, 10))), + new PartitionOffsets(new PartitionSession(2, 2, TOPIC1), Arrays.asList(OffsetsRange.of(0, 1), OffsetsRange.of(2, 3))) ); - reader.updateOffsetsInTransaction(active, Collections.singletonMap(TOPIC, offsets), updateSettings); + reader.updateOffsetsInTransaction(active, Collections.singletonMap(TOPIC1, offsets), updateSettings); txStatus.complete(Status.of(StatusCode.ABORTED)); mock.assertIsClosed(); @@ -138,7 +181,7 @@ public void updateOffsetsInTxErrorTest() { ReadStreamMock mock = new ReadStreamMock(); ReaderSettings settings = ReaderSettings.newBuilder() - .addTopic(TopicReadSettings.newBuilder().setPath(TOPIC).build()) + .addTopic(TopicReadSettings.newBuilder().setPath(TOPIC1).build()) .setConsumerName("consumer") .setRetryConfig(TopicRetryConfig.NEVER) .setMaxMemoryUsageBytes(1000) @@ -155,11 +198,11 @@ public void updateOffsetsInTxErrorTest() { CompletableFuture txStatus = new CompletableFuture<>(); TxMock active = new TxMock(txStatus); List offsets = Arrays.asList( - new PartitionOffsets(new PartitionSession(1, 1, TOPIC), Arrays.asList(OffsetsRange.of(0, 10))), - new PartitionOffsets(new PartitionSession(2, 2, TOPIC), + new PartitionOffsets(new PartitionSession(1, 1, TOPIC1), Arrays.asList(OffsetsRange.of(0, 10))), + new PartitionOffsets(new PartitionSession(2, 2, TOPIC1), Arrays.asList(OffsetsRange.of(0, 1), OffsetsRange.of(2, 3))) ); - reader.updateOffsetsInTransaction(active, Collections.singletonMap(TOPIC, offsets), updateSettings); + reader.updateOffsetsInTransaction(active, Collections.singletonMap(TOPIC1, offsets), updateSettings); txStatus.completeExceptionally(new RuntimeException("tx problem")); mock.assertIsClosed(); 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 ec07402e9..11405690c 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 @@ -88,7 +88,7 @@ public void initAndShutdownTest() throws InterruptedException { reader.shutdown(); // double shutdow is allowed Exception ex = Assert.assertThrows(RuntimeException.class, () -> reader.receive(0, TimeUnit.MILLISECONDS)); - Assert.assertEquals("Reader was stopped", ex.getMessage()); + Assert.assertEquals("Reader was stopped with Status{code = SUCCESS}", ex.getMessage()); mock.closeStream(Status.SUCCESS); } @@ -112,8 +112,34 @@ public void shutdownBeforeInitTest() throws InterruptedException { mock.assertIsClosed(); mock.closeStream(Status.SUCCESS); - Exception ex = Assert.assertThrows(RuntimeException.class, () -> reader.receive(0, TimeUnit.MILLISECONDS)); - Assert.assertEquals("Reader was stopped", ex.getMessage()); + Exception ex1 = Assert.assertThrows(RuntimeException.class, () -> reader.initAndWait()); + Assert.assertEquals("Reader was closed with Status{code = SUCCESS}", ex1.getMessage()); + + Exception ex2 = Assert.assertThrows(RuntimeException.class, () -> reader.receive(0, TimeUnit.MILLISECONDS)); + Assert.assertEquals("Reader was stopped with Status{code = SUCCESS}", ex2.getMessage()); + } + + @Test + public void shutdownWithoutInitTest() throws InterruptedException { + ReadStreamMock mock = new ReadStreamMock(); + + ReaderSettings settings = ReaderSettings.newBuilder() + .addTopic(TopicReadSettings.newBuilder().setPath("/test-topic").build()) + .setConsumerName("consumer") + .build(); + + SyncReader reader = new SyncReaderImpl(mockRpc(mock), settings, REGISTRY); + + reader.shutdown(); // shutdown without init + mock.assertIsNotStarted(); + + Exception ex1 = Assert.assertThrows(RuntimeException.class, () -> reader.initAndWait()); + Assert.assertEquals("Reader was closed with Status{code = SUCCESS, issues = [Closed by client (S_INFO)]}", + ex1.getMessage()); + + Exception ex2 = Assert.assertThrows(RuntimeException.class, () -> reader.receive(0, TimeUnit.MILLISECONDS)); + Assert.assertEquals("Reader was stopped with Status{code = SUCCESS, issues = [Closed by client (S_INFO)]}", + ex2.getMessage()); } @Test diff --git a/topic/src/test/java/tech/ydb/topic/settings/ReaderSettingsTest.java b/topic/src/test/java/tech/ydb/topic/settings/ReaderSettingsTest.java new file mode 100644 index 000000000..038728b5e --- /dev/null +++ b/topic/src/test/java/tech/ydb/topic/settings/ReaderSettingsTest.java @@ -0,0 +1,51 @@ +package tech.ydb.topic.settings; + +import org.junit.Assert; +import org.junit.Test; + +/** + * + * @author Aleksandr Gorshenin {@literal } + */ +public class ReaderSettingsTest { + + @Test + public void validateTopicsListTest() { + Exception ex = Assert.assertThrows( + IllegalArgumentException.class, + () -> ReaderSettings.newBuilder().setConsumerName("consumer").build() + ); + Assert.assertEquals("Missing topics for read settings. At least one should be set", ex.getMessage()); + } + + @Test + public void validateConsumerNameTest() { + Exception ex = Assert.assertThrows( + IllegalArgumentException.class, + () -> ReaderSettings.newBuilder().addTopic("/topic").build() + ); + Assert.assertEquals("Missing consumer name for read settings. Use withoutConsumer option explicitly if you " + + "want to read without a consumer", ex.getMessage()); + } + + @Test + public void validateWithoutConsumerTest() { + Exception ex = Assert.assertThrows( + IllegalArgumentException.class, + () -> ReaderSettings.newBuilder().addTopic("/topic").setConsumerName("c").withoutConsumer().build() + ); + Assert.assertEquals( + "Both mutually exclusive options consumerName and withoutConsumer are set for read settings", + ex.getMessage() + ); + } + + @Test + public void validateRetryConfigTest() { + Exception ex = Assert.assertThrows( + NullPointerException.class, + () -> ReaderSettings.newBuilder().addTopic("/topic").setConsumerName("c").setRetryConfig(null).build() + ); + Assert.assertEquals("RetryConfig must not be null", ex.getMessage()); + } +} From f6e6234f6951e4a6e7d276f5ab11e6b8595974d6 Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Thu, 17 Sep 2026 11:31:53 +0100 Subject: [PATCH 09/12] Fixes by copylot --- .../tech/ydb/topic/impl/TopicStreamBase.java | 7 ++++++ .../ydb/topic/read/impl/AsyncReaderImpl.java | 2 +- .../ydb/topic/read/impl/LazyExecutor.java | 13 ++-------- .../tech/ydb/topic/read/impl/ReadSession.java | 24 +++++++++++-------- .../tech/ydb/topic/read/impl/ReaderImpl.java | 4 +++- .../ydb/topic/read/impl/SyncReaderImpl.java | 1 + .../topic/TopicReadersIntegrationTest.java | 1 + 7 files changed, 29 insertions(+), 23 deletions(-) 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 f091afe7f..9791e5e14 100644 --- a/topic/src/main/java/tech/ydb/topic/impl/TopicStreamBase.java +++ b/topic/src/main/java/tech/ydb/topic/impl/TopicStreamBase.java @@ -64,6 +64,13 @@ public void close() { } } + protected void fail(Status status) { + logger.warn("[{}] stopped by fail {}", debugId, status); + if (streamStatus.complete(status)) { + stream.close(); + } + } + @Override public void send(W req) { if (streamStatus.isDone()) { 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 b24d275d2..2aea0c394 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 @@ -127,7 +127,7 @@ public void handleReaderClosed(Status status) { try { eventHandler.onReaderClosed(new ReaderClosedEvent()); } catch (Throwable th) { - failSession(th, "onSessionStarted"); + failSession(th, "onReaderClosed"); throw th; } }); 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 index c5d575d26..8589a652f 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/LazyExecutor.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/LazyExecutor.java @@ -4,7 +4,6 @@ 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; @@ -69,15 +68,7 @@ private void shutdown(ExecutorService service) { return; } - try { - service.shutdown(); - if (!service.awaitTermination(5, TimeUnit.SECONDS)) { - logger.warn("executor {} shutdown timeout exceeded, interrupt all tasks", name); - service.shutdownNow(); - } - } catch (InterruptedException e) { - logger.warn("executor {} shutdown interrupted", name, e); - Thread.currentThread().interrupt(); - } + // do not wait for termination + service.shutdown(); } } 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 3fe0dd4d0..21923fc84 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 @@ -136,12 +136,18 @@ public StartPartitionSessionEvent onStartPartition(YdbTopic.StreamReadMessage.St req.getPartitionOffsets().getEnd() ); - String traceID = debugId + '/' + psid + "-p" + pid; + String tid = debugId + '/' + psid + "-p" + pid; + if (partitions.putIfAbsent(psid, partition) != null) { + logger.error("[{}] Received second StartPartitionSessionRequest for the already active {}", debugId, + partition); + Issue issue = Issue.of("Restarting read session due to receiving second StartPartitionSessionRequest with " + + partition, Issue.Severity.FATAL); + fail(Status.of(StatusCode.CLIENT_INTERNAL_ERROR, issue)); + return null; + } logger.info("[{}] Received StartPartitionSessionRequest for {} and consumer \"{}\" with committedOffset {}" - + " and partitionOffsets {}", traceID, partition, config.getConsumerName(), committed, offsets); - - partitions.put(psid, partition); - return new StartPartitionRequest(traceID, partition, committed, offsets); + + " and partitionOffsets {}", tid, partition, config.getConsumerName(), committed, offsets); + return new StartPartitionRequest(tid, partition, committed, offsets); } public PartitionSession onClosePartition(long partitionSessionId) { @@ -169,11 +175,9 @@ public StopPartitionSessionEvent onStopPartition(YdbTopic.StreamReadMessage.Stop if (partition == null) { logger.error("[{}] Received graceful StopPartitionSessionRequest for partition session {}, " + "but have no such partition session active", debugId, psid); - send(YdbTopic.StreamReadMessage.FromClient.newBuilder().setStopPartitionSessionResponse( - YdbTopic.StreamReadMessage.StopPartitionSessionResponse.newBuilder() - .setPartitionSessionId(psid) - .build()) - .build()); + Issue issue = Issue.of("Restarting read session due to receiving StopPartitionSessionRequest with " + + "PartitionSessionId " + psid + " that SDK knows nothing about", Issue.Severity.FATAL); + fail(Status.of(StatusCode.CLIENT_INTERNAL_ERROR, issue)); return 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 99fb68f07..c9603ba49 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 @@ -119,7 +119,9 @@ protected void onNext(ReadSession stream, FromServer message) { stream.onInit(message.getInitResponse()); } else if (message.hasStartPartitionSessionRequest()) { StartPartitionSessionEvent event = stream.onStartPartition(message.getStartPartitionSessionRequest()); - handler.handleStartPartitionSessionRequest(event); + if (event != null) { + handler.handleStartPartitionSessionRequest(event); + } } else if (message.hasStopPartitionSessionRequest()) { YdbTopic.StreamReadMessage.StopPartitionSessionRequest req = message.getStopPartitionSessionRequest(); if (req.getGraceful()) { diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java b/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java index ccf968ef1..53e877d69 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java @@ -184,6 +184,7 @@ public Message receiveInternal(ReceiveSettings receiveSettings, long timeout, Ti } if (!next.isActive()) { + next.confirm(); continue; } diff --git a/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java b/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java index 607cbeb7e..0b7b8f78b 100644 --- a/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java +++ b/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java @@ -174,6 +174,7 @@ private static void writeToTopic(String topicPath, int partitionID, int count) { private static void writeToTopic(String topicPath, String producerId, int startFrom, int count) { writeToTopic(startFrom, count, WriterSettings.newBuilder() + .setLogPrefix("writers-test-" + producerId) .setTopicPath(topicPath) .setProducerId(producerId) .build()); From 42e7b0ec370d33731a602a3971d5fa5a5789fe02 Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Thu, 17 Sep 2026 11:57:54 +0100 Subject: [PATCH 10/12] Updates tests --- .../topic/TopicReadersIntegrationTest.java | 3 +- .../ydb/topic/read/impl/ReadStreamMock.java | 6 +- .../ydb/topic/read/impl/ReaderImplTest.java | 161 +++++++++++++----- 3 files changed, 130 insertions(+), 40 deletions(-) diff --git a/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java b/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java index 0b7b8f78b..b29b68f16 100644 --- a/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java +++ b/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java @@ -230,7 +230,8 @@ public void singleThreadExecutorTest() throws Exception { processing.completeExceptionally(new RuntimeException("shutdown")); shutdown.get(5, TimeUnit.SECONDS); - executor.shutdownNow(); + executor.shutdown(); + executor.awaitTermination(5, TimeUnit.SECONDS); } @Test 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 cc2b4fec5..0d9d0dd6c 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 @@ -92,13 +92,17 @@ public void responseUpdateToken() { } public void responseStartPartition(String topicPath, long partitionID, long committedOffset) { + responseStartPartition(topicPath, partitionID, committedOffset, partCounter.incrementAndGet()); + } + + public void responseStartPartition(String topicPath, long partitionID, long committedOffset, long psid) { FromServer msg = FromServer.newBuilder() .setStatus(StatusCodesProtos.StatusIds.StatusCode.SUCCESS) .setStartPartitionSessionRequest(YdbTopic.StreamReadMessage.StartPartitionSessionRequest.newBuilder() .setPartitionSession(YdbTopic.StreamReadMessage.PartitionSession.newBuilder() .setPath(topicPath) .setPartitionId(partitionID) - .setPartitionSessionId(partCounter.incrementAndGet()) + .setPartitionSessionId(psid) .build()) .setCommittedOffset(committedOffset) .build()) 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 11a47f20f..4fdc89122 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 @@ -10,8 +10,10 @@ import java.util.concurrent.ScheduledExecutorService; import org.junit.Assert; +import org.junit.Rule; import org.junit.Test; import org.junit.function.ThrowingRunnable; +import org.mockito.ArgumentCaptor; import org.mockito.Mockito; import tech.ydb.common.transaction.TxMode; @@ -23,10 +25,14 @@ import tech.ydb.topic.description.OffsetsRange; import tech.ydb.topic.read.PartitionOffsets; import tech.ydb.topic.read.PartitionSession; +import tech.ydb.topic.read.events.StartPartitionSessionEvent; +import tech.ydb.topic.read.events.StopPartitionSessionEvent; import tech.ydb.topic.settings.ReaderSettings; import tech.ydb.topic.settings.TopicReadSettings; import tech.ydb.topic.settings.TopicRetryConfig; import tech.ydb.topic.settings.UpdateOffsetsInTransactionSettings; +import tech.ydb.topic.utils.HideLoggers; +import tech.ydb.topic.utils.HideLoggersRule; public class ReaderImplTest { private static final CodecRegistry REGISTRY = new CodecRegistry(); @@ -34,6 +40,9 @@ public class ReaderImplTest { private static final String TOPIC1 = "/test/topic"; private static final String TOPIC2 = "/test/topic2"; + @Rule + public final HideLoggersRule hideLogger = new HideLoggersRule(); + private static TopicRpc mockRpc(ReadStreamMock first, ReadStreamMock... rest) { TopicRpc rpc = Mockito.mock(TopicRpc.class); Mockito.when(rpc.getScheduler()).thenReturn(Mockito.mock(ScheduledExecutorService.class)); @@ -49,10 +58,118 @@ private static void assertIllegalArgument(String msg, ThrowingRunnable runnable) Assert.assertEquals(msg, ex.getMessage()); } + private static ReaderImpl startReader(ReaderImpl.Handler handler, ReadStreamMock mock) { + ReaderSettings settings = ReaderSettings.newBuilder() + .addTopic(TopicReadSettings.newBuilder().setPath(TOPIC1).build()) + .setConsumerName("consumer") + .setRetryConfig(TopicRetryConfig.NEVER) + .setMaxMemoryUsageBytes(1000) + .build(); + ReadConfig config = new ReadConfig(REGISTRY, Runnable::run, Runnable::run, settings); + ReaderImpl reader = new ReaderImpl(mockRpc(mock), "test-reader", settings, config, handler); + reader.start(); + mock.responseInit("read-session-1"); + mock.assertSentMessagesCount(2); + mock.assertLastMessage().isReadRequest(1000); + return reader; + } + @Test - public void updateTokenTest() { + @HideLoggers({ ReadSession.class }) + public void duplicateStartPartitionRequestTest() { + ReadStreamMock mock = new ReadStreamMock(); + ReaderImpl.Handler handler = Mockito.mock(ReaderImpl.Handler.class); + ReaderImpl reader = startReader(handler, mock); + + mock.responseStartPartition(TOPIC1, 123, 0, 1); + + ArgumentCaptor start = ArgumentCaptor.forClass(StartPartitionSessionEvent.class); + Mockito.verify(handler).handleStartPartitionSessionRequest(start.capture()); + + Assert.assertEquals(new PartitionSession(1, 123, TOPIC1), start.getValue().getPartitionSession()); + + // Once recived, a second start partition request is a protocol error. + mock.responseStartPartition(TOPIC1, 123, 0, 1); + + // partiton is not started + Mockito.verify(handler, Mockito.never()).handleStopPartitionSession(Mockito.any()); + + ArgumentCaptor closed = ArgumentCaptor.forClass(Status.class); + Mockito.verify(handler).handleReaderClosed(closed.capture()); + Assert.assertEquals(StatusCode.CLIENT_INTERNAL_ERROR, closed.getValue().getCode()); + mock.assertSentMessagesCount(2); + mock.assertIsClosed(); + reader.close(); + mock.assertIsClosed(); + } + + @Test + @HideLoggers({ ReadSession.class }) + public void duplicateStopPartitionRequestTest() { + ReadStreamMock mock = new ReadStreamMock(); + ReaderImpl.Handler handler = Mockito.mock(ReaderImpl.Handler.class); + ReaderImpl reader = startReader(handler, mock); + + mock.responseStartPartition(TOPIC1, 123, 0); + ArgumentCaptor start = ArgumentCaptor.forClass(StartPartitionSessionEvent.class); + Mockito.verify(handler).handleStartPartitionSessionRequest(start.capture()); + start.getValue().confirm(); + + // Both requests arrive before the application confirms the stop. + mock.responseStopPartition(1, true); + mock.responseStopPartition(1, true); + ArgumentCaptor stop = ArgumentCaptor.forClass(StopPartitionSessionEvent.class); + Mockito.verify(handler, Mockito.times(2)).handleStopPartitionSession(stop.capture()); + for (StopPartitionSessionEvent event : stop.getAllValues()) { + Assert.assertSame(start.getValue().getPartitionSession(), event.getPartitionSession()); + } + mock.assertSentMessagesCount(3); + + stop.getAllValues().get(0).confirm(); + mock.assertSentMessagesCount(4); + mock.assertLastMessage().isStopPartition(1); + stop.getAllValues().get(1).confirm(); + stop.getAllValues().get(0).confirm(); + mock.assertSentMessagesCount(4); + mock.assertIsActive(); + + // Once confirmed, a graceful stop for the removed session is a protocol error. + mock.responseStopPartition(1, true); + Mockito.verify(handler, Mockito.times(2)).handleStopPartitionSession(Mockito.any()); + ArgumentCaptor closed = ArgumentCaptor.forClass(Status.class); + Mockito.verify(handler).handleReaderClosed(closed.capture()); + Assert.assertEquals(StatusCode.CLIENT_INTERNAL_ERROR, closed.getValue().getCode()); + mock.assertSentMessagesCount(4); + mock.assertIsClosed(); + reader.close(); + mock.assertIsClosed(); + } + + @Test + public void duplicateForcedStopPartitionRequestTest() { ReadStreamMock mock = new ReadStreamMock(); + ReaderImpl.Handler handler = Mockito.mock(ReaderImpl.Handler.class); + ReaderImpl reader = startReader(handler, mock); + + mock.responseStartPartition(TOPIC1, 123, 0); + ArgumentCaptor start = ArgumentCaptor.forClass(StartPartitionSessionEvent.class); + Mockito.verify(handler).handleStartPartitionSessionRequest(start.capture()); + start.getValue().confirm(); + + mock.responseStopPartition(1, false); + mock.responseStopPartition(1, false); + Mockito.verify(handler).handleClosePartitionSession(start.getValue().getPartitionSession()); + Mockito.verify(handler, Mockito.never()).handleStopPartitionSession(Mockito.any()); + Mockito.verify(handler, Mockito.never()).handleReaderClosed(Mockito.any()); + mock.assertSentMessagesCount(3); // forced stops do not require acknowledgement + mock.assertIsActive(); + reader.close(); + mock.assertIsClosed(); + } + + @Test + public void updateTokenTest() { ReaderSettings settings = ReaderSettings.newBuilder() .addTopic(TopicReadSettings.newBuilder() .setPath(TOPIC1) @@ -65,12 +182,13 @@ public void updateTokenTest() { .withoutConsumer() .build(); + ReadStreamMock mock = new ReadStreamMock(); ReadConfig config = new ReadConfig(REGISTRY, Runnable::run, Runnable::run, settings); ReaderImpl.Handler handler = Mockito.mock(ReaderImpl.Handler.class); ReaderImpl reader = new ReaderImpl(mockRpc(mock), "test-reader", settings, config, handler); - reader.start(); + reader.start(); mock.assertSentMessagesCount(1); mock.assertLastMessage().isInitRequest(null, TOPIC1, TOPIC2); @@ -92,19 +210,8 @@ public void updateTokenTest() { @Test public void updateOffsetsInTxValidationTest() { ReadStreamMock mock = new ReadStreamMock(); - - ReaderSettings settings = ReaderSettings.newBuilder() - .addTopic(TopicReadSettings.newBuilder().setPath(TOPIC1).build()) - .setConsumerName("consumer") - .setRetryConfig(TopicRetryConfig.NEVER) - .setMaxMemoryUsageBytes(1000) - .build(); - - ReadConfig config = new ReadConfig(REGISTRY, Runnable::run, Runnable::run, settings); ReaderImpl.Handler handler = Mockito.mock(ReaderImpl.Handler.class); - - ReaderImpl reader = new ReaderImpl(mockRpc(mock), "test-reader", settings, config, handler); - reader.start(); + ReaderImpl reader = startReader(handler, mock); UpdateOffsetsInTransactionSettings updateSettings = UpdateOffsetsInTransactionSettings.newBuilder() .withTraceId("test-trace").build(); @@ -145,19 +252,8 @@ public void updateOffsetsInTxValidationTest() { @Test public void updateOffsetsInTxFailTest() { ReadStreamMock mock = new ReadStreamMock(); - - ReaderSettings settings = ReaderSettings.newBuilder() - .addTopic(TopicReadSettings.newBuilder().setPath(TOPIC1).build()) - .setConsumerName("consumer") - .setRetryConfig(TopicRetryConfig.NEVER) - .setMaxMemoryUsageBytes(1000) - .build(); - - ReadConfig config = new ReadConfig(REGISTRY, Runnable::run, Runnable::run, settings); ReaderImpl.Handler handler = Mockito.mock(ReaderImpl.Handler.class); - - ReaderImpl reader = new ReaderImpl(mockRpc(mock), "test-reader", settings, config, handler); - reader.start(); + ReaderImpl reader = startReader(handler, mock); UpdateOffsetsInTransactionSettings updateSettings = UpdateOffsetsInTransactionSettings.newBuilder().build(); @@ -179,19 +275,8 @@ public void updateOffsetsInTxFailTest() { @Test public void updateOffsetsInTxErrorTest() { ReadStreamMock mock = new ReadStreamMock(); - - ReaderSettings settings = ReaderSettings.newBuilder() - .addTopic(TopicReadSettings.newBuilder().setPath(TOPIC1).build()) - .setConsumerName("consumer") - .setRetryConfig(TopicRetryConfig.NEVER) - .setMaxMemoryUsageBytes(1000) - .build(); - - ReadConfig config = new ReadConfig(REGISTRY, Runnable::run, Runnable::run, settings); ReaderImpl.Handler handler = Mockito.mock(ReaderImpl.Handler.class); - - ReaderImpl reader = new ReaderImpl(mockRpc(mock), "test-reader", settings, config, handler); - reader.start(); + ReaderImpl reader = startReader(handler, mock); UpdateOffsetsInTransactionSettings updateSettings = UpdateOffsetsInTransactionSettings.newBuilder().build(); From 92755a1df94e2cd7e6ff5806ed67758c6158022e Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Thu, 17 Sep 2026 14:34:04 +0100 Subject: [PATCH 11/12] Added integration tests for reader retries --- .../tech/ydb/topic/impl/TopicStreamBase.java | 6 +- .../ydb/topic/read/impl/AsyncReaderImpl.java | 4 +- .../ydb/topic/read/impl/LazyExecutor.java | 12 +- .../tech/ydb/topic/read/impl/ReaderImpl.java | 2 +- .../ydb/topic/read/impl/SyncReaderImpl.java | 10 +- .../ydb/topic/FailableReaderInterceptor.java | 255 ++++++++++++++++++ .../topic/TopicReadersIntegrationTest.java | 191 ++++++++++++- .../topic/read/impl/AsyncReaderImplTest.java | 3 +- .../topic/read/impl/SyncReaderImplTest.java | 3 +- .../tech/ydb/topic/utils/ErrorsHandler.java | 7 +- 10 files changed, 473 insertions(+), 20 deletions(-) create mode 100644 topic/src/test/java/tech/ydb/topic/FailableReaderInterceptor.java 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 9791e5e14..522143003 100644 --- a/topic/src/main/java/tech/ydb/topic/impl/TopicStreamBase.java +++ b/topic/src/main/java/tech/ydb/topic/impl/TopicStreamBase.java @@ -18,6 +18,7 @@ public abstract class TopicStreamBase impl private final W initRequest; private final CompletableFuture streamStatus = new CompletableFuture<>(); private volatile String token; + private volatile boolean isStopped = false; public TopicStreamBase(Logger logger, String debugId, GrpcReadWriteStream stream, W initRequest) { this.logger = logger; @@ -44,6 +45,7 @@ public CompletableFuture start(Consumer messageHandler) { } } }).whenComplete((st, th) -> { + isStopped = true; Status status = st != null ? st : Status.of(StatusCode.CLIENT_INTERNAL_ERROR, th); logger.debug("[{}] finished with status {}", debugId, status); streamStatus.complete(status); @@ -60,6 +62,7 @@ public CompletableFuture start(Consumer messageHandler) { public void close() { logger.debug("[{}] closed by app", debugId); if (!streamStatus.isDone()) { + isStopped = true; stream.close(); } } @@ -67,13 +70,14 @@ public void close() { protected void fail(Status status) { logger.warn("[{}] stopped by fail {}", debugId, status); if (streamStatus.complete(status)) { + isStopped = true; stream.close(); } } @Override public void send(W req) { - if (streamStatus.isDone()) { + if (isStopped) { 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 2aea0c394..3c0a217c7 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 @@ -65,9 +65,9 @@ public AsyncReaderImpl(TopicRpc topicRpc, String readerName = settings.getReaderName(); String consumerName = settings.getConsumerName(); - logger.info("Reader{} (generated id {}) created for topic(s) {} and {}", - readerName != null ? (" '" + readerName + "'") : "", + logger.info("[{}] AsyncReader{} created for topic(s) {} and {}", debugId, + readerName != null ? (" '" + readerName + "'") : "", settings.getTopics().stream().map(t -> "\"" + t.getPath() + "\"").collect(Collectors.joining(", ")), consumerName != null ? (" consumer \"" + consumerName + "\"") : "without a consumer" ); 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 index 8589a652f..406d981e6 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/LazyExecutor.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/LazyExecutor.java @@ -4,6 +4,7 @@ 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; @@ -68,7 +69,14 @@ private void shutdown(ExecutorService service) { return; } - // do not wait for termination - service.shutdown(); + 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/main/java/tech/ydb/topic/read/impl/ReaderImpl.java b/topic/src/main/java/tech/ydb/topic/read/impl/ReaderImpl.java index c9603ba49..29c1aef42 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 @@ -80,7 +80,7 @@ protected ReadSession createNewStream(String id) { @Override protected void onRetry(ReadSession stream, Status status) { - logger.warn("[{}] paused by status {}", debugId, status); + logger.warn("[{}] stopped by status {}", debugId, status); if (errorHandler != null) { try { errorHandler.accept(status, null); diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java b/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java index 53e877d69..1d4f082d9 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/SyncReaderImpl.java @@ -67,9 +67,9 @@ public SyncReaderImpl(TopicRpc topicRpc, ReaderSettings settings, @Nonnull Codec String readerName = settings.getReaderName(); String consumerName = settings.getConsumerName(); - logger.info("Reader{} (generated id {}) created for topic(s) {} and {}", - readerName != null ? (" '" + readerName + "'") : "", + logger.info("[{}] SyncReader{} created for topic(s) {} and {}", debugId, + readerName != null ? (" '" + readerName + "'") : "", settings.getTopics().stream().map(t -> "\"" + t.getPath() + "\"").collect(Collectors.joining(", ")), consumerName != null ? (" consumer \"" + consumerName + "\"") : "without a consumer" ); @@ -237,7 +237,7 @@ public void handleDataReceivedEvent(ReaderImpl.PartitionControl control, DataRec int messagesCount = event.getMessages().size(); long offsetStart = event.getMessages().get(0).getOffset(); long offsetEnd = event.getMessages().get(event.getMessages().size() - 1).getOffset(); - logger.debug("{} Putting a batch into queueData with {} message(s) (offsets {}-{}) from {}", + logger.debug("[{}] Putting a batch into read queue with {} message(s) (offsets {}-{}) from {}", debugId, messagesCount, offsetStart, offsetEnd, ps); for (Message msg: event.getMessages()) { @@ -252,8 +252,8 @@ public void handleDataReceivedEvent(ReaderImpl.PartitionControl control, DataRec } @Override - public void handleCommitResponse(long committedOffset, PartitionSession partitionSession) { - logger.debug("CommitResponse received for{} with committedOffset {}", partitionSession, committedOffset); + public void handleCommitResponse(long committedOffset, PartitionSession ps) { + logger.debug("[{}] commit response received for {} with committedOffset {}", debugId, ps, committedOffset); } @Override diff --git a/topic/src/test/java/tech/ydb/topic/FailableReaderInterceptor.java b/topic/src/test/java/tech/ydb/topic/FailableReaderInterceptor.java new file mode 100644 index 000000000..831f62345 --- /dev/null +++ b/topic/src/test/java/tech/ydb/topic/FailableReaderInterceptor.java @@ -0,0 +1,255 @@ +package tech.ydb.topic; + +import java.util.HashMap; +import java.util.Map; +import java.util.NavigableMap; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentSkipListMap; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Consumer; + +import io.grpc.CallOptions; +import io.grpc.Channel; +import io.grpc.ClientCall; +import io.grpc.ClientInterceptor; +import io.grpc.ManagedChannelBuilder; +import io.grpc.Metadata; +import io.grpc.MethodDescriptor; +import io.grpc.Status; + +import tech.ydb.proto.StatusCodesProtos; +import tech.ydb.proto.topic.YdbTopic.StreamReadMessage.CommitOffsetRequest.PartitionCommitOffset; +import tech.ydb.proto.topic.YdbTopic.StreamReadMessage.CommitOffsetResponse.PartitionCommittedOffset; +import tech.ydb.proto.topic.YdbTopic.StreamReadMessage.FromClient; +import tech.ydb.proto.topic.YdbTopic.StreamReadMessage.FromServer; + + +/** + * + * @author Aleksandr Gorshenin + */ +public class FailableReaderInterceptor implements Consumer>, ClientInterceptor { + private final AtomicInteger initCounter = new AtomicInteger(); + private final AtomicInteger readCounter = new AtomicInteger(); + + private final Map initErrors = new HashMap<>(); + private final Map readErrors = new HashMap<>(); + private final Map> ackErrors = new ConcurrentHashMap<>(); + private final Map> sendErrors = new ConcurrentHashMap<>(); + + public void reset() { + initErrors.clear(); + readErrors.clear(); + ackErrors.clear(); + sendErrors.clear(); + initCounter.set(0); + readCounter.set(0); + } + + @Override + public void accept(ManagedChannelBuilder t) { + t.intercept(this); + } + + public void unavailableOnInit(int number) { + initErrors.put(number, closeStream(Status.UNAVAILABLE)); + } + + public void badRequestOnInit(int number) { + initErrors.put(number, sendError(StatusCodesProtos.StatusIds.StatusCode.BAD_REQUEST)); + } + + public void unavailableOnReadResponse(int number) { + readErrors.put(number, closeStream(Status.UNAVAILABLE)); + } + + public void badRequestOnReadResponse(int number) { + readErrors.put(number, sendError(StatusCodesProtos.StatusIds.StatusCode.BAD_REQUEST)); + } + + public void unavailableOnCommitAck(long partitionID, long offset) { + ackErrors.computeIfAbsent(partitionID, id -> new ConcurrentSkipListMap<>()) + .put(offset, closeStream(Status.UNAVAILABLE)); + } + + public void badRequestOnCommitAck(long partitionID, long offset) { + ackErrors.computeIfAbsent(partitionID, id -> new ConcurrentSkipListMap<>()) + .put(offset, sendError(StatusCodesProtos.StatusIds.StatusCode.BAD_REQUEST)); + } + + public void unavailableOnCommitWithOffset(long partitionID, long offset) { + sendErrors.computeIfAbsent(partitionID, id -> new ConcurrentSkipListMap<>()) + .put(offset, closeStream(Status.UNAVAILABLE)); + } + + public void badSessionOnCommitWithOffset(long partitionID, long offset) { + sendErrors.computeIfAbsent(partitionID, id -> new ConcurrentSkipListMap<>()) + .put(offset, sendError(StatusCodesProtos.StatusIds.StatusCode.BAD_SESSION)); + } + + + @Override + public ClientCall interceptCall(MethodDescriptor method, CallOptions callOptions, Channel next) { + return new ProxyCall<>(next.newCall(method, callOptions)); + } + + interface Error { + boolean fail(ClientCall.Listener listener); + } + + private class ProxyCall extends ClientCall { + + private final ClientCall realCall; + private final Map partitions = new ConcurrentHashMap<>(); + private volatile ProxyListener proxyListener; + private volatile boolean isClosed = false; + + ProxyCall(ClientCall delegate) { + this.realCall = delegate; + } + + @Override + public void start(Listener listener, Metadata headers) { + proxyListener = new ProxyListener<>(listener); + realCall.start(proxyListener, headers); + } + + @Override + public void request(int numMessages) { + realCall.request(numMessages); + } + + @Override + public void cancel(String message, Throwable cause) { + realCall.cancel(message, cause); + } + + @Override + public void halfClose() { + realCall.halfClose(); + } + + @Override + @SuppressWarnings("unchecked") + public void sendMessage(W message) { + if (isClosed) { + return; + } + + Error error = null; + if (message instanceof FromClient) { + FromClient msg = (FromClient) message; + if (msg.hasCommitOffsetRequest()) { + for (PartitionCommitOffset p: msg.getCommitOffsetRequest().getCommitOffsetsList()) { + Long pid = partitions.getOrDefault(p.getPartitionSessionId(), -1L); + NavigableMap local = sendErrors.get(pid); + if (local == null || local.isEmpty()) { + continue; + } + long offset = p.getOffsetsList().get(p.getOffsetsCount() - 1).getEnd(); + Map.Entry nextError = local.firstEntry(); + if (nextError != null && nextError.getKey() <= offset) { + error = local.remove(nextError.getKey()); + break; + } + } + } + } + + if (error == null) { + realCall.sendMessage(message); + return; + } + + isClosed = error.fail((Listener) proxyListener); + if (isClosed) { + realCall.halfClose(); + } + } + + private class ProxyListener extends Listener { + private final Listener realListener; + + ProxyListener(Listener realListener) { + this.realListener = realListener; + } + + @Override + public void onClose(Status status, Metadata trailers) { + if (!isClosed) { + realListener.onClose(status, trailers); + } + } + + @Override + public void onHeaders(Metadata headers) { + if (!isClosed) { + realListener.onHeaders(headers); + } + } + + @Override + @SuppressWarnings("unchecked") + public void onMessage(R message) { + if (isClosed) { + return; + } + + Error error = null; + if (message instanceof FromServer) { + FromServer msg = (FromServer) message; + if (msg.hasStartPartitionSessionRequest()) { + tech.ydb.proto.topic.YdbTopic.StreamReadMessage.PartitionSession partition = + msg.getStartPartitionSessionRequest().getPartitionSession(); + partitions.put(partition.getPartitionSessionId(), partition.getPartitionId()); + } + if (msg.hasInitResponse()) { + error = initErrors.get(initCounter.incrementAndGet()); + } + if (msg.hasReadResponse()) { + error = readErrors.get(readCounter.incrementAndGet()); + } + if (msg.hasCommitOffsetResponse()) { + for (PartitionCommittedOffset p: msg.getCommitOffsetResponse().getPartitionsCommittedOffsetsList()) { + Long pid = partitions.getOrDefault(p.getPartitionSessionId(), -1L); + NavigableMap local = ackErrors.get(pid); + if (local == null || local.isEmpty()) { + continue; + } + long lastAck = p.getCommittedOffset(); + Map.Entry nextError = local.firstEntry(); + if (nextError != null && nextError.getKey() <= lastAck) { + error = local.remove(nextError.getKey()); + break; + } + } + } + } + if (error == null) { + realListener.onMessage(message); + return; + } + + isClosed = error.fail((Listener) realListener); + if (isClosed) { + realCall.halfClose(); + } + } + } + } + + private static Error closeStream(Status grpcStatus) { + return (ClientCall.Listener listener) -> { + listener.onClose(grpcStatus, new Metadata()); + return true; + }; + } + + private static Error sendError(StatusCodesProtos.StatusIds.StatusCode ydbStatus) { + return (ClientCall.Listener listener) -> { + listener.onMessage(FromServer.newBuilder().setStatus(ydbStatus).build()); + return false; + }; + } +} + diff --git a/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java b/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java index b29b68f16..a89127c72 100644 --- a/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java +++ b/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java @@ -10,6 +10,7 @@ import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.concurrent.Semaphore; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicLong; @@ -28,6 +29,7 @@ import tech.ydb.common.transaction.TxMode; import tech.ydb.core.Status; +import tech.ydb.core.StatusCode; import tech.ydb.table.SessionRetryContext; import tech.ydb.table.TableClient; import tech.ydb.table.transaction.TableTransaction; @@ -38,6 +40,7 @@ import tech.ydb.topic.impl.SerialExecutor; import tech.ydb.topic.read.AsyncReader; import tech.ydb.topic.read.Message; +import tech.ydb.topic.read.SyncReader; import tech.ydb.topic.read.events.DataReceivedEvent; import tech.ydb.topic.read.events.ReadEventHandler; import tech.ydb.topic.read.events.StartPartitionSessionEvent; @@ -55,6 +58,7 @@ import tech.ydb.topic.settings.TopicReadSettings; import tech.ydb.topic.settings.UpdateOffsetsInTransactionSettings; import tech.ydb.topic.settings.WriterSettings; +import tech.ydb.topic.utils.ErrorsHandler; import tech.ydb.topic.utils.HideLoggers; import tech.ydb.topic.utils.HideLoggersRule; import tech.ydb.topic.write.SyncWriter; @@ -66,8 +70,11 @@ public class TopicReadersIntegrationTest { private static final Logger logger = LoggerFactory.getLogger(YdbTopicsIntegrationTest.class); + private static final FailableReaderInterceptor PROXY = new FailableReaderInterceptor(); + @ClassRule - public final static GrpcTransportRule ydbTransport = new GrpcTransportRule(); + public final static GrpcTransportRule ydbTransport = new GrpcTransportRule() + .withGrpcTransportCustomizer(b -> b.addChannelInitializer(PROXY)); @Rule public final HideLoggersRule hideLogger = new HideLoggersRule(); @@ -82,8 +89,9 @@ public class TopicReadersIntegrationTest { @BeforeClass public static void initClient() { client = TopicClient.newClient(ydbTransport).build(); - logger.info("Create test topic {} ...", TEST_TOPIC); + client.dropTopic(TEST_TOPIC).join(); + client.dropTopic(SPLITTED_TOPIC).join(); client.createTopic(TEST_TOPIC, CreateTopicSettings.newBuilder() .addConsumer(Consumer.newBuilder().setName(TEST_CONSUMER1).build()) .setPartitioningSettings(PartitioningSettings.newBuilder() @@ -130,12 +138,14 @@ public static void initClient() { public static void closeClient() { logger.info("Drop test topic {} ...", TEST_TOPIC); client.dropTopic(TEST_TOPIC).join(); + logger.info("Drop test topic {} ...", SPLITTED_TOPIC); client.dropTopic(SPLITTED_TOPIC).join(); client.close(); } @Before public void resetConsumer() { + PROXY.reset(); List> resets = new ArrayList<>(); DescribeConsumerSettings dc = DescribeConsumerSettings.newBuilder().withIncludeStats(true).build(); @@ -299,6 +309,183 @@ public void readAllSplittedTest() throws InterruptedException { } } + @Test(timeout = 120000) + public void syncReadAllWithDefaultRetryPolicyTest() throws Exception { + // Fail initialization, reading, and both sides of committing, in this order. + PROXY.unavailableOnInit(1); + PROXY.badRequestOnInit(2); + PROXY.unavailableOnReadResponse(1); + PROXY.unavailableOnInit(4); + PROXY.unavailableOnInit(5); + PROXY.unavailableOnInit(6); + PROXY.badRequestOnReadResponse(2); + + PROXY.badSessionOnCommitWithOffset(1, 60); + PROXY.unavailableOnCommitAck(1, 200); + + AtomicLong[] committed = new AtomicLong[] { new AtomicLong(), new AtomicLong(), new AtomicLong() }; + CountDownLatch totalCommitted = new CountDownLatch(3600); + + ErrorsHandler errors = new ErrorsHandler(); + ReaderSettings settings = ReaderSettings.newBuilder() + .addTopic(TopicReadSettings.newBuilder().setPath(TEST_TOPIC).build()) + .setConsumerName(TEST_CONSUMER1) + .setErrorsHandler((st, th) -> { // restore partition states + DescribeConsumerSettings s = DescribeConsumerSettings.newBuilder() + .withIncludeStats(true).build(); + ConsumerDescription desc = client.describeConsumer(TEST_TOPIC, TEST_CONSUMER1, s).join().getValue(); + for (ConsumerPartitionInfo partition: desc.getPartitions()) { + int pid = (int) partition.getPartitionId(); + long lastCommit = partition.getConsumerStats().getCommittedOffset(); + long diff = lastCommit - committed[pid].getAndSet(lastCommit); + for (int idx = 0; idx < diff; idx++) { + totalCommitted.countDown(); + } + } + errors.accept(st, th); + }) + .build(); + + Thread worker = new Thread(() -> { + Semaphore commitInflyLimit = new Semaphore(100); + SyncReader reader = client.createSyncReader(settings); + reader.init(); + try { + while (!Thread.interrupted() && totalCommitted.getCount() > 0) { + Message msg = reader.receive(10, TimeUnit.MILLISECONDS); + if (msg == null) { + continue; + } + + int pid = (int) msg.getPartitionSession().getPartitionId(); + AtomicLong lastCommit = committed[pid]; + long messageCommit = msg.getOffset() + 1; + // limit commit infly to avoid last message committing before test errors + commitInflyLimit.acquire(); + msg.commit().whenComplete((res, th) -> { + commitInflyLimit.release(); + if (th == null) { // commit is successful + long diff = messageCommit - lastCommit.getAndSet(messageCommit); + for (int idx = 0; idx < diff; idx++) { + totalCommitted.countDown(); + } + } + }); + } + } catch (InterruptedException ex) { + // nothing + } finally { + reader.shutdown(); + } + }); + + worker.start(); + try { + Assert.assertTrue("All messages must be committed", totalCommitted.await(30, TimeUnit.SECONDS)); + Assert.assertEquals(1000, committed[0].get()); + Assert.assertEquals(500, committed[1].get()); + Assert.assertEquals(2100, committed[2].get()); + } finally { + worker.interrupt(); + worker.join(1_000); + } + + errors.assertCodes( + StatusCode.TRANSPORT_UNAVAILABLE, + StatusCode.BAD_REQUEST, + StatusCode.TRANSPORT_UNAVAILABLE, + StatusCode.TRANSPORT_UNAVAILABLE, + StatusCode.TRANSPORT_UNAVAILABLE, + StatusCode.TRANSPORT_UNAVAILABLE, + StatusCode.BAD_REQUEST, + StatusCode.BAD_SESSION, + StatusCode.TRANSPORT_UNAVAILABLE + ); + } + + @Test(timeout = 120000) + public void asyncReadAllWithDefaultRetryPolicyTest() throws Exception { + // Fail initialization, reading, and both sides of committing, in this order. + PROXY.unavailableOnInit(1); + PROXY.badRequestOnInit(2); + PROXY.unavailableOnReadResponse(1); + PROXY.unavailableOnInit(4); + PROXY.unavailableOnInit(5); + + PROXY.badSessionOnCommitWithOffset(2, 100); + PROXY.unavailableOnCommitAck(2, 500); + + ErrorsHandler errors = new ErrorsHandler(); + ReaderSettings settings = ReaderSettings.newBuilder() + .addTopic(TopicReadSettings.newBuilder().setPath(TEST_TOPIC).build()) + .setReaderName("async-read-all-with-default-retry-policy") + .setConsumerName(TEST_CONSUMER1) + .setMaxBatchSize(100) // commits by 100 messages + .setErrorsHandler(errors) + .build(); + + AtomicLong[] offsets = new AtomicLong[] { new AtomicLong(), new AtomicLong(), new AtomicLong() }; + AtomicLong[] committed = new AtomicLong[] { new AtomicLong(), new AtomicLong(), new AtomicLong() }; + CountDownLatch totalCommitted = new CountDownLatch(3600); + + AsyncReader reader = client.createAsyncReader(settings, ReadEventHandlersSettings.newBuilder() + .setEventHandler(new ReadEventHandler() { + @Override + public void onStartPartitionSession(StartPartitionSessionEvent event) { + int pid = (int) event.getPartitionSession().getPartitionId(); + // restore offset position + offsets[pid].set(event.getCommittedOffset()); + // restore lost commits + long diff = event.getCommittedOffset() - committed[pid].getAndSet(event.getCommittedOffset()); + for (int idx = 0; idx < diff; idx++) { + totalCommitted.countDown(); + } + event.confirm(); + } + + @Override + public void onMessages(DataReceivedEvent event) { + int pid = (int) event.getPartitionSession().getPartitionId(); + AtomicLong offset = offsets[pid]; + AtomicLong lastCommit = committed[pid]; + for (Message msg : event.getMessages()) { + Assert.assertEquals(offset.getAndIncrement(), msg.getOffset()); + } + + long eventCommit = event.getRangeToCommit().getEnd(); + event.commit().thenRun(() -> { + long diff = eventCommit - lastCommit.getAndSet(eventCommit); + for (int idx = 0; idx < diff; idx++) { + totalCommitted.countDown(); + } + }); + } + }).build()); + + reader.init().join(); + try { + Assert.assertTrue(totalCommitted.await(30, TimeUnit.SECONDS)); + Assert.assertEquals(1000, offsets[0].get()); + Assert.assertEquals(500, offsets[1].get()); + Assert.assertEquals(2100, offsets[2].get()); + Assert.assertEquals(1000, committed[0].get()); + Assert.assertEquals(500, committed[1].get()); + Assert.assertEquals(2100, committed[2].get()); + } finally { + reader.shutdown().join(); + } + + errors.assertCodes( + StatusCode.TRANSPORT_UNAVAILABLE, + StatusCode.BAD_REQUEST, + StatusCode.TRANSPORT_UNAVAILABLE, + StatusCode.TRANSPORT_UNAVAILABLE, + StatusCode.TRANSPORT_UNAVAILABLE, + StatusCode.BAD_SESSION, + StatusCode.TRANSPORT_UNAVAILABLE + ); + } + @Test @Ignore // requires auto-partitioning supporr public void readAllSplittedWithoutCommitTest() throws InterruptedException { 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 e474f1d12..12fe2cc6e 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 @@ -230,8 +230,7 @@ public void invalidBatchesTest() { reader.shutdown(); mock.responseData(1200).partition(1, 1000).batch(Codec.RAW, MSG1, MSG2, MSG3, MSG4, MSG5).and().send(); - mock.assertSentMessagesCount(5); - mock.assertLastMessage().isReadRequest(1200); + mock.assertSentMessagesCount(4); Mockito.verify(handler, Mockito.never()).onMessages(Mockito.any(DataReceivedEvent.class)); mock.closeStream(Status.of(StatusCode.INTERNAL_ERROR)); } 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 11405690c..6ca412fe2 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 @@ -223,8 +223,7 @@ public void invalidBatchesTest() throws InterruptedException { reader.shutdown(); // batch after shutdown is just skipped mock.responseData(1200).partition(1, 1000).batch(Codec.RAW, MSG1, MSG2, MSG3, MSG4, MSG5).and().send(); - mock.assertSentMessagesCount(5); - mock.assertLastMessage().isReadRequest(1200); + mock.assertSentMessagesCount(4); mock.closeStream(Status.of(StatusCode.INTERNAL_ERROR)); } diff --git a/topic/src/test/java/tech/ydb/topic/utils/ErrorsHandler.java b/topic/src/test/java/tech/ydb/topic/utils/ErrorsHandler.java index bb137bd43..8f2db06b4 100644 --- a/topic/src/test/java/tech/ydb/topic/utils/ErrorsHandler.java +++ b/topic/src/test/java/tech/ydb/topic/utils/ErrorsHandler.java @@ -33,9 +33,10 @@ public void assertEmpty() { public void assertCodes(StatusCode... codes) { Iterator it = problems.iterator(); - for (StatusCode code: codes) { - Assert.assertTrue("Expected " + code + ", but has nothing", it.hasNext()); - Assert.assertEquals(code, it.next()); + for (int idx = 0; idx < codes.length; idx++) { + StatusCode code = codes[idx]; + Assert.assertTrue("Expected " + code + " on position " + idx + ", but has nothing", it.hasNext()); + Assert.assertEquals("Unexpected code on position " + idx, code, it.next()); } Assert.assertFalse("Unexpected error code", it.hasNext()); } From c8d73e4ff4b73792cf4bd4470b415d3287f5df2d Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Fri, 18 Sep 2026 14:07:26 +0100 Subject: [PATCH 12/12] Small fixes --- .../ydb/topic/read/impl/AsyncReaderImpl.java | 28 +++++++++++-------- .../ydb/topic/read/impl/LazyExecutor.java | 4 +-- .../tech/ydb/topic/read/impl/ReaderImpl.java | 16 +++++++++-- .../topic/read/impl/AsyncReaderImplTest.java | 2 +- 4 files changed, 33 insertions(+), 17 deletions(-) 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 3c0a217c7..b5bfb9af5 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 @@ -95,7 +95,22 @@ public CompletableFuture shutdown() { } private void close(Status status) { + if (shutdownFuture.isDone()) { + return; + } + + controlEventsExecutor.execute(() -> { + try { + eventHandler.onReaderClosed(new ReaderClosedEvent()); + } catch (Throwable th) { + logger.error("[{}] onReaderClosed finished with exception", th); + throw th; + } + }); + + // stop decompressong decompressor.close(); + // wait while processer finished all tasks processor.close(); initFuture.completeExceptionally(new RuntimeException("Reader closed with " + status)); shutdownFuture.complete(null); @@ -111,7 +126,7 @@ private class AsyncHandler implements ReaderImpl.Handler { @Override public void handleSessionStarted(String sessionId) { initFuture.complete(null); - processor.execute(() -> { + controlEventsExecutor.execute(() -> { try { eventHandler.onSessionStarted(new SessionStartedEvent(sessionId)); } catch (Throwable th) { @@ -123,16 +138,7 @@ public void handleSessionStarted(String sessionId) { @Override public void handleReaderClosed(Status status) { - processor.execute(() -> { - try { - eventHandler.onReaderClosed(new ReaderClosedEvent()); - } catch (Throwable th) { - failSession(th, "onReaderClosed"); - throw th; - } - }); - - close(status); // wait while processer finished all tasks + close(status); } @Override 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 index 406d981e6..da8271c74 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/LazyExecutor.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/LazyExecutor.java @@ -71,8 +71,8 @@ private void shutdown(ExecutorService service) { try { service.shutdown(); - if (!service.awaitTermination(100, TimeUnit.MILLISECONDS)) { - service.shutdownNow(); + if (!service.awaitTermination(1, TimeUnit.SECONDS)) { + logger.warn("executor {} doesn't compelete all tasks", name); } } catch (InterruptedException e) { logger.warn("executor {} shutdown interrupted", name, e); 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 29c1aef42..02ce252d4 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 @@ -64,6 +64,8 @@ public interface Handler { private final FromClient initRequest; + private volatile String currentSessionId = null; + public ReaderImpl(TopicRpc rpc, String id, ReaderSettings settings, ReadConfig config, Handler handler) { super(logger, id, settings.getRetryConfig(), rpc.getScheduler()); this.rpc = rpc; @@ -81,6 +83,7 @@ protected ReadSession createNewStream(String id) { @Override protected void onRetry(ReadSession stream, Status status) { logger.warn("[{}] stopped by status {}", debugId, status); + currentSessionId = null; if (errorHandler != null) { try { errorHandler.accept(status, null); @@ -93,6 +96,7 @@ protected void onRetry(ReadSession stream, Status status) { @Override protected void onClose(ReadSession stream, Status status) { + currentSessionId = null; if (!status.isSuccess()) { logger.warn("[{}] closed by status {}", debugId, status); } else { @@ -115,6 +119,7 @@ protected void onNext(ReadSession stream, FromServer message) { if (message.hasInitResponse()) { resetRetries(); + currentSessionId = message.getInitResponse().getSessionId(); handler.handleSessionStarted(message.getInitResponse().getSessionId()); stream.onInit(message.getInitResponse()); } else if (message.hasStartPartitionSessionRequest()) { @@ -184,18 +189,23 @@ public CompletableFuture updateOffsetsInTransaction(YdbTransaction trans logger.debug(str.toString()); } + Object sessionId = currentSessionId; // store current session id to fail it for tx errors transaction.getStatusFuture().whenComplete((status, error) -> { if (status != null && !status.isSuccess()) { String msg = "Restarting read session due to transaction " + transaction.getId() + " with partition offsets from read session " + debugId + " was not committed with status: " + status; - fail(Status.of(StatusCode.CLIENT_INTERNAL_ERROR, Issue.of(msg, Issue.Severity.ERROR))); + if (sessionId == currentSessionId) { + fail(Status.of(StatusCode.CLIENT_INTERNAL_ERROR, Issue.of(msg, Issue.Severity.ERROR))); + } } if (error != null) { String msg = "Restarting read session due to transaction " + transaction.getId() + " with partition offsets from read session " + debugId + " was not committed with reason: " + error.getMessage(); - fail(Status.of(StatusCode.CLIENT_INTERNAL_ERROR, error, Issue.of(msg, Issue.Severity.ERROR))); + if (sessionId == currentSessionId) { + fail(Status.of(StatusCode.CLIENT_INTERNAL_ERROR, Issue.of(msg, Issue.Severity.ERROR))); + } } }); @@ -240,7 +250,7 @@ private static YdbTopic.UpdateOffsetsInTransactionRequest.TopicOffsets buildTopi .build(); } - public static YdbTopic.OffsetsRange buildOffsetRange(OffsetsRange range) { + static YdbTopic.OffsetsRange buildOffsetRange(OffsetsRange range) { return YdbTopic.OffsetsRange.newBuilder() .setStart(range.getStart()) .setEnd(range.getEnd()) 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 12fe2cc6e..221905268 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 @@ -159,7 +159,7 @@ public void shutdownWithoutInitTest() { ex.getCause().getMessage()); mock.assertIsNotStarted(); mock.assertSentMessagesCount(0); - Mockito.verifyNoInteractions(handler); + Mockito.verify(handler).onReaderClosed(Mockito.any()); } @Test