├── .editorconfig ├── .gitignore ├── LICENSE ├── Package.resolved ├── Package.swift ├── README.md ├── Sources └── Limiter │ └── Limiter.swift └── Tests └── LimiterTests └── LimiterTests.swift /.editorconfig: -------------------------------------------------------------------------------- 1 | # editorconfig.org 2 | [*] 3 | indent_style = space 4 | indent_size = 2 5 | trim_trailing_whitespace = true 6 | insert_final_newline = true 7 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .DS_Store 2 | /.build 3 | /Packages 4 | xcuserdata/ 5 | DerivedData/ 6 | .swiftpm/configuration/registries.json 7 | .swiftpm/xcode/package.xcworkspace/contents.xcworkspacedata 8 | .netrc 9 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2025 Philipp Gabriel 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /Package.resolved: -------------------------------------------------------------------------------- 1 | { 2 | "originHash" : "5c85006e223cbe628406d142aaaec2f081d76f3232c5f265f8d8ce6352ca3b9c", 3 | "pins" : [ 4 | { 5 | "identity" : "swift-clocks", 6 | "kind" : "remoteSourceControl", 7 | "location" : "https://github.com/pointfreeco/swift-clocks", 8 | "state" : { 9 | "revision" : "cc46202b53476d64e824e0b6612da09d84ffde8e", 10 | "version" : "1.0.6" 11 | } 12 | }, 13 | { 14 | "identity" : "swift-collections", 15 | "kind" : "remoteSourceControl", 16 | "location" : "https://github.com/apple/swift-collections", 17 | "state" : { 18 | "revision" : "671108c96644956dddcd89dd59c203dcdb36cec7", 19 | "version" : "1.1.4" 20 | } 21 | }, 22 | { 23 | "identity" : "swift-concurrency-extras", 24 | "kind" : "remoteSourceControl", 25 | "location" : "https://github.com/pointfreeco/swift-concurrency-extras", 26 | "state" : { 27 | "revision" : "82a4ae7170d98d8538ec77238b7eb8e7199ef2e8", 28 | "version" : "1.3.1" 29 | } 30 | }, 31 | { 32 | "identity" : "xctest-dynamic-overlay", 33 | "kind" : "remoteSourceControl", 34 | "location" : "https://github.com/pointfreeco/xctest-dynamic-overlay", 35 | "state" : { 36 | "revision" : "39de59b2d47f7ef3ca88a039dff3084688fe27f4", 37 | "version" : "1.5.2" 38 | } 39 | } 40 | ], 41 | "version" : 3 42 | } 43 | -------------------------------------------------------------------------------- /Package.swift: -------------------------------------------------------------------------------- 1 | // swift-tools-version: 6.0 2 | import PackageDescription 3 | 4 | let package = Package( 5 | name: "swift-concurrency-limiter", 6 | platforms: [.macOS(.v15), .iOS(.v18)], 7 | products: [ 8 | .library(name: "Limiter", targets: ["Limiter"]) 9 | ], 10 | dependencies: [ 11 | .package(url: "https://github.com/apple/swift-collections", from: "1.0.0"), 12 | .package(url: "https://github.com/pointfreeco/swift-clocks", from: "1.0.0") 13 | ], 14 | targets: [ 15 | .target( 16 | name: "Limiter", 17 | dependencies: [.product(name: "OrderedCollections", package: "swift-collections")] 18 | ), 19 | .testTarget( 20 | name: "LimiterTests", 21 | dependencies: [ 22 | "Limiter", 23 | .product(name: "Clocks", package: "swift-clocks"), 24 | ] 25 | ) 26 | ] 27 | ) 28 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Limiter 2 | 3 | > [!CAUTION] 4 | > This package uses unreleased Swift features and should only be used for demonstration or education purposes, for now. 5 | -------------------------------------------------------------------------------- /Sources/Limiter/Limiter.swift: -------------------------------------------------------------------------------- 1 | import Synchronization 2 | import OrderedCollections 3 | 4 | @available(macOS 9999, *) 5 | public final class AsyncLimiter: Sendable { 6 | 7 | typealias Continuation = UnsafeContinuation 8 | 9 | struct SuspensionID: Hashable { 10 | let id: Int 11 | } 12 | 13 | struct State { 14 | 15 | enum EnterAction { 16 | case resume(Continuation) 17 | case cancel(Continuation) 18 | case escalate([UnsafeCurrentTask]) 19 | } 20 | 21 | enum CancelAction { 22 | case cancel(Continuation) 23 | } 24 | 25 | enum LeaveAction { 26 | case resume(Continuation) 27 | } 28 | 29 | enum EscalationAction { 30 | case escalate([UnsafeCurrentTask]) 31 | } 32 | 33 | enum Suspension { 34 | case active(UnsafeCurrentTask) 35 | case suspended(Continuation, UnsafeCurrentTask) 36 | case cancelled 37 | } 38 | 39 | let limit: Int 40 | var suspensions: OrderedDictionary = [:] 41 | 42 | mutating func enter(id: SuspensionID, continuation: Continuation, task: UnsafeCurrentTask) -> EnterAction? { 43 | switch suspensions[id] { 44 | case .none: 45 | if suspensions.count >= limit { 46 | var tasks: [UnsafeCurrentTask] = [] 47 | for (_, suspension) in suspensions { 48 | switch suspension { 49 | case .active(let task): 50 | tasks.append(task) 51 | case .suspended(_, let task): 52 | tasks.append(task) 53 | case .cancelled: 54 | break 55 | } 56 | } 57 | suspensions[id] = .suspended(continuation, task) 58 | if tasks.isEmpty { 59 | return nil 60 | } else { 61 | return .escalate(tasks) 62 | } 63 | } else { 64 | suspensions[id] = .active(task) 65 | return .resume(continuation) 66 | } 67 | case .active, .suspended: 68 | preconditionFailure("Invalid state") 69 | case .cancelled: 70 | suspensions[id] = nil 71 | return .cancel(continuation) 72 | } 73 | } 74 | 75 | mutating func cancel(id: SuspensionID) -> CancelAction? { 76 | switch suspensions[id] { 77 | case .none: 78 | suspensions[id] = .cancelled 79 | return nil 80 | case .active: 81 | return nil 82 | case .suspended(let continuation, _): 83 | suspensions[id] = nil 84 | return .cancel(continuation) 85 | case .cancelled: 86 | preconditionFailure("Invalid state") 87 | } 88 | } 89 | 90 | mutating func leave(id: SuspensionID) -> LeaveAction? { 91 | switch suspensions[id] { 92 | case .none, .suspended, .cancelled: 93 | preconditionFailure("Invalid state") 94 | case .active: 95 | suspensions[id] = nil 96 | } 97 | for suspension in suspensions { 98 | switch suspension.value { 99 | case .active, .cancelled: 100 | continue 101 | case .suspended(let continuation, let task): 102 | suspensions[suspension.key] = .active(task) 103 | return .resume(continuation) 104 | } 105 | } 106 | return nil 107 | } 108 | 109 | func escalate(id: SuspensionID) -> EscalationAction? { 110 | switch suspensions[id] { 111 | case .none, .active, .cancelled: 112 | return nil 113 | case .suspended: 114 | break 115 | } 116 | var tasks: [UnsafeCurrentTask] = [] 117 | for suspension in suspensions { 118 | if suspension.key == id { 119 | break 120 | } 121 | switch suspension.value { 122 | case .active(let task): 123 | tasks.append(task) 124 | case .suspended(_, let task): 125 | tasks.append(task) 126 | case .cancelled: 127 | break 128 | } 129 | } 130 | if tasks.isEmpty { 131 | return nil 132 | } else { 133 | return .escalate(tasks) 134 | } 135 | } 136 | } 137 | 138 | let counter: Atomic 139 | let state: Mutex 140 | 141 | public init(limit: Int) { 142 | precondition(limit >= 1, "Needs at least one concurrent operation") 143 | counter = .init(0) 144 | state = .init(.init(limit: limit)) 145 | } 146 | 147 | public func withControl( 148 | isolation: isolated (any Actor)? = #isolation, 149 | operation: () async throws(E) -> sending R 150 | ) async throws -> R { 151 | 152 | let id = nextID() 153 | 154 | try await withTaskPriorityEscalationHandler(operation: { 155 | try await withTaskCancellationHandler(operation: { 156 | try await withUnsafeThrowingContinuation(isolation: isolation) { continuation in 157 | 158 | let action = state.withLock { state in 159 | withUnsafeCurrentTask { task in 160 | state.enter(id: id, continuation: continuation, task: task!) 161 | } 162 | } 163 | 164 | switch action { 165 | case .cancel(let continuation): 166 | continuation.resume(throwing: CancellationError()) 167 | case .resume(let continuation): 168 | continuation.resume() 169 | case .escalate(let tasks): 170 | for task in tasks { 171 | UnsafeCurrentTask.escalatePriority(task, to: Task.currentPriority) 172 | } 173 | case .none: 174 | break 175 | } 176 | } 177 | }, onCancel: { 178 | let action = state.withLock { state in 179 | state.cancel(id: id) 180 | } 181 | 182 | switch action { 183 | case .cancel(let continuation): 184 | continuation.resume(throwing: CancellationError()) 185 | case .none: 186 | break 187 | } 188 | }, isolation: isolation) 189 | }, onPriorityEscalated: { priority in 190 | let action = state.withLock { state in 191 | state.escalate(id: id) 192 | } 193 | 194 | switch action { 195 | case .escalate(let tasks): 196 | for task in tasks { 197 | UnsafeCurrentTask.escalatePriority(task, to: priority) 198 | } 199 | case .none: 200 | break 201 | } 202 | }, isolation: isolation) 203 | 204 | defer { 205 | 206 | let action = state.withLock { state in 207 | state.leave(id: id) 208 | } 209 | 210 | switch action { 211 | case .resume(let continuation): 212 | continuation.resume() 213 | case .none: 214 | break 215 | } 216 | } 217 | 218 | return try await operation() 219 | } 220 | 221 | func nextID() -> SuspensionID { 222 | return SuspensionID(id: counter.wrappingAdd(1, ordering: .relaxed).newValue) 223 | } 224 | } 225 | -------------------------------------------------------------------------------- /Tests/LimiterTests/LimiterTests.swift: -------------------------------------------------------------------------------- 1 | import Clocks 2 | import Foundation 3 | import Limiter 4 | import Synchronization 5 | import Testing 6 | 7 | // This is required, because the default executor schedules the high prio task before the low prio. 8 | final class NaiveExecutor: TaskExecutor { 9 | 10 | let queue: DispatchQueue 11 | 12 | init(queue: DispatchQueue) { 13 | self.queue = queue 14 | } 15 | 16 | func enqueue(_ job: consuming ExecutorJob) { 17 | let job = UnownedJob(job) 18 | let executor = asUnownedTaskExecutor() 19 | queue.async { 20 | job.runSynchronously(on: executor) 21 | } 22 | } 23 | } 24 | 25 | struct CustomError: Error { } 26 | 27 | func never() async throws { 28 | let never = AsyncStream { _ in } 29 | for await _ in never { } 30 | throw CancellationError() 31 | } 32 | 33 | @available(macOS 9999, *) 34 | @Test 35 | func testAsyncMutex() async throws { 36 | 37 | @globalActor actor CustomActor { static let shared = CustomActor() } 38 | 39 | try await withThrowingTaskGroup(of: Int.self) { taskGroup in 40 | 41 | let clock = TestClock() 42 | let limiter = AsyncLimiter(limit: 1) 43 | 44 | let now = clock.now 45 | let firstDeadline = now.advanced(by: .seconds(2)) 46 | let secondDeadline = now.advanced(by: .seconds(1)) 47 | 48 | taskGroup.addTask { @Sendable @CustomActor in 49 | return try await limiter.withControl { 50 | try await clock.sleep(until: firstDeadline) 51 | return 1 52 | } 53 | } 54 | 55 | taskGroup.addTask { @Sendable @CustomActor in 56 | return try await limiter.withControl { 57 | try await clock.sleep(until: secondDeadline) 58 | return 2 59 | } 60 | } 61 | 62 | await clock.advance(by: .seconds(2)) 63 | let first = try await taskGroup.next()! 64 | let second = try await taskGroup.next()! 65 | 66 | #expect(first == 1) 67 | #expect(second == 2) 68 | } 69 | } 70 | 71 | @available(macOS 9999, *) 72 | @Test 73 | func testImmediateCancel() async { 74 | 75 | let limiter = AsyncLimiter(limit: 1) 76 | let task = Task { 77 | try await limiter.withControl { 78 | do { 79 | try await never() 80 | } catch { 81 | throw CustomError() 82 | } 83 | } 84 | } 85 | 86 | task.cancel() 87 | 88 | await #expect { 89 | try await task.value 90 | } throws: { error in 91 | // NB: If cancellation happened when job was already scheduled, the operation has to handle cancellation. 92 | // In this case, this will result in CustomError. 93 | error is CancellationError || error is CustomError 94 | } 95 | } 96 | 97 | @available(macOS 9999, *) 98 | @Test 99 | func testDelayedCancel() async { 100 | 101 | let executor = NaiveExecutor(queue: .init(label: #function)) 102 | await withTaskExecutorPreference(executor) { 103 | let limiter = AsyncLimiter(limit: 1) 104 | 105 | let task1 = Task(executorPreference: executor) { 106 | try await limiter.withControl { 107 | try await never() 108 | } 109 | } 110 | 111 | let task2 = Task(executorPreference: executor) { 112 | try await limiter.withControl { 113 | do { 114 | try await never() 115 | } catch { 116 | throw CustomError() 117 | } 118 | } 119 | } 120 | 121 | await Task.yield() 122 | task2.cancel() 123 | task1.cancel() 124 | 125 | await #expect(throws: CancellationError.self) { 126 | try await task2.value 127 | } 128 | } 129 | } 130 | 131 | @available(macOS 9999, *) 132 | @Test 133 | func testEscalation() async throws { 134 | 135 | struct Priorities { 136 | var task1Before: TaskPriority? 137 | var task1Inside: TaskPriority? 138 | var task1After: TaskPriority? 139 | var task2Before: TaskPriority? 140 | var task2Inside: TaskPriority? 141 | var task2After: TaskPriority? 142 | var task3Before: TaskPriority? 143 | var task3Inside: TaskPriority? 144 | var task3After: TaskPriority? 145 | } 146 | 147 | let executor = NaiveExecutor(queue: .init(label: #function)) 148 | let clock = TestClock() 149 | let deadline = clock.now.advanced(by: .seconds(1)) 150 | let limiter = AsyncLimiter(limit: 1) 151 | let priorities = Mutex(Priorities()) 152 | let (stream, continuation) = AsyncStream.makeStream(of: Void.self) 153 | 154 | Task(executorPreference: executor, priority: .low) { @Sendable in 155 | priorities.withLock { $0.task1Before = Task.currentPriority } 156 | try await limiter.withControl { 157 | priorities.withLock { $0.task1Inside = Task.currentPriority } 158 | try await clock.sleep(until: deadline) 159 | priorities.withLock { $0.task1After = Task.currentPriority } 160 | } 161 | continuation.yield() 162 | } 163 | 164 | Task(executorPreference: executor, priority: .medium) { @Sendable in 165 | priorities.withLock { $0.task2Before = Task.currentPriority } 166 | try await limiter.withControl { 167 | priorities.withLock { $0.task2Inside = Task.currentPriority } 168 | try await clock.sleep(until: deadline) 169 | priorities.withLock { $0.task2After = Task.currentPriority } 170 | } 171 | continuation.yield() 172 | } 173 | 174 | Task(executorPreference: executor, priority: .high) { @Sendable in 175 | priorities.withLock { $0.task3Before = Task.currentPriority } 176 | try await limiter.withControl { 177 | priorities.withLock { $0.task3Inside = Task.currentPriority } 178 | try await clock.sleep(until: deadline) 179 | priorities.withLock { $0.task3After = Task.currentPriority } 180 | } 181 | continuation.yield() 182 | } 183 | 184 | await clock.advance(by: .seconds(1)) 185 | var iterator = stream.makeAsyncIterator() 186 | await iterator.next()! 187 | await iterator.next()! 188 | await iterator.next()! 189 | continuation.finish() 190 | 191 | priorities.withLock { priorities in 192 | #expect(priorities.task1Before == .low) 193 | #expect(priorities.task1Inside == .low) 194 | #expect(priorities.task1After == .high) 195 | #expect(priorities.task2Before == .medium) 196 | #expect(priorities.task2Inside == .high) 197 | #expect(priorities.task2After == .high) 198 | #expect(priorities.task3Before == .high) 199 | #expect(priorities.task3Inside == .high) 200 | #expect(priorities.task3After == .high) 201 | } 202 | } 203 | --------------------------------------------------------------------------------