From 862d4e21163a7b91bdadbb8c40000cc9814b06c9 Mon Sep 17 00:00:00 2001 From: michael crosby Date: Wed, 4 Jun 2025 14:14:58 -0400 Subject: [PATCH] add rotating allocator for UInt32 This creates an allocator based on a FIFO that will allocate through the range before reusing previously released allocations. Signed-off-by: michael crosby --- Package.swift | 7 + .../NetworkAddress+Allocator.swift | 18 +++ .../RotatingAddressAllocator.swift | 125 ++++++++++++++++++ .../TestNetworkAddress+Allocator.swift | 64 +++++++++ 4 files changed, 214 insertions(+) create mode 100644 Sources/ContainerizationExtras/RotatingAddressAllocator.swift diff --git a/Package.swift b/Package.swift index 76521c5c..df305c67 100644 --- a/Package.swift +++ b/Package.swift @@ -244,6 +244,13 @@ let package = Package( .product(name: "Logging", package: "swift-log"), ] ), + .testTarget( + name: "ContainerizationExtrasTests", + dependencies: [ + "ContainerizationExtras", + "CShim", + ] + ), .target( name: "CShim" ), diff --git a/Sources/ContainerizationExtras/NetworkAddress+Allocator.swift b/Sources/ContainerizationExtras/NetworkAddress+Allocator.swift index fce1b57b..a9eb28e4 100644 --- a/Sources/ContainerizationExtras/NetworkAddress+Allocator.swift +++ b/Sources/ContainerizationExtras/NetworkAddress+Allocator.swift @@ -73,6 +73,24 @@ extension UInt32 { indexToAddress: { lower + UInt32($0) } ) } + + /// Creates a rotating allocator for vsock ports, or any UInt32 values. + public static func rotatingAllocator(lower: UInt32, size: UInt32) throws -> any AddressAllocator { + guard 0xffff_ffff - lower + 1 >= size else { + throw AllocatorError.rangeExceeded + } + + return RotatingAddressAllocator( + size: size, + addressToIndex: { address in + guard address >= lower && address <= lower + UInt32(size) else { + return nil + } + return Int(address - lower) + }, + indexToAddress: { lower + UInt32($0) } + ) + } } extension Character { diff --git a/Sources/ContainerizationExtras/RotatingAddressAllocator.swift b/Sources/ContainerizationExtras/RotatingAddressAllocator.swift new file mode 100644 index 00000000..dce2b67f --- /dev/null +++ b/Sources/ContainerizationExtras/RotatingAddressAllocator.swift @@ -0,0 +1,125 @@ +//===----------------------------------------------------------------------===// +// Copyright © 2025 Apple Inc. and the Containerization project authors. +// All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +//===----------------------------------------------------------------------===// + +import Synchronization + +package final class RotatingAddressAllocator: AddressAllocator { + package typealias AddressType = UInt32 + + private struct State { + var allocations: [AddressType] + var enabled: Bool + var allocationCount: Int + let addressToIndex: AddressToIndexTransform + let indexToAddress: IndexToAddressTransform + + init( + size: UInt32, + addressToIndex: @escaping AddressToIndexTransform, + indexToAddress: @escaping IndexToAddressTransform + ) { + self.allocations = [UInt32](0.. + + /// Create an allocator with specified size and index mappings. + package init( + size: UInt32, + addressToIndex: @escaping AddressToIndexTransform, + indexToAddress: @escaping IndexToAddressTransform + ) { + let state = State( + size: size, + addressToIndex: addressToIndex, + indexToAddress: indexToAddress + ) + self.stateGuard = Mutex(state) + } + + public func allocate() throws -> AddressType { + try self.stateGuard.withLock { state in + guard state.enabled else { + throw AllocatorError.allocatorDisabled + } + + guard state.allocations.count > 0 else { + throw AllocatorError.allocatorFull + } + + let value = state.allocations.removeFirst() + + guard let address = state.indexToAddress(Int(value)) else { + throw AllocatorError.invalidIndex(Int(value)) + } + + state.allocationCount += 1 + return address + } + } + + package func reserve(_ address: AddressType) throws { + try self.stateGuard.withLock { state in + guard state.enabled else { + throw AllocatorError.allocatorDisabled + } + + guard let index = state.addressToIndex(address) else { + throw AllocatorError.invalidAddress(address.description) + } + + let i = state.allocations.firstIndex(of: UInt32(index)) + guard let i else { + throw AllocatorError.alreadyAllocated("\(address.description)") + } + + _ = state.allocations.remove(at: i) + state.allocationCount += 1 + } + } + + package func release(_ address: AddressType) throws { + try self.stateGuard.withLock { state in + guard let index = (state.addressToIndex(address)) else { + throw AllocatorError.invalidAddress(address.description) + } + let value = UInt32(index) + + guard !state.allocations.contains(value) else { + throw AllocatorError.notAllocated("\(address.description)") + } + + state.allocations.append(value) + state.allocationCount -= 1 + } + } + + package func disableAllocator() -> Bool { + self.stateGuard.withLock { state in + guard state.allocationCount == 0 else { + return false + } + state.enabled = false + return true + } + } +} diff --git a/Tests/ContainerizationExtrasTests/TestNetworkAddress+Allocator.swift b/Tests/ContainerizationExtrasTests/TestNetworkAddress+Allocator.swift index 8e94cee3..1a5294ed 100644 --- a/Tests/ContainerizationExtrasTests/TestNetworkAddress+Allocator.swift +++ b/Tests/ContainerizationExtrasTests/TestNetworkAddress+Allocator.swift @@ -182,4 +182,68 @@ final class TestAddressAllocators { let value = try allocator.allocate() #expect(value == address) } + + @Test + func testRotatingUInt32PortAllocator() throws { + var allocations = Set() + let lower = UInt32(5000) + let allocator = try UInt32.rotatingAllocator(lower: lower, size: 3) + allocations.insert(try allocator.allocate()) + allocations.insert(try allocator.allocate()) + allocations.insert(try allocator.allocate()) + do { + _ = try allocator.allocate() + #expect(Bool(false), "Expected AllocatorError.allocatorFull to be thrown") + } catch { + #expect(error as? AllocatorError == .allocatorFull, "Unexpected error thrown: \(error)") + } + + let address = UInt32(5001) + try allocator.release(address) + let value = try allocator.allocate() + #expect(value == address) + } + + @Test + func testRotatingFIFOUInt32PortAllocator() throws { + let lower = UInt32(5000) + let allocator = try UInt32.rotatingAllocator(lower: lower, size: 3) + let first = try allocator.allocate() + #expect(first == 5000) + let second = try allocator.allocate() + #expect(second == 5001) + + try allocator.release(first) + let third = try allocator.allocate() + // even after a release, it should continue to allocate in the range + // before reusing an previous allocation on the stack. + #expect(third == 5002) + + // now the next allocation should be our first port + let reused = try allocator.allocate() + #expect(reused == first) + + try allocator.release(third) + let thirdReused = try allocator.allocate() + #expect(thirdReused == third) + } + + @Test + func testRotatingReservedUInt32PortAllocator() throws { + let lower = UInt32(5000) + let allocator = try UInt32.rotatingAllocator(lower: lower, size: 3) + + try allocator.reserve(5001) + let first = try allocator.allocate() + #expect(first == 5000) + // this should skip the reserved 5001 + let second = try allocator.allocate() + #expect(second == 5002) + + // no release our reserved + try allocator.release(5001) + + let third = try allocator.allocate() + #expect(third == 5001) + } }