Permalink
Browse files

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
  • Loading branch information...
weissi authored and Lukasa committed Jan 3, 2019
1 parent 695b1eb commit 08eb3259ce6041f0fa8ba6adee46a002ce66eacf
@@ -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) {
@@ -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
@@ -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.
@@ -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),
]
}
}
@@ -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 {

0 comments on commit 08eb325

Please sign in to comment.