diff --git a/Sources/NIO/ByteBuffer-core.swift b/Sources/NIO/ByteBuffer-core.swift index b5f6331b5f..dcc89425bc 100644 --- a/Sources/NIO/ByteBuffer-core.swift +++ b/Sources/NIO/ByteBuffer-core.swift @@ -218,12 +218,13 @@ public struct ByteBuffer { return new } - public func reallocStorage(capacity: Capacity) { - let ptr = self.allocator.realloc(self.bytes, Int(capacity)) + public func reallocStorage(capacity minimumNeededCapacity: Capacity) { + let newCapacity = minimumNeededCapacity.nextPowerOf2ClampedToMax() + let ptr = self.allocator.realloc(self.bytes, Int(newCapacity)) /* bind the memory so we can assume it elsewhere to be bound to UInt8 */ - ptr.bindMemory(to: UInt8.self, capacity: Int(capacity)) + ptr.bindMemory(to: UInt8.self, capacity: Int(newCapacity)) self.bytes = ptr - self.capacity = capacity + self.capacity = newCapacity self.fullSlice = 0.. self._slice.upperBound { - // double the capacity, we may want to use different strategies depending on the actual current capacity later on. - var newCapacity = max(1, toCapacity(self.capacity)) - - // double the capacity until the requested capacity can be full-filled - repeat { - precondition(newCapacity != Capacity.max, "cannot make ByteBuffers larger than \(newCapacity)") - if newCapacity < (Capacity.max >> 1) { - newCapacity = newCapacity << 1 + let totalNeededCapacityWhenKeepingSlice = self._slice.lowerBound + index + capacity + if totalNeededCapacityWhenKeepingSlice > self._slice.upperBound { + // we need to at least adjust the slice's upper bound which we can do as we're the unique owner of the storage, + // let's see if adjusting the slice's upper bound buys us enough storage + if totalNeededCapacityWhenKeepingSlice > self._storage.capacity { + let newStorageMinCapacity = index + capacity + // nope, we need to actually re-allocate again. If our slice does not start at 0, let's also rebase + if self._slice.lowerBound == 0 { + self._storage.reallocStorage(capacity: newStorageMinCapacity) } else { - newCapacity = Capacity.max + self._storage = self._storage.reallocSlice(self._slice.lowerBound ..< self._slice.upperBound, + capacity: newStorageMinCapacity) } - } while newCapacity < index || newCapacity - index < capacity - - self._storage.reallocStorage(capacity: newCapacity) - self._slice = _slice.lowerBound..<_slice.lowerBound + newCapacity + self._slice = self._storage.fullSlice + } else { + // yes, let's just extend the slice until the end of the buffer + self._slice = _slice.lowerBound ..< self._storage.capacity + } } + assert(self._slice.lowerBound + index + capacity <= self._slice.upperBound) + assert(self._slice.lowerBound >= 0, "illegal slice: negative lower bound: \(self._slice.lowerBound)") + assert(self._slice.upperBound <= self._storage.capacity, "illegal slice: upper bound (\(self._slice.upperBound)) exceeds capacity: \(self._storage.capacity)") } // MARK: Internal API diff --git a/Tests/NIOTests/ByteBufferTest+XCTest.swift b/Tests/NIOTests/ByteBufferTest+XCTest.swift index d4bb52e611..66b5a2ad54 100644 --- a/Tests/NIOTests/ByteBufferTest+XCTest.swift +++ b/Tests/NIOTests/ByteBufferTest+XCTest.swift @@ -30,6 +30,8 @@ extension ByteBufferTest { ("testEqualsComparesReadBuffersOnly", testEqualsComparesReadBuffersOnly), ("testSimpleReadTest", testSimpleReadTest), ("testSimpleWrites", testSimpleWrites), + ("testMakeSureUniquelyOwnedSliceDoesNotGetReallocatedOnWrite", testMakeSureUniquelyOwnedSliceDoesNotGetReallocatedOnWrite), + ("testWriteToUniquelyOwnedSliceWhichTriggersAReallocation", testWriteToUniquelyOwnedSliceWhichTriggersAReallocation), ("testReadWrite", testReadWrite), ("testStaticStringReadTests", testStaticStringReadTests), ("testString", testString), diff --git a/Tests/NIOTests/ByteBufferTest.swift b/Tests/NIOTests/ByteBufferTest.swift index f675cc3d1c..b791c9b2c6 100644 --- a/Tests/NIOTests/ByteBufferTest.swift +++ b/Tests/NIOTests/ByteBufferTest.swift @@ -94,6 +94,39 @@ class ByteBufferTest: XCTestCase { XCTAssertEqual(6, buf.readableBytes) } + func makeSliceToBufferWhichIsDeallocated() -> ByteBuffer { + var buf = self.allocator.buffer(capacity: 16) + let oldCapacity = buf.capacity + buf.write(bytes: 0..<16) + XCTAssertEqual(oldCapacity, buf.capacity) + return buf.getSlice(at: 15, length: 1)! + } + + func testMakeSureUniquelyOwnedSliceDoesNotGetReallocatedOnWrite() { + var slice = self.makeSliceToBufferWhichIsDeallocated() + XCTAssertEqual(1, slice.capacity) + let oldStorageBegin = slice.withUnsafeReadableBytes { ptr in + return UInt(bitPattern: ptr.baseAddress!) + } + slice.set(integer: 1 as UInt8, at: 0) + let newStorageBegin = slice.withUnsafeReadableBytes { ptr in + return UInt(bitPattern: ptr.baseAddress!) + } + XCTAssertEqual(oldStorageBegin, newStorageBegin) + } + + func testWriteToUniquelyOwnedSliceWhichTriggersAReallocation() { + var slice = self.makeSliceToBufferWhichIsDeallocated() + XCTAssertEqual(1, slice.capacity) + // this will cause a re-allocation, the whole buffer should be 32 bytes then, the slice having 17 of that. + // this fills 16 bytes so will still fit + slice.write(bytes: Array(16..<32)) + XCTAssertEqual(Array(15..<32), slice.readBytes(length: slice.readableBytes)!) + + // and this will need another re-allocation + slice.write(bytes: Array(32..<47)) + } + func testReadWrite() { buf.write(string: "X") buf.write(string: "Y")