diff --git a/topic/src/main/java/tech/ydb/topic/description/Codec.java b/topic/src/main/java/tech/ydb/topic/description/Codec.java index 893613dc8..7f771a2ca 100644 --- a/topic/src/main/java/tech/ydb/topic/description/Codec.java +++ b/topic/src/main/java/tech/ydb/topic/description/Codec.java @@ -5,7 +5,6 @@ import java.io.OutputStream; /** - * * Interface for custom codec implementation. *

@@ -50,6 +49,17 @@ public interface Codec { */ int getId(); + /** + * Returns a conservative upper bound for the encoded size of an input with the specified size. + * Codecs whose output may be larger than their input should override this method. + * + * @param inputSizeBytes input size in bytes + * @return encoded size upper bound in bytes + */ + default long getMaxEncodedSize(int inputSizeBytes) { + return inputSizeBytes; + } + /** * Decode data * diff --git a/topic/src/main/java/tech/ydb/topic/impl/GzipCodec.java b/topic/src/main/java/tech/ydb/topic/impl/GzipCodec.java index 314c047ee..2ad3ed3a0 100644 --- a/topic/src/main/java/tech/ydb/topic/impl/GzipCodec.java +++ b/topic/src/main/java/tech/ydb/topic/impl/GzipCodec.java @@ -36,6 +36,13 @@ public int getId() { return Codec.GZIP; } + @Override + public long getMaxEncodedSize(int inputSizeBytes) { + long size = inputSizeBytes; + // zlib's compressBound formula plus the GZIP header and trailer + return size + (size >>> 12) + (size >>> 14) + (size >>> 25) + 31; + } + @Override public InputStream decode(InputStream byteArrayInputStream) throws IOException { return new GZIPInputStream(byteArrayInputStream); diff --git a/topic/src/main/java/tech/ydb/topic/impl/LzopCodec.java b/topic/src/main/java/tech/ydb/topic/impl/LzopCodec.java index 63c9eb386..f31254679 100644 --- a/topic/src/main/java/tech/ydb/topic/impl/LzopCodec.java +++ b/topic/src/main/java/tech/ydb/topic/impl/LzopCodec.java @@ -16,6 +16,10 @@ * Compression codec which implements the LZO algorithm */ public class LzopCodec implements Codec { + private static final int BLOCK_SIZE = 256 * 1024; + private static final int HEADER_SIZE = 38; + private static final int BLOCK_OVERHEAD = 8; + private static final int END_MARKER_SIZE = 4; private static final LzopCodec INSTANCE = new LzopCodec(); @@ -40,6 +44,14 @@ public int getId() { return Codec.LZOP; } + @Override + public long getMaxEncodedSize(int inputSizeBytes) { + long size = inputSizeBytes; + // LzopOutputStream header and end marker, plus two length fields for each non-empty block + long blockCount = (size + BLOCK_SIZE - 1) / BLOCK_SIZE; + return size + HEADER_SIZE + blockCount * BLOCK_OVERHEAD + END_MARKER_SIZE; + } + @Override public InputStream decode(InputStream byteArrayInputStream) throws IOException { return new LzopInputStream(byteArrayInputStream); diff --git a/topic/src/main/java/tech/ydb/topic/impl/ZstdBackwardCodec.java b/topic/src/main/java/tech/ydb/topic/impl/ZstdBackwardCodec.java index f3b31141f..480fb00a3 100644 --- a/topic/src/main/java/tech/ydb/topic/impl/ZstdBackwardCodec.java +++ b/topic/src/main/java/tech/ydb/topic/impl/ZstdBackwardCodec.java @@ -4,6 +4,7 @@ import java.io.InputStream; import java.io.OutputStream; +import com.github.luben.zstd.Zstd; import com.github.luben.zstd.ZstdInputStream; import com.github.luben.zstd.ZstdOutputStream; @@ -37,6 +38,11 @@ public int getId() { return Codec.ZSTD; } + @Override + public long getMaxEncodedSize(int inputSizeBytes) { + return Zstd.compressBound(inputSizeBytes); + } + @Override public InputStream decode(InputStream byteArrayInputStream) throws IOException { return new ZstdInputStream(byteArrayInputStream); diff --git a/topic/src/main/java/tech/ydb/topic/impl/ZstdCodec.java b/topic/src/main/java/tech/ydb/topic/impl/ZstdCodec.java index 67c8181ed..d7b8b9d7f 100644 --- a/topic/src/main/java/tech/ydb/topic/impl/ZstdCodec.java +++ b/topic/src/main/java/tech/ydb/topic/impl/ZstdCodec.java @@ -4,6 +4,7 @@ import java.io.InputStream; import java.io.OutputStream; +import com.github.luben.zstd.Zstd; import com.github.luben.zstd.ZstdInputStreamNoFinalizer; import com.github.luben.zstd.ZstdOutputStreamNoFinalizer; @@ -37,6 +38,11 @@ public int getId() { return Codec.ZSTD; } + @Override + public long getMaxEncodedSize(int inputSizeBytes) { + return Zstd.compressBound(inputSizeBytes); + } + @Override public InputStream decode(InputStream byteArrayInputStream) throws IOException { return new ZstdInputStreamNoFinalizer(byteArrayInputStream); diff --git a/topic/src/main/java/tech/ydb/topic/write/impl/BufferManager.java b/topic/src/main/java/tech/ydb/topic/write/impl/BufferManager.java index e3f056d9d..fe48f96b8 100644 --- a/topic/src/main/java/tech/ydb/topic/write/impl/BufferManager.java +++ b/topic/src/main/java/tech/ydb/topic/write/impl/BufferManager.java @@ -168,10 +168,43 @@ public void releaseMessage(long messageSize) { countAvailable.release(); } - public void updateMessageSize(long oldSize, long newSize) { + /** + * Tries to update the buffer reservation for an already acquired message. + * Releases unused blocks when the message becomes smaller. When it grows, reserves the requested additional + * blocks if possible, or all currently available blocks otherwise. + * + * @param oldSize currently reserved message size in bytes + * @param newSize required message size in bytes + * @return effective reservation size in bytes + */ + public long updateMessageSize(long oldSize, long newSize) { int oldBlocks = calculateBlocksCount(oldSize, blockBitsCount); int newBlocks = calculateBlocksCount(newSize, blockBitsCount); - blocksAvailable.release(oldBlocks - newBlocks); + int difference = oldBlocks - newBlocks; + + if (difference >= 0) { + blocksAvailable.release(difference); + return newSize; + } + + int requiredBlocks = -difference; + + if (blocksAvailable.tryAcquire(requiredBlocks)) { + return newSize; + } + + int acquiredBlocks = blocksAvailable.drainPermits(); + + if (acquiredBlocks >= requiredBlocks) { + blocksAvailable.release(acquiredBlocks - requiredBlocks); + return newSize; + } + + if (acquiredBlocks == 0) { + return oldSize; + } + + return (oldBlocks + (long) acquiredBlocks) << blockBitsCount; } private static int calculateBlockSize(long maxBufferSize) { diff --git a/topic/src/main/java/tech/ydb/topic/write/impl/WriterQueue.java b/topic/src/main/java/tech/ydb/topic/write/impl/WriterQueue.java index 28cd8df98..1a577da14 100644 --- a/topic/src/main/java/tech/ydb/topic/write/impl/WriterQueue.java +++ b/topic/src/main/java/tech/ydb/topic/write/impl/WriterQueue.java @@ -34,9 +34,12 @@ public class WriterQueue { interface EncodedMsg { SentMessage getSentMessage(); + void confirm(WriteAck ack); + void close(RuntimeException ex); } + private static final Logger logger = LoggerFactory.getLogger(WriterImpl.class); private final String debugId; @@ -56,7 +59,7 @@ interface EncodedMsg { private volatile EnqueuedMessage lastAcceptedMessage = null; public WriterQueue(String debugId, WriterSettings settings, CodecRegistry codecRegistry, - Executor compressionExecutor, Runnable readyNotify) { + Executor compressionExecutor, Runnable readyNotify) { this.debugId = debugId; this.buffer = new BufferManager(debugId, settings); @@ -216,40 +219,46 @@ List updateSeqNo(long newSeqNo) { CompletableFuture enqueue(Message message, YdbTransaction tx) throws QueueOverflowException, InterruptedException { - long msgSize = Math.min(message.getData().length, buffer.getMaxSize()); - buffer.acquire(msgSize); - return accept(message, tx, msgSize); + long reservedSizeBytes = reservationSizeBytes(message.getData().length); + buffer.acquire(reservedSizeBytes); + return accept(message, tx, reservedSizeBytes); } CompletableFuture tryEnqueue(Message message, YdbTransaction tx) throws QueueOverflowException { - long msgSize = Math.min(message.getData().length, buffer.getMaxSize()); - buffer.tryAcquire(msgSize); - return accept(message, tx, msgSize); + long reservedSizeBytes = reservationSizeBytes(message.getData().length); + buffer.tryAcquire(reservedSizeBytes); + return accept(message, tx, reservedSizeBytes); } CompletableFuture tryEnqueue(Message message, YdbTransaction tx, long timeout, TimeUnit unit) throws QueueOverflowException, InterruptedException, TimeoutException { - long msgSize = Math.min(message.getData().length, buffer.getMaxSize()); - buffer.tryAcquire(msgSize, timeout, unit); - return accept(message, tx, msgSize); + long reservedSizeBytes = reservationSizeBytes(message.getData().length); + buffer.tryAcquire(reservedSizeBytes, timeout, unit); + return accept(message, tx, reservedSizeBytes); } + /** + * Calculates the buffer reservation required before encoding a message using the codec-provided size bound. + * The reservation is capped at the full buffer size to preserve support for a single oversized message. + */ + private long reservationSizeBytes(int inputSizeBytes) { + return Math.min(codec.getMaxEncodedSize(inputSizeBytes), buffer.getMaxSize()); + } - private CompletableFuture accept(Message message, YdbTransaction tx, long msgSize) { - EnqueuedMessage msg = new EnqueuedMessage(new MessageMeta(message, tx), msgSize); + private CompletableFuture accept(Message message, YdbTransaction tx, long reservedSizeBytes) { + EnqueuedMessage msg = new EnqueuedMessage(new MessageMeta(message, tx), reservedSizeBytes); lastAcceptedMessage = msg; queue.add(msg); if (codec.getId() == Codec.RAW) { // fast track without compression - msg.completeWithData(UnsafeByteOperations.unsafeWrap(message.getData()), msgSize); + msg.completeWithData(UnsafeByteOperations.unsafeWrap(message.getData()), reservedSizeBytes); readyNotify.run(); return msg.getAckFuture(); } - // encode message try { - compressionExecutor.execute(() -> encode(message.getData(), msgSize, msg)); + compressionExecutor.execute(() -> encode(message.getData(), reservedSizeBytes, msg)); } catch (Throwable ex) { logger.warn("[{}] Message wasn't sent because of processing error", debugId, ex); msg.completeWithProblem(ex); @@ -259,7 +268,7 @@ private CompletableFuture accept(Message message, YdbTransaction tx, l return msg.getAckFuture(); } - private void encode(byte[] data, long msgSize, EnqueuedMessage msg) { + private void encode(byte[] data, long reservedSize, EnqueuedMessage msg) { if (msg.isReady()) { return; } @@ -270,15 +279,27 @@ private void encode(byte[] data, long msgSize, EnqueuedMessage msg) { os.write(data, 0, data.length); } - logger.trace("[{}] Message compressed from {} to {} bytes", debugId, msgSize, encoded.size()); + if (logger.isTraceEnabled()) { + logger.trace("[{}] Message compressed from {} to {} bytes", debugId, data.length, encoded.size()); + } - long bufferSize = msgSize; - if (msgSize > encoded.size()) { // if compressed lenght is less than uncompression - update buffer size - bufferSize = encoded.size(); - buffer.updateMessageSize(msgSize, bufferSize); + long bufferSizeBytes = Math.min(encoded.size(), buffer.getMaxSize()); + ByteString encodedData = encoded.toByteString(); + long updatedBufferSizeBytes = buffer.updateMessageSize(reservedSize, bufferSizeBytes); + + if (updatedBufferSizeBytes < bufferSizeBytes) { + logger.warn( + "[{}] Cannot fully reserve {} bytes for encoded message; reserving {} bytes. " + + "Writer buffer size {} may be temporarily exceeded", + debugId, + bufferSizeBytes, + updatedBufferSizeBytes, + buffer.getMaxSize() + ); } - msg.completeWithData(encoded.toByteString(), bufferSize); + bufferSizeBytes = updatedBufferSizeBytes; + msg.completeWithData(encodedData, bufferSizeBytes); } catch (Throwable ex) { logger.warn("[{}] Message wasn't sent because of encoding error", debugId, ex); msg.completeWithProblem(ex); @@ -286,7 +307,7 @@ private void encode(byte[] data, long msgSize, EnqueuedMessage msg) { readyNotify.run(); } - private class SkippedMsg implements EncodedMsg { + private static class SkippedMsg implements EncodedMsg { private final CompletableFuture ackFuture; private final WriteAck ack; @@ -312,7 +333,7 @@ public void close(RuntimeException ex) { } } - private class ProblemMsg implements EncodedMsg { + private static class ProblemMsg implements EncodedMsg { private final CompletableFuture ackFuture; private final Throwable problem; diff --git a/topic/src/test/java/tech/ydb/topic/write/impl/BufferManagerTest.java b/topic/src/test/java/tech/ydb/topic/write/impl/BufferManagerTest.java index af76bfc1d..4ba5488b8 100644 --- a/topic/src/test/java/tech/ydb/topic/write/impl/BufferManagerTest.java +++ b/topic/src/test/java/tech/ydb/topic/write/impl/BufferManagerTest.java @@ -139,7 +139,7 @@ public void testTryAcquireWithTimeoutBytesTimeoutThrows() throws Exception { + "Buffer currently has 1 messages with 5 / 20 bytes available", () -> bm.tryAcquire(15, 1, TimeUnit.MILLISECONDS)); - bm.updateMessageSize(15, 10); + Assert.assertEquals(10, bm.updateMessageSize(15, 10)); bm.tryAcquire(6, 1, TimeUnit.MILLISECONDS); // success assertTimeout("[test] Rejecting a message due to reaching message queue in-flight limit of 2", () -> bm.tryAcquire(1, 1, TimeUnit.MILLISECONDS)); @@ -158,6 +158,26 @@ public void testTryAcquireWithTimeoutBytesTimeoutThrows() throws Exception { bm.releaseMessage(10); } + @Test + public void testUpdateMessageSizeAcquiresAvailableCapacity() throws Exception { + BufferManager bufferManager = manager(10, 3); + bufferManager.tryAcquire(4); + bufferManager.tryAcquire(4); + + long updatedSize = bufferManager.updateMessageSize(4, 8); + Assert.assertEquals(6, updatedSize); + assertOverflow( + "[test] Rejecting a message of 1 bytes: not enough space in message queue. " + + "Buffer currently has 2 messages with 0 / 10 bytes available", + () -> bufferManager.tryAcquire(1) + ); + + bufferManager.releaseMessage(4); + bufferManager.releaseMessage(updatedSize); + bufferManager.tryAcquire(10); + bufferManager.releaseMessage(10); + } + @Test public void testLargeBufferConfiguration() throws QueueOverflowException { BufferManager bm = manager(Integer.MAX_VALUE, 10); diff --git a/topic/src/test/java/tech/ydb/topic/write/impl/WriterQueueTest.java b/topic/src/test/java/tech/ydb/topic/write/impl/WriterQueueTest.java index 313793463..1ce4b0c08 100644 --- a/topic/src/test/java/tech/ydb/topic/write/impl/WriterQueueTest.java +++ b/topic/src/test/java/tech/ydb/topic/write/impl/WriterQueueTest.java @@ -1,10 +1,13 @@ package tech.ydb.topic.write.impl; +import java.io.FilterOutputStream; import java.io.IOException; import java.io.InputStream; import java.io.OutputStream; +import java.util.ArrayDeque; import java.util.Arrays; import java.util.List; +import java.util.Queue; import java.util.concurrent.CompletableFuture; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; @@ -23,11 +26,13 @@ import tech.ydb.topic.write.QueueOverflowException; import tech.ydb.topic.write.WriteAck; +import static java.util.Collections.singletonList; + /** * @author Aleksandr Gorshenin */ public class WriterQueueTest { - private static final Message SMALL_MSG = Message.of(new byte[] { 0x00, 0x01, 0x02, 0x03, 0x05 }); + private static final Message SMALL_MSG = Message.of(new byte[]{0x00, 0x01, 0x02, 0x03, 0x05}); @Rule public final HideLoggersRule hideLogger = new HideLoggersRule(); @@ -35,6 +40,7 @@ public class WriterQueueTest { private static Message smallMsg(int seqNo) { return Message.newBuilder().setData(SMALL_MSG.getData()).setSeqNo(seqNo).build(); } + private static WriterSettings rawSettings() { return WriterSettings.newBuilder() .setTopicPath("/test") @@ -117,6 +123,116 @@ public void testRawCompressor() throws Exception { Assert.assertTrue(f3.isDone()); } + @Test + public void testCompressedExpansionIsAccounted() throws Exception { + Codec codec = expandingCodec(Codec.GZIP, 1); + WriterSettings settings = WriterSettings.newBuilder() + .setTopicPath("/test") + .setCodec(codec.getId()) + .setMaxSendBufferMemorySize(10) + .build(); + WriterQueue writerQueue = new WriterQueue("test", settings, new CodecRegistry(singletonList(codec)), + Runnable::run, () -> { + }); + + writerQueue.tryEnqueue(SMALL_MSG, null); + assertOverflow( + "[test] Rejecting a message of 5 bytes: not enough space in message queue. " + + "Buffer currently has 1 messages with 4 / 10 bytes available", + () -> writerQueue.tryEnqueue(SMALL_MSG, null) + ); + } + + @Test + public void testLzopFramingIsAccounted() throws Exception { + Queue encodingTasks = new ArrayDeque<>(); + WriterSettings settings = WriterSettings.newBuilder() + .setTopicPath("/test") + .setCodec(Codec.LZOP) + .setMaxSendBufferMemorySize(109) + .build(); + WriterQueue writerQueue = new WriterQueue("test", settings, new CodecRegistry(), encodingTasks::add, () -> { + }); + + writerQueue.tryEnqueue(SMALL_MSG, null); + assertOverflow( + "[test] Rejecting a message of 55 bytes: not enough space in message queue. " + + "Buffer currently has 1 messages with 54 / 109 bytes available", + () -> writerQueue.tryEnqueue(SMALL_MSG, null) + ); + } + + @Test + public void testBuiltInCodecSizeBoundsDoNotOverflow() { + CodecRegistry registry = new CodecRegistry(); + + Assert.assertTrue(registry.getCodec(Codec.GZIP).getMaxEncodedSize(Integer.MAX_VALUE) > Integer.MAX_VALUE); + Assert.assertTrue(registry.getCodec(Codec.LZOP).getMaxEncodedSize(Integer.MAX_VALUE) > Integer.MAX_VALUE); + Assert.assertTrue(registry.getCodec(Codec.ZSTD).getMaxEncodedSize(Integer.MAX_VALUE) > Integer.MAX_VALUE); + } + + @Test + @HideLoggers({WriterImpl.class}) + public void testEncodedMessageLargerThanBufferIsProcessed() throws Exception { + Codec codec = expandingCodec(9004, 6); + Queue encodingTasks = new ArrayDeque<>(); + WriterSettings settings = WriterSettings.newBuilder() + .setTopicPath("/test") + .setCodec(codec.getId()) + .setMaxSendBufferMemorySize(10) + .build(); + WriterQueue writerQueue = new WriterQueue("test", settings, new CodecRegistry(singletonList(codec)), + encodingTasks::add, () -> { + }); + + writerQueue.tryEnqueue(SMALL_MSG, null); + writerQueue.tryEnqueue(SMALL_MSG, null); + + encodingTasks.remove().run(); + SentMessage first = writerQueue.nextMessageToSend(); + Assert.assertNotNull(first); + Assert.assertEquals(11, first.getPb().getData().size()); + Assert.assertEquals(5, first.getBufferSize()); + writerQueue.confirmAck(new WriteAck(first.getSeqNo(), WriteAck.State.WRITTEN, null, null)); + + encodingTasks.remove().run(); + SentMessage second = writerQueue.nextMessageToSend(); + Assert.assertNotNull(second); + Assert.assertEquals(11, second.getPb().getData().size()); + Assert.assertEquals(10, second.getBufferSize()); + } + + @Test + @HideLoggers({WriterImpl.class}) + public void testExpansionUsesAvailableCapacity() throws Exception { + Codec codec = expandingCodec(9005, 4); + Queue encodingTasks = new ArrayDeque<>(); + WriterSettings settings = WriterSettings.newBuilder() + .setTopicPath("/test") + .setCodec(codec.getId()) + .setMaxSendBufferMemorySize(12) + .build(); + WriterQueue writerQueue = new WriterQueue("test", settings, new CodecRegistry(singletonList(codec)), + encodingTasks::add, () -> { + }); + + writerQueue.tryEnqueue(SMALL_MSG, null); + writerQueue.tryEnqueue(SMALL_MSG, null); + + encodingTasks.remove().run(); + SentMessage first = writerQueue.nextMessageToSend(); + Assert.assertNotNull(first); + Assert.assertEquals(9, first.getPb().getData().size()); + Assert.assertEquals(7, first.getBufferSize()); + writerQueue.confirmAck(new WriteAck(first.getSeqNo(), WriteAck.State.WRITTEN, null, null)); + + encodingTasks.remove().run(); + SentMessage second = writerQueue.nextMessageToSend(); + Assert.assertNotNull(second); + Assert.assertEquals(9, second.getPb().getData().size()); + Assert.assertEquals(9, second.getBufferSize()); + } + @Test @HideLoggers({ WriterImpl.class }) public void testGzipNullCompressor() throws Exception { @@ -138,7 +254,7 @@ public void testGzipNullCompressor() throws Exception { } @Test - @HideLoggers({ WriterImpl.class }) + @HideLoggers({WriterImpl.class}) public void testWrongCodec() throws Exception { // Codec that always throws on encode Codec failingCodec = new Codec() { @@ -216,7 +332,7 @@ public void testSmallBufferWriting() throws QueueOverflowException { q.tryEnqueue(smallMsg(10), null); // success q.tryEnqueue(smallMsg(20), null); // success assertOverflow("[test] Rejecting a message of 5 bytes: not enough space in message queue. " - + "Buffer currently has 2 messages with 2 / 12 bytes available", + + "Buffer currently has 2 messages with 2 / 12 bytes available", () -> q.tryEnqueue(smallMsg(30), null)); Assert.assertEquals(20, assertSendAll(q, 2)); @@ -224,7 +340,7 @@ public void testSmallBufferWriting() throws QueueOverflowException { q.tryEnqueue(smallMsg(30), null); // success assertOverflow("[test] Rejecting a message of 5 bytes: not enough space in message queue. " - + "Buffer currently has 2 messages with 2 / 12 bytes available", + + "Buffer currently has 2 messages with 2 / 12 bytes available", () -> q.tryEnqueue(smallMsg(40), null)); Assert.assertEquals(30, assertSendAll(q, 1)); @@ -341,4 +457,29 @@ public void testLostAcks() throws Exception { Assert.assertEquals(WriteAck.State.WRITTEN, f4.join().getState()); Assert.assertEquals(WriteAck.State.WRITTEN, f5.join().getState()); } + + private static Codec expandingCodec(int id, int additionalBytes) { + return new Codec() { + @Override + public int getId() { + return id; + } + + @Override + public InputStream decode(InputStream inputStream) { + return inputStream; + } + + @Override + public OutputStream encode(OutputStream outputStream) { + return new FilterOutputStream(outputStream) { + @Override + public void close() throws IOException { + out.write(new byte[additionalBytes]); + super.close(); + } + }; + } + }; + } }