add rotating allocator for UInt32 (#11)

This creates an allocator based on a FIFO that will allocate through the
range before reusing previously released allocations.

Signed-off-by: michael crosby <[email protected]>

Author: crosbymichael

Package.swift

@@ -244,6 +244,13 @@ let package = Package(
                 .product(name: "Logging", package: "swift-log"),
             ]
         ),
+        .testTarget(
+            name: "ContainerizationExtrasTests",
+            dependencies: [
+                "ContainerizationExtras",
+                "CShim",
+            ]
+        ),
         .target(
             name: "CShim"
         ),

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<UInt32> {
+        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 {

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<AddressType>
+        let indexToAddress: IndexToAddressTransform<AddressType>
+
+        init(
+            size: UInt32,
+            addressToIndex: @escaping AddressToIndexTransform<AddressType>,
+            indexToAddress: @escaping IndexToAddressTransform<AddressType>
+        ) {
+            self.allocations = [UInt32](0..<size)
+            self.enabled = true
+            self.allocationCount = 0
+            self.addressToIndex = addressToIndex
+            self.indexToAddress = indexToAddress
+        }
+    }
+
+    private let stateGuard: Mutex<State>
+
+    /// Create an allocator with specified size and index mappings.
+    package init(
+        size: UInt32,
+        addressToIndex: @escaping AddressToIndexTransform<AddressType>,
+        indexToAddress: @escaping IndexToAddressTransform<AddressType>
+    ) {
+        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
+        }
+    }
+}

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<UInt32>()
+        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)
+    }
 }