Improving topic write buffer capacity accounting

topic/src/main/java/tech/ydb/topic/description/Codec.java

@@ -5,7 +5,6 @@
 import java.io.OutputStream;
 
 /**
-
  *
  * Interface for custom codec implementation.
  * <p>
@@ -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
      *

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);

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);

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);

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);

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) {

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<SentMessage> updateSeqNo(long newSeqNo) {
 
     CompletableFuture<WriteAck> 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<WriteAck> 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<WriteAck> 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<WriteAck> accept(Message message, YdbTransaction tx, long msgSize) {
-        EnqueuedMessage msg = new EnqueuedMessage(new MessageMeta(message, tx), msgSize);
+    private CompletableFuture<WriteAck> 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<WriteAck> 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,23 +279,35 @@ 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);
         }
         readyNotify.run();
     }
 
-    private class SkippedMsg implements EncodedMsg {
+    private static class SkippedMsg implements EncodedMsg {
         private final CompletableFuture<WriteAck> 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<WriteAck> ackFuture;
         private final Throwable problem;
 

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);