Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

byte-buffer: use malloc_good_size on Darwin to allocate memory #3066

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 16 additions & 3 deletions Sources/NIOCore/ByteBuffer-core.swift
Original file line number Diff line number Diff line change
Expand Up @@ -299,6 +299,19 @@ public struct ByteBuffer {
_ByteBufferSlice(0..<self.capacity)
}

/// Returns the size that this storage would allocate for a given buffer capacity.
///
/// On Darwin, use `malloc_good_size` to get optimal size (capacity rounded up, but hopefully not allocating extra memory pages).
/// On other platforms, use the next power of 2 clamped to `UInt32.max`.
@inlinable
static func mallocSize(capacity: _Capacity) -> _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))!
Expand All @@ -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,
Expand All @@ -332,7 +345,7 @@ public struct ByteBuffer {

@inlinable
func reallocStorage(capacity minimumNeededCapacity: _Capacity) {
let newCapacity = minimumNeededCapacity.nextPowerOf2ClampedToMax()
let newCapacity = minimumNeededCapacity == 0 ? 0 : _Storage.mallocSize(capacity: minimumNeededCapacity)
natikgadzhi marked this conversation as resolved.
Show resolved Hide resolved
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))
Expand All @@ -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),
Expand Down
57 changes: 57 additions & 0 deletions Tests/NIOCoreTests/ByteBufferStorageMallocTest.swift
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
//===----------------------------------------------------------------------===//
//
// 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

#if canImport(Darwin)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's make these tests cross-platform by redefining the logic here. We can just have an abstraction function that is essentially equivalent to mallocSize, but allows us to validate that that's what we're actually using.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Lukasa apologies for the delay — I think this gets it to a good spot. Once that gets in, I'll get a bit more experience around ByteBuffer and then see if it's useful to figure out better Linux allocation strategies.

import Darwin

// 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 = malloc_good_size(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 = malloc_good_size(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)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why are we asking for capacity differently here than in the rest of the tests?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same approach as the rest of the tests would be fine, BUT zero being explicit feels nice.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it might be nice to keep the tests discussing capacity the same way.

}

}
#endif