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; - } -} 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..6f5fb2ce2 --- /dev/null +++ b/topic/src/main/java/tech/ydb/topic/read/impl/BufferManager.java @@ -0,0 +1,233 @@ +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(BufferManager.class); + private static final OffsetsRange ALL = OffsetsRange.of(Long.MIN_VALUE, Long.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<>(); + + 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; + 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(long bufferSize, List dataList) { + 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; + 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) { // threshold + long request = released.getAndSet(0); + if (request > 0) { + totalReleased.addAndGet(request); + requestFunc.accept(request); + } + } + } + + private static void recalcBuffer(int[] buffer, long buffSize) { + // buffSize guard + if (buffSize == 0) { + Arrays.fill(buffer, 0); + return; + } + + long total = 0; + for (int v: buffer) { + 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) { + currSum += buffer[idx]; + long newBuff = currSum * buffSize / total; + buffer[idx] = (int) (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/main/java/tech/ydb/topic/read/impl/MessageDecoder.java b/topic/src/main/java/tech/ydb/topic/read/impl/MessageDecoder.java index d1fd88143..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 @@ -1,6 +1,5 @@ package tech.ydb.topic.read.impl; -import java.io.IOException; import java.util.Queue; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.Executor; @@ -10,7 +9,7 @@ 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. @@ -18,102 +17,65 @@ */ 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; + long getTotalAvailable() { + return totalAvailable.get(); } - private class DecodeTask implements Runnable { - private final String traceID; - private final Batch batch; - private final Runnable readyHandler; + void add(ReadPartitionDecoder.EncodedMessage task) { + decodingQueue.add(task); + } - DecodeTask(String traceID, Batch batch, Runnable readyHandler) { - this.traceID = traceID; - this.batch = batch; - this.readyHandler = readyHandler; + void free(long bufferSize) { + if (isStopped) { + return; } - public Batch getBatch() { - return batch; + 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); - } + while (!isStopped && totalAvailable.get() > 0) { + ReadPartitionDecoder.EncodedMessage next = decodingQueue.poll(); + if (next == null) { + return; + } - batch.getMessages().forEach(message -> { + long size = next.allocate(); + totalAvailable.addAndGet(-size); 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); + decompressionExecutor.execute(() -> next.decode(codecRegistry)); + } catch (Throwable ex) { + logger.error("Cannot execute decompression ", ex); + next.setError(ex); } - }); - batch.markAsReady(); - - if (logger.isTraceEnabled()) { - logger.trace("[{}] Finished decoding batch", traceID); } - - 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..f2959c4fd --- /dev/null +++ b/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionDecoder.java @@ -0,0 +1,176 @@ +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(MessageDecoder.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() { + decoder.free(allocatedTotal.getAndSet(0)); + } + + 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 void setError(Throwable th) { + problem = new IOException("Decompression for " + getPartitionSession() + " error", th); + releaseRange(OffsetsRange.of(getOffset())); + isReady = true; + readyHandler.run(); + } + + public long allocate() { + if (isStopped) { + problem = new IOException("" + getPartitionSession() + " is already closed"); + 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); + } 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 0a055a0e3..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 @@ -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; @@ -23,44 +24,48 @@ * @author Nikolay Perfilov */ public abstract class ReadPartitionSession { + private static final Logger logger = LoggerFactory.getLogger(ReaderImpl.class); private final String traceID; private final ReadSession session; private final PartitionSession partition; private final int maxBatchSize; - private final 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; } + public PartitionSession getPartition() { + return partition; + } + @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 +79,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) { + 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<>(); - for (YdbTopic.StreamReadMessage.ReadResponse.MessageData msg: batch.getMessageDataList()) { + 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 +120,17 @@ public CompletableFuture addBatches(List[0])); + public void releaseRange(OffsetsRange range) { + decoder.releaseRange(range); + bufferManager.releaseRange(partition.getId(), range); + sendDataToReadersIfNeeded(); } public void sendDataToReadersIfNeeded() { @@ -137,60 +139,52 @@ public void sendDataToReadersIfNeeded() { } if (isReadingNow.compareAndSet(false, true)) { - List batchesToRead = new ArrayList<>(); - - Batch next = readingQueue.peek(); - if (next == null || !next.isReady()) { + Iterator it = readingQueue.iterator(); + if (!it.hasNext()) { isReadingNow.set(false); - if (next != null && next.isReady()) { + if (!readingQueue.isEmpty()) { 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; + MessageImpl next = it.next(); + if (!next.isReady()) { + isReadingNow.set(false); + if (next.isReady()) { + sendDataToReadersIfNeeded(); } + return; + } - next = readingQueue.poll(); - - 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 +193,24 @@ 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..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 @@ -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() @@ -106,6 +104,8 @@ public void startAndInitialize() { protected void onStop() { logger.debug("[{}] Session onStop called", streamId); + decoder.stop(); + partSessions.values().forEach(ReadPartitionSession::stop); partSessions.clear(); @@ -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(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) { diff --git a/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java b/topic/src/test/java/tech/ydb/topic/TopicReadersIntegrationTest.java index 3c031aa6e..9d92fc873 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 messages 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(30, TimeUnit.SECONDS)); + Assert.assertEquals(1000, offsets[0].get()); + Assert.assertEquals(500, offsets[1].get()); + Assert.assertEquals(2100, offsets[2].get()); + } finally { + reader.shutdown().join(); + } } } diff --git a/topic/src/test/java/tech/ydb/topic/read/impl/BufferManagerTest.java b/topic/src/test/java/tech/ydb/topic/read/impl/BufferManagerTest.java new file mode 100644 index 000000000..ef9f4af73 --- /dev/null +++ b/topic/src/test/java/tech/ydb/topic/read/impl/BufferManagerTest.java @@ -0,0 +1,185 @@ +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); + + // partial 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(BufferManager.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 have no effect + 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()); + } + + @Test + public void intOverflowTest() { + AtomicLong requested = new AtomicLong(0); + BufferManager bm = new BufferManager("trace-5", 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()); + } + + @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()); + } +} 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..d7824c5d9 --- /dev/null +++ b/topic/src/test/java/tech/ydb/topic/read/impl/MessageDecoderTest.java @@ -0,0 +1,434 @@ +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 hideLogger = 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()); + + Assert.assertEquals(9800, decoder.getTotalAvailable()); + partition.releaseRange(OffsetsRange.of(0, 10)); + Assert.assertEquals(10000, decoder.getTotalAvailable()); + } + + @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()); + Assert.assertEquals(100, decoder.getTotalAvailable()); + + decoder.decodeNext(); + + Assert.assertEquals(-50, decoder.getTotalAvailable()); + 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 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()); + + Assert.assertEquals(70, m4.getData().length); + 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 + 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()); + 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)); + 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.assertEquals(-50, decoder.getTotalAvailable()); + + 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.assertEquals(0, decoder.getTotalAvailable()); + + 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.assertEquals(40, decoder.getTotalAvailable()); + + 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()); + 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 + 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); // 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])); + + 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)); + + Assert.assertEquals(200, decoder.getTotalAvailable()); + + decoder.decodeNext(); + + Assert.assertEquals(0, decoder.getTotalAvailable()); + 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(200, decoder.getTotalAvailable()); + + 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)); + Assert.assertEquals(200, decoder.getTotalAvailable()); + + decoder.decodeNext(); + + Assert.assertEquals(200, decoder.getTotalAvailable()); + 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); + + Assert.assertEquals(200, decoder.getTotalAvailable()); + p1.close(); + Assert.assertEquals(200, decoder.getTotalAvailable()); + } +}