Moving PooledRecvBufferAllocator from NIOPosix to NIOCore

Sources/NIOPosix/PooledRecvBufferAllocator.swift → Sources/NIOCore/NIOPooledRecvBufferAllocator.swift

@@ -12,147 +12,170 @@
 //
 //===----------------------------------------------------------------------===//
 
-import NIOCore
-
 /// A receive buffer allocator which cycles through a pool of buffers.
-internal struct PooledRecvBufferAllocator {
+///
+/// Channels can read multiple times per cycle (based on `ChannelOptions.maxMessagesPerRead`), and they reuse
+/// the inbound buffer for each read. If a `ChannelHandler` holds onto this buffer, then CoWing will be needed.
+/// A `NIOPooledRecvBufferAllocator` cycles through preallocated buffers to avoid CoWs during the same read cycle.
+public struct NIOPooledRecvBufferAllocator: Sendable {
     // The pool will either use a single buffer (i.e. `buffer`) OR store multiple buffers
     // in `buffers`. If `buffers` is non-empty then `buffer` MUST be `nil`. If `buffer`
     // is non-nil then `buffers` MUST be empty.
     //
     // The backing storage is changed from `buffer` to `buffers` when a second buffer is
     // needed (and if capacity allows).
-    private var buffer: Optional<ByteBuffer>
-    private var buffers: [ByteBuffer]
+    @usableFromInline
+    internal var _buffer: Optional<ByteBuffer>
+    @usableFromInline
+    internal var _buffers: [ByteBuffer]
     /// The index into `buffers` of the index which was last used.
-    private var lastUsedIndex: Int
+    @usableFromInline
+    internal var _lastUsedIndex: Int
 
     /// Maximum number of buffers to store in the pool.
-    internal private(set) var capacity: Int
+    public private(set) var capacity: Int
     /// The receive allocator providing hints for the next buffer size to use.
-    internal var recvAllocator: RecvByteBufferAllocator
+    public var recvAllocator: RecvByteBufferAllocator
 
     /// The return value from the last call to `recvAllocator.record(actualReadBytes:)`.
-    private var mayGrow: Bool
+    @usableFromInline
+    internal var _mayGrow: Bool
 
-    init(capacity: Int, recvAllocator: RecvByteBufferAllocator) {
+    /// Builds a new instance of `NIOPooledRecvBufferAllocator`
+    ///
+    /// - Parameters:
+    ///   - capacity: Maximum number of buffers to store in the pool.
+    ///   - recvAllocator: The receive allocator providing hints for the next buffer size to use.
+    public init(capacity: Int, recvAllocator: RecvByteBufferAllocator) {
         precondition(capacity > 0)
         self.capacity = capacity
-        self.buffer = nil
-        self.buffers = []
-        self.lastUsedIndex = 0
+        self._buffer = nil
+        self._buffers = []
+        self._lastUsedIndex = 0
         self.recvAllocator = recvAllocator
-        self.mayGrow = false
+        self._mayGrow = false
     }
 
     /// Returns the number of buffers in the pool.
-    var count: Int {
-        if self.buffer == nil {
+    public var count: Int {
+        if self._buffer == nil {
             // Empty or switched to `buffers` for storage.
-            return self.buffers.count
+            return self._buffers.count
         } else {
             // `buffer` is non-nil; `buffers` must be empty and the count must be 1.
-            assert(self.buffers.isEmpty)
+            assert(self._buffers.isEmpty)
             return 1
         }
     }
 
     /// Update the capacity of the underlying buffer pool.
-    mutating func updateCapacity(to newCapacity: Int) {
+    ///
+    /// - Parameters:
+    ///   - newCapacity: The new capacity for the underlying buffer pool.
+    public mutating func updateCapacity(to newCapacity: Int) {
         precondition(newCapacity > 0)
 
         if newCapacity > self.capacity {
             self.capacity = newCapacity
-            if !self.buffers.isEmpty {
-                self.buffers.reserveCapacity(newCapacity)
+            if !self._buffers.isEmpty {
+                self._buffers.reserveCapacity(newCapacity)
             }
         } else if newCapacity < self.capacity {
             self.capacity = newCapacity
             // Drop buffers if over capacity.
-            while self.buffers.count > self.capacity {
-                self.buffers.removeLast()
+            while self._buffers.count > self.capacity {
+                self._buffers.removeLast()
             }
             // Reset the last used index.
-            if self.lastUsedIndex >= self.capacity {
-                self.lastUsedIndex = 0
+            if self._lastUsedIndex >= self.capacity {
+                self._lastUsedIndex = 0
             }
         }
     }
 
     /// Record the number of bytes which were read.
     ///
-    /// Returns whether the next buffer will be larger than the last.
-    mutating func record(actualReadBytes: Int) {
-        self.mayGrow = self.recvAllocator.record(actualReadBytes: actualReadBytes)
+    /// - Parameters:
+    ///   - actualReadBytes: Number of bytes being recorded
+    public mutating func record(actualReadBytes: Int) {
+        self._mayGrow = self.recvAllocator.record(actualReadBytes: actualReadBytes)
     }
 
     /// Provides a buffer with enough writable capacity as determined by the underlying
     /// receive allocator to the given closure.
-    mutating func buffer<Result>(
+    ///
+    /// - Parameters:
+    ///    - allocator: `ByteBufferAllocator` used to construct a new buffer if needed
+    ///    - body: Closure where the caller can use the new or existing buffer
+    /// - Returns: A tuple containing the `ByteBuffer` used and the `Result` yielded by the closure provided.
+    @inlinable
+    public mutating func buffer<Result>(
         allocator: ByteBufferAllocator,
         _ body: (inout ByteBuffer) throws -> Result
     ) rethrows -> (ByteBuffer, Result) {
         // Reuse an existing buffer if we can do so without CoWing.
-        if let bufferAndResult = try self.reuseExistingBuffer(body) {
+        if let bufferAndResult = try self._reuseExistingBuffer(body) {
             return bufferAndResult
         } else {
             // No available buffers or the allocator does not offer up buffer sizes; directly
             // allocate a new one.
-            return try self.allocateNewBuffer(using: allocator, body)
+            return try self._allocateNewBuffer(using: allocator, body)
         }
     }
 
-    private mutating func reuseExistingBuffer<Result>(
+    @inlinable
+    internal mutating func _reuseExistingBuffer<Result>(
         _ body: (inout ByteBuffer) throws -> Result
     ) rethrows -> (ByteBuffer, Result)? {
         if let nextBufferSize = self.recvAllocator.nextBufferSize() {
-            if let result = try self.buffer?.modifyIfUniquelyOwned(minimumCapacity: nextBufferSize, body) {
+            if let result = try self._buffer?._modifyIfUniquelyOwned(minimumCapacity: nextBufferSize, body) {
                 // `result` can only be non-nil if `buffer` is non-nil.
-                return (self.buffer!, result)
+                return (self._buffer!, result)
             } else {
                 // Cycle through the buffers starting at the last used buffer.
-                let resultAndIndex = try self.buffers.loopingFirstIndexWithResult(startingAt: self.lastUsedIndex) {
+                let resultAndIndex = try self._buffers._loopingFirstIndexWithResult(startingAt: self._lastUsedIndex) {
                     buffer in
-                    try buffer.modifyIfUniquelyOwned(minimumCapacity: nextBufferSize, body)
+                    try buffer._modifyIfUniquelyOwned(minimumCapacity: nextBufferSize, body)
                 }
 
                 if let (result, index) = resultAndIndex {
-                    self.lastUsedIndex = index
-                    return (self.buffers[index], result)
+                    self._lastUsedIndex = index
+                    return (self._buffers[index], result)
                 }
             }
-        } else if self.buffer != nil, !self.mayGrow {
+        } else if self._buffer != nil, !self._mayGrow {
             // No hint about the buffer size (so pooling is not being used) and the allocator
             // indicated that the next buffer will not grow in size so reuse the existing stored
             // buffer.
-            self.buffer!.clear()
-            let result = try body(&self.buffer!)
-            return (self.buffer!, result)
+            self._buffer!.clear()
+            let result = try body(&self._buffer!)
+            return (self._buffer!, result)
         }
 
         // Couldn't reuse an existing buffer.
         return nil
     }
 
-    private mutating func allocateNewBuffer<Result>(
+    @inlinable
+    internal mutating func _allocateNewBuffer<Result>(
         using allocator: ByteBufferAllocator,
         _ body: (inout ByteBuffer) throws -> Result
     ) rethrows -> (ByteBuffer, Result) {
         // Couldn't reuse a buffer; create a new one and store it if there's capacity.
         var newBuffer = self.recvAllocator.buffer(allocator: allocator)
 
-        if let buffer = self.buffer {
-            assert(self.buffers.isEmpty)
+        if let buffer = self._buffer {
+            assert(self._buffers.isEmpty)
             // We have a stored buffer, either:
             // 1. We have capacity to add more and use `buffers` for storage, or
             // 2. Our capacity is 1; we can't use `buffers` for storage.
             if self.capacity > 1 {
-                self.buffer = nil
-                self.buffers.reserveCapacity(self.capacity)
-                self.buffers.append(buffer)
-                self.buffers.append(newBuffer)
-                self.lastUsedIndex = self.buffers.index(before: self.buffers.endIndex)
-                return try self.modifyBuffer(atIndex: self.lastUsedIndex, body)
+                self._buffer = nil
+                self._buffers.reserveCapacity(self.capacity)
+                self._buffers.append(buffer)
+                self._buffers.append(newBuffer)
+                self._lastUsedIndex = self._buffers.index(before: self._buffers.endIndex)
+                return try self._modifyBuffer(atIndex: self._lastUsedIndex, body)
             } else {
                 let result = try body(&newBuffer)
                 return (newBuffer, result)
@@ -163,32 +186,34 @@ internal struct PooledRecvBufferAllocator {
             //    buffer is nil), or
             // 2. we've already switched to using buffers for storage and it's not yet full, or
             // 3. we've already switched to using buffers for storage and it's full.
-            if self.buffers.isEmpty {
-                self.buffer = newBuffer
-                let result = try body(&self.buffer!)
-                return (self.buffer!, result)
-            } else if self.buffers.count < self.capacity {
-                self.buffers.append(newBuffer)
-                self.lastUsedIndex = self.buffers.index(before: self.buffers.endIndex)
-                return try self.modifyBuffer(atIndex: self.lastUsedIndex, body)
+            if self._buffers.isEmpty {
+                self._buffer = newBuffer
+                let result = try body(&self._buffer!)
+                return (self._buffer!, result)
+            } else if self._buffers.count < self.capacity {
+                self._buffers.append(newBuffer)
+                self._lastUsedIndex = self._buffers.index(before: self._buffers.endIndex)
+                return try self._modifyBuffer(atIndex: self._lastUsedIndex, body)
             } else {
                 let result = try body(&newBuffer)
                 return (newBuffer, result)
             }
         }
     }
 
-    private mutating func modifyBuffer<Result>(
+    @inlinable
+    internal mutating func _modifyBuffer<Result>(
         atIndex index: Int,
         _ body: (inout ByteBuffer) throws -> Result
     ) rethrows -> (ByteBuffer, Result) {
-        let result = try body(&self.buffers[index])
-        return (self.buffers[index], result)
+        let result = try body(&self._buffers[index])
+        return (self._buffers[index], result)
     }
 }
 
 extension ByteBuffer {
-    fileprivate mutating func modifyIfUniquelyOwned<Result>(
+    @inlinable
+    internal mutating func _modifyIfUniquelyOwned<Result>(
         minimumCapacity: Int,
         _ body: (inout ByteBuffer) throws -> Result
     ) rethrows -> Result? {
@@ -206,18 +231,20 @@ extension Array {
     ///
     /// - Returns: The result and index of the first element passed to `body` which returned
     ///   non-nil, or `nil` if no such element exists.
-    fileprivate mutating func loopingFirstIndexWithResult<Result>(
+    @inlinable
+    internal mutating func _loopingFirstIndexWithResult<Result>(
         startingAt middleIndex: Index,
         whereNonNil body: (inout Element) throws -> Result?
     ) rethrows -> (Result, Index)? {
-        if let result = try self.firstIndexWithResult(in: middleIndex..<self.endIndex, whereNonNil: body) {
+        if let result = try self._firstIndexWithResult(in: middleIndex..<self.endIndex, whereNonNil: body) {
             return result
         }
 
-        return try self.firstIndexWithResult(in: self.startIndex..<middleIndex, whereNonNil: body)
+        return try self._firstIndexWithResult(in: self.startIndex..<middleIndex, whereNonNil: body)
     }
 
-    private mutating func firstIndexWithResult<Result>(
+    @inlinable
+    internal mutating func _firstIndexWithResult<Result>(
         in indices: Range<Index>,
         whereNonNil body: (inout Element) throws -> Result?
     ) rethrows -> (Result, Index)? {

Sources/NIOPosix/BaseSocketChannel.swift

@@ -251,7 +251,7 @@ class BaseSocketChannel<SocketType: BaseSocketProtocol>: SelectableChannel, Chan
     // MARK: Variables, on EventLoop thread only
     var readPending = false
     var pendingConnect: Optional<EventLoopPromise<Void>>
-    var recvBufferPool: PooledRecvBufferAllocator
+    var recvBufferPool: NIOPooledRecvBufferAllocator
     var maxMessagesPerRead: UInt = 4
     private var inFlushNow: Bool = false  // Guard against re-entrance of flushNow() method.
     private var autoRead: Bool = true

Tests/NIOPosixTests/PooledRecvBufferAllocatorTests.swift → Tests/NIOCoreTests/NIOPooledRecvBufferAllocatorTests.swift

@@ -17,10 +17,10 @@ import XCTest
 
 @testable import NIOPosix
 
-internal final class PooledRecvBufferAllocatorTests: XCTestCase {
+internal final class NIOPooledRecvBufferAllocatorTests: XCTestCase {
     func testPoolFillsToCapacity() {
         let allocator = ByteBufferAllocator()
-        var pool = PooledRecvBufferAllocator(
+        var pool = NIOPooledRecvBufferAllocator(
             capacity: 3,
             recvAllocator: FixedSizeRecvByteBufferAllocator(capacity: 1024)
         )
@@ -55,7 +55,7 @@ internal final class PooledRecvBufferAllocatorTests: XCTestCase {
 
     func testBuffersAreRecycled() {
         let allocator = ByteBufferAllocator()
-        var pool = PooledRecvBufferAllocator(
+        var pool = NIOPooledRecvBufferAllocator(
             capacity: 5,
             recvAllocator: FixedSizeRecvByteBufferAllocator(capacity: 1024)
         )
@@ -76,7 +76,7 @@ internal final class PooledRecvBufferAllocatorTests: XCTestCase {
 
     func testFirstAvailableBufferUsed() {
         let allocator = ByteBufferAllocator()
-        var pool = PooledRecvBufferAllocator(
+        var pool = NIOPooledRecvBufferAllocator(
             capacity: 3,
             recvAllocator: FixedSizeRecvByteBufferAllocator(capacity: 1024)
         )
@@ -99,7 +99,7 @@ internal final class PooledRecvBufferAllocatorTests: XCTestCase {
 
     func testBuffersAreClearedBetweenCalls() {
         let allocator = ByteBufferAllocator()
-        var pool = PooledRecvBufferAllocator(
+        var pool = NIOPooledRecvBufferAllocator(
             capacity: 3,
             recvAllocator: FixedSizeRecvByteBufferAllocator(capacity: 1024)
         )
@@ -135,7 +135,7 @@ internal final class PooledRecvBufferAllocatorTests: XCTestCase {
 
     func testPoolCapacityIncrease() {
         let allocator = ByteBufferAllocator()
-        var pool = PooledRecvBufferAllocator(
+        var pool = NIOPooledRecvBufferAllocator(
             capacity: 3,
             recvAllocator: FixedSizeRecvByteBufferAllocator(capacity: 1024)
         )
@@ -171,7 +171,7 @@ internal final class PooledRecvBufferAllocatorTests: XCTestCase {
 
     func testPoolCapacityDecrease() {
         let allocator = ByteBufferAllocator()
-        var pool = PooledRecvBufferAllocator(
+        var pool = NIOPooledRecvBufferAllocator(
             capacity: 5,
             recvAllocator: FixedSizeRecvByteBufferAllocator(capacity: 1024)
         )
@@ -192,17 +192,6 @@ internal final class PooledRecvBufferAllocatorTests: XCTestCase {
     }
 }
 
-extension ByteBuffer {
-    // Copied from NIOCoreTests/ByteBufferTest.swift
-    fileprivate func storagePointerIntegerValue() -> UInt {
-        var pointer: UInt = 0
-        self.withVeryUnsafeBytes { ptr in
-            pointer = UInt(bitPattern: ptr.baseAddress!)
-        }
-        return pointer
-    }
-}
-
 extension Array where Element == ByteBuffer {
     fileprivate func allHaveUniqueStorage() -> Bool {
         self.count == Set(self.map { $0.storagePointerIntegerValue() }).count