diff --git a/Sources/NIOCore/ByteBuffer-core.swift b/Sources/NIOCore/ByteBuffer-core.swift index 812a8cc386..514f43d160 100644 --- a/Sources/NIOCore/ByteBuffer-core.swift +++ b/Sources/NIOCore/ByteBuffer-core.swift @@ -299,6 +299,19 @@ public struct ByteBuffer { _ByteBufferSlice(0.. _Capacity { + #if canImport(Darwin) + return _Capacity(malloc_good_size(Int(capacity))) + #else + return capacity.nextPowerOf2ClampedToMax() + #endif + } + @inlinable static func _allocateAndPrepareRawMemory(bytes: _Capacity, allocator: Allocator) -> UnsafeMutableRawPointer { let ptr = allocator.malloc(size_t(bytes))! @@ -314,7 +327,7 @@ public struct ByteBuffer { @inlinable func allocateStorage(capacity: _Capacity) -> _Storage { - let newCapacity = capacity == 0 ? 0 : capacity.nextPowerOf2ClampedToMax() + let newCapacity = capacity == 0 ? 0 : _Storage.mallocSize(capacity: capacity) return _Storage( bytesNoCopy: _Storage._allocateAndPrepareRawMemory(bytes: newCapacity, allocator: self.allocator), capacity: newCapacity, @@ -332,7 +345,7 @@ public struct ByteBuffer { @inlinable func reallocStorage(capacity minimumNeededCapacity: _Capacity) { - let newCapacity = minimumNeededCapacity.nextPowerOf2ClampedToMax() + let newCapacity = _Storage.mallocSize(capacity: minimumNeededCapacity) let ptr = self.allocator.realloc(self.bytes, size_t(newCapacity))! // bind the memory so we can assume it elsewhere to be bound to UInt8 ptr.bindMemory(to: UInt8.self, capacity: Int(newCapacity)) @@ -346,7 +359,7 @@ public struct ByteBuffer { @inlinable static func reallocated(minimumCapacity: _Capacity, allocator: Allocator) -> _Storage { - let newCapacity = minimumCapacity == 0 ? 0 : minimumCapacity.nextPowerOf2ClampedToMax() + let newCapacity = minimumCapacity == 0 ? 0 : _Storage.mallocSize(capacity: minimumCapacity) // TODO: Use realloc if possible return _Storage( bytesNoCopy: _Storage._allocateAndPrepareRawMemory(bytes: newCapacity, allocator: allocator), diff --git a/Tests/NIOCoreTests/ByteBufferStorageMallocTest.swift b/Tests/NIOCoreTests/ByteBufferStorageMallocTest.swift new file mode 100644 index 0000000000..e07e2637b1 --- /dev/null +++ b/Tests/NIOCoreTests/ByteBufferStorageMallocTest.swift @@ -0,0 +1,64 @@ +//===----------------------------------------------------------------------===// +// +// This source file is part of the SwiftNIO open source project +// +// Copyright (c) 2025 Apple Inc. and the SwiftNIO project authors +// Licensed under Apache License v2.0 +// +// See LICENSE.txt for license information +// See CONTRIBUTORS.txt for the list of SwiftNIO project authors +// +// SPDX-License-Identifier: Apache-2.0 +// +//===----------------------------------------------------------------------===// + +import XCTest + +@testable import NIOCore + +/// Returns the malloc size that we expect to be allocated for a given requested capacity size. +/// On Darwin, we want to use `malloc_good_size` to get the optimal size, but on all other platforms +/// we use the next power of 2, clamped. +private func expectedMallocSize(_ size: Int) -> Int { + #if canImport(Darwin) + return Darwin.malloc_good_size(size) + #else + return UInt32(size).nextPowerOf2ClampedToMax() + #endif +} + +// Tests that ByteBuffer allocates memory in an optimal way depending on the host platform. + +final class ByteBufferStorageMallocTest: XCTestCase { + + func testInitialAllocationUsesGoodSize() { + let allocator = ByteBufferAllocator() + let requestedCapacity = 1000 + let expectedCapacity = expectedMallocSize(requestedCapacity) + + let buffer = allocator.buffer(capacity: requestedCapacity) + XCTAssertEqual(Int(buffer._storage.capacity), expectedCapacity) + } + + func testReallocationUsesGoodSize() { + let allocator = ByteBufferAllocator() + var buffer = allocator.buffer(capacity: 16) + let initialCapacity = buffer.capacity + + // Write more bytes than the current capacity to trigger reallocation + let newSize = initialCapacity + 100 + let expectedCapacity = expectedMallocSize(Int(newSize)) + + // This will trigger reallocation + buffer.writeBytes(Array(repeating: UInt8(0), count: Int(newSize))) + + XCTAssertEqual(Int(buffer._storage.capacity), expectedCapacity) + } + + func testZeroCapacity() { + let allocator = ByteBufferAllocator() + let buffer = allocator.buffer(capacity: 0) + XCTAssertEqual(buffer.capacity, 0) + } + +}