From 5696175b6be97022e0ca82956c1d307d01f331f9 Mon Sep 17 00:00:00 2001 From: sukru tikves Date: Sun, 13 Sep 2026 14:30:53 -0700 Subject: [PATCH 1/2] Add async request queue to llm-server Replace the binary busy flag with a bounded async RequestQueue: concurrent requests wait FIFO and get a 429 only when the queue is full. - acquire() returns a QueuePermit that frees the slot exactly once, on scope exit or deinit, so a dropped streaming response can't leak it and wedge the server - Cancel queued waiters cleanly: remove the waiter and throw CancellationError - --max-queue-depth (default 16); reject negative values - hasRecurrentState skips prefix-reuse accounting for recurrent (SSM/hybrid) models - RequestQueue, QueuePermit and ServerError live in CoreAILMCommon, unit-tested via CoreAILMCommonTests --- .../Sources/CoreAILMCommon/RequestQueue.swift | 162 ++++++++++++++ .../CoreAIPipelinedEngine.swift | 2 + .../CoreAISequentialEngine.swift | 1 + .../InferenceEngines/InferenceEngine.swift | 13 ++ .../Tools/llm-server/ChatHandler.swift | 34 ++- .../Tools/llm-server/CompletionHandler.swift | 11 +- .../Tools/llm-server/LLMServerMain.swift | 13 +- .../Tools/llm-server/ServerState.swift | 43 ++-- .../RequestQueueTests.swift | 204 ++++++++++++++++++ 9 files changed, 439 insertions(+), 44 deletions(-) create mode 100644 swift/Sources/CoreAILMCommon/RequestQueue.swift create mode 100644 swift/Tests/CoreAILMCommonTests/RequestQueueTests.swift diff --git a/swift/Sources/CoreAILMCommon/RequestQueue.swift b/swift/Sources/CoreAILMCommon/RequestQueue.swift new file mode 100644 index 00000000..5b118c73 --- /dev/null +++ b/swift/Sources/CoreAILMCommon/RequestQueue.swift @@ -0,0 +1,162 @@ +// Copyright 2026 Apple Inc. +// +// Use of this source code is governed by a BSD-3-clause license that can +// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +import Foundation +import Synchronization + +// MARK: - Request Queue + +/// Async semaphore gating concurrency at one in-flight request. A second request +/// waits in a FIFO queue instead of getting an immediate 429; it is rejected only +/// once `maxDepth` requests are already waiting. +/// +/// `acquire()` returns a `QueuePermit` that frees the slot exactly once. Callers +/// hold the permit for the lifetime of the request and let it release on scope +/// exit (or hand it to a streaming body closure); the permit also releases on +/// deallocation, so the slot is never leaked if a caller drops it. +public final class RequestQueue: Sendable { + public let maxDepth: Int + private let state = Mutex(QueueState()) + + private struct QueueState { + var isActive: Bool = false + // Waiters carry a token so the cancellation handler can find and remove + // its own continuation (CheckedContinuation is not Equatable). + var waiters: [(id: UInt64, cont: CheckedContinuation)] = [] + var nextID: UInt64 = 0 + } + + public init(maxDepth: Int) { + self.maxDepth = max(0, maxDepth) + } + + public var depth: Int { + state.withLock { $0.waiters.count + ($0.isActive ? 1 : 0) } + } + + public var queuedCount: Int { + state.withLock { $0.waiters.count } + } + + public var isActive: Bool { + state.withLock { $0.isActive } + } + + /// Acquire the exclusive slot, waiting in FIFO order if it is held. Throws + /// `ServerError.queueFull` when `maxDepth` waiters are already queued, or + /// `CancellationError` if the awaiting task is cancelled while queued. + public func acquire() async throws -> QueuePermit { + let needsWait: Bool = state.withLock { s in + if !s.isActive { + s.isActive = true + return false + } + return true + } + if !needsWait { return QueuePermit(self) } + + // Reserve a token before suspending so the cancellation handler below can + // match this waiter even if it runs before the continuation is installed. + let id = state.withLock { s -> UInt64 in + defer { s.nextID &+= 1 } + return s.nextID + } + + try await withTaskCancellationHandler { + try await withCheckedThrowingContinuation { (cont: CheckedContinuation) in + enum Action { case acquired, rejected, queued, cancelled } + let action: Action = state.withLock { s -> Action in + if Task.isCancelled { return .cancelled } + if !s.isActive { + s.isActive = true + return .acquired + } + if s.waiters.count >= maxDepth { + return .rejected + } + s.waiters.append((id: id, cont: cont)) + return .queued + } + switch action { + case .acquired: + cont.resume() + case .rejected: + cont.resume(throwing: ServerError.queueFull(depth: maxDepth)) + case .cancelled: + cont.resume(throwing: CancellationError()) + case .queued: + break + } + } + } onCancel: { + // Remove our still-queued waiter and resume it throwing. The shared + // Mutex serializes this against release(), so a given waiter is popped + // by exactly one of the two and resumed exactly once. + let cont: CheckedContinuation? = state.withLock { s in + guard let idx = s.waiters.firstIndex(where: { $0.id == id }) else { return nil } + return s.waiters.remove(at: idx).cont + } + cont?.resume(throwing: CancellationError()) + } + return QueuePermit(self) + } + + /// Hand the slot to the next FIFO waiter, or clear the active flag. Internal; + /// callers release through `QueuePermit`. + fileprivate func release() { + let next: CheckedContinuation? = state.withLock { s in + if !s.waiters.isEmpty { + return s.waiters.removeFirst().cont + } + s.isActive = false + return nil + } + next?.resume() + } +} + +// MARK: - Queue Permit + +/// A held `RequestQueue` slot. Releases exactly once — on the first `release()` +/// call or, as a safety net, on deallocation (e.g. if a streaming response body +/// is dropped by the server before its writer closure ever runs). +public final class QueuePermit: Sendable { + private let queue: RequestQueue + private let released = Mutex(false) + + fileprivate init(_ queue: RequestQueue) { + self.queue = queue + } + + public func release() { + let shouldRelease = released.withLock { r -> Bool in + if r { return false } + r = true + return true + } + if shouldRelease { queue.release() } + } + + deinit { release() } +} + +// MARK: - Server Errors + +public enum ServerError: Error, LocalizedError { + case badRequest(String) + case queueFull(depth: Int) + + public var isBadRequest: Bool { + if case .badRequest = self { return true } + return false + } + + public var errorDescription: String? { + switch self { + case .badRequest(let msg): return msg + case .queueFull(let depth): return "Queue full (depth: \(depth))" + } + } +} diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift index 64fcb27f..e63c8ff5 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift @@ -68,6 +68,8 @@ final class CoreAIPipelinedEngine: InferenceEngine, ConstrainedGenerationCapable var processedTokenCount: Int { engine.processedTokenCount } + var hasRecurrentState: Bool { engine.hasNonTruncatableStates } + init( config: ModelConfig, preparedModel: PreparedModel, diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift index 643c3008..f4ffe7b7 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift @@ -35,6 +35,7 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable public var supportsLogits: Bool { true } public var vocabSize: Int { config.vocabSize } + public var hasRecurrentState: Bool { hasNonTruncatableStates } public let config: ModelConfig // Core AI function handle diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift index 5b0a3f1e..739d0944 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift @@ -143,6 +143,15 @@ public protocol InferenceEngine: Sendable { /// Useful for debugging multi-turn efficiency and verifying prefix caching behavior. var lastPrefixHitCount: Int { get } + /// Whether the engine carries recurrent state (SSM / hybrid models). + /// + /// Recurrent/conv states summarize the entire prefix and cannot be truncated + /// by moving a KV cursor, so token-prefix reuse is impossible: on any rewind + /// the engine full-resets and replays the whole prompt. Callers (e.g. the + /// server's prefix-reuse accounting) should short-circuit when this is true. + /// Defaults to `false`. + var hasRecurrentState: Bool { get } + // MARK: - Configuration associatedtype ConfigType: Codable, InferenceConfiguration @@ -198,6 +207,10 @@ extension InferenceEngine { public var lastPrefixHitCount: Int { 0 } } +extension InferenceEngine { + public var hasRecurrentState: Bool { false } +} + extension InferenceEngine { /// Default: engine is not busy. public var isBusy: Bool { false } diff --git a/swift/Sources/Tools/llm-server/ChatHandler.swift b/swift/Sources/Tools/llm-server/ChatHandler.swift index 5dbd76bd..2820c4f3 100644 --- a/swift/Sources/Tools/llm-server/ChatHandler.swift +++ b/swift/Sources/Tools/llm-server/ChatHandler.swift @@ -81,6 +81,9 @@ func startServer(state: ServerState, port: Int) async throws { configuration: .init( address: .hostname("127.0.0.1", port: port) )) + + // Process lifecycle is left to the supervising process. TODO: any in-process + // idle-exit must first drain in-flight engine work (no drain hook today). try await app.run() } @@ -100,8 +103,15 @@ private func handleAutoRoute(request: Request, state: ServerState) async throws private func handleChatCompletionsFromBody(body: ByteBuffer, state: ServerState, sessionID: String? = nil) async throws -> Response { - guard state.tryAcquire() else { - let err = ErrorResponse(error: .init(message: "Server is busy.", type: "server_error", code: "busy")) + // Non-streaming paths release the permit at each return; the streaming path + // hands it to the body closure. Release-on-deinit backstops a dropped response. + let permit: QueuePermit + do { + permit = try await state.queue.acquire() + } catch let error as ServerError { + let err = ErrorResponse( + error: .init( + message: error.localizedDescription, type: "server_error", code: "queue_full")) let data = try JSONEncoder().encode(err) return Response( status: .tooManyRequests, headers: [.contentType: "application/json"], @@ -111,7 +121,7 @@ private func handleChatCompletionsFromBody(body: ByteBuffer, state: ServerState, do { chatRequest = try JSONDecoder().decode(ChatCompletionRequest.self, from: body) } catch { - state.release() + permit.release() let err = ErrorResponse(error: .init(message: "\(error)", type: "invalid_request_error", code: nil)) let data = try JSONEncoder().encode(err) return Response( @@ -121,15 +131,17 @@ private func handleChatCompletionsFromBody(body: ByteBuffer, state: ServerState, do { let shouldStream = chatRequest.stream ?? false if shouldStream { - return try await handleStreamingRequest(chatRequest: chatRequest, state: state, sessionID: sessionID) + // Ownership of the permit passes to the streaming body closure. + return try await handleStreamingRequest( + chatRequest: chatRequest, state: state, sessionID: sessionID, permit: permit) } else { let response = try await handleNonStreamingRequest( chatRequest: chatRequest, state: state, sessionID: sessionID) - state.release() + permit.release() return response } } catch let error as ServerError { - state.release() + permit.release() let status: HTTPResponse.Status = error.isBadRequest ? .badRequest : .internalServerError let err = ErrorResponse(error: .init(message: "\(error)", type: "invalid_request_error", code: nil)) let data = try JSONEncoder().encode(err) @@ -137,7 +149,7 @@ private func handleChatCompletionsFromBody(body: ByteBuffer, state: ServerState, status: status, headers: [.contentType: "application/json"], body: .init(byteBuffer: ByteBuffer(data: data)) ) } catch { - state.release() + permit.release() let err = ErrorResponse(error: .init(message: "\(error)", type: "server_error", code: nil)) let data = try JSONEncoder().encode(err) return Response( @@ -327,7 +339,9 @@ private func handleNonStreamingRequest(chatRequest: ChatCompletionRequest, state // MARK: - Streaming (SSE) -private func handleStreamingRequest(chatRequest: ChatCompletionRequest, state: ServerState, sessionID: String? = nil) +private func handleStreamingRequest( + chatRequest: ChatCompletionRequest, state: ServerState, sessionID: String? = nil, permit: QueuePermit +) async throws -> Response { let requestMaxTokens = chatRequest.maxCompletionTokens ?? chatRequest.maxTokens ?? state.config.defaultMaxTokens @@ -364,7 +378,9 @@ private func handleStreamingRequest(chatRequest: ChatCompletionRequest, state: S } let responseBody = ResponseBody { writer in - defer { state.release() } + // The closure owns the permit for the stream's lifetime; release-on-deinit + // covers a response the framework drops before entering this closure. + defer { permit.release() } do { let encoder = JSONEncoder() let genStart = SuspendingClock().now diff --git a/swift/Sources/Tools/llm-server/CompletionHandler.swift b/swift/Sources/Tools/llm-server/CompletionHandler.swift index e845d7ad..7962148f 100644 --- a/swift/Sources/Tools/llm-server/CompletionHandler.swift +++ b/swift/Sources/Tools/llm-server/CompletionHandler.swift @@ -30,14 +30,19 @@ func handleCompletionsFromBody(body: ByteBuffer, state: ServerState) async throw status: .notImplemented, headers: [.contentType: "application/json"], body: .init(byteBuffer: ByteBuffer(data: data))) } - guard state.tryAcquire() else { - let err = ErrorResponse(error: .init(message: "Server is busy.", type: "server_error", code: "busy")) + let permit: QueuePermit + do { + permit = try await state.queue.acquire() + } catch let error as ServerError { + let err = ErrorResponse( + error: .init( + message: error.localizedDescription, type: "server_error", code: "queue_full")) let data = try JSONEncoder().encode(err) return Response( status: .tooManyRequests, headers: [.contentType: "application/json"], body: .init(byteBuffer: ByteBuffer(data: data))) } - defer { state.release() } + defer { permit.release() } let req: CompletionRequest do { req = try JSONDecoder().decode(CompletionRequest.self, from: body) diff --git a/swift/Sources/Tools/llm-server/LLMServerMain.swift b/swift/Sources/Tools/llm-server/LLMServerMain.swift index ef4242da..41cd72fb 100644 --- a/swift/Sources/Tools/llm-server/LLMServerMain.swift +++ b/swift/Sources/Tools/llm-server/LLMServerMain.swift @@ -65,6 +65,9 @@ struct LLMServer: AsyncParsableCommand { ) var chunkThreshold: Int? + @Option(name: .customLong("max-queue-depth"), help: "Max requests queued before returning 429 (default: 16)") + var maxQueueDepth: Int = 16 + @Flag(name: .customLong("no-thinking"), help: "Disable thinking/reasoning (appends /no_think or sets template)") var noThinking: Bool = false @@ -77,6 +80,12 @@ struct LLMServer: AsyncParsableCommand { @Flag(help: "Enable verbose logging") var verbose: Bool = false + func validate() throws { + guard maxQueueDepth >= 0 else { + throw ValidationError("--max-queue-depth must be >= 0 (got \(maxQueueDepth))") + } + } + func run() async throws { CLILogger.level = verbose ? 1 : 0 @@ -188,7 +197,8 @@ struct LLMServer: AsyncParsableCommand { supportsLogprobs: supportsLogprobs, maxContextLength: bundle.maxContextLength, vocabSize: bundle.vocabSize, - additionalEosTokenIds: additionalEosTokenIds + additionalEosTokenIds: additionalEosTokenIds, + maxQueueDepth: maxQueueDepth ) let state = ServerState( @@ -203,6 +213,7 @@ struct LLMServer: AsyncParsableCommand { print(" Logprobs: \(supportsLogprobs ? "supported" : "not supported (use --variant coreai-sequential)")") print(" Context: \(bundle.maxContextLength) tokens") print(" No-thinking: \(noThinking)") + print(" Max queue depth: \(maxQueueDepth)") let topKStr = topK.map { "\($0)" } ?? "nil" let topPStr = topP.map { "\($0)" } ?? "nil" print(" Sampling: temperature=\(temperature), topK=\(topKStr), topP=\(topPStr)") diff --git a/swift/Sources/Tools/llm-server/ServerState.swift b/swift/Sources/Tools/llm-server/ServerState.swift index 2023bd5b..8b0cf1ca 100644 --- a/swift/Sources/Tools/llm-server/ServerState.swift +++ b/swift/Sources/Tools/llm-server/ServerState.swift @@ -24,6 +24,7 @@ struct ServerConfig: Sendable { let maxContextLength: Int let vocabSize: Int? let additionalEosTokenIds: [Int32] + let maxQueueDepth: Int } // MARK: - Server Stats @@ -117,12 +118,12 @@ final class ServerState: @unchecked Sendable { let tokenizer: any Tokenizer let config: ServerConfig let stats = ServerStats() + let queue: RequestQueue let toolCallDetection: ToolCallDetection? let thinkingFormat: ThinkTagParser.Format private let _state = Mutex(InternalState()) private struct InternalState { - var generating: Bool = false var lastSessionID: String? = nil var lastPromptTokens: [Int32] = [] var prefixHits: Int = 0 @@ -135,6 +136,7 @@ final class ServerState: @unchecked Sendable { self.engine = engine self.tokenizer = tokenizer self.config = config + self.queue = RequestQueue(maxDepth: config.maxQueueDepth) self.toolCallDetection = detectToolCallFormat(using: tokenizer) self.thinkingFormat = detectThinkingFormat(using: tokenizer) } @@ -150,20 +152,16 @@ final class ServerState: @unchecked Sendable { ) } - func tryAcquire() -> Bool { - _state.withLock { s in - guard !s.generating else { return false } - s.generating = true - return true - } - } - - func release() { - _state.withLock { $0.generating = false } - } - /// Prepare engine for a new request. Returns the number of prefix tokens reused. func prepareForRequest(sessionID: String?, promptTokens: [Int32]) async -> Int { + // Recurrent-state models (SSM/hybrid) cannot reuse a token prefix: the + // engine full-resets and replays the whole prompt on any rewind. Skip the + // prefix-reuse fast path entirely and count it as a miss for honest stats. + if engine.hasRecurrentState { + _state.withLock { $0.prefixMisses += 1 } + return 0 + } + let action = _state.withLock { s -> PrepareAction in guard let sid = sessionID, sid == s.lastSessionID else { s.lastSessionID = sessionID @@ -235,7 +233,7 @@ final class ServerState: @unchecked Sendable { /// Readiness snapshot for /ready endpoint. func readySnapshot() -> ReadyResponse { - let busy = _state.withLock { $0.generating } + let busy = queue.isActive let usedTokens = engine.processedTokenCount let maxTokens = config.maxContextLength let utilization = maxTokens > 0 ? Double(usedTokens) / Double(maxTokens) : 0 @@ -285,20 +283,3 @@ enum RequestID { return "coreai-\(n)" } } - -// MARK: - Server Errors - -enum ServerError: Error, LocalizedError { - case badRequest(String) - - var isBadRequest: Bool { - if case .badRequest = self { return true } - return false - } - - var errorDescription: String? { - switch self { - case .badRequest(let msg): return msg - } - } -} diff --git a/swift/Tests/CoreAILMCommonTests/RequestQueueTests.swift b/swift/Tests/CoreAILMCommonTests/RequestQueueTests.swift new file mode 100644 index 00000000..a0673d6e --- /dev/null +++ b/swift/Tests/CoreAILMCommonTests/RequestQueueTests.swift @@ -0,0 +1,204 @@ +// Copyright 2026 Apple Inc. +// +// Use of this source code is governed by a BSD-3-clause license that can +// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +import CoreAILMCommon +import Testing + +/// Unit tests for the async `RequestQueue` semaphore that replaced the binary +/// tryAcquire/release busy-gate in the llm-server. +struct RequestQueueTests { + // MARK: - Helpers + + /// Poll `condition` until it becomes true or `timeout` elapses. Used to wait + /// for waiters to enqueue without relying on fixed sleeps. + private func waitUntil( + timeout: Duration = .seconds(2), + _ condition: @Sendable () -> Bool + ) async throws { + let deadline = ContinuousClock.now + timeout + while ContinuousClock.now < deadline { + if condition() { return } + try await Task.sleep(for: .milliseconds(2)) + } + Issue.record("waitUntil timed out") + } + + /// Records the order in which waiters wake up. + private actor Recorder { + private(set) var values: [Int] = [] + func record(_ value: Int) { values.append(value) } + } + + // MARK: - Tests + + @Test("First acquire succeeds immediately and marks the queue active") + func acquireImmediate() async throws { + let queue = RequestQueue(maxDepth: 4) + #expect(queue.isActive == false) + + let permit = try await queue.acquire() + #expect(queue.isActive == true) + #expect(queue.queuedCount == 0) + + permit.release() + #expect(queue.isActive == false) + } + + @Test("release with no waiters clears the active flag") + func releaseWithoutWaiters() async throws { + let queue = RequestQueue(maxDepth: 4) + let permit = try await queue.acquire() + #expect(queue.isActive == true) + + permit.release() + #expect(queue.isActive == false) + #expect(queue.queuedCount == 0) + } + + @Test("A second concurrent acquire waits, then proceeds after release") + func secondAcquireWaits() async throws { + let queue = RequestQueue(maxDepth: 4) + let permit0 = try await queue.acquire() // holds the single active slot + + let waiter = Task { try await queue.acquire() } + try await waitUntil { queue.queuedCount == 1 } // confirm it blocked + + permit0.release() // hand the slot to the waiter + let permit1 = try await waiter.value + #expect(queue.isActive == true) + #expect(queue.queuedCount == 0) + + permit1.release() + #expect(queue.isActive == false) + } + + @Test("Waiters wake in FIFO order") + func fifoWakeupOrder() async throws { + let queue = RequestQueue(maxDepth: 8) + let recorder = Recorder() + let permit0 = try await queue.acquire() // active slot held by the test + + // Enqueue three waiters in a deterministic order: launch each only after + // the previous one is observed in the queue. Each waiter records its id + // and releases so the next in line wakes. + var tasks: [Task] = [] + for index in 0..<3 { + let task = Task { + let permit = try await queue.acquire() + await recorder.record(index) + permit.release() + } + tasks.append(task) + try await waitUntil { queue.queuedCount == index + 1 } + } + + permit0.release() // trigger the cascade: waiter 0 -> 1 -> 2 + for task in tasks { try await task.value } + + #expect(await recorder.values == [0, 1, 2]) + #expect(queue.isActive == false) + } + + @Test("acquire rejects with queueFull once the queue reaches maxDepth") + func rejectsAtMaxDepth() async throws { + let queue = RequestQueue(maxDepth: 2) + let permit0 = try await queue.acquire() // active slot + + let w1 = Task { try await queue.acquire() } + try await waitUntil { queue.queuedCount == 1 } + let w2 = Task { try await queue.acquire() } + try await waitUntil { queue.queuedCount == 2 } + + // Queue is now at maxDepth (2 waiters); the next acquire must reject. + do { + _ = try await queue.acquire() + Issue.record("expected queueFull to be thrown") + } catch let error as ServerError { + guard case .queueFull(let depth) = error else { + Issue.record("expected .queueFull, got \(error)") + return + } + #expect(depth == 2) + } + + // Drain the two queued waiters so no continuation is leaked. + permit0.release() + let permit1 = try await w1.value + permit1.release() + let permit2 = try await w2.value + permit2.release() + #expect(queue.isActive == false) + } + + @Test("maxDepth 0 rejects the second concurrent request immediately") + func maxDepthZeroRejects() async throws { + let queue = RequestQueue(maxDepth: 0) + let permit0 = try await queue.acquire() // active slot, no queuing allowed + + await #expect(throws: ServerError.self) { + _ = try await queue.acquire() + } + #expect(queue.queuedCount == 0) + + permit0.release() + let permit1 = try await queue.acquire() // slot free again + #expect(queue.isActive == true) + permit1.release() + } + + @Test("Negative maxDepth is clamped to zero") + func negativeMaxDepthClamped() async throws { + let queue = RequestQueue(maxDepth: -5) + #expect(queue.maxDepth == 0) + + let permit0 = try await queue.acquire() + await #expect(throws: ServerError.self) { + _ = try await queue.acquire() + } + permit0.release() + } + + @Test("Permit release is idempotent") + func permitReleaseIdempotent() async throws { + let queue = RequestQueue(maxDepth: 4) + let permit = try await queue.acquire() + + permit.release() + permit.release() // second release is a no-op + #expect(queue.isActive == false) + + let permit2 = try await queue.acquire() + #expect(queue.isActive == true) + permit2.release() + } + + @Test( + "Cancelling a queued waiter removes it and frees the slot", + .timeLimit(.minutes(1))) + func cancelledWaiterIsRemoved() async throws { + let queue = RequestQueue(maxDepth: 4) + let permit0 = try await queue.acquire() // test holds the active slot + + let waiter = Task { try await queue.acquire() } // must suspend + try await waitUntil { queue.queuedCount == 1 } + + waiter.cancel() // cancel while suspended in acquire() + + // The waiter is removed promptly and its acquire throws CancellationError. + try await waitUntil { queue.queuedCount == 0 } + #expect(queue.queuedCount == 0) + await #expect(throws: CancellationError.self) { + _ = try await waiter.value + } + + // The slot was not handed to the cancelled waiter: releasing frees it and + // a fresh acquire still succeeds. + permit0.release() + #expect(queue.isActive == false) + let permit1 = try await queue.acquire() + #expect(queue.isActive == true) + permit1.release() + } +} From 1b97dcb7e22bf0611e7ac8927313a7a3986110df Mon Sep 17 00:00:00 2001 From: sukru tikves Date: Mon, 14 Sep 2026 10:32:13 -0700 Subject: [PATCH 2/2] Use Atomic for QueuePermit release guard Replace the Mutex once-guard with a lock-free Atomic.compareExchange; the release path takes the RequestQueue Mutex, so relaxed ordering is sufficient. Add a concurrent-release unit test that races 50 release() calls against a single queued waiter. --- .../Sources/CoreAILMCommon/RequestQueue.swift | 15 ++++---- .../RequestQueueTests.swift | 36 +++++++++++++++++++ 2 files changed, 44 insertions(+), 7 deletions(-) diff --git a/swift/Sources/CoreAILMCommon/RequestQueue.swift b/swift/Sources/CoreAILMCommon/RequestQueue.swift index 5b118c73..4b2329b7 100644 --- a/swift/Sources/CoreAILMCommon/RequestQueue.swift +++ b/swift/Sources/CoreAILMCommon/RequestQueue.swift @@ -124,19 +124,20 @@ public final class RequestQueue: Sendable { /// is dropped by the server before its writer closure ever runs). public final class QueuePermit: Sendable { private let queue: RequestQueue - private let released = Mutex(false) + private let released = Atomic(false) fileprivate init(_ queue: RequestQueue) { self.queue = queue } public func release() { - let shouldRelease = released.withLock { r -> Bool in - if r { return false } - r = true - return true - } - if shouldRelease { queue.release() } + // Flip false -> true atomically; exactly one caller wins the exchange and + // hands the slot back. `.relaxed` suffices: the exchange's atomicity (not + // its ordering) picks the single winner, and queue.release() takes the + // RequestQueue Mutex, which provides the actual cross-thread handoff. + let (won, _) = released.compareExchange( + expected: false, desired: true, ordering: .relaxed) + if won { queue.release() } } deinit { release() } diff --git a/swift/Tests/CoreAILMCommonTests/RequestQueueTests.swift b/swift/Tests/CoreAILMCommonTests/RequestQueueTests.swift index a0673d6e..61e4b0b3 100644 --- a/swift/Tests/CoreAILMCommonTests/RequestQueueTests.swift +++ b/swift/Tests/CoreAILMCommonTests/RequestQueueTests.swift @@ -174,6 +174,42 @@ struct RequestQueueTests { permit2.release() } + @Test( + "Concurrent releases hand the slot back exactly once", + .timeLimit(.minutes(1))) + func concurrentReleasesReturnSlotOnce() async throws { + let queue = RequestQueue(maxDepth: 4) + let permit = try await queue.acquire() + + // Enqueue exactly one waiter so a double-return would be observable: if + // the slot is handed back more than once, the second hand-back would wake + // this waiter (or corrupt isActive) even though nobody re-acquired. + let waiter = Task { try await queue.acquire() } + try await waitUntil { queue.queuedCount == 1 } + + // Race many release() calls on the same permit. The atomic once-guard must + // let exactly one win, so the slot is handed to the single waiter once. + await withTaskGroup(of: Void.self) { group in + for _ in 0..<50 { + group.addTask { permit.release() } + } + } + + // The single waiter got the slot exactly once; nothing is left queued. + let permit1 = try await waiter.value + #expect(queue.isActive == true) + #expect(queue.queuedCount == 0) + + permit1.release() + #expect(queue.isActive == false) + + // Slot is usable again — accounting was not corrupted by the race. + let permit2 = try await queue.acquire() + #expect(queue.isActive == true) + permit2.release() + #expect(queue.isActive == false) + } + @Test( "Cancelling a queued waiter removes it and frees the slot", .timeLimit(.minutes(1)))