Issue-2734 - Clamp buffer to maximum upon large write operation

Sources/NIOCore/ByteBuffer-core.swift

@@ -825,6 +825,23 @@ public struct ByteBuffer {
         }
         return true
     }
+
+    /// Shrinks the buffers capacity to a `desiredCapacity` without altering readably bytes.
+    ///
+    /// - Parameter desiredCapacity: The desired capacity for the buffers capacity to be shrunken to
+    /// - Returns: Bool indicating whether the buffer capacity has been shrunk to the desiredCapacity.
+    @inlinable
+    @discardableResult public mutating func shrinkBufferCapacity(to desiredCapacity: Int) -> Bool {
+        guard desiredCapacity < capacity, desiredCapacity > readableBytes else {
+            return false
+        }
+
+        let shrunkenCapacity = _toCapacity(desiredCapacity)
+        self._storage = self._storage.reallocSlice(self._slice.lowerBound..<self._slice.lowerBound + shrunkenCapacity, capacity: shrunkenCapacity)
+        self._slice = self._storage.fullSlice
+
+        return true
+    }
 
     /// The reader index or the number of bytes previously read from this `ByteBuffer`. `readerIndex` is `0` for a
     /// newly allocated `ByteBuffer`.

Sources/NIOCore/Codec.swift

@@ -768,9 +768,16 @@ public final class MessageToByteHandler<Encoder: MessageToByteEncoder>: ChannelO
     private var state: State = .notInChannelYet
     private let encoder: Encoder
     private var buffer: ByteBuffer? = nil
+    private let desiredBufferCapacity: Int?
 
+    public init(_ encoder: Encoder, desiredBufferCapacity: Int) {
+        self.encoder = encoder
+        self.desiredBufferCapacity = desiredBufferCapacity.nextPowerOf2()
+    }
+
     public init(_ encoder: Encoder) {
         self.encoder = encoder
+        self.desiredBufferCapacity = nil
     }
 }
 
@@ -817,4 +824,11 @@ extension MessageToByteHandler {
             context.fireErrorCaught(error)
         }
     }
+
+    public func flush(context: ChannelHandlerContext) {
+        context.flush()
+        if let desiredBufferCapacity {
+            self.buffer?.shrinkBufferCapacity(to: desiredBufferCapacity)
+        }
+    }
 }

Tests/NIOCoreTests/ByteBufferTest.swift

@@ -1837,6 +1837,39 @@ class ByteBufferTest: XCTestCase {
         XCTAssertEqual(0xaa, buffer.getInteger(at: 0, as: UInt8.self))
         XCTAssertEqual(0, buffer2.readableBytes)
     }
+
+    func testShrinkBufferCapacity() {
+        let desiredCapacity = 1024
+        var buffer = self.allocator.buffer(capacity: 512)
+
+        // For any item, it should not shrink buffer capacity to a value larger than the current buffer capacity
+        buffer.clear()
+        buffer.writeString("Any item")
+        XCTAssertFalse(buffer.shrinkBufferCapacity(to: 2048))
+        XCTAssertEqual(buffer.capacity, 512)
+
+        // If desired capacity are less than or equal to buffer capacity, should not shrink
+        buffer.clear()
+        buffer.writeString(String(repeating: "x", count: desiredCapacity))
+        XCTAssertEqual(buffer.capacity, 1024) // Before
+        XCTAssertFalse(buffer.shrinkBufferCapacity(to: desiredCapacity))
+        XCTAssertEqual(buffer.capacity, 1024) // After
+
+        // If desiredCapacity is less than readable bytes, do not shrink
+        buffer.clear()
+        buffer.writeString(String(repeating: "x", count: desiredCapacity + 1))
+        XCTAssertEqual(buffer.capacity, 2048)
+        XCTAssertFalse(buffer.shrinkBufferCapacity(to: desiredCapacity))
+        XCTAssertEqual(buffer.capacity, 2048)
+
+        // If desired capacity is greater than the readable bytes and less than buffer capacity, should shrink to desiredCapacity.nextPowerOf2Clamp
+        // Where desiredCapacity.nextPowerOf2Clamp == desiredCapacity as desiredCapacity should already be a power of 2 value
+        buffer.clear()
+        buffer.writeString(String(repeating: "x", count: desiredCapacity - 1))
+        XCTAssertEqual(buffer.capacity, 2048)
+        XCTAssertTrue(buffer.shrinkBufferCapacity(to: desiredCapacity))
+        XCTAssertEqual(buffer.capacity, 1024)
+    }
 
     func testDumpBytesFormat() throws {
         self.buf.clear()