Add option to reserve writable capacity to writeWithUnsafeMutableBytes (

#1105) (#1175)

Motivation:

writeWithUnsafeMutableBytes cannot tolerate arbitrarily large writes because it is not possible to resize the buffer by the time body receives the buffer pointer. This change provides an option for the user to reserve writable capacity before the pointer is passed to body.

Modifications:

Added a backward compatible optional parameter to writeWithUnsafeMutableBytes which can reserve writable capacity before the pointer is passed to body.

Result:

Additive change only. Users can optionally specify a minimum writable capacity when calling writeWithUnsafeMutableBytes.

Sources/NIO/ByteBuffer-core.swift

@@ -492,14 +492,32 @@ public struct ByteBuffer {
         return try body(.init(rebasing: self._slicedStorageBuffer.dropFirst(self.writerIndex)))
     }
 
+    /// This vends a pointer of the `ByteBuffer` at the `writerIndex` after ensuring that the buffer has at least `minimumWritableBytes` of writable bytes available.
+    ///
+    /// - warning: Do not escape the pointer from the closure for later use.
+    ///
+    /// - parameters:
+    ///     - minimumWritableBytes: The number of writable bytes to reserve capacity for before vending the `ByteBuffer` pointer to `body`.
+    ///     - body: The closure that will accept the yielded bytes and return the number of bytes written.
+    /// - returns: The number of bytes written.
     @discardableResult
     @inlinable
-    public mutating func writeWithUnsafeMutableBytes(_ body: (UnsafeMutableRawBufferPointer) throws -> Int) rethrows -> Int {
-        let bytesWritten = try withUnsafeMutableWritableBytes(body)
+    public mutating func writeWithUnsafeMutableBytes(minimumWritableBytes: Int, _ body: (UnsafeMutableRawBufferPointer) throws -> Int) rethrows -> Int {
+        if minimumWritableBytes > 0 {
+            self.reserveCapacity(self.writerIndex + minimumWritableBytes)
+        }
+        let bytesWritten = try self.withUnsafeMutableWritableBytes(body)
         self._moveWriterIndex(to: self._writerIndex + _toIndex(bytesWritten))
         return bytesWritten
     }
 
+    @available(*, deprecated, message: "please use writeWithUnsafeMutableBytes(minimumWritableBytes:_:) instead to ensure sufficient write capacity.")
+    @discardableResult
+    @inlinable
+    public mutating func writeWithUnsafeMutableBytes(_ body: (UnsafeMutableRawBufferPointer) throws -> Int) rethrows -> Int {
+        return try self.writeWithUnsafeMutableBytes(minimumWritableBytes: 0, body)
+    }
+
     /// This vends a pointer to the storage of the `ByteBuffer`. It's marked as _very unsafe_ because it might contain
     /// uninitialised memory and it's undefined behaviour to read it. In most cases you should use `withUnsafeReadableBytes`.
     ///

Sources/NIO/NonBlockingFileIO.swift

@@ -188,7 +188,7 @@ public struct NonBlockingFileIO {
         return self.threadPool.runIfActive(eventLoop: eventLoop) { () -> ByteBuffer in
             var bytesRead = 0
             while bytesRead < byteCount {
-                let n = try buf.writeWithUnsafeMutableBytes { ptr in
+                let n = try buf.writeWithUnsafeMutableBytes(minimumWritableBytes: byteCount - bytesRead) { ptr in
                     let res = try fileHandle.withUnsafeFileDescriptor { descriptor in
                         try Posix.read(descriptor: descriptor,
                                        pointer: ptr.baseAddress!,

Sources/NIO/SocketChannel.swift

@@ -15,7 +15,7 @@
 extension ByteBuffer {
     mutating func withMutableWritePointer(body: (UnsafeMutableRawBufferPointer) throws -> IOResult<Int>) rethrows -> IOResult<Int> {
         var singleResult: IOResult<Int>!
-        _ = try self.writeWithUnsafeMutableBytes { ptr in
+        _ = try self.writeWithUnsafeMutableBytes(minimumWritableBytes: 0) { ptr in
             let localWriteResult = try body(ptr)
             singleResult = localWriteResult
             switch localWriteResult {

Tests/NIOHTTP1Tests/HTTPServerClientTest.swift

@@ -191,7 +191,7 @@ class HTTPServerClientTest : XCTestCase {
 
                 case "/massive-response":
                     var buf = context.channel.allocator.buffer(capacity: HTTPServerClientTest.massiveResponseLength)
-                    buf.writeWithUnsafeMutableBytes { targetPtr in
+                    buf.writeWithUnsafeMutableBytes(minimumWritableBytes: HTTPServerClientTest.massiveResponseLength) { targetPtr in
                         return HTTPServerClientTest.massiveResponseBytes.withUnsafeBytes { srcPtr in
                             precondition(targetPtr.count >= srcPtr.count)
                             targetPtr.copyMemory(from: srcPtr)

Tests/NIOTests/ByteBufferTest+XCTest.swift

@@ -41,6 +41,8 @@ extension ByteBufferTest {
                 ("testCoWWorks", testCoWWorks),
                 ("testWithMutableReadPointerMovesReaderIndexAndReturnsNumBytesConsumed", testWithMutableReadPointerMovesReaderIndexAndReturnsNumBytesConsumed),
                 ("testWithMutableWritePointerMovesWriterIndexAndReturnsNumBytesWritten", testWithMutableWritePointerMovesWriterIndexAndReturnsNumBytesWritten),
+                ("testWithMutableWritePointerWithMinimumSpecifiedAdjustsCapacity", testWithMutableWritePointerWithMinimumSpecifiedAdjustsCapacity),
+                ("testWithMutableWritePointerWithMinimumSpecifiedWhileAtMaxCapacity", testWithMutableWritePointerWithMinimumSpecifiedWhileAtMaxCapacity),
                 ("testSetGetInt8", testSetGetInt8),
                 ("testSetGetInt16", testSetGetInt16),
                 ("testSetGetInt32", testSetGetInt32),

Tests/NIOTests/ByteBufferTest.swift

@@ -228,11 +228,58 @@ class ByteBufferTest: XCTestCase {
     func testWithMutableWritePointerMovesWriterIndexAndReturnsNumBytesWritten() {
         XCTAssertEqual(0, buf.writerIndex)
 
-        let bytesWritten = buf.writeWithUnsafeMutableBytes { (_: UnsafeMutableRawBufferPointer) in return 5 }
+        let bytesWritten = buf.writeWithUnsafeMutableBytes(minimumWritableBytes: 5) {
+            XCTAssertTrue($0.count >= 5)
+            return 5
+        }
         XCTAssertEqual(5, bytesWritten)
         XCTAssertEqual(5, buf.writerIndex)
     }
 
+    func testWithMutableWritePointerWithMinimumSpecifiedAdjustsCapacity() {
+        XCTAssertEqual(0, buf.writerIndex)
+        XCTAssertEqual(512, buf.capacity)
+
+        var bytesWritten = buf.writeWithUnsafeMutableBytes(minimumWritableBytes: 256) {
+            XCTAssertTrue($0.count >= 256)
+            return 256
+        }
+        XCTAssertEqual(256, bytesWritten)
+        XCTAssertEqual(256, buf.writerIndex)
+        XCTAssertEqual(512, buf.capacity)
+
+        bytesWritten += buf.writeWithUnsafeMutableBytes(minimumWritableBytes: 1024) {
+            XCTAssertTrue($0.count >= 1024)
+            return 1024
+        }
+        let expectedBytesWritten = 256 + 1024
+        XCTAssertEqual(expectedBytesWritten, bytesWritten)
+        XCTAssertEqual(expectedBytesWritten, buf.writerIndex)
+        XCTAssertTrue(buf.capacity >= expectedBytesWritten)
+    }
+
+    func testWithMutableWritePointerWithMinimumSpecifiedWhileAtMaxCapacity() {
+        XCTAssertEqual(0, buf.writerIndex)
+        XCTAssertEqual(512, buf.capacity)
+
+        var bytesWritten = buf.writeWithUnsafeMutableBytes(minimumWritableBytes: 512) {
+            XCTAssertTrue($0.count >= 512)
+            return 512
+        }
+        XCTAssertEqual(512, bytesWritten)
+        XCTAssertEqual(512, buf.writerIndex)
+        XCTAssertEqual(512, buf.capacity)
+
+        bytesWritten += buf.writeWithUnsafeMutableBytes(minimumWritableBytes: 1) {
+            XCTAssertTrue($0.count >= 1)
+            return 1
+        }
+        let expectedBytesWritten = 512 + 1
+        XCTAssertEqual(expectedBytesWritten, bytesWritten)
+        XCTAssertEqual(expectedBytesWritten, buf.writerIndex)
+        XCTAssertTrue(buf.capacity >= expectedBytesWritten)
+    }
+
     func testSetGetInt8() throws {
         try setGetInt(index: 0, v: Int8.max)
     }
@@ -598,7 +645,7 @@ class ByteBufferTest: XCTestCase {
         let cap = buf.capacity
         var otherBuf = buf
         XCTAssertEqual(otherBuf, buf)
-        otherBuf?.writeWithUnsafeMutableBytes { ptr in
+        otherBuf?.writeWithUnsafeMutableBytes(minimumWritableBytes: 0) { ptr in
             XCTAssertEqual(cap, ptr.count)
             let intPtr = ptr.baseAddress!.bindMemory(to: UInt8.self, capacity: ptr.count)
             for i in 0..<ptr.count {

Tests/NIOTests/DatagramChannelTests.swift

@@ -258,7 +258,7 @@ final class DatagramChannelTests: XCTestCase {
             // will cause EMSGSIZE.
             let bufferSize = 1024 * 5
             var buffer = self.firstChannel.allocator.buffer(capacity: bufferSize)
-            buffer.writeWithUnsafeMutableBytes {
+            buffer.writeWithUnsafeMutableBytes(minimumWritableBytes: bufferSize) {
                 _ = memset($0.baseAddress!, 4, $0.count)
                 return $0.count
             }
@@ -282,7 +282,7 @@ final class DatagramChannelTests: XCTestCase {
         // We want to try to trigger EMSGSIZE. To be safe, we're going to allocate a 10MB buffer here and fill it.
         let bufferSize = 1024 * 1024 * 10
         var buffer = self.firstChannel.allocator.buffer(capacity: bufferSize)
-        buffer.writeWithUnsafeMutableBytes {
+        buffer.writeWithUnsafeMutableBytes(minimumWritableBytes: bufferSize) {
             _ = memset($0.baseAddress!, 4, $0.count)
             return $0.count
         }
@@ -305,7 +305,7 @@ final class DatagramChannelTests: XCTestCase {
         // We want to try to trigger EMSGSIZE. To be safe, we're going to allocate a 10MB buffer here and fill it.
         let bufferSize = 1024 * 1024 * 10
         var buffer = self.firstChannel.allocator.buffer(capacity: bufferSize)
-        buffer.writeWithUnsafeMutableBytes {
+        buffer.writeWithUnsafeMutableBytes(minimumWritableBytes: bufferSize) {
             _ = memset($0.baseAddress!, 4, $0.count)
             return $0.count
         }
@@ -338,7 +338,7 @@ final class DatagramChannelTests: XCTestCase {
         // We want to try to trigger EMSGSIZE. To be safe, we're going to allocate a 10MB buffer here and fill it.
         let bufferSize = 1024 * 1024 * 10
         var buffer = self.firstChannel.allocator.buffer(capacity: bufferSize)
-        buffer.writeWithUnsafeMutableBytes {
+        buffer.writeWithUnsafeMutableBytes(minimumWritableBytes: bufferSize) {
             _ = memset($0.baseAddress!, 4, $0.count)
             return $0.count
         }