From 7bac8e41394d7689ab0cf3d6db56aa576e61dc40 Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Thu, 10 Sep 2026 14:42:54 +0100 Subject: [PATCH 1/2] 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 2/2] 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(); - } - } -}