From d7feca6eaa6630d2fcde37b478ae496ff7ac65ad Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Mon, 3 Aug 2026 11:22:01 +0100 Subject: [PATCH 01/12] Added BufferManager to control every message lifecycle --- .../ydb/topic/read/impl/BufferManager.java | 220 ++++++++++++++++++ .../topic/read/impl/BufferManagerTest.java | 157 +++++++++++++ 2 files changed, 377 insertions(+) create mode 100644 topic/src/main/java/tech/ydb/topic/read/impl/BufferManager.java create mode 100644 topic/src/test/java/tech/ydb/topic/read/impl/BufferManagerTest.java diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/BufferManager.java b/topic/src/main/java/tech/ydb/topic/read/impl/BufferManager.java new file mode 100644 index 000000000..c554785ee --- /dev/null +++ b/topic/src/main/java/tech/ydb/topic/read/impl/BufferManager.java @@ -0,0 +1,220 @@ +package tech.ydb.topic.read.impl; + +import java.util.Arrays; +import java.util.Iterator; +import java.util.List; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.function.Consumer; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import tech.ydb.proto.topic.YdbTopic; +import tech.ydb.topic.description.OffsetsRange; + +/** + * + * @author Aleksandr Gorshenin {@literal } + */ +public class BufferManager { + private static final Logger logger = LoggerFactory.getLogger(ReaderImpl.class); + private static final OffsetsRange ALL = OffsetsRange.of(Integer.MIN_VALUE, Integer.MAX_VALUE); + + private final String traceID; + private final long maxBufferSize; + private final Consumer requestFunc; + + private final AtomicLong released = new AtomicLong(0); + private final ConcurrentHashMap partitions = new ConcurrentHashMap<>(); + + public BufferManager(String traceID, long maxBufferSize, Consumer requestFunc) { + this.traceID = traceID; + this.maxBufferSize = maxBufferSize; + this.requestFunc = requestFunc; + } + + public void init(String sessionId) { + logger.info("[{}] Session {} initialized. Requesting {} bytes...", traceID, sessionId, maxBufferSize); + requestFunc.accept(maxBufferSize); + } + + // Has no reentrant thread safety + public void allocate(int bufferSize, List dataList) { + logger.trace("[{}] Received ReadResponse of {} bytes", traceID, bufferSize); + + // calculate message count + int messagesCount = 0; + for (YdbTopic.StreamReadMessage.ReadResponse.PartitionData data: dataList) { + for (YdbTopic.StreamReadMessage.ReadResponse.Batch batch: data.getBatchesList()) { + messagesCount += batch.getMessageDataCount(); + } + } + + if (messagesCount == 0) { + logger.error("[{}] Received empty ReadResponse of {} bytes", traceID, bufferSize); + release(bufferSize); + return; + } + + // get real size for every message + int[] msgSize = new int[messagesCount]; + int msgIdx = 0; + for (YdbTopic.StreamReadMessage.ReadResponse.PartitionData data: dataList) { + for (YdbTopic.StreamReadMessage.ReadResponse.Batch batch: data.getBatchesList()) { + for (YdbTopic.StreamReadMessage.ReadResponse.MessageData msg: batch.getMessageDataList()) { + msgSize[msgIdx] = msg.getData().size(); + msgIdx++; + } + } + } + + // recalculate real messages size to expected buffer size + recalcBuffer(msgSize, bufferSize); + + // build batch and messages with calculated buffer size + msgIdx = 0; + for (YdbTopic.StreamReadMessage.ReadResponse.PartitionData data: dataList) { + PartitionBuffer part = partitions.computeIfAbsent(data.getPartitionSessionId(), PartitionBuffer::new); + for (YdbTopic.StreamReadMessage.ReadResponse.Batch batch: data.getBatchesList()) { + if (batch.getMessageDataCount() <= 0) { + continue; + } + + long startOffset = batch.getMessageData(0).getOffset(); + int[] batchSizes = new int[batch.getMessageDataCount()]; + for (int idx = 0; idx < batch.getMessageDataCount(); idx++) { + batchSizes[idx] = msgSize[msgIdx++]; + } + + part.add(new BatchBuffer(startOffset, batchSizes)); + } + + if (!partitions.containsKey(data.getPartitionSessionId())) { + release(part.release(ALL)); + } + } + } + + // Thread safe + public void releasePartition(Long id) { + PartitionBuffer part = partitions.remove(id); + if (part != null) { + release(part.release(ALL)); + } + } + + // Thread safe + public void releaseRange(Long id, OffsetsRange range) { + PartitionBuffer part = partitions.get(id); + if (part != null) { + release(part.release(range)); + } + } + + private void release(long total) { + long now = released.addAndGet(total); + if (now >= maxBufferSize / 10) { // threshhold + long request = released.getAndSet(0); + if (request > 0) { + requestFunc.accept(request); + } + } + } + + private static void recalcBuffer(int[] buffer, int buffSize) { + // corner cases + if (buffSize == 0) { + Arrays.fill(buffer, 0); + return; + } + + int total = 0; + for (int v: buffer) { + total += v; + } + + int currBuff = 0; + int currSum = 0; + for (int idx = 0; idx < buffer.length; idx += 1) { + currSum += buffer[idx]; + int newBuff = currSum * buffSize / total; + buffer[idx] = newBuff - currBuff; + currBuff = newBuff; + } + } + + private static class PartitionBuffer { + private final ConcurrentLinkedQueue batches = new ConcurrentLinkedQueue<>(); + + PartitionBuffer(Long id) { + } + + public void add(BatchBuffer range) { + batches.add(range); + } + + public long release(OffsetsRange range) { + long released = 0; + + Iterator it = batches.iterator(); + while (it.hasNext()) { + BatchBuffer next = it.next(); + if (next.getStartOffset() > range.getEnd()) { // fast path + break; + } + + released += next.release(range); + if (!next.isActive()) { + it.remove(); + } + } + + return released; + } + } + + private static class BatchBuffer { + private final long startOffset; + private final AtomicInteger[] messages; + private final AtomicLong total; + + BatchBuffer(long startOffset, int[] messageSizes) { + this.startOffset = startOffset; + this.messages = new AtomicInteger[messageSizes.length]; + long totalSize = 0; + for (int idx = 0; idx < messageSizes.length; idx++) { + this.messages[idx] = new AtomicInteger(messageSizes[idx]); + totalSize += messageSizes[idx]; + } + this.total = new AtomicLong(totalSize); + } + + public long getStartOffset() { + return startOffset; + } + + public boolean isActive() { + return total.get() > 0; + } + + public long release(OffsetsRange range) { + int first = (int) (Math.max(startOffset, range.getStart()) - startOffset); + int last = (int) Math.min(messages.length, range.getEnd() - startOffset); + + if (last <= first) { + return 0; + } + + long released = 0; + for (int idx = first; idx < last; idx++) { + released += messages[idx].getAndSet(0); + } + total.addAndGet(-released); + + return released; + } + } +} diff --git a/topic/src/test/java/tech/ydb/topic/read/impl/BufferManagerTest.java b/topic/src/test/java/tech/ydb/topic/read/impl/BufferManagerTest.java new file mode 100644 index 000000000..779f25a0c --- /dev/null +++ b/topic/src/test/java/tech/ydb/topic/read/impl/BufferManagerTest.java @@ -0,0 +1,157 @@ +package tech.ydb.topic.read.impl; + +import java.util.Arrays; +import java.util.Collections; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.function.Consumer; + +import com.google.protobuf.ByteString; +import org.junit.Assert; +import org.junit.Rule; +import org.junit.Test; + +import tech.ydb.proto.topic.YdbTopic.StreamReadMessage.ReadResponse; +import tech.ydb.topic.description.OffsetsRange; +import tech.ydb.topic.utils.HideLoggers; +import tech.ydb.topic.utils.HideLoggersRule; + +/** + * @author Aleksandr Gorshenin + */ +public class BufferManagerTest { + @Rule + public final HideLoggersRule hideLogger = new HideLoggersRule(); + + private static final class Resource implements Consumer { + private final AtomicInteger count = new AtomicInteger(0); + private final AtomicLong total = new AtomicLong(0); + + @Override + public void accept(Long value) { + count.incrementAndGet(); + total.addAndGet(value); + } + + public void assertCalls(int expectedCount, long expectedTotal) { + Assert.assertEquals("Unexpected resource requests count", expectedCount, count.get()); + Assert.assertEquals("Unexpected resource requests total value", expectedTotal, total.get()); + } + } + + private static ReadResponse.MessageData msg(long offset, int size) { + return ReadResponse.MessageData.newBuilder() + .setOffset(offset) + .setData(ByteString.copyFrom(new byte[size])) + .build(); + } + + private static ReadResponse.Batch batch(long startOffset, int... sizes) { + ReadResponse.Batch.Builder builder = ReadResponse.Batch.newBuilder(); + long offset = startOffset; + for (int size : sizes) { + builder.addMessageData(msg(offset++, size)); + } + return builder.build(); + } + + private static ReadResponse.PartitionData partition(long partition, ReadResponse.Batch... batches) { + ReadResponse.PartitionData.Builder builder = ReadResponse.PartitionData.newBuilder() + .setPartitionSessionId(partition); + for (ReadResponse.Batch batch : batches) { + builder.addBatches(batch); + } + return builder.build(); + } + + @Test + public void simpleUseTest() { + Resource r = new Resource(); + BufferManager bm = new BufferManager("trace-1", 1234, r); + + r.assertCalls(0, 0); + + bm.init("s1"); + r.assertCalls(1, 1234); + + bm.allocate(500, Arrays.asList( // 500 = 100 + 100 + 250 + 40 + 5 + 5 + partition(1, batch(0, 100, 100), batch(2, 250), batch(6, 40, 5, 5), batch(20)), // empty batch + partition(2) // empty partition + )); + r.assertCalls(1, 1234); + + // release unknown partitions doesn't affect anything + bm.releasePartition(3L); + bm.releaseRange(4L, OffsetsRange.of(0, 100)); + r.assertCalls(1, 1234); + + // partitial release more than 10% of buffer size - a new request + bm.releaseRange(1L, OffsetsRange.of(1, 3)); + r.assertCalls(2, 1234 + 350); + + // double release - no more request + bm.releaseRange(1L, OffsetsRange.of(1, 3)); + r.assertCalls(2, 1234 + 350); + + // partitial release less than 10% of buffer size - no additional request + bm.releaseRange(1L, OffsetsRange.of(6, 10)); + r.assertCalls(2, 1234 + 350); + + bm.releasePartition(1L); + r.assertCalls(3, 1234 + 500); + + // double release - no more request + bm.releasePartition(1L); + r.assertCalls(3, 1234 + 500); + } + + @Test + @HideLoggers(ReaderImpl.class) + public void incorrectAllocateTest() { + Resource r = new Resource(); + BufferManager bm = new BufferManager("trace-2", 20000, r); + r.assertCalls(0, 0); + + bm.allocate(5000, Collections.emptyList()); + r.assertCalls(1, 5000); + + bm.allocate(4000, Arrays.asList(partition(1), partition(2))); + r.assertCalls(2, 9000); + } + + @Test + public void zeroAllocateTest() { + Resource r = new Resource(); + BufferManager bm = new BufferManager("trace-3", 20000, r); + r.assertCalls(0, 0); + + // zero allocation is allowed + bm.allocate(0, Arrays.asList(partition(1, batch(0, 100, 100), batch(2, 250), batch(6, 40, 5, 5), batch(20)))); + r.assertCalls(0, 0); + + // releases has no effects + bm.releaseRange(1L, OffsetsRange.of(1, 3)); + bm.releasePartition(1L); + r.assertCalls(0, 0); + } + + @Test + public void concurrentReleaseTest() { + AtomicLong requested = new AtomicLong(0); + BufferManager bm = new BufferManager("trace-4", 1 /* no threshold */, requested::addAndGet); + + Thread t = new Thread(() -> { + for (int idx = 0; idx < 1000; idx += 1) { + bm.allocate(100, Arrays.asList(partition(1, batch(idx, 100)))); + } + }); + t.start(); + + while (t.isAlive()) { + bm.releasePartition(1L); + } + + bm.releasePartition(1L); + Assert.assertEquals(100000L, requested.get()); + } +} From 17302a185eeeca9e3d1c231ef7e7511a3c6bd5b6 Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Mon, 10 Aug 2026 11:02:07 +0100 Subject: [PATCH 02/12] Added per-message MessageDecoder --- .../ydb/topic/read/impl/MessageDecoder.java | 104 +++-------- .../tech/ydb/topic/read/impl/MessageImpl.java | 41 +---- .../topic/read/impl/ReadPartitionDecoder.java | 167 ++++++++++++++++++ .../topic/read/impl/ReadPartitionSession.java | 143 +++++++-------- .../tech/ydb/topic/read/impl/ReadSession.java | 72 +++----- 5 files changed, 292 insertions(+), 235 deletions(-) create mode 100644 topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionDecoder.java 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 d1fd88143..79fc28295 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 @@ -1,119 +1,63 @@ package tech.ydb.topic.read.impl; -import java.io.IOException; import java.util.Queue; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.Executor; import java.util.concurrent.atomic.AtomicLong; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - import tech.ydb.topic.description.CodecRegistry; -import tech.ydb.topic.utils.Encoder; +import tech.ydb.topic.impl.SerialRunnable; /** * Decodes message batches while limiting memory consumption for uncompressed data. * @author Aleksandr Gorshenin */ public class MessageDecoder { - private static final Logger logger = LoggerFactory.getLogger(MessageDecoder.class); + private final AtomicLong totalAvailable; - private final AtomicLong availableBufferSize; private final Executor decompressionExecutor; private final CodecRegistry codecRegistry; - private final Queue decodingQueue = new ConcurrentLinkedQueue<>(); + private final Queue decodingQueue = new ConcurrentLinkedQueue<>(); + private final SerialRunnable decodeNext = new SerialRunnable(new DecodeNext()); + private volatile boolean isStopped = false; public MessageDecoder(long maxBufferSize, Executor decompressionExecutor, CodecRegistry codecRegistry) { - this.availableBufferSize = new AtomicLong(maxBufferSize); + this.totalAvailable = new AtomicLong(maxBufferSize); this.decompressionExecutor = decompressionExecutor; this.codecRegistry = codecRegistry; } - public void decode(String traceID, Batch batch, Runnable readyHandler) { - decodingQueue.offer(new DecodeTask(traceID, batch, readyHandler)); - tryToDecodeNextBatch(); + public void decodeNext() { + decodeNext.run(); } - private void tryToDecodeNextBatch() { - if (availableBufferSize.get() <= 0) { - return; - } - - while (availableBufferSize.get() > 0) { - DecodeTask task = decodingQueue.poll(); - if (task == null) { - return; - } - - Batch batch = task.getBatch(); - if (batch.getReadFuture().isDone()) { // session was closed, just skip decoding - continue; - } - - long bufferSize = getUncompressedSize(batch); - availableBufferSize.addAndGet(-bufferSize); - batch.getReadFuture().whenComplete((v, th) -> { - availableBufferSize.addAndGet(bufferSize); - tryToDecodeNextBatch(); - }); - - decompressionExecutor.execute(task); - } + public void stop() { + this.isStopped = true; } - private long getUncompressedSize(Batch batch) { - long uncompressed = 0; - long compressed = 0; - for (MessageImpl msg: batch.getMessages()) { - uncompressed += msg.getUncompressedSize(); - compressed += msg.getData().length; - } - - if (uncompressed > 0) { - return uncompressed; - } - - // TODO: Implement moving average for compression level - return 2 * compressed; + void add(ReadPartitionDecoder.EncodedMessage task) { + decodingQueue.add(task); } - private class DecodeTask implements Runnable { - private final String traceID; - private final Batch batch; - private final Runnable readyHandler; - - DecodeTask(String traceID, Batch batch, Runnable readyHandler) { - this.traceID = traceID; - this.batch = batch; - this.readyHandler = readyHandler; - } - - public Batch getBatch() { - return batch; + void free(long bufferSize) { + if (bufferSize > 0) { + totalAvailable.addAndGet(bufferSize); + decodeNext.run(); } + } + private final class DecodeNext implements Runnable { @Override public void run() { - if (logger.isTraceEnabled()) { - logger.trace("[{}] Started decoding batch", traceID); - } - - batch.getMessages().forEach(message -> { - try { - message.setData(Encoder.decode(batch.getCodec(), message.getData(), codecRegistry)); - } catch (IOException exception) { - message.setException(exception); - logger.warn("[{}] Exception was thrown while decoding a message: ", traceID, exception); + while (!isStopped && totalAvailable.get() > 0) { + ReadPartitionDecoder.EncodedMessage next = decodingQueue.poll(); + if (next == null) { + return; } - }); - batch.markAsReady(); - if (logger.isTraceEnabled()) { - logger.trace("[{}] Finished decoding batch", traceID); + totalAvailable.addAndGet(-next.allocate()); + decompressionExecutor.execute(() -> next.decode(codecRegistry)); } - - readyHandler.run(); } } } diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/MessageImpl.java b/topic/src/main/java/tech/ydb/topic/read/impl/MessageImpl.java index 08d2ef969..0fb94e9d6 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/MessageImpl.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/MessageImpl.java @@ -1,6 +1,5 @@ package tech.ydb.topic.read.impl; -import java.io.IOException; import java.time.Instant; import java.util.Collections; import java.util.List; @@ -11,7 +10,6 @@ import tech.ydb.proto.topic.YdbTopic; import tech.ydb.topic.description.MetadataItem; import tech.ydb.topic.description.OffsetsRange; -import tech.ydb.topic.read.DecompressionException; import tech.ydb.topic.read.Message; import tech.ydb.topic.read.MessageCommitter; import tech.ydb.topic.read.PartitionOffsets; @@ -20,60 +18,35 @@ /** * @author Nikolay Perfilov */ -public class MessageImpl implements Message { +public abstract class MessageImpl implements Message { private final PartitionSession session; private final MessageCommitter committer; - private final long offset; + private final BatchMeta batchMeta; private final OffsetsRange commitRange; - private final long uncompressedSize; private final long seqNo; + private final long offset; private final Instant createdAt; private final String messageGroupId; - private final BatchMeta batchMeta; private final List metadataItems; - private byte[] data; - private IOException exception = null; - public MessageImpl(PartitionSession session, MessageCommitter committer, BatchMeta meta, OffsetsRange commitRange, YdbTopic.StreamReadMessage.ReadResponse.MessageData msg) { this.session = session; this.committer = committer; - this.uncompressedSize = msg.getUncompressedSize(); - this.offset = msg.getOffset(); + this.batchMeta = meta; this.commitRange = commitRange; + + this.offset = msg.getOffset(); this.seqNo = msg.getSeqNo(); this.createdAt = ProtobufUtils.protoToInstant(msg.getCreatedAt()); this.messageGroupId = msg.getMessageGroupId(); this.metadataItems = msg.getMetadataItemsList().stream() .map(metadataItem -> new MetadataItem(metadataItem.getKey(), metadataItem.getValue().toByteArray())) .collect(Collectors.toList()); - this.batchMeta = meta; - - this.data = msg.getData().toByteArray(); } - @Override - public byte[] getData() { - if (exception != null) { - throw new DecompressionException("Error occurred while decoding a message", - exception, data, batchMeta.getCodec()); - } - return data; - } - - public long getUncompressedSize() { - return uncompressedSize; - } - - public void setData(byte[] data) { - this.data = data; - } - - public void setException(IOException exception) { - this.exception = exception; - } + public abstract boolean isReady(); @Override public long getOffset() { diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionDecoder.java b/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionDecoder.java new file mode 100644 index 000000000..896c55e06 --- /dev/null +++ b/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionDecoder.java @@ -0,0 +1,167 @@ +package tech.ydb.topic.read.impl; + +import java.io.IOException; +import java.io.InputStream; +import java.util.Iterator; +import java.util.Queue; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.atomic.AtomicLong; + +import com.google.common.io.ByteStreams; +import com.google.protobuf.ByteString; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import tech.ydb.proto.topic.YdbTopic; +import tech.ydb.topic.description.Codec; +import tech.ydb.topic.description.CodecRegistry; +import tech.ydb.topic.description.OffsetsRange; +import tech.ydb.topic.read.DecompressionException; +import tech.ydb.topic.read.MessageCommitter; +import tech.ydb.topic.read.PartitionSession; + +/** + * + * @author Aleksandr Gorshenin {@literal } + */ +public class ReadPartitionDecoder { + private static final Logger logger = LoggerFactory.getLogger(ReaderImpl.class); + + private final String traceID; + private final PartitionSession partition; + private final MessageCommitter committer; + private final MessageDecoder decoder; + private final Runnable readyHandler; + + private final AtomicLong allocatedTotal = new AtomicLong(0); + private final Queue allocatedMessages = new ConcurrentLinkedQueue<>(); + private volatile boolean isStopped = false; + + public ReadPartitionDecoder(String traceId, MessageDecoder decoder, PartitionSession partition, + MessageCommitter committer, Runnable readyHandler) { + this.traceID = traceId; + this.decoder = decoder; + this.partition = partition; + this.committer = committer; + this.readyHandler = readyHandler; + } + + public MessageImpl decode(BatchMeta m, OffsetsRange r, YdbTopic.StreamReadMessage.ReadResponse.MessageData msg) { + EncodedMessage encoded = new EncodedMessage(m, r, msg); + decoder.add(encoded); + return encoded; + } + + public void releaseRange(OffsetsRange range) { + long released = 0; + Iterator it = allocatedMessages.iterator(); + while (it.hasNext()) { + EncodedMessage msg = it.next(); + if (msg.getOffset() < range.getStart()) { + continue; + } + + if (msg.getOffset() < range.getEnd()) { + released += msg.uncompressedSize; + it.remove(); + } else { + break; + } + } + + while (true) { + long curr = allocatedTotal.get(); + long free = Math.min(released, curr); + if (free <= 0) { + return; + } + if (allocatedTotal.compareAndSet(curr, curr - free)) { + decoder.free(free); + return; + } + } + } + + public void close() { + isStopped = true; + release(); + } + + private void release() { + long total = allocatedTotal.getAndSet(0); + if (total > 0) { + decoder.free(total); + } + } + + public class EncodedMessage extends MessageImpl { + private final int codecCode; + private final long uncompressedSize; + private ByteString origin; + private byte[] data = null; + private IOException problem = null; + + private volatile boolean isReady = false; + + EncodedMessage(BatchMeta meta, OffsetsRange range, YdbTopic.StreamReadMessage.ReadResponse.MessageData msg) { + super(partition, committer, meta, range, msg); + this.origin = msg.getData(); + this.codecCode = meta.getCodec(); + this.uncompressedSize = msg.getUncompressedSize() > 0 ? msg.getUncompressedSize() : 2 * origin.size(); + } + + @Override + public byte[] getData() { + if (problem != null) { + throw new DecompressionException("Error occurred while decoding a message", problem, + origin.toByteArray(), codecCode); + } + return data; + } + + @Override + public boolean isReady() { + return isReady; + } + + public long allocate() { + if (isStopped) { + isReady = true; + return 0; + } + + allocatedTotal.addAndGet(uncompressedSize); + allocatedMessages.add(this); + return uncompressedSize; + } + + public void decode(CodecRegistry registry) { + if (isStopped) { + release(); + return; + } + + try { + Codec codec = registry.getCodec(codecCode); + if (codec == null) { + logger.warn("[{}] Cannot decode a message because codec {} is not registered", traceID, codecCode); + problem = new IOException("Codec " + codecCode + " is not registered"); + return; + } + + logger.trace("[{}] Started decoding batch", traceID); + try (InputStream encoded = origin.newInput(); InputStream decoded = codec.decode(encoded)) { + data = ByteStreams.toByteArray(decoded); + logger.trace("[{}] Finished decoding batch", traceID); + origin = null; + } catch (IOException ex) { + logger.warn("[{}] Exception was thrown while decoding a message: ", traceID, ex); + problem = ex; + } + } finally { + isReady = true; + readyHandler.run(); + } + } + } +} 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 0a055a0e3..9de528abc 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 @@ -1,7 +1,7 @@ package tech.ydb.topic.read.impl; import java.util.ArrayList; -import java.util.LinkedList; +import java.util.Iterator; import java.util.List; import java.util.Queue; import java.util.concurrent.CompletableFuture; @@ -13,6 +13,7 @@ import org.slf4j.LoggerFactory; import tech.ydb.proto.topic.YdbTopic; +import tech.ydb.topic.description.Codec; import tech.ydb.topic.description.OffsetsRange; import tech.ydb.topic.read.Message; import tech.ydb.topic.read.PartitionSession; @@ -29,38 +30,36 @@ public abstract class ReadPartitionSession { private final ReadSession session; private final PartitionSession partition; private final int maxBatchSize; - private final MessageDecoder decoder; + private final BufferManager bufferManager; private final MessageCommitterImpl committer; + private final ReadPartitionDecoder decoder; private volatile long lastReadOffset; private volatile boolean isStopped = false; - private final Queue readingQueue = new ConcurrentLinkedQueue<>(); - private final AtomicBoolean isReadingNow = new AtomicBoolean(); + private final Queue readingQueue = new ConcurrentLinkedQueue<>(); + private final AtomicBoolean isReadingNow = new AtomicBoolean(false); ReadPartitionSession(String traceID, ReadSession session, PartitionSession partition, long lastCommittedOffset) { this.traceID = traceID; this.session = session; this.partition = partition; this.maxBatchSize = session.getMaxBatchSize(); - this.decoder = session.getMessageDecoder(); + this.bufferManager = session.getBufferManager(); this.committer = new MessageCommitterImpl(this, lastCommittedOffset); + this.decoder = new ReadPartitionDecoder(traceID, session.getMessageDecoder(), partition, committer, + this::sendDataToReadersIfNeeded); this.lastReadOffset = lastCommittedOffset; } - @Override - public String toString() { - return "[" + traceID + "]"; - } - public PartitionSession getPartition() { return partition; } boolean commitOffsets(List ranges) { if (isStopped) { - logger.info("{} Offset ranges {} are requested to be committed, but partition session is already closed", - this, ranges.stream().map(OffsetsRange::toString).collect(Collectors.joining(","))); + 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); @@ -74,38 +73,39 @@ void confirmCommit(long committedOffset) { public void stop() { isStopped = true; committer.failPendingCommits(); - // complete all read futures - for (Batch batch: readingQueue) { - batch.complete(); - } - logger.info("{} stopped", this); + decoder.close(); + logger.info("[{}] stopped", traceID); } abstract CompletableFuture handleDataReceivedEvent(DataReceivedEvent event); - public CompletableFuture addBatches(List batchList) { + public boolean addBatches(List batchList) { if (isStopped) { - return CompletableFuture.completedFuture(null); + return false; } - List> batchFutures = new LinkedList<>(); for (YdbTopic.StreamReadMessage.ReadResponse.Batch batch: batchList) { if (batch.getMessageDataCount() == 0) { - logger.error("{} Received empty batch. This shouldn't happen", this); - return CompletableFuture.completedFuture(null); + logger.error("[{}] Received empty batch. This shouldn't happen", traceID); + continue; } BatchMeta meta = new BatchMeta(batch); - List messages = new ArrayList<>(); + List messages = new ArrayList<>(batch.getMessageDataCount()); for (YdbTopic.StreamReadMessage.ReadResponse.MessageData msg: batch.getMessageDataList()) { if (lastReadOffset > msg.getOffset()) { - logger.error("{} Received a message with offset {} which is less than last read offset {} ", - this, msg.getOffset(), lastReadOffset); + logger.error("[{}] Received a message with offset {} which is less than last read offset {} ", + traceID, msg.getOffset(), lastReadOffset); lastReadOffset = msg.getOffset(); } OffsetsRange commitRange = OffsetsRange.of(lastReadOffset, msg.getOffset() + 1); - messages.add(new MessageImpl(partition, committer, meta, commitRange, msg)); + if (meta.getCodec() == Codec.RAW) { + messages.add(new RawMessage(meta, commitRange, msg)); + } else { + messages.add(decoder.decode(meta, commitRange, msg)); + } + lastReadOffset = commitRange.getEnd(); } @@ -114,21 +114,17 @@ public CompletableFuture addBatches(List[0])); + public void releaseRange(OffsetsRange range) { + decoder.releaseRange(range); + bufferManager.releaseRange(partition.getId(), range); + sendDataToReadersIfNeeded(); } public void sendDataToReadersIfNeeded() { @@ -136,61 +132,47 @@ public void sendDataToReadersIfNeeded() { return; } - if (isReadingNow.compareAndSet(false, true)) { - List batchesToRead = new ArrayList<>(); - - Batch next = readingQueue.peek(); - if (next == null || !next.isReady()) { + if (isReadingNow.compareAndSet(false, true)) { + Iterator it = readingQueue.iterator(); + if (!it.hasNext()) { isReadingNow.set(false); - if (next != null && next.isReady()) { - sendDataToReadersIfNeeded(); - } return; } - next = readingQueue.poll(); - - batchesToRead.add(next); - List messagesToRead = new ArrayList<>(next.getMessages()); - - int batchSize = messagesToRead.size(); - while (maxBatchSize <= 0 || batchSize < maxBatchSize) { - next = readingQueue.peek(); - if (next == null || !next.isReady()) { - break; - } - if (maxBatchSize > 0 && next.getMessages().size() + batchSize > maxBatchSize) { - break; - } - next = readingQueue.poll(); + MessageImpl next = it.next(); + if (!next.isReady()) { + isReadingNow.set(false); + return; + } - batchesToRead.add(next); - messagesToRead.addAll(next.getMessages()); - batchSize += next.getMessages().size(); + List messagesToRead = new ArrayList<>(); + while (next != null && next.isReady() && (maxBatchSize <= 0 || messagesToRead.size() < maxBatchSize)) { + messagesToRead.add(next); + it.remove(); + next = it.hasNext() ? it.next() : null; } // Should be called maximum in 1 thread at a time DataReceivedEvent event = new DataReceivedEventImpl(partition, committer, messagesToRead); - if (logger.isDebugEnabled()) { - logger.debug("[{}] DataReceivedEvent callback with {} message(s) (offsets {}-{}) is about " + - "to be called...", traceID, messagesToRead.size(), - messagesToRead.get(0).getOffset(), - messagesToRead.get(messagesToRead.size() - 1).getOffset()); - } + logger.debug("[{}] DataReceivedEvent callback with {} message(s) (offsets {}-{}) is about " + + "to be called...", traceID, messagesToRead.size(), + messagesToRead.get(0).getOffset(), + messagesToRead.get(messagesToRead.size() - 1).getOffset()); handleDataReceivedEvent(event).whenComplete((res, th) -> { if (th != null) { logger.error("[{}] DataReceivedEvent callback with {} message(s) (offsets {}-{}) finished" + " with error: ", traceID, messagesToRead.size(), messagesToRead.get(0).getOffset(), messagesToRead.get(messagesToRead.size() - 1).getOffset(), th); - } else if (logger.isDebugEnabled()) { + } else { logger.debug("[{}] DataReceivedEvent callback with {} message(s) (offsets {}-{}) " + "successfully finished", traceID, messagesToRead.size(), messagesToRead.get(0).getOffset(), messagesToRead.get(messagesToRead.size() - 1).getOffset()); } + + releaseRange(event.getRangeToCommit()); isReadingNow.set(false); - batchesToRead.forEach(Batch::complete); sendDataToReadersIfNeeded(); }); } else { @@ -199,4 +181,23 @@ public void sendDataToReadersIfNeeded() { } } } + + private class RawMessage extends MessageImpl { + private final byte[] data; + + RawMessage(BatchMeta meta, OffsetsRange range, YdbTopic.StreamReadMessage.ReadResponse.MessageData msg) { + super(partition, committer, meta, range, msg); + this.data = msg.getData().toByteArray(); + } + + @Override + public byte[] getData() { + return data; + } + + @Override + public boolean isReady() { + return true; + } + } } 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 ec7bc6c56..99ea562ea 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 @@ -2,13 +2,11 @@ import java.time.Duration; import java.time.Instant; -import java.util.ArrayList; import java.util.List; import java.util.Map; import java.util.UUID; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.atomic.AtomicLong; import java.util.stream.Collectors; import org.slf4j.Logger; @@ -48,13 +46,9 @@ public final class ReadSession extends SessionBase partitions = new ConcurrentHashMap<>(); private final Map partSessions = new ConcurrentHashMap<>(); @@ -64,9 +58,9 @@ public ReadSession(TopicRpc rpc, ReaderImpl reader, MessageDecoder decoder, Stri this.reader = reader; this.rpc = rpc; this.decoder = decoder; + this.bufferManager = new BufferManager(id, settings.getMaxMemoryUsageBytes(), this::sendReadRequest); this.consumerName = settings.getConsumerName(); - this.maxMemoryUsageBytes = settings.getMaxMemoryUsageBytes(); this.maxBatchSize = settings.getMaxBatchSize(); this.initRequest = buildInitRequest(settings); } @@ -84,6 +78,10 @@ MessageDecoder getMessageDecoder() { return decoder; } + BufferManager getBufferManager() { + return bufferManager; + } + @Override protected void sendUpdateTokenRequest(String token) { streamConnection.sendNext(YdbTopic.StreamReadMessage.FromClient.newBuilder() @@ -111,6 +109,8 @@ protected void onStop() { partitions.values().forEach(reader::handleClosePartitionSession); partitions.clear(); + + decoder.stop(); } protected void closeDueToError(Status status, Throwable th) { @@ -121,18 +121,12 @@ protected void closeDueToError(Status status, Throwable th) { } } - private void sendReadRequest() { - long currentSizeBytesToRequest = sizeBytesToRequest.getAndSet(0); - if (currentSizeBytesToRequest <= 0) { - logger.debug("[{}] Nothing to request in DataRequest. sizeBytesToRequest == {}", streamId, - currentSizeBytesToRequest); - return; - } - logger.debug("[{}] Sending DataRequest with {} bytes", streamId, currentSizeBytesToRequest); + private void sendReadRequest(long sizeToRequest) { + logger.debug("[{}] Sending DataRequest with {} bytes", streamId, sizeToRequest); send(YdbTopic.StreamReadMessage.FromClient.newBuilder() .setReadRequest(YdbTopic.StreamReadMessage.ReadRequest.newBuilder() - .setBytesSize(currentSizeBytesToRequest) + .setBytesSize(sizeToRequest) .build()) .build()); } @@ -163,12 +157,8 @@ void sendCommitOffsetRequest(PartitionSession session, List ranges } private void onInitResponse(YdbTopic.StreamReadMessage.InitResponse response) { - String sessionId = response.getSessionId(); - - sizeBytesToRequest.set(maxMemoryUsageBytes); - logger.info("[{}] Session {} initialized. Requesting {} bytes...", streamId, sessionId, maxMemoryUsageBytes); - reader.onSessionStarted(sessionId); - sendReadRequest(); + reader.onSessionStarted(response.getSessionId()); + bufferManager.init(response.getSessionId()); } private void onStartPartitionSessionRequest(YdbTopic.StreamReadMessage.StartPartitionSessionRequest req) { @@ -252,6 +242,7 @@ protected void onStopPartitionSessionRequest(YdbTopic.StreamReadMessage.StopPart if (rps != null) { logger.info("[{}] Received force StopPartitionSessionRequest for {} ", streamId, rps.getPartition()); rps.stop(); + bufferManager.releasePartition(psid); } reader.handleClosePartitionSession(partition); @@ -292,46 +283,27 @@ public void confirm() { session.stop(); } } + + bufferManager.releasePartition(psid); } }); } private void onReadResponse(YdbTopic.StreamReadMessage.ReadResponse response) { - final long responseBytesSize = response.getBytesSize(); - logger.debug("[{}] Received ReadResponse of {} bytes", streamId, responseBytesSize); - List> batchReadFutures = new ArrayList<>(); + logger.debug("[{}] Received ReadResponse of {} bytes", streamId, response.getBytesSize()); + bufferManager.allocate((int) response.getBytesSize(), response.getPartitionDataList()); for (YdbTopic.StreamReadMessage.ReadResponse.PartitionData data: response.getPartitionDataList()) { long psid = data.getPartitionSessionId(); - ReadPartitionSession session = partSessions.get(data.getPartitionSessionId()); - if (session == null) { + ReadPartitionSession session = partSessions.get(psid); + if (session == null || !session.addBatches(data.getBatchesList())) { logger.warn("[{}] Received PartitionData for unknown(most likely already closed) PartitionSessionId={}", streamId, psid); - // TODO: release memory buffer - continue; + bufferManager.releasePartition(psid); } - - // Completes when all messages from a batch are read by user - batchReadFutures.add(session.addBatches(data.getBatchesList())); } - CompletableFuture.allOf(batchReadFutures.toArray(new CompletableFuture[0])) - .whenComplete((res, th) -> { - if (th != null) { - logger.error("[{}] Exception while waiting for batches to be read:", streamId, th); - return; - } - if (isStopped()) { - logger.trace("[{}] Finished handling ReadResponse of {} bytes. Read session is already " + - "closed -- no need to send ReadRequest", streamId, responseBytesSize); - return; - } - - logger.trace("[{}] Finished handling ReadResponse of {} bytes. Sending ReadRequest...", - streamId, responseBytesSize); - this.sizeBytesToRequest.addAndGet(responseBytesSize); - sendReadRequest(); - }); + decoder.decodeNext(); } protected void onCommitOffsetResponse(YdbTopic.StreamReadMessage.CommitOffsetResponse response) { From 520f48e2cb9b81fda11b8da237562a59e8024c2c Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Mon, 10 Aug 2026 15:25:15 +0100 Subject: [PATCH 03/12] Added unit tests for MessageDecoder & ReadPartitionDecoder --- .../ydb/topic/read/impl/MessageDecoder.java | 15 +- .../topic/read/impl/ReadPartitionDecoder.java | 14 +- .../topic/read/impl/MessageDecoderTest.java | 394 ++++++++++++++++++ 3 files changed, 416 insertions(+), 7 deletions(-) create mode 100644 topic/src/test/java/tech/ydb/topic/read/impl/MessageDecoderTest.java 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 79fc28295..5d06bae5e 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 @@ -5,6 +5,9 @@ import java.util.concurrent.Executor; import java.util.concurrent.atomic.AtomicLong; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + import tech.ydb.topic.description.CodecRegistry; import tech.ydb.topic.impl.SerialRunnable; @@ -13,6 +16,7 @@ * @author Aleksandr Gorshenin */ public class MessageDecoder { + private static final Logger logger = LoggerFactory.getLogger(MessageDecoder.class); private final AtomicLong totalAvailable; private final Executor decompressionExecutor; @@ -55,8 +59,15 @@ public void run() { return; } - totalAvailable.addAndGet(-next.allocate()); - decompressionExecutor.execute(() -> next.decode(codecRegistry)); + long size = next.allocate(); + totalAvailable.addAndGet(-size); + try { + decompressionExecutor.execute(() -> next.decode(codecRegistry)); + } catch (Throwable ex) { + logger.error("Cannot execute decompression ", ex); + next.setError(ex); + totalAvailable.addAndGet(size); + } } } } diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionDecoder.java b/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionDecoder.java index 896c55e06..3f31eeb41 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionDecoder.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionDecoder.java @@ -25,7 +25,7 @@ * @author Aleksandr Gorshenin {@literal } */ public class ReadPartitionDecoder { - private static final Logger logger = LoggerFactory.getLogger(ReaderImpl.class); + private static final Logger logger = LoggerFactory.getLogger(MessageDecoder.class); private final String traceID; private final PartitionSession partition; @@ -88,10 +88,7 @@ public void close() { } private void release() { - long total = allocatedTotal.getAndSet(0); - if (total > 0) { - decoder.free(total); - } + decoder.free(allocatedTotal.getAndSet(0)); } public class EncodedMessage extends MessageImpl { @@ -124,8 +121,15 @@ public boolean isReady() { return isReady; } + public void setError(Throwable th) { + problem = new IOException("Decompression for " + getPartitionSession() + " error", th); + isReady = true; + readyHandler.run(); + } + public long allocate() { if (isStopped) { + problem = new IOException("" + getPartitionSession() + " is already closed"); isReady = true; return 0; } diff --git a/topic/src/test/java/tech/ydb/topic/read/impl/MessageDecoderTest.java b/topic/src/test/java/tech/ydb/topic/read/impl/MessageDecoderTest.java new file mode 100644 index 000000000..cb3bc7a15 --- /dev/null +++ b/topic/src/test/java/tech/ydb/topic/read/impl/MessageDecoderTest.java @@ -0,0 +1,394 @@ +package tech.ydb.topic.read.impl; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.util.Arrays; +import java.util.Queue; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.Executor; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.zip.GZIPOutputStream; + +import com.google.protobuf.ByteString; +import org.junit.Assert; +import org.junit.Rule; +import org.junit.Test; +import org.junit.function.ThrowingRunnable; + +import tech.ydb.proto.topic.YdbTopic.StreamReadMessage.ReadResponse; +import tech.ydb.topic.description.Codec; +import tech.ydb.topic.description.CodecRegistry; +import tech.ydb.topic.description.OffsetsRange; +import tech.ydb.topic.read.DecompressionException; +import tech.ydb.topic.read.PartitionSession; +import tech.ydb.topic.utils.HideLoggers; +import tech.ydb.topic.utils.HideLoggersRule; + +/** + * Unit tests for {@link MessageDecoder} flow control and scheduling. + */ +public class MessageDecoderTest { + private static final CodecRegistry REGISTRY = new CodecRegistry(); + private static final PartitionSession PS1 = new PartitionSession(1, 1, "/topic"); + private static final PartitionSession PS2 = new PartitionSession(2, 2, "/topic"); + + private static void assertDecompressionException(String msg, ThrowingRunnable runnable) { + DecompressionException ex = Assert.assertThrows(DecompressionException.class, runnable); + Assert.assertEquals(msg, ex.getCause().getMessage()); + } + + private static byte[] gzip(byte[] data) { + try { + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + try (GZIPOutputStream gz = new GZIPOutputStream(baos)) { + gz.write(data); + } + return baos.toByteArray(); + } catch (IOException ex) { + throw new RuntimeException(ex); + } + } + + private static ReadResponse.MessageData rawMsg(long offset, long uncompressedSize, byte[] data) { + return ReadResponse.MessageData.newBuilder() + .setOffset(offset) + .setUncompressedSize(uncompressedSize) + .setData(ByteString.copyFrom(data)) + .build(); + } + + private static ReadResponse.MessageData gzipMsg(int offset, int uncompressedSize) { + byte[] data = new byte[uncompressedSize]; + Arrays.fill(data, (byte) offset); + return ReadResponse.MessageData.newBuilder() + .setOffset(offset) + .setUncompressedSize(uncompressedSize) + .setData(ByteString.copyFrom(gzip(data))) + .build(); + } + + private static BatchMeta meta(int codec) { + return new BatchMeta(ReadResponse.Batch.newBuilder().setCodec(codec).build()); + } + + @Rule + public final HideLoggersRule hideLoggerRule = new HideLoggersRule(); + + @Test + public void rawDecodeTest() { + MessageDecoder decoder = new MessageDecoder(10000, Runnable::run, REGISTRY); + AtomicInteger ready = new AtomicInteger(0); + ReadPartitionDecoder partition = new ReadPartitionDecoder("t1", decoder, PS1, null, ready::incrementAndGet); + + BatchMeta meta = meta(Codec.RAW); + MessageImpl m1 = partition.decode(meta, OffsetsRange.of(1), rawMsg(1, 100, new byte[] {1, 2})); + MessageImpl m2 = partition.decode(meta, OffsetsRange.of(2), rawMsg(2, 100, new byte[] {3})); + + Assert.assertFalse(m1.isReady()); + Assert.assertFalse(m2.isReady()); + + decoder.decodeNext(); + + Assert.assertTrue(m1.isReady()); + Assert.assertTrue(m2.isReady()); + Assert.assertArrayEquals(new byte[] {1, 2}, m1.getData()); + Assert.assertArrayEquals(new byte[] {3}, m2.getData()); + Assert.assertEquals(2, ready.get()); + } + + @Test + public void flowControlByBudgetTest() { + MessageDecoder decoder = new MessageDecoder(100, Runnable::run, REGISTRY); + + AtomicInteger ready = new AtomicInteger(0); + ReadPartitionDecoder p1 = new ReadPartitionDecoder("p1", decoder, PS1, null, ready::incrementAndGet); + + BatchMeta meta = meta(Codec.GZIP); + MessageImpl m1 = p1.decode(meta, OffsetsRange.of(1), gzipMsg(1, 40)); + MessageImpl m2 = p1.decode(meta, OffsetsRange.of(2), gzipMsg(2, 50)); + MessageImpl m3 = p1.decode(meta, OffsetsRange.of(3), gzipMsg(3, 60)); + MessageImpl m4 = p1.decode(meta, OffsetsRange.of(4), gzipMsg(4, 70)); + MessageImpl m5 = p1.decode(meta, OffsetsRange.of(5), gzipMsg(5, 120)); + + Assert.assertFalse(m1.isReady()); + Assert.assertFalse(m2.isReady()); + Assert.assertFalse(m3.isReady()); + Assert.assertFalse(m4.isReady()); + Assert.assertFalse(m5.isReady()); + Assert.assertEquals(0, ready.get()); + + decoder.decodeNext(); + + Assert.assertTrue(m1.isReady()); + Assert.assertTrue(m2.isReady()); + Assert.assertTrue(m3.isReady()); + Assert.assertFalse(m4.isReady()); + Assert.assertFalse(m5.isReady()); + Assert.assertEquals(3, ready.get()); + + Assert.assertEquals(40, m1.getData().length); + Assert.assertEquals(50, m2.getData().length); + Assert.assertEquals(60, m3.getData().length); + Assert.assertEquals(1, m1.getData()[0]); + Assert.assertEquals(2, m2.getData()[0]); + Assert.assertEquals(3, m3.getData()[0]); + + p1.releaseRange(OffsetsRange.of(1)); // 40 is not enough to resume decoding + p1.releaseRange(OffsetsRange.of(4)); // that offset is not decodec yet + Assert.assertFalse(m4.isReady()); + Assert.assertFalse(m5.isReady()); + Assert.assertEquals(3, ready.get()); + + p1.releaseRange(OffsetsRange.of(2)); + Assert.assertTrue(m4.isReady()); + Assert.assertFalse(m5.isReady()); + Assert.assertEquals(4, ready.get()); + + Assert.assertEquals(70, m4.getData().length); + Assert.assertEquals(4, m4.getData()[0]); + + p1.releaseRange(OffsetsRange.of(0, 3)); // double release + Assert.assertFalse(m5.isReady()); + Assert.assertEquals(4, ready.get()); + + p1.releaseRange(OffsetsRange.of(0, 5)); + Assert.assertTrue(m5.isReady()); + Assert.assertEquals(5, ready.get()); + } + + @Test + public void partitionFlowTest() { + MessageDecoder decoder = new MessageDecoder(100, Runnable::run, REGISTRY); + + AtomicInteger r1 = new AtomicInteger(0); + ReadPartitionDecoder p1 = new ReadPartitionDecoder("p1", decoder, PS1, null, r1::incrementAndGet); + AtomicInteger r2 = new AtomicInteger(0); + ReadPartitionDecoder p2 = new ReadPartitionDecoder("p2", decoder, PS2, null, r2::incrementAndGet); + + BatchMeta meta = meta(Codec.GZIP); + MessageImpl m1 = p1.decode(meta, OffsetsRange.of(1), gzipMsg(1, 40)); + MessageImpl m2 = p1.decode(meta, OffsetsRange.of(2), gzipMsg(2, 50)); + MessageImpl m3 = p2.decode(meta, OffsetsRange.of(10), gzipMsg(1, 60)); + MessageImpl m4 = p2.decode(meta, OffsetsRange.of(11), gzipMsg(2, 40)); + MessageImpl m5 = p2.decode(meta, OffsetsRange.of(12), gzipMsg(3, 30)); + + Assert.assertEquals(0, r1.get()); + Assert.assertEquals(0, r2.get()); + + decoder.decodeNext(); + + MessageImpl m6 = p1.decode(meta, OffsetsRange.of(4), gzipMsg(4, 10)); + MessageImpl m7 = p1.decode(meta, OffsetsRange.of(5), gzipMsg(5, 20)); + MessageImpl m8 = p2.decode(meta, OffsetsRange.of(14), gzipMsg(14, 10)); + MessageImpl m9 = p2.decode(meta, OffsetsRange.of(15), gzipMsg(15, 20)); + + decoder.decodeNext(); + + Assert.assertTrue(m1.isReady()); + Assert.assertTrue(m2.isReady()); + Assert.assertTrue(m3.isReady()); + Assert.assertFalse(m4.isReady()); + Assert.assertFalse(m5.isReady()); + Assert.assertFalse(m6.isReady()); + Assert.assertFalse(m7.isReady()); + Assert.assertFalse(m8.isReady()); + Assert.assertFalse(m9.isReady()); + + Assert.assertEquals(2, r1.get()); + Assert.assertEquals(1, r2.get()); + + p1.close(); + + Assert.assertTrue(m1.isReady()); + Assert.assertTrue(m2.isReady()); + Assert.assertTrue(m3.isReady()); + Assert.assertTrue(m4.isReady()); + Assert.assertFalse(m5.isReady()); + Assert.assertFalse(m6.isReady()); + Assert.assertFalse(m7.isReady()); + Assert.assertFalse(m8.isReady()); + Assert.assertFalse(m9.isReady()); + + Assert.assertEquals(2, r1.get()); + Assert.assertEquals(2, r2.get()); + + p1.releaseRange(OffsetsRange.of(0, 100)); + p2.releaseRange(OffsetsRange.of(0, 100)); + + Assert.assertTrue(m5.isReady()); + Assert.assertTrue(m6.isReady()); + Assert.assertTrue(m7.isReady()); + Assert.assertTrue(m8.isReady()); + Assert.assertTrue(m9.isReady()); + + Assert.assertEquals(2, r1.get()); // p1 is already stopped + Assert.assertEquals(5, r2.get()); +} + + @Test + public void decodeStopTest() { + MessageDecoder decoder = new MessageDecoder(70, Runnable::run, REGISTRY); + + AtomicInteger ready = new AtomicInteger(0); + ReadPartitionDecoder partition = new ReadPartitionDecoder("t3", decoder, PS1, null, ready::incrementAndGet); + + BatchMeta meta = meta(Codec.GZIP); + MessageImpl m1 = partition.decode(meta, OffsetsRange.of(1), gzipMsg(1, 40)); + MessageImpl m2 = partition.decode(meta, OffsetsRange.of(2), gzipMsg(2, 50)); + MessageImpl m3 = partition.decode(meta, OffsetsRange.of(3), gzipMsg(3, 60)); + + Assert.assertEquals(0, ready.get()); + + decoder.decodeNext(); + Assert.assertEquals(2, ready.get()); + Assert.assertTrue(m1.isReady()); + Assert.assertTrue(m2.isReady()); + Assert.assertFalse(m3.isReady()); + + decoder.stop(); + Assert.assertEquals(2, ready.get()); + Assert.assertFalse(m3.isReady()); + + partition.releaseRange(OffsetsRange.of(0, 10)); + Assert.assertEquals(2, ready.get()); + Assert.assertFalse(m3.isReady()); + } + + @Test + public void decodesOnProvidedExecutorTest() { + Queue decodeTasks = new ConcurrentLinkedQueue<>(); + MessageDecoder decoder = new MessageDecoder(1000, decodeTasks::add, REGISTRY); + + AtomicInteger p1ready = new AtomicInteger(); + AtomicInteger p2ready = new AtomicInteger(); + ReadPartitionDecoder p1 = new ReadPartitionDecoder("p1", decoder, PS1, null, p1ready::incrementAndGet); + ReadPartitionDecoder p2 = new ReadPartitionDecoder("p2", decoder, PS2, null, p2ready::incrementAndGet); + + BatchMeta meta = meta(Codec.GZIP); + MessageImpl m1 = p1.decode(meta, OffsetsRange.of(1), gzipMsg(1, 400)); + MessageImpl m2 = p2.decode(meta, OffsetsRange.of(1), gzipMsg(1, 500)); + MessageImpl m3 = p2.decode(meta, OffsetsRange.of(2), gzipMsg(2, 600)); + MessageImpl m4 = p1.decode(meta, OffsetsRange.of(2), gzipMsg(2, 400)); + MessageImpl m5 = p1.decode(meta, OffsetsRange.of(3), gzipMsg(3, 400)); + + Assert.assertFalse(m1.isReady()); + Assert.assertFalse(m2.isReady()); + Assert.assertFalse(m3.isReady()); + Assert.assertFalse(m4.isReady()); + Assert.assertFalse(m5.isReady()); + Assert.assertEquals(0, decodeTasks.size()); + Assert.assertEquals(0, p1ready.get()); + Assert.assertEquals(0, p2ready.get()); + + decoder.decodeNext(); + + Assert.assertFalse(m1.isReady()); + Assert.assertFalse(m2.isReady()); + Assert.assertEquals(3, decodeTasks.size()); + Assert.assertEquals(0, p1ready.get()); + Assert.assertEquals(0, p2ready.get()); + + decodeTasks.poll().run(); + + Assert.assertTrue(m1.isReady()); + Assert.assertFalse(m2.isReady()); + Assert.assertEquals(2, decodeTasks.size()); + Assert.assertEquals(1, p1ready.get()); + Assert.assertEquals(0, p2ready.get()); + + decodeTasks.poll().run(); + + Assert.assertTrue(m1.isReady()); + Assert.assertTrue(m2.isReady()); + Assert.assertEquals(1, decodeTasks.size()); + Assert.assertEquals(1, p1ready.get()); + Assert.assertEquals(1, p2ready.get()); + + decoder.stop(); + p1.close(); + p2.close(); + + Assert.assertEquals(1, decodeTasks.size()); + decodeTasks.poll().run(); + Assert.assertEquals(0, decodeTasks.size()); + + Assert.assertEquals(1, p1ready.get()); + Assert.assertEquals(1, p2ready.get()); + } + + @Test + @HideLoggers(MessageDecoder.class) + public void decodeProblemsTest() { + MessageDecoder decoder = new MessageDecoder(200, Runnable::run, REGISTRY); + + AtomicInteger ready = new AtomicInteger(); + ReadPartitionDecoder p1 = new ReadPartitionDecoder("p1", decoder, PS1, null, ready::incrementAndGet); + + BatchMeta meta1 = meta(1244); // Uknown codec + MessageImpl m1 = p1.decode(meta1, OffsetsRange.of(1), rawMsg(1, 50, new byte[0])); + MessageImpl m2 = p1.decode(meta1, OffsetsRange.of(2), rawMsg(2, 50, new byte[0])); + + BatchMeta meta2 = meta(Codec.GZIP); + MessageImpl m3 = p1.decode(meta2, OffsetsRange.of(3), rawMsg(3, 50, new byte[] { 0x1, 0x2 })); + MessageImpl m4 = p1.decode(meta2, OffsetsRange.of(4), rawMsg(4, 50, new byte[] { 0x1, 0x2 })); + + MessageImpl m5 = p1.decode(meta2, OffsetsRange.of(5), gzipMsg(5, 50)); + MessageImpl m6 = p1.decode(meta2, OffsetsRange.of(6), gzipMsg(6, 50)); + + decoder.decodeNext(); + + Assert.assertEquals(4, ready.get()); + + Assert.assertTrue(m1.isReady()); + Assert.assertTrue(m2.isReady()); + Assert.assertTrue(m3.isReady()); + Assert.assertTrue(m4.isReady()); + Assert.assertFalse(m5.isReady()); + Assert.assertFalse(m6.isReady()); + + assertDecompressionException("Codec 1244 is not registered", m1::getData); + assertDecompressionException("Codec 1244 is not registered", m2::getData); + assertDecompressionException("Not in GZIP format", m3::getData); + assertDecompressionException("Not in GZIP format", m4::getData); + + p1.close(); + + Assert.assertEquals(4, ready.get()); + Assert.assertTrue(m5.isReady()); + Assert.assertTrue(m6.isReady()); + assertDecompressionException("Partition session 1 (partition 1) for topic \"/topic\" is already closed", + m5::getData); + assertDecompressionException("Partition session 1 (partition 1) for topic \"/topic\" is already closed", + m6::getData); + } + + @Test + @HideLoggers(MessageDecoder.class) + public void wrongDecoderTest() { + Executor decompressor = (Runnable command) -> { + throw new RejectedExecutionException("rejected"); + }; + + MessageDecoder decoder = new MessageDecoder(200, decompressor, REGISTRY); + + AtomicInteger ready = new AtomicInteger(); + ReadPartitionDecoder p1 = new ReadPartitionDecoder("p1", decoder, PS1, null, ready::incrementAndGet); + + BatchMeta meta = meta(Codec.GZIP); + MessageImpl m1 = p1.decode(meta, OffsetsRange.of(5), gzipMsg(5, 50)); + MessageImpl m2 = p1.decode(meta, OffsetsRange.of(6), gzipMsg(6, 50)); + + decoder.decodeNext(); + + Assert.assertEquals(2, ready.get()); + + Assert.assertTrue(m1.isReady()); + Assert.assertTrue(m2.isReady()); + + assertDecompressionException("Decompression for Partition session 1 (partition 1) for topic \"/topic\" error", + m1::getData); + assertDecompressionException("Decompression for Partition session 1 (partition 1) for topic \"/topic\" error", + m2::getData); + } +} From 81472d2b7a7b38dc2abeaaa83ed9c872ba4320cb Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Mon, 10 Aug 2026 20:57:33 +0100 Subject: [PATCH 04/12] Fixed integer overflow in BufferManager --- .../tech/ydb/topic/read/impl/BufferManager.java | 10 +++++----- .../topic/read/impl/ReadPartitionSession.java | 5 +++++ .../ydb/topic/read/impl/BufferManagerTest.java | 17 +++++++++++++++++ 3 files changed, 27 insertions(+), 5 deletions(-) diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/BufferManager.java b/topic/src/main/java/tech/ydb/topic/read/impl/BufferManager.java index c554785ee..e1bc8e9c5 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/BufferManager.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/BufferManager.java @@ -131,17 +131,17 @@ private static void recalcBuffer(int[] buffer, int buffSize) { return; } - int total = 0; + long total = 0; for (int v: buffer) { total += v; } - int currBuff = 0; - int currSum = 0; + long currBuff = 0; + long currSum = 0; for (int idx = 0; idx < buffer.length; idx += 1) { currSum += buffer[idx]; - int newBuff = currSum * buffSize / total; - buffer[idx] = newBuff - currBuff; + long newBuff = currSum * buffSize / total; + buffer[idx] = (int) (newBuff - currBuff); currBuff = newBuff; } } 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 9de528abc..6897a7cf4 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 @@ -56,6 +56,11 @@ public PartitionSession getPartition() { return partition; } + @Override + public String toString() { + return "[" + traceID + "]"; + } + boolean commitOffsets(List ranges) { if (isStopped) { logger.info("[{}] Offset ranges {} are requested to be committed, but partition session is already closed", diff --git a/topic/src/test/java/tech/ydb/topic/read/impl/BufferManagerTest.java b/topic/src/test/java/tech/ydb/topic/read/impl/BufferManagerTest.java index 779f25a0c..9868a21b1 100644 --- a/topic/src/test/java/tech/ydb/topic/read/impl/BufferManagerTest.java +++ b/topic/src/test/java/tech/ydb/topic/read/impl/BufferManagerTest.java @@ -154,4 +154,21 @@ public void concurrentReleaseTest() { bm.releasePartition(1L); Assert.assertEquals(100000L, requested.get()); } + + @Test + public void intOverflowTest() { + AtomicLong requested = new AtomicLong(0); + BufferManager bm = new BufferManager("trace-4", 1000, requested::addAndGet); + + bm.allocate(100000, Arrays.asList( + partition(1, batch(1, 10000, 10000, 10000)), + partition(2, batch(1, 20000, 30000, 10000)), + partition(3, batch(1, 40000, 40000, 40000)) + )); + + bm.releasePartition(1L); + bm.releasePartition(2L); + bm.releasePartition(3L); + Assert.assertEquals(100000, requested.get()); + } } From 8810b049eedfa42cd7aa37caa32f016fef49024b Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Mon, 10 Aug 2026 20:58:45 +0100 Subject: [PATCH 05/12] Extends TopicReadersIntegrationTest --- .../topic/TopicReadersIntegrationTest.java | 295 +++++++++++++----- 1 file changed, 209 insertions(+), 86 deletions(-) diff --git a/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java b/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java index 3c031aa6e..d57f3161e 100644 --- a/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java +++ b/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java @@ -3,12 +3,13 @@ import java.util.Arrays; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicLong; -import org.junit.After; import org.junit.AfterClass; import org.junit.Assert; import org.junit.Before; @@ -16,10 +17,10 @@ import org.junit.ClassRule; import org.junit.Rule; import org.junit.Test; -import org.junit.rules.Timeout; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import tech.ydb.core.Status; import tech.ydb.test.junit4.GrpcTransportRule; import tech.ydb.topic.description.Consumer; import tech.ydb.topic.read.AsyncReader; @@ -29,7 +30,9 @@ import tech.ydb.topic.read.events.StartPartitionSessionEvent; import tech.ydb.topic.read.impl.AsyncReaderImpl; import tech.ydb.topic.read.impl.ReaderImpl; +import tech.ydb.topic.settings.CommitOffsetSettings; import tech.ydb.topic.settings.CreateTopicSettings; +import tech.ydb.topic.settings.PartitioningSettings; import tech.ydb.topic.settings.ReadEventHandlersSettings; import tech.ydb.topic.settings.ReaderSettings; import tech.ydb.topic.settings.StartPartitionSessionSettings; @@ -49,9 +52,6 @@ public class TopicReadersIntegrationTest { @ClassRule public final static GrpcTransportRule ydbTransport = new GrpcTransportRule(); - @Rule - public final Timeout timeout = new Timeout(10, TimeUnit.SECONDS); - @Rule public final HideLoggersRule hideLogger = new HideLoggersRule(); @@ -64,42 +64,71 @@ public class TopicReadersIntegrationTest { @BeforeClass public static void initClient() { client = TopicClient.newClient(ydbTransport).build(); + + logger.info("Create test topic {} ...", TEST_TOPIC); + client.createTopic(TEST_TOPIC, CreateTopicSettings.newBuilder() + .addConsumer(Consumer.newBuilder().setName(TEST_CONSUMER1).build()) + .setPartitioningSettings(PartitioningSettings.newBuilder() + .setMinActivePartitions(3) + .setMaxActivePartitions(3) + .build()) + .build() + ).join().expectSuccess("can't create a new topic"); + + CompletableFuture f1 = CompletableFuture.runAsync(() -> writeToTopic(0, 1000)); + CompletableFuture f2 = CompletableFuture.runAsync(() -> writeToTopic(1, 500)); + CompletableFuture f3 = CompletableFuture.runAsync(() -> writeToTopic(2, 2100)); + + CompletableFuture.allOf(f1, f2, f3).join(); } @AfterClass public static void closeClient() { + logger.info("Drop test topic {} ...", TEST_TOPIC); + client.dropTopic(TEST_TOPIC).join(); client.close(); } @Before - public void initTopic() throws Exception { - logger.info("Create test topic {} ...", TEST_TOPIC); - client.createTopic(TEST_TOPIC, CreateTopicSettings.newBuilder() - .addConsumer(Consumer.newBuilder().setName(TEST_CONSUMER1).build()) - .build() - ).join().expectSuccess("can't create a new topic"); + public void resetConsumer() { + CompletableFuture r1 = resetPartition(0); + CompletableFuture r2 = resetPartition(1); + CompletableFuture r3 = resetPartition(2); + r1.join().expectSuccess(); + r2.join().expectSuccess(); + r3.join().expectSuccess(); + } - // send 6 messages with offsets 0-5 - String[] messages = new String[] { "test1", "test2", "test3", "test4", "test5", "stop" }; + private static CompletableFuture resetPartition(int partitionID) { + return client.commitOffset(TEST_TOPIC, CommitOffsetSettings.newBuilder() + .setConsumer(TEST_CONSUMER1) + .setOffset(0) + .setPartitionId(partitionID) + .build()); + } + private static void writeToTopic(int partitionID, int count) { WriterSettings settings = WriterSettings.newBuilder() .setTopicPath(TEST_TOPIC) - .setProducerId("helper") + .setProducerId("p" + partitionID) + .setPartitionId(partitionID) .build(); + SyncWriter writer = client.createSyncWriter(settings); writer.initAndWait(); - for (String message : messages) { - writer.send(tech.ydb.topic.write.Message.of(message.getBytes())); + for (int idx = 1; idx <= count; idx++) { + byte[] msg = ("p" + partitionID + "_msg" + idx).getBytes(); + byte[] data = new byte[100]; + System.arraycopy(msg, 0, data, 0, msg.length); + writer.send(tech.ydb.topic.write.Message.of(data)); } - writer.flush(); - writer.shutdown(10, TimeUnit.SECONDS); - } - - @After - public void dropTopic() { - logger.info("Drop test topic {} ...", TEST_TOPIC); - client.dropTopic(TEST_TOPIC).join().expectSuccess("can't drop test topic"); + try { + writer.flush(); + writer.shutdown(10, TimeUnit.SECONDS); + } catch (InterruptedException | ExecutionException | TimeoutException ex) { + throw new AssertionError("cannot write", ex); + } } @Test @@ -144,20 +173,26 @@ public void readAllTest() throws InterruptedException { .setConsumerName(TEST_CONSUMER1) .build(); - AtomicLong offset = new AtomicLong(); - CountDownLatch read = new CountDownLatch(6); + AtomicLong[] offsets = new AtomicLong[] { new AtomicLong(), new AtomicLong(), new AtomicLong() }; + CountDownLatch read = new CountDownLatch(3600); AsyncReader reader = client.createAsyncReader(readerSettings, ReadEventHandlersSettings.newBuilder() .setEventHandler((DataReceivedEvent event) -> { - for (Message msg: event.getMessages()) { - Assert.assertEquals(offset.getAndIncrement(), msg.getOffset()); - read.countDown(); - } - }).build()); + AtomicLong offset = offsets[(int) event.getPartitionSession().getPartitionId()]; + for (Message msg : event.getMessages()) { + Assert.assertEquals(offset.getAndIncrement(), msg.getOffset()); + read.countDown(); + } + }).build()); reader.init().join(); - Assert.assertTrue(read.await(5, TimeUnit.SECONDS)); - Assert.assertEquals(6, offset.get()); - reader.shutdown().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(); + } } @Test @@ -171,7 +206,7 @@ public void readAllByPartitionIdTest() throws InterruptedException { .build(); AtomicLong offset = new AtomicLong(); - CountDownLatch read = new CountDownLatch(6); + CountDownLatch read = new CountDownLatch(1000); AsyncReader reader = client.createAsyncReader(readerSettings, ReadEventHandlersSettings.newBuilder() .setEventHandler((DataReceivedEvent event) -> { for (Message msg: event.getMessages()) { @@ -181,9 +216,12 @@ public void readAllByPartitionIdTest() throws InterruptedException { }).build()); reader.init().join(); - Assert.assertTrue(read.await(5, TimeUnit.SECONDS)); - Assert.assertEquals(6, offset.get()); - reader.shutdown().join(); + try { + Assert.assertTrue(read.await(30, TimeUnit.SECONDS)); + Assert.assertEquals(1000, offset.get()); + } finally { + reader.shutdown().join(); + } } @Test @@ -197,7 +235,7 @@ public void readAllWithoutConsumerTest() throws InterruptedException { .build(); AtomicLong offset = new AtomicLong(); - CountDownLatch read = new CountDownLatch(6); + CountDownLatch read = new CountDownLatch(1000); AsyncReader reader = client.createAsyncReader(readerSettings, ReadEventHandlersSettings.newBuilder() .setEventHandler((DataReceivedEvent event) -> { for (Message msg: event.getMessages()) { @@ -206,130 +244,215 @@ public void readAllWithoutConsumerTest() throws InterruptedException { } }).build()); + reader.init().join(); - Assert.assertTrue(read.await(5, TimeUnit.SECONDS)); - Assert.assertEquals(6, offset.get()); - reader.shutdown().join(); + try { + Assert.assertTrue(read.await(30, TimeUnit.SECONDS)); + Assert.assertEquals(1000, offset.get()); + } finally { + reader.shutdown().join(); + } } @Test public void readFromTest() throws Exception { ReaderSettings readerSettings = ReaderSettings.newBuilder() - .addTopic(TopicReadSettings.newBuilder().setPath(TEST_TOPIC).build()) + .addTopic(TopicReadSettings.newBuilder() + .setPath(TEST_TOPIC) + .setPartitionIds(Arrays.asList(1L)) + .build()) .setConsumerName(TEST_CONSUMER1) .build(); - AtomicLong counter = new AtomicLong(); - CompletableFuture result = new CompletableFuture<>(); + AtomicLong offset = new AtomicLong(123L); + CountDownLatch read = new CountDownLatch(500-123); AsyncReader reader = client.createAsyncReader(readerSettings, ReadEventHandlersSettings.newBuilder() .setEventHandler(new ReadEventHandler() { @Override public void onStartPartitionSession(StartPartitionSessionEvent event) { Assert.assertEquals(0, event.getCommittedOffset()); Assert.assertEquals(0, event.getPartitionOffsets().getStart()); - Assert.assertEquals(6, event.getPartitionOffsets().getEnd()); + Assert.assertEquals(500, event.getPartitionOffsets().getEnd()); - // read only from offset 2 - event.confirm(StartPartitionSessionSettings.newBuilder().setReadOffset(2L).build()); + // read only from offset 123 + event.confirm(StartPartitionSessionSettings.newBuilder().setReadOffset(123L).build()); } @Override public void onMessages(DataReceivedEvent event) { for (Message msg : event.getMessages()) { - Assert.assertEquals(msg.getOffset(), 2 + counter.get()); - long read = counter.incrementAndGet(); - if (new String(msg.getData()).equals("stop")) { - result.complete(read); - } + Assert.assertEquals(offset.getAndIncrement(), msg.getOffset()); + read.countDown(); } } }).build()); reader.init().join(); - Assert.assertEquals(Long.valueOf(4), result.join()); - reader.shutdown().join(); + try { + Assert.assertTrue(read.await(30, TimeUnit.SECONDS)); + Assert.assertEquals(500, offset.get()); + } finally { + reader.shutdown().join(); + } } @Test public void readFromWithCommitTest() throws Exception { ReaderSettings readerSettings = ReaderSettings.newBuilder() - .addTopic(TopicReadSettings.newBuilder().setPath(TEST_TOPIC).build()) + .addTopic(TopicReadSettings.newBuilder() + .setPath(TEST_TOPIC) + .setPartitionIds(Arrays.asList(1L)) + .build()) .setConsumerName(TEST_CONSUMER1) .build(); - AtomicLong counter = new AtomicLong(); - CompletableFuture result = new CompletableFuture<>(); + AtomicLong offset = new AtomicLong(200L); + CountDownLatch committed = new CountDownLatch(500-200); AsyncReader reader = client.createAsyncReader(readerSettings, ReadEventHandlersSettings.newBuilder() .setEventHandler(new ReadEventHandler() { @Override public void onStartPartitionSession(StartPartitionSessionEvent event) { Assert.assertEquals(0, event.getCommittedOffset()); Assert.assertEquals(0, event.getPartitionOffsets().getStart()); - Assert.assertEquals(6, event.getPartitionOffsets().getEnd()); + Assert.assertEquals(500, event.getPartitionOffsets().getEnd()); - // read only from offset 2 + // read only from offset 200 event.confirm(StartPartitionSessionSettings.newBuilder() - .setReadOffset(2L) - .setCommitOffset(2L) + .setReadOffset(200L) + .setCommitOffset(200L) .build()); } @Override public void onMessages(DataReceivedEvent event) { - event.commit().join(); - for (Message msg : event.getMessages()) { - Assert.assertEquals(msg.getOffset(), 2 + counter.get()); - long read = counter.incrementAndGet(); - if (new String(msg.getData()).equals("stop")) { - result.complete(read); - } + for (Message msg: event.getMessages()) { + Assert.assertEquals(msg.getOffset(), offset.getAndIncrement()); + msg.commit().whenComplete((r, th) -> { + Assert.assertNull(th); + committed.countDown(); + }); } } }).build()); reader.init().join(); - Assert.assertEquals(Long.valueOf(4), result.join()); - reader.shutdown().join(); + try { + Assert.assertTrue(committed.await(30, TimeUnit.SECONDS)); + Assert.assertEquals(500, offset.get()); + } finally { + reader.shutdown().join(); + } } @Test public void readRetentionedTopicTest() throws Exception { ReaderSettings readerSettings = ReaderSettings.newBuilder() - .addTopic(TopicReadSettings.newBuilder().setPath(TEST_TOPIC).build()) + .addTopic(TopicReadSettings.newBuilder() + .setPath(TEST_TOPIC) + .setPartitionIds(Arrays.asList(1L)) + .build()) .setConsumerName(TEST_CONSUMER1) .build(); - AtomicLong counter = new AtomicLong(); - CompletableFuture result = new CompletableFuture<>(); + AtomicLong offset = new AtomicLong(150L); + CountDownLatch lastCommitted = new CountDownLatch(1); AsyncReader reader = client.createAsyncReader(readerSettings, ReadEventHandlersSettings.newBuilder() .setEventHandler(new ReadEventHandler() { @Override public void onStartPartitionSession(StartPartitionSessionEvent event) { Assert.assertEquals(0, event.getCommittedOffset()); Assert.assertEquals(0, event.getPartitionOffsets().getStart()); - Assert.assertEquals(6, event.getPartitionOffsets().getEnd()); + Assert.assertEquals(500, event.getPartitionOffsets().getEnd()); - // emulate topic retention - skip first 2 message but don't commit them + // emulate topic retention - skip first 150 message but don't commit them event.confirm(StartPartitionSessionSettings.newBuilder() - .setReadOffset(2L) + .setReadOffset(150L) .build()); } @Override public void onMessages(DataReceivedEvent event) { - event.commit().join(); - for (Message msg : event.getMessages()) { - Assert.assertEquals(msg.getOffset(), 2 + counter.get()); - long read = counter.incrementAndGet(); - if (new String(msg.getData()).equals("stop")) { - result.complete(read); - } + for (Message msg: event.getMessages()) { + Assert.assertEquals(msg.getOffset(), offset.getAndIncrement()); } + + event.commit().whenComplete((r, th) -> { + Assert.assertNull(th); + if (event.getRangeToCommit().getEnd() >= 500) { + lastCommitted.countDown(); + } + }); } }).build()); reader.init().join(); - Assert.assertEquals(Long.valueOf(4), result.join()); - reader.shutdown().join(); + try { + Assert.assertTrue(lastCommitted.await(30, TimeUnit.SECONDS)); + Assert.assertEquals(500, offset.get()); + } finally { + reader.shutdown().join(); + } + } + + @Test + public void smallBufferTest() throws InterruptedException { + ReaderSettings readerSettings = ReaderSettings.newBuilder() + .addTopic(TopicReadSettings.newBuilder().setPath(TEST_TOPIC).build()) + .setConsumerName(TEST_CONSUMER1) + .setMaxMemoryUsageBytes(1000) + .build(); + + AtomicLong[] offsets = new AtomicLong[] { new AtomicLong(), new AtomicLong(), new AtomicLong() }; + CountDownLatch read = new CountDownLatch(3600); + 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()); + read.countDown(); + } + }).build()); + + 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(); + } + } + + @Test + public void directDecompressorTest() throws InterruptedException { + ReaderSettings readerSettings = ReaderSettings.newBuilder() + .addTopic(TopicReadSettings.newBuilder().setPath(TEST_TOPIC).build()) + .setConsumerName(TEST_CONSUMER1) + .setDecompressionExecutor(Runnable::run) + .setMaxMemoryUsageBytes(1000) + .build(); + + AtomicLong[] offsets = new AtomicLong[] { new AtomicLong(), new AtomicLong(), new AtomicLong() }; + CountDownLatch read = new CountDownLatch(3600); + AsyncReader reader = client.createAsyncReader(readerSettings, ReadEventHandlersSettings.newBuilder() + .setExecutor(Runnable::run) + .setEventHandler((DataReceivedEvent event) -> { + AtomicLong offset = offsets[(int) event.getPartitionSession().getPartitionId()]; + for (Message msg : event.getMessages()) { + Assert.assertEquals(offset.getAndIncrement(), msg.getOffset()); + read.countDown(); + } + }).build()); + + reader.init().join(); + try { + Assert.assertTrue(read.await(10, TimeUnit.MINUTES)); + Assert.assertEquals(1000, offsets[0].get()); + Assert.assertEquals(500, offsets[1].get()); + Assert.assertEquals(2100, offsets[2].get()); + } finally { + reader.shutdown().join(); + } } } From d336c0dd65a18d67e7c1119fe7d350667f25a4b9 Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Tue, 11 Aug 2026 08:58:03 +0100 Subject: [PATCH 06/12] Fixed data race in ReadPartitionSession --- .../main/java/tech/ydb/topic/read/impl/BufferManager.java | 4 ++-- .../java/tech/ydb/topic/read/impl/ReadPartitionSession.java | 6 ++++++ 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/BufferManager.java b/topic/src/main/java/tech/ydb/topic/read/impl/BufferManager.java index e1bc8e9c5..f3c17a864 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/BufferManager.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/BufferManager.java @@ -20,8 +20,8 @@ * @author Aleksandr Gorshenin {@literal } */ public class BufferManager { - private static final Logger logger = LoggerFactory.getLogger(ReaderImpl.class); - private static final OffsetsRange ALL = OffsetsRange.of(Integer.MIN_VALUE, Integer.MAX_VALUE); + private static final Logger logger = LoggerFactory.getLogger(BufferManager.class); + private static final OffsetsRange ALL = OffsetsRange.of(Long.MIN_VALUE, Long.MAX_VALUE); private final String traceID; private final long maxBufferSize; 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 6897a7cf4..b200f071f 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 @@ -141,12 +141,18 @@ public void sendDataToReadersIfNeeded() { Iterator it = readingQueue.iterator(); if (!it.hasNext()) { isReadingNow.set(false); + if (!readingQueue.isEmpty()) { + sendDataToReadersIfNeeded(); + } return; } MessageImpl next = it.next(); if (!next.isReady()) { isReadingNow.set(false); + if (next.isReady()) { + sendDataToReadersIfNeeded(); + } return; } From 62fd7070141d8ca5d9d506779b67083ec1be6400 Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Tue, 11 Aug 2026 09:06:08 +0100 Subject: [PATCH 07/12] Added empty messages guard to BufferManager --- .../tech/ydb/topic/read/impl/BufferManager.java | 13 ++++++++++--- .../java/tech/ydb/topic/read/impl/ReadSession.java | 2 +- .../tech/ydb/topic/read/impl/BufferManagerTest.java | 13 ++++++++++++- 3 files changed, 23 insertions(+), 5 deletions(-) diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/BufferManager.java b/topic/src/main/java/tech/ydb/topic/read/impl/BufferManager.java index f3c17a864..b19176f24 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/BufferManager.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/BufferManager.java @@ -42,7 +42,7 @@ public void init(String sessionId) { } // Has no reentrant thread safety - public void allocate(int bufferSize, List dataList) { + public void allocate(long bufferSize, List dataList) { logger.trace("[{}] Received ReadResponse of {} bytes", traceID, bufferSize); // calculate message count @@ -124,8 +124,8 @@ private void release(long total) { } } - private static void recalcBuffer(int[] buffer, int buffSize) { - // corner cases + private static void recalcBuffer(int[] buffer, long buffSize) { + // buffSize guard if (buffSize == 0) { Arrays.fill(buffer, 0); return; @@ -136,6 +136,13 @@ private static void recalcBuffer(int[] buffer, int buffSize) { total += v; } + // empty messages guard + if (total == 0) { + Arrays.fill(buffer, 0); + buffer[buffer.length - 1] = (int) buffSize; + return; + } + long currBuff = 0; long currSum = 0; for (int idx = 0; idx < buffer.length; idx += 1) { 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 99ea562ea..0e9bac841 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 @@ -291,7 +291,7 @@ public void confirm() { private void onReadResponse(YdbTopic.StreamReadMessage.ReadResponse response) { logger.debug("[{}] Received ReadResponse of {} bytes", streamId, response.getBytesSize()); - bufferManager.allocate((int) response.getBytesSize(), response.getPartitionDataList()); + bufferManager.allocate(response.getBytesSize(), response.getPartitionDataList()); for (YdbTopic.StreamReadMessage.ReadResponse.PartitionData data: response.getPartitionDataList()) { long psid = data.getPartitionSessionId(); diff --git a/topic/src/test/java/tech/ydb/topic/read/impl/BufferManagerTest.java b/topic/src/test/java/tech/ydb/topic/read/impl/BufferManagerTest.java index 9868a21b1..ebb6a1fe2 100644 --- a/topic/src/test/java/tech/ydb/topic/read/impl/BufferManagerTest.java +++ b/topic/src/test/java/tech/ydb/topic/read/impl/BufferManagerTest.java @@ -158,7 +158,7 @@ public void concurrentReleaseTest() { @Test public void intOverflowTest() { AtomicLong requested = new AtomicLong(0); - BufferManager bm = new BufferManager("trace-4", 1000, requested::addAndGet); + BufferManager bm = new BufferManager("trace-5", 1000, requested::addAndGet); bm.allocate(100000, Arrays.asList( partition(1, batch(1, 10000, 10000, 10000)), @@ -171,4 +171,15 @@ public void intOverflowTest() { bm.releasePartition(3L); Assert.assertEquals(100000, requested.get()); } + + @Test + public void zeroMessagesTest() { + AtomicLong requested = new AtomicLong(0); + BufferManager bm = new BufferManager("trace-6", 1000, requested::addAndGet); + + bm.allocate(100000, Arrays.asList(partition(1, batch(1, 0, 0, 0)))); + + bm.releasePartition(1L); + Assert.assertEquals(100000, requested.get()); + } } From dbae2cf5d1ce4245679f94f99c899ae7b6036636 Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Tue, 11 Aug 2026 11:17:43 +0100 Subject: [PATCH 08/12] style fixes --- .../topic/read/impl/ReadPartitionDecoder.java | 6 +++++- .../topic/read/impl/ReadPartitionSession.java | 16 +++++++++------- .../ydb/topic/TopicReadersIntegrationTest.java | 2 +- 3 files changed, 15 insertions(+), 9 deletions(-) diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionDecoder.java b/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionDecoder.java index 3f31eeb41..a09440a56 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionDecoder.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionDecoder.java @@ -157,12 +157,16 @@ public void decode(CodecRegistry registry) { try (InputStream encoded = origin.newInput(); InputStream decoded = codec.decode(encoded)) { data = ByteStreams.toByteArray(decoded); logger.trace("[{}] Finished decoding batch", traceID); - origin = null; } catch (IOException ex) { logger.warn("[{}] Exception was thrown while decoding a message: ", traceID, ex); problem = ex; } } finally { + if (problem == null) { + origin = null; + } else { + data = null; + } isReady = true; readyHandler.run(); } 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 b200f071f..564c878bc 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 @@ -24,6 +24,7 @@ * @author Nikolay Perfilov */ public abstract class ReadPartitionSession { + private static final Logger logger = LoggerFactory.getLogger(ReaderImpl.class); private final String traceID; @@ -89,7 +90,7 @@ public boolean addBatches(List ba return false; } - for (YdbTopic.StreamReadMessage.ReadResponse.Batch batch: batchList) { + for (YdbTopic.StreamReadMessage.ReadResponse.Batch batch : batchList) { if (batch.getMessageDataCount() == 0) { logger.error("[{}] Received empty batch. This shouldn't happen", traceID); continue; @@ -97,7 +98,7 @@ public boolean addBatches(List ba BatchMeta meta = new BatchMeta(batch); List messages = new ArrayList<>(batch.getMessageDataCount()); - for (YdbTopic.StreamReadMessage.ReadResponse.MessageData msg: batch.getMessageDataList()) { + for (YdbTopic.StreamReadMessage.ReadResponse.MessageData msg : batch.getMessageDataList()) { if (lastReadOffset > msg.getOffset()) { logger.error("[{}] Received a message with offset {} which is less than last read offset {} ", traceID, msg.getOffset(), lastReadOffset); @@ -137,7 +138,7 @@ public void sendDataToReadersIfNeeded() { return; } - if (isReadingNow.compareAndSet(false, true)) { + if (isReadingNow.compareAndSet(false, true)) { Iterator it = readingQueue.iterator(); if (!it.hasNext()) { isReadingNow.set(false); @@ -165,10 +166,10 @@ public void sendDataToReadersIfNeeded() { // Should be called maximum in 1 thread at a time DataReceivedEvent event = new DataReceivedEventImpl(partition, committer, messagesToRead); - logger.debug("[{}] DataReceivedEvent callback with {} message(s) (offsets {}-{}) is about " + - "to be called...", traceID, messagesToRead.size(), - messagesToRead.get(0).getOffset(), - messagesToRead.get(messagesToRead.size() - 1).getOffset()); + logger.debug("[{}] DataReceivedEvent callback with {} message(s) (offsets {}-{}) is about " + + "to be called...", traceID, messagesToRead.size(), + messagesToRead.get(0).getOffset(), + messagesToRead.get(messagesToRead.size() - 1).getOffset()); handleDataReceivedEvent(event).whenComplete((res, th) -> { if (th != null) { logger.error("[{}] DataReceivedEvent callback with {} message(s) (offsets {}-{}) finished" @@ -194,6 +195,7 @@ public void sendDataToReadersIfNeeded() { } private class RawMessage extends MessageImpl { + private final byte[] data; RawMessage(BatchMeta meta, OffsetsRange range, YdbTopic.StreamReadMessage.ReadResponse.MessageData msg) { diff --git a/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java b/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java index d57f3161e..6d71968ec 100644 --- a/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java +++ b/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java @@ -447,7 +447,7 @@ public void directDecompressorTest() throws InterruptedException { reader.init().join(); try { - Assert.assertTrue(read.await(10, TimeUnit.MINUTES)); + 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()); From d86e2e9d094c88f9f968a0661e5eab65a1f5b041 Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Tue, 11 Aug 2026 13:48:28 +0100 Subject: [PATCH 09/12] Typo fixes --- .../src/main/java/tech/ydb/topic/read/impl/BufferManager.java | 2 +- topic/src/main/java/tech/ydb/topic/read/impl/ReadSession.java | 4 ++-- .../test/java/tech/ydb/topic/TopicReadersIntegrationTest.java | 2 +- .../test/java/tech/ydb/topic/read/impl/BufferManagerTest.java | 4 ++-- .../java/tech/ydb/topic/read/impl/MessageDecoderTest.java | 4 ++-- 5 files changed, 8 insertions(+), 8 deletions(-) diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/BufferManager.java b/topic/src/main/java/tech/ydb/topic/read/impl/BufferManager.java index b19176f24..d3259d68a 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/BufferManager.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/BufferManager.java @@ -116,7 +116,7 @@ public void releaseRange(Long id, OffsetsRange range) { private void release(long total) { long now = released.addAndGet(total); - if (now >= maxBufferSize / 10) { // threshhold + if (now >= maxBufferSize / 10) { // threshold long request = released.getAndSet(0); if (request > 0) { requestFunc.accept(request); 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 0e9bac841..f66ac9725 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 @@ -104,13 +104,13 @@ public void startAndInitialize() { protected void onStop() { logger.debug("[{}] Session onStop called", streamId); + decoder.stop(); + partSessions.values().forEach(ReadPartitionSession::stop); partSessions.clear(); partitions.values().forEach(reader::handleClosePartitionSession); partitions.clear(); - - decoder.stop(); } protected void closeDueToError(Status status, Throwable th) { diff --git a/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java b/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java index 6d71968ec..9d92fc873 100644 --- a/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java +++ b/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java @@ -364,7 +364,7 @@ public void onStartPartitionSession(StartPartitionSessionEvent event) { Assert.assertEquals(0, event.getPartitionOffsets().getStart()); Assert.assertEquals(500, event.getPartitionOffsets().getEnd()); - // emulate topic retention - skip first 150 message but don't commit them + // emulate topic retention - skip first 150 messages but don't commit them event.confirm(StartPartitionSessionSettings.newBuilder() .setReadOffset(150L) .build()); diff --git a/topic/src/test/java/tech/ydb/topic/read/impl/BufferManagerTest.java b/topic/src/test/java/tech/ydb/topic/read/impl/BufferManagerTest.java index ebb6a1fe2..6ec683f66 100644 --- a/topic/src/test/java/tech/ydb/topic/read/impl/BufferManagerTest.java +++ b/topic/src/test/java/tech/ydb/topic/read/impl/BufferManagerTest.java @@ -85,7 +85,7 @@ public void simpleUseTest() { bm.releaseRange(4L, OffsetsRange.of(0, 100)); r.assertCalls(1, 1234); - // partitial release more than 10% of buffer size - a new request + // partial release more than 10% of buffer size - a new request bm.releaseRange(1L, OffsetsRange.of(1, 3)); r.assertCalls(2, 1234 + 350); @@ -129,7 +129,7 @@ public void zeroAllocateTest() { bm.allocate(0, Arrays.asList(partition(1, batch(0, 100, 100), batch(2, 250), batch(6, 40, 5, 5), batch(20)))); r.assertCalls(0, 0); - // releases has no effects + // releases have no effect bm.releaseRange(1L, OffsetsRange.of(1, 3)); bm.releasePartition(1L); r.assertCalls(0, 0); diff --git a/topic/src/test/java/tech/ydb/topic/read/impl/MessageDecoderTest.java b/topic/src/test/java/tech/ydb/topic/read/impl/MessageDecoderTest.java index cb3bc7a15..b07a66968 100644 --- a/topic/src/test/java/tech/ydb/topic/read/impl/MessageDecoderTest.java +++ b/topic/src/test/java/tech/ydb/topic/read/impl/MessageDecoderTest.java @@ -135,7 +135,7 @@ public void flowControlByBudgetTest() { Assert.assertEquals(3, m3.getData()[0]); p1.releaseRange(OffsetsRange.of(1)); // 40 is not enough to resume decoding - p1.releaseRange(OffsetsRange.of(4)); // that offset is not decodec yet + p1.releaseRange(OffsetsRange.of(4)); // that offset is not decoded yet Assert.assertFalse(m4.isReady()); Assert.assertFalse(m5.isReady()); Assert.assertEquals(3, ready.get()); @@ -325,7 +325,7 @@ public void decodeProblemsTest() { AtomicInteger ready = new AtomicInteger(); ReadPartitionDecoder p1 = new ReadPartitionDecoder("p1", decoder, PS1, null, ready::incrementAndGet); - BatchMeta meta1 = meta(1244); // Uknown codec + BatchMeta meta1 = meta(1244); // Unknown codec MessageImpl m1 = p1.decode(meta1, OffsetsRange.of(1), rawMsg(1, 50, new byte[0])); MessageImpl m2 = p1.decode(meta1, OffsetsRange.of(2), rawMsg(2, 50, new byte[0])); From bd17fd172008bd1d701f30879fc4ed77f6ad613c Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Tue, 11 Aug 2026 14:17:52 +0100 Subject: [PATCH 10/12] Fixed corner case in MessageDecoder --- .../ydb/topic/read/impl/MessageDecoder.java | 9 +++- .../topic/read/impl/ReadPartitionDecoder.java | 1 + .../topic/read/impl/MessageDecoderTest.java | 42 ++++++++++++++++++- 3 files changed, 50 insertions(+), 2 deletions(-) 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 5d06bae5e..8357af332 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 @@ -39,11 +39,19 @@ public void stop() { this.isStopped = true; } + long getTotalAvailable() { + return totalAvailable.get(); + } + void add(ReadPartitionDecoder.EncodedMessage task) { decodingQueue.add(task); } void free(long bufferSize) { + if (isStopped) { + return; + } + if (bufferSize > 0) { totalAvailable.addAndGet(bufferSize); decodeNext.run(); @@ -66,7 +74,6 @@ public void run() { } catch (Throwable ex) { logger.error("Cannot execute decompression ", ex); next.setError(ex); - totalAvailable.addAndGet(size); } } } diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionDecoder.java b/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionDecoder.java index a09440a56..f2959c4fd 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionDecoder.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionDecoder.java @@ -123,6 +123,7 @@ public boolean isReady() { public void setError(Throwable th) { problem = new IOException("Decompression for " + getPartitionSession() + " error", th); + releaseRange(OffsetsRange.of(getOffset())); isReady = true; readyHandler.run(); } diff --git a/topic/src/test/java/tech/ydb/topic/read/impl/MessageDecoderTest.java b/topic/src/test/java/tech/ydb/topic/read/impl/MessageDecoderTest.java index b07a66968..4834023c5 100644 --- a/topic/src/test/java/tech/ydb/topic/read/impl/MessageDecoderTest.java +++ b/topic/src/test/java/tech/ydb/topic/read/impl/MessageDecoderTest.java @@ -95,6 +95,10 @@ public void rawDecodeTest() { Assert.assertArrayEquals(new byte[] {1, 2}, m1.getData()); Assert.assertArrayEquals(new byte[] {3}, m2.getData()); Assert.assertEquals(2, ready.get()); + + Assert.assertEquals(9800, decoder.getTotalAvailable()); + partition.releaseRange(OffsetsRange.of(0, 10)); + Assert.assertEquals(10000, decoder.getTotalAvailable()); } @Test @@ -117,9 +121,11 @@ public void flowControlByBudgetTest() { Assert.assertFalse(m4.isReady()); Assert.assertFalse(m5.isReady()); Assert.assertEquals(0, ready.get()); + Assert.assertEquals(100, decoder.getTotalAvailable()); decoder.decodeNext(); + Assert.assertEquals(-50, decoder.getTotalAvailable()); Assert.assertTrue(m1.isReady()); Assert.assertTrue(m2.isReady()); Assert.assertTrue(m3.isReady()); @@ -136,11 +142,14 @@ public void flowControlByBudgetTest() { p1.releaseRange(OffsetsRange.of(1)); // 40 is not enough to resume decoding p1.releaseRange(OffsetsRange.of(4)); // that offset is not decoded yet + + Assert.assertEquals(-10, decoder.getTotalAvailable()); Assert.assertFalse(m4.isReady()); Assert.assertFalse(m5.isReady()); Assert.assertEquals(3, ready.get()); p1.releaseRange(OffsetsRange.of(2)); + Assert.assertEquals(-30, decoder.getTotalAvailable()); Assert.assertTrue(m4.isReady()); Assert.assertFalse(m5.isReady()); Assert.assertEquals(4, ready.get()); @@ -149,12 +158,19 @@ public void flowControlByBudgetTest() { Assert.assertEquals(4, m4.getData()[0]); p1.releaseRange(OffsetsRange.of(0, 3)); // double release + + Assert.assertEquals(-30, decoder.getTotalAvailable()); Assert.assertFalse(m5.isReady()); Assert.assertEquals(4, ready.get()); p1.releaseRange(OffsetsRange.of(0, 5)); Assert.assertTrue(m5.isReady()); Assert.assertEquals(5, ready.get()); + + Assert.assertEquals(-20, decoder.getTotalAvailable()); + + p1.close(); + Assert.assertEquals(100, decoder.getTotalAvailable()); } @Test @@ -175,8 +191,10 @@ public void partitionFlowTest() { Assert.assertEquals(0, r1.get()); Assert.assertEquals(0, r2.get()); + Assert.assertEquals(100, decoder.getTotalAvailable()); decoder.decodeNext(); + Assert.assertEquals(-50, decoder.getTotalAvailable()); MessageImpl m6 = p1.decode(meta, OffsetsRange.of(4), gzipMsg(4, 10)); MessageImpl m7 = p1.decode(meta, OffsetsRange.of(5), gzipMsg(5, 20)); @@ -184,6 +202,7 @@ public void partitionFlowTest() { MessageImpl m9 = p2.decode(meta, OffsetsRange.of(15), gzipMsg(15, 20)); decoder.decodeNext(); + Assert.assertEquals(-50, decoder.getTotalAvailable()); Assert.assertTrue(m1.isReady()); Assert.assertTrue(m2.isReady()); @@ -199,6 +218,7 @@ public void partitionFlowTest() { Assert.assertEquals(1, r2.get()); p1.close(); + Assert.assertEquals(0, decoder.getTotalAvailable()); Assert.assertTrue(m1.isReady()); Assert.assertTrue(m2.isReady()); @@ -216,6 +236,8 @@ public void partitionFlowTest() { p1.releaseRange(OffsetsRange.of(0, 100)); p2.releaseRange(OffsetsRange.of(0, 100)); + Assert.assertEquals(40, decoder.getTotalAvailable()); + Assert.assertTrue(m5.isReady()); Assert.assertTrue(m6.isReady()); Assert.assertTrue(m7.isReady()); @@ -224,7 +246,7 @@ public void partitionFlowTest() { Assert.assertEquals(2, r1.get()); // p1 is already stopped Assert.assertEquals(5, r2.get()); -} + } @Test public void decodeStopTest() { @@ -239,20 +261,27 @@ public void decodeStopTest() { MessageImpl m3 = partition.decode(meta, OffsetsRange.of(3), gzipMsg(3, 60)); Assert.assertEquals(0, ready.get()); + Assert.assertEquals(70, decoder.getTotalAvailable()); decoder.decodeNext(); + + Assert.assertEquals(-20, decoder.getTotalAvailable()); Assert.assertEquals(2, ready.get()); Assert.assertTrue(m1.isReady()); Assert.assertTrue(m2.isReady()); Assert.assertFalse(m3.isReady()); decoder.stop(); + Assert.assertEquals(-20, decoder.getTotalAvailable()); + Assert.assertEquals(2, ready.get()); Assert.assertFalse(m3.isReady()); partition.releaseRange(OffsetsRange.of(0, 10)); Assert.assertEquals(2, ready.get()); Assert.assertFalse(m3.isReady()); + + Assert.assertEquals(-20, decoder.getTotalAvailable()); } @Test @@ -336,8 +365,11 @@ public void decodeProblemsTest() { MessageImpl m5 = p1.decode(meta2, OffsetsRange.of(5), gzipMsg(5, 50)); MessageImpl m6 = p1.decode(meta2, OffsetsRange.of(6), gzipMsg(6, 50)); + Assert.assertEquals(200, decoder.getTotalAvailable()); + decoder.decodeNext(); + Assert.assertEquals(0, decoder.getTotalAvailable()); Assert.assertEquals(4, ready.get()); Assert.assertTrue(m1.isReady()); @@ -354,6 +386,8 @@ public void decodeProblemsTest() { p1.close(); + Assert.assertEquals(200, decoder.getTotalAvailable()); + Assert.assertEquals(4, ready.get()); Assert.assertTrue(m5.isReady()); Assert.assertTrue(m6.isReady()); @@ -378,9 +412,11 @@ public void wrongDecoderTest() { BatchMeta meta = meta(Codec.GZIP); MessageImpl m1 = p1.decode(meta, OffsetsRange.of(5), gzipMsg(5, 50)); MessageImpl m2 = p1.decode(meta, OffsetsRange.of(6), gzipMsg(6, 50)); + Assert.assertEquals(200, decoder.getTotalAvailable()); decoder.decodeNext(); + Assert.assertEquals(200, decoder.getTotalAvailable()); Assert.assertEquals(2, ready.get()); Assert.assertTrue(m1.isReady()); @@ -390,5 +426,9 @@ public void wrongDecoderTest() { m1::getData); assertDecompressionException("Decompression for Partition session 1 (partition 1) for topic \"/topic\" error", m2::getData); + + Assert.assertEquals(200, decoder.getTotalAvailable()); + p1.close(); + Assert.assertEquals(200, decoder.getTotalAvailable()); } } From 521dfa8cd76cb7433ca2e324f4c33fe62177a2b7 Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Tue, 11 Aug 2026 14:23:57 +0100 Subject: [PATCH 11/12] Extended BufferManager logs --- .../main/java/tech/ydb/topic/read/impl/BufferManager.java | 8 +++++++- .../java/tech/ydb/topic/read/impl/BufferManagerTest.java | 2 +- .../java/tech/ydb/topic/read/impl/MessageDecoderTest.java | 2 +- 3 files changed, 9 insertions(+), 3 deletions(-) diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/BufferManager.java b/topic/src/main/java/tech/ydb/topic/read/impl/BufferManager.java index d3259d68a..6f5fb2ce2 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/BufferManager.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/BufferManager.java @@ -30,6 +30,9 @@ public class BufferManager { private final AtomicLong released = new AtomicLong(0); private final ConcurrentHashMap partitions = new ConcurrentHashMap<>(); + private final AtomicLong totalAllocated = new AtomicLong(0); + private final AtomicLong totalReleased = new AtomicLong(0); + public BufferManager(String traceID, long maxBufferSize, Consumer requestFunc) { this.traceID = traceID; this.maxBufferSize = maxBufferSize; @@ -43,7 +46,9 @@ public void init(String sessionId) { // Has no reentrant thread safety public void allocate(long bufferSize, List dataList) { - logger.trace("[{}] Received ReadResponse of {} bytes", traceID, bufferSize); + logger.debug("[{}] Received ReadResponse of {} bytes, {} allocated and {} released before", + traceID, bufferSize, totalAllocated.get(), totalReleased.get()); + totalAllocated.addAndGet(bufferSize); // calculate message count int messagesCount = 0; @@ -119,6 +124,7 @@ private void release(long total) { if (now >= maxBufferSize / 10) { // threshold long request = released.getAndSet(0); if (request > 0) { + totalReleased.addAndGet(request); requestFunc.accept(request); } } diff --git a/topic/src/test/java/tech/ydb/topic/read/impl/BufferManagerTest.java b/topic/src/test/java/tech/ydb/topic/read/impl/BufferManagerTest.java index 6ec683f66..ef9f4af73 100644 --- a/topic/src/test/java/tech/ydb/topic/read/impl/BufferManagerTest.java +++ b/topic/src/test/java/tech/ydb/topic/read/impl/BufferManagerTest.java @@ -106,7 +106,7 @@ public void simpleUseTest() { } @Test - @HideLoggers(ReaderImpl.class) + @HideLoggers(BufferManager.class) public void incorrectAllocateTest() { Resource r = new Resource(); BufferManager bm = new BufferManager("trace-2", 20000, r); diff --git a/topic/src/test/java/tech/ydb/topic/read/impl/MessageDecoderTest.java b/topic/src/test/java/tech/ydb/topic/read/impl/MessageDecoderTest.java index 4834023c5..d7824c5d9 100644 --- a/topic/src/test/java/tech/ydb/topic/read/impl/MessageDecoderTest.java +++ b/topic/src/test/java/tech/ydb/topic/read/impl/MessageDecoderTest.java @@ -73,7 +73,7 @@ private static BatchMeta meta(int codec) { } @Rule - public final HideLoggersRule hideLoggerRule = new HideLoggersRule(); + public final HideLoggersRule hideLogger = new HideLoggersRule(); @Test public void rawDecodeTest() { From af96f011c7441560a1242f7fe15c4bc290f2c76a Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Thu, 13 Aug 2026 10:48:46 +0100 Subject: [PATCH 12/12] Removed class Batch --- .../java/tech/ydb/topic/read/impl/Batch.java | 48 ------------------- 1 file changed, 48 deletions(-) delete mode 100644 topic/src/main/java/tech/ydb/topic/read/impl/Batch.java diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/Batch.java b/topic/src/main/java/tech/ydb/topic/read/impl/Batch.java deleted file mode 100644 index 0024c4786..000000000 --- a/topic/src/main/java/tech/ydb/topic/read/impl/Batch.java +++ /dev/null @@ -1,48 +0,0 @@ -package tech.ydb.topic.read.impl; - -import java.util.List; -import java.util.concurrent.CompletableFuture; - -import tech.ydb.topic.description.Codec; - -/** - * @author Nikolay Perfilov - */ -public class Batch { - private final BatchMeta meta; - private final List messages; - - // Completes when batch is read - private final CompletableFuture readFuture = new CompletableFuture<>(); - private volatile boolean isReady = false; - - public Batch(BatchMeta meta, List messages) { - this.meta = meta; - this.messages = messages; - this.isReady = meta.getCodec() == Codec.RAW; - } - - public List getMessages() { - return messages; - } - - public void complete() { - readFuture.complete(null); - } - - public CompletableFuture getReadFuture() { - return readFuture; - } - - public int getCodec() { - return meta.getCodec(); - } - - public boolean isReady() { - return isReady; - } - - public void markAsReady() { - this.isReady = true; - } -}