better DispatchData support (#732)

Motivation:

DispatchData's APIs are less than optimal and the current one we used
did allocate at least on Linux to copy bytes out.

Modifications:

introduce ByteBuffer.write/read/set/getDispatchData

Result:

when Dispatch gets fixed, lower allocations

IntegrationTests/tests_04_performance/test_01_resources/template/Sources/SwiftBootstrap/SwiftMain.swift

@@ -308,15 +308,16 @@ public func swiftMain() -> Int {
         let foundationData = "A".data(using: .utf8)!
         @inline(never)
         func doWrites(buffer: inout ByteBuffer) {
-            /* all of those should be 0 allocations */
-
+            /* these ones are zero allocations */
             // buffer.write(bytes: foundationData) // see SR-7542
             buffer.write(bytes: [0x41])
-            buffer.write(bytes: dispatchData)
             buffer.write(bytes: "A".utf8)
             buffer.write(string: "A")
             buffer.write(staticString: "A")
             buffer.write(integer: 0x41, as: UInt8.self)
+
+            /* those down here should be one allocation each (on Linux) */
+            buffer.write(bytes: dispatchData) // see https://bugs.swift.org/browse/SR-9597
         }
         @inline(never)
         func doReads(buffer: inout ByteBuffer) {

Sources/NIO/ByteBuffer-aux.swift

@@ -12,6 +12,8 @@
 //
 //===----------------------------------------------------------------------===//
 
+import Dispatch
+
 extension ByteBuffer {
 
     // MARK: Bytes ([UInt8]) APIs
@@ -169,6 +171,78 @@ extension ByteBuffer {
         return self.getString(at: self.readerIndex, length: length)! /* must work, enough readable bytes */
     }
 
+    // MARK: DispatchData APIs
+    /// Write `dispatchData` into this `ByteBuffer`, moving the writer index forward appropriately.
+    ///
+    /// - parameters:
+    ///     - dispatchData: The `DispatchData` instance to write to the `ByteBuffer`.
+    /// - returns: The number of bytes written.
+    @discardableResult
+    public mutating func write(dispatchData: DispatchData) -> Int {
+        let written = self.set(dispatchData: dispatchData, at: self.writerIndex)
+        self._moveWriterIndex(forwardBy: written)
+        return written
+    }
+
+    /// Write `dispatchData` into this `ByteBuffer` at `index`. Does not move the writer index.
+    ///
+    /// - parameters:
+    ///     - dispatchData: The `DispatchData` to write.
+    ///     - index: The index for the first serialized byte.
+    /// - returns: The number of bytes written.
+    @discardableResult
+    public mutating func set(dispatchData: DispatchData, at index: Int) -> Int {
+        let allBytesCount = dispatchData.count
+        self.reserveCapacity(index + allBytesCount)
+        self.withVeryUnsafeMutableBytes { destCompleteStorage in
+            assert(destCompleteStorage.count >= index + allBytesCount)
+            let dest = destCompleteStorage[index ..< index + allBytesCount]
+            dispatchData.copyBytes(to: .init(rebasing: dest), count: dest.count)
+        }
+        return allBytesCount
+    }
+
+    /// Get the bytes at `index` from this `ByteBuffer` as a `DispatchData`. Does not move the reader index.
+    ///
+    /// - note: Please consider using `readDispatchData` which is a safer alternative that automatically maintains the
+    ///         `readerIndex` and won't allow you to read uninitialized memory.
+    /// - warning: This method allows the user to read any of the bytes in the `ByteBuffer`'s storage, including
+    ///           _uninitialized_ ones. To use this API in a safe way the user needs to make sure all the requested
+    ///           bytes have been written before and are therefore initialized. Note that bytes between (including)
+    ///           `readerIndex` and (excluding) `writerIndex` are always initialized by contract and therefore must be
+    ///           safe to read.
+    /// - parameters:
+    ///     - index: The starting index into `ByteBuffer` containing the string of interest.
+    ///     - length: The number of bytes.
+    /// - returns: A `DispatchData` value deserialized from this `ByteBuffer` or `nil` if the requested bytes aren't contained in this `ByteBuffer`.
+    public func getDispatchData(at index: Int, length: Int) -> DispatchData? {
+        precondition(index >= 0, "index must not be negative")
+        precondition(length >= 0, "length must not be negative")
+        return self.withVeryUnsafeBytes { pointer in
+            guard index <= pointer.count - length else {
+                return nil
+            }
+            return DispatchData(bytes: UnsafeRawBufferPointer(rebasing: pointer[index..<(index+length)]))
+        }
+    }
+
+    /// Read `length` bytes off this `ByteBuffer` and return them as a `DispatchData`. Move the reader index forward by `length`.
+    ///
+    /// - parameters:
+    ///     - length: The number of bytes.
+    /// - returns: A `DispatchData` value containing the bytes from this `ByteBuffer` or `nil` if there aren't at least `length` bytes readable.
+    public mutating func readDispatchData(length: Int) -> DispatchData? {
+        precondition(length >= 0, "length must not be negative")
+        guard self.readableBytes >= length else {
+            return nil
+        }
+        defer {
+            self._moveReaderIndex(forwardBy: length)
+        }
+        return self.getDispatchData(at: self.readerIndex, length: length)! /* must work, enough readable bytes */
+    }
+
+
     // MARK: Other APIs
 
     /// Yields an immutable buffer pointer containing this `ByteBuffer`'s readable bytes. Will move the reader index

Sources/NIO/ByteBuffer-core.swift

@@ -548,6 +548,16 @@ public struct ByteBuffer {
         return try body(.init(self._slicedStorageBuffer))
     }
 
+    /// 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 `withUnsafeMutableWritableBytes`.
+    ///
+    /// - warning: Do not escape the pointer from the closure for later use.
+    @inlinable
+    public mutating func withVeryUnsafeMutableBytes<T>(_ body: (UnsafeMutableRawBufferPointer) throws -> T) rethrows -> T {
+        self._copyStorageAndRebaseIfNeeded() // this will trigger a CoW if necessary
+        return try body(.init(self._slicedStorageBuffer))
+    }
+
     /// Yields a buffer pointer containing this `ByteBuffer`'s readable bytes.
     ///
     /// - warning: Do not escape the pointer from the closure for later use.

Tests/NIOTests/ByteBufferTest+XCTest.swift

@@ -127,6 +127,13 @@ extension ByteBufferTest {
                 ("testReadWithFunctionsThatReturnNumberOfReadBytesAreDiscardable", testReadWithFunctionsThatReturnNumberOfReadBytesAreDiscardable),
                 ("testWriteAndSetAndGetAndReadEncoding", testWriteAndSetAndGetAndReadEncoding),
                 ("testPossiblyLazilyBridgedString", testPossiblyLazilyBridgedString),
+                ("testWithVeryUnsafeMutableBytesWorksOnEmptyByteBuffer", testWithVeryUnsafeMutableBytesWorksOnEmptyByteBuffer),
+                ("testWithVeryUnsafeMutableBytesYieldsPointerToWholeStorage", testWithVeryUnsafeMutableBytesYieldsPointerToWholeStorage),
+                ("testWithVeryUnsafeMutableBytesYieldsPointerToWholeStorageAndCanBeWritenTo", testWithVeryUnsafeMutableBytesYieldsPointerToWholeStorageAndCanBeWritenTo),
+                ("testWithVeryUnsafeMutableBytesDoesCoW", testWithVeryUnsafeMutableBytesDoesCoW),
+                ("testWithVeryUnsafeMutableBytesDoesCoWonSlices", testWithVeryUnsafeMutableBytesDoesCoWonSlices),
+                ("testGetDispatchDataWorks", testGetDispatchDataWorks),
+                ("testGetDispatchDataReadWrite", testGetDispatchDataReadWrite),
            ]
    }
 }

Tests/NIOTests/ByteBufferTest.swift

@@ -940,11 +940,13 @@ class ByteBufferTest: XCTestCase {
         testIndexOrLengthFunc({ buf.readSlice(length: $0) })
         testIndexOrLengthFunc({ buf.readBytes(length: $0) })
         testIndexOrLengthFunc({ buf.readData(length: $0) })
+        testIndexOrLengthFunc({ buf.readDispatchData(length: $0) })
 
         testIndexAndLengthFunc(buf.getBytes)
         testIndexAndLengthFunc(buf.getData)
         testIndexAndLengthFunc(buf.getSlice)
         testIndexAndLengthFunc(buf.getString)
+        testIndexAndLengthFunc(buf.getDispatchData)
     }
 
     func testWriteForContiguousCollections() throws {
@@ -1579,6 +1581,111 @@ class ByteBufferTest: XCTestCase {
         XCTAssertEqual(10, written)
         XCTAssertEqual("abcäää\n", String(decoding: self.buf.readableBytesView, as: Unicode.UTF8.self))
     }
+
+    func testWithVeryUnsafeMutableBytesWorksOnEmptyByteBuffer() {
+        var buf = self.allocator.buffer(capacity: 0)
+        XCTAssertEqual(0, buf.capacity)
+        buf.withVeryUnsafeMutableBytes { ptr in
+            XCTAssertEqual(0, ptr.count)
+        }
+    }
+
+    func testWithVeryUnsafeMutableBytesYieldsPointerToWholeStorage() {
+        var buf = self.allocator.buffer(capacity: 16)
+        let capacity = buf.capacity
+        XCTAssertGreaterThanOrEqual(capacity, 16)
+        buf.write(string: "1234")
+        XCTAssertEqual(capacity, buf.capacity)
+        buf.withVeryUnsafeMutableBytes { ptr in
+            XCTAssertEqual(capacity, ptr.count)
+            XCTAssertEqual("1234", String(decoding: ptr[0..<4], as: Unicode.UTF8.self))
+        }
+    }
+
+    func testWithVeryUnsafeMutableBytesYieldsPointerToWholeStorageAndCanBeWritenTo() {
+        var buf = self.allocator.buffer(capacity: 16)
+        let capacity = buf.capacity
+        XCTAssertGreaterThanOrEqual(capacity, 16)
+        buf.write(string: "1234")
+        XCTAssertEqual(capacity, buf.capacity)
+        buf.withVeryUnsafeMutableBytes { ptr in
+            XCTAssertEqual(capacity, ptr.count)
+            XCTAssertEqual("1234", String(decoding: ptr[0..<4], as: Unicode.UTF8.self))
+            UnsafeMutableRawBufferPointer(rebasing: ptr[4..<8]).copyBytes(from: "5678".utf8)
+        }
+        buf.moveWriterIndex(forwardBy: 4)
+        XCTAssertEqual("12345678", buf.readString(length: buf.readableBytes))
+
+        buf.withVeryUnsafeMutableBytes { ptr in
+            XCTAssertEqual(capacity, ptr.count)
+            XCTAssertEqual("12345678", String(decoding: ptr[0..<8], as: Unicode.UTF8.self))
+            ptr[0] = "X".utf8.first!
+            UnsafeMutableRawBufferPointer(rebasing: ptr[8..<16]).copyBytes(from: "abcdefgh".utf8)
+        }
+        buf.moveWriterIndex(forwardBy: 8)
+        XCTAssertEqual("abcdefgh", buf.readString(length: buf.readableBytes))
+        XCTAssertEqual("X", buf.getString(at: 0, length: 1))
+    }
+
+    func testWithVeryUnsafeMutableBytesDoesCoW() {
+        var buf = self.allocator.buffer(capacity: 16)
+        let capacity = buf.capacity
+        XCTAssertGreaterThanOrEqual(capacity, 16)
+        buf.write(string: "1234")
+        let bufCopy = buf
+        XCTAssertEqual(capacity, buf.capacity)
+        buf.withVeryUnsafeMutableBytes { ptr in
+            XCTAssertEqual(capacity, ptr.count)
+            XCTAssertEqual("1234", String(decoding: ptr[0..<4], as: Unicode.UTF8.self))
+            UnsafeMutableRawBufferPointer(rebasing: ptr[0..<8]).copyBytes(from: "abcdefgh".utf8)
+        }
+        buf.moveWriterIndex(forwardBy: 4)
+        XCTAssertEqual("1234", String(decoding: bufCopy.readableBytesView, as: Unicode.UTF8.self))
+        XCTAssertEqual("abcdefgh", String(decoding: buf.readableBytesView, as: Unicode.UTF8.self))
+    }
+
+    func testWithVeryUnsafeMutableBytesDoesCoWonSlices() {
+        var buf = self.allocator.buffer(capacity: 16)
+        let capacity = buf.capacity
+        XCTAssertGreaterThanOrEqual(capacity, 16)
+        buf.write(string: "1234567890")
+        var buf2 = buf.getSlice(at: 4, length: 4)!
+        XCTAssertEqual(capacity, buf.capacity)
+        let capacity2 = buf2.capacity
+        buf2.withVeryUnsafeMutableBytes { ptr in
+            XCTAssertEqual(capacity2, ptr.count)
+            XCTAssertEqual("5678", String(decoding: ptr[0..<4], as: Unicode.UTF8.self))
+            UnsafeMutableRawBufferPointer(rebasing: ptr[0..<4]).copyBytes(from: "QWER".utf8)
+        }
+        XCTAssertEqual("QWER", String(decoding: buf2.readableBytesView, as: Unicode.UTF8.self))
+        XCTAssertEqual("1234567890", String(decoding: buf.readableBytesView, as: Unicode.UTF8.self))
+    }
+
+    func testGetDispatchDataWorks() {
+        self.buf.clear()
+        self.buf.write(string: "abcdefgh")
+
+        XCTAssertEqual(0, self.buf.getDispatchData(at: 7, length: 0)!.count)
+        XCTAssertNil(self.buf.getDispatchData(at: self.buf.capacity, length: 1))
+        XCTAssertEqual("abcdefgh", String(decoding: self.buf.getDispatchData(at: 0, length: 8)!, as: Unicode.UTF8.self))
+        XCTAssertEqual("ef", String(decoding: self.buf.getDispatchData(at: 4, length: 2)!, as: Unicode.UTF8.self))
+    }
+
+    func testGetDispatchDataReadWrite() {
+        var buffer = UnsafeMutableRawBufferPointer.allocate(byteCount: 4, alignment: 0)
+        buffer.copyBytes(from: "1234".utf8)
+        defer {
+            buffer.deallocate()
+        }
+        self.buf.clear()
+        self.buf.write(string: "abcdefgh")
+        self.buf.write(dispatchData: DispatchData.empty)
+        self.buf.write(dispatchData: DispatchData(bytes: UnsafeRawBufferPointer(buffer)))
+        XCTAssertEqual(12, self.buf.readableBytes)
+        XCTAssertEqual("abcdefgh1234", String(decoding: self.buf.readDispatchData(length: 12)!, as: Unicode.UTF8.self))
+        XCTAssertNil(self.buf.readDispatchData(length: 1))
+        XCTAssertEqual(0, self.buf.readDispatchData(length: 0)?.count ?? 12)
+    }
 }
 
 private enum AllocationExpectationState: Int {