diff --git a/Documentation/Daemon.md b/Documentation/Daemon.md index ca2cf9c..6d8a11d 100644 --- a/Documentation/Daemon.md +++ b/Documentation/Daemon.md @@ -84,6 +84,16 @@ Each key is a **server ID** — a short name you'll use with `--server` in the C |-------|------|---------|-------------| | `httpPort` | `Int` | — | Enable HTTP+SSE transport on this port (in addition to Bonjour) | +## IMAP Connections + +Per server the daemon keeps: + +- **One primary connection** for listing, fetching, moving, flagging, drafts and everything else. It runs one request at a time: IMAP keeps the selected mailbox per connection, so a request selects its mailbox and finishes its commands before the next request may select another. +- **Up to two extra connections** for searches, counts and raw `.eml` downloads (`post search`, `post fetch --eml`). Each is leased to one request at a time, so these can run while the primary connection is busy. They open on first use; one-at-a-time use only ever opens the first. If the server refuses an extra connection, the request runs on the primary connection instead. +- **IDLE connections** for watched mailboxes (see below). + +Concurrent requests are safe: each returns the same result it would return on its own. + ## IMAP IDLE IDLE is optional and can be enabled per server. When enabled, the daemon maintains a dedicated IMAP connection that listens for real-time mailbox changes — no polling needed. diff --git a/Sources/PostServer/AsyncSerialLock.swift b/Sources/PostServer/AsyncSerialLock.swift new file mode 100644 index 0000000..a6227e9 --- /dev/null +++ b/Sources/PostServer/AsyncSerialLock.swift @@ -0,0 +1,33 @@ +import Foundation + +/// A lock for async code: one holder at a time, waiters resume in arrival order. +/// +/// IMAP keeps the selected mailbox per connection, and every Post operation selects +/// a mailbox and then awaits further commands. Actors don't serialize work across +/// those suspension points, so without this lock a concurrent request's SELECT can +/// land in between and the commands run in the wrong mailbox. The lock is held for +/// the whole operation, not per command. +actor AsyncSerialLock { + private var isLocked = false + private var waiters: [CheckedContinuation] = [] + + /// Number of callers waiting for the lock; for tests. + var waiterCount: Int { waiters.count } + + func lock() async { + guard isLocked else { + isLocked = true + return + } + await withCheckedContinuation { waiters.append($0) } + } + + func unlock() { + if waiters.isEmpty { + isLocked = false + } else { + // Hand the lock straight to the next waiter; it stays held. + waiters.removeFirst().resume() + } + } +} diff --git a/Sources/PostServer/ConnectionSlotPool.swift b/Sources/PostServer/ConnectionSlotPool.swift new file mode 100644 index 0000000..3618a6d --- /dev/null +++ b/Sources/PostServer/ConnectionSlotPool.swift @@ -0,0 +1,39 @@ +import Foundation + +/// Leases a fixed number of slots, each to one holder at a time. +/// +/// Each slot stands for one of an account's extra IMAP connections. A search or raw +/// download leases a slot for its whole duration, so no two operations ever share an +/// extra connection and its selected mailbox. The lowest free slot is handed out first, +/// so one-at-a-time use keeps a single extra connection open. +actor ConnectionSlotPool { + let size: Int + private var free: Set + private var waiters: [CheckedContinuation] = [] + + init(size: Int) { + precondition(size > 0, "A connection slot pool needs at least one slot") + self.size = size + self.free = Set(0.. Int { + if let slot = free.min() { + free.remove(slot) + return slot + } + return await withCheckedContinuation { waiters.append($0) } + } + + func release(_ slot: Int) { + if waiters.isEmpty { + free.insert(slot) + } else { + // Hand the slot straight to the next waiter. + waiters.removeFirst().resume(returning: slot) + } + } +} diff --git a/Sources/PostServer/IMAPConnectionManager.swift b/Sources/PostServer/IMAPConnectionManager.swift index df0250a..364496e 100644 --- a/Sources/PostServer/IMAPConnectionManager.swift +++ b/Sources/PostServer/IMAPConnectionManager.swift @@ -3,6 +3,11 @@ import Logging import SwiftMail public actor IMAPConnectionManager { + /// Extra connections per account for searches and raw downloads, on top of the + /// primary connection and any IDLE connections. Kept small: servers cap the + /// connections per user (Dovecot 10 per IP, Gmail 15). + public static let extraConnectionsPerServer = 2 + private let configuration: PostConfiguration private var connections: [String: IMAPServer] = [:] /// In-flight connection tasks so concurrent callers share a single attempt. diff --git a/Sources/PostServer/MailboxCommandRunning.swift b/Sources/PostServer/MailboxCommandRunning.swift new file mode 100644 index 0000000..4eb0150 --- /dev/null +++ b/Sources/PostServer/MailboxCommandRunning.swift @@ -0,0 +1,59 @@ +import Foundation +import NIOIMAPCore +import SwiftMail + +// NIOIMAPCore also declares IMAPServer, MessageIdentifier and MessageIdentifierSet, so the +// SwiftMail types are qualified here. + +/// The commands searches and raw downloads need. Both the shared primary connection +/// (`IMAPServer`) and an extra connection (`IMAPNamedConnection`) provide them, so the +/// same code runs on an extra connection or, as a fallback, on the primary one. +protocol MailboxCommandRunning: Actor { + func selectMailbox(_ mailboxName: String) async throws -> SwiftMail.Mailbox.Selection + + func extendedSearch( + identifierSet: SwiftMail.MessageIdentifierSet?, + criteria: [SwiftMail.SearchCriteria], + sortCriteria: [SortCriterion], + sortCharset: String, + calendar: Calendar, + partialRange: PartialRange? + ) async throws -> SwiftMail.ExtendedSearchResult + + func fetchMessageInfosBulk( + using identifierSet: SwiftMail.MessageIdentifierSet, + options: SwiftMail.FetchMessageInfoOptions, + headerFields: [String]? + ) async throws -> [SwiftMail.MessageInfo] + + func fetchRawMessage(identifier: T) async throws -> Data +} + +extension SwiftMail.IMAPServer: MailboxCommandRunning {} +extension SwiftMail.IMAPNamedConnection: MailboxCommandRunning {} + +extension MailboxCommandRunning { + /// Unsorted extended search; named apart from `extendedSearch` so calls on a concrete + /// `IMAPServer` stay unambiguous. + func searchUIDs( + identifierSet: SwiftMail.MessageIdentifierSet?, + criteria: [SwiftMail.SearchCriteria], + partialRange: PartialRange? + ) async throws -> SwiftMail.ExtendedSearchResult { + try await extendedSearch( + identifierSet: identifierSet, + criteria: criteria, + sortCriteria: [], + sortCharset: "UTF-8", + calendar: Calendar(identifier: .gregorian), + partialRange: partialRange + ) + } + + /// Envelope, flags and dates only, no message bodies. + func fetchHeaderInfos( + using identifierSet: SwiftMail.MessageIdentifierSet + ) async throws -> [SwiftMail.MessageInfo] { + try await fetchMessageInfosBulk(using: identifierSet, options: .default, headerFields: nil) + } +} diff --git a/Sources/PostServer/PostServer.swift b/Sources/PostServer/PostServer.swift index 02adfde..3d56e95 100644 --- a/Sources/PostServer/PostServer.swift +++ b/Sources/PostServer/PostServer.swift @@ -166,6 +166,12 @@ public actor PostServer { } internal var connectionManager: IMAPConnectionManager + /// One lock per account, held for a whole operation on its primary connection, and the + /// leases for its extra connections. They live here, not in the connection manager, so a + /// configuration reload that replaces the manager keeps them: an operation that started + /// before the reload and one that starts after it still exclude each other. + private var primaryLocks: [String: AsyncSerialLock] = [:] + private var extraConnectionPools: [String: ConnectionSlotPool] = [:] private let apiKeyStore = APIKeyStore() private var scopedServerIDsByToken: [String: Set] = [:] private var scopesByToken: [String: Set] = [:] @@ -229,11 +235,8 @@ public actor PostServer { // Stop all IDLE watches stopIdleWatches() - // Shut down existing connections - await connectionManager.shutdown() - - // Replace connection manager with new config - connectionManager = IMAPConnectionManager(configuration: newConfig) + // Shut down existing connections and replace the connection manager + await replaceConnectionManager(configuration: newConfig) do { let keyCount = try primeAPIKeyScopes() logger.info("API key scopes refreshed (\(keyCount) key(s)).") @@ -250,6 +253,35 @@ public actor PostServer { } } + /// Shuts down the current connections and switches to a connection manager for + /// `configuration`. The per-account locks and extra-connection leases are kept. + internal func replaceConnectionManager(configuration: PostConfiguration) async { + await connectionManager.shutdown() + connectionManager = IMAPConnectionManager(configuration: configuration) + } + + /// The lock that keeps "select a mailbox, then run commands in it" together on the + /// account's shared primary connection. + internal func primaryLock(for serverId: String) -> AsyncSerialLock { + if let lock = primaryLocks[serverId] { + return lock + } + let lock = AsyncSerialLock() + primaryLocks[serverId] = lock + return lock + } + + /// The leases for the account's extra connections + /// (`IMAPConnectionManager.extraConnectionsPerServer` slots). + internal func extraConnectionPool(for serverId: String) -> ConnectionSlotPool { + if let pool = extraConnectionPools[serverId] { + return pool + } + let pool = ConnectionSlotPool(size: IMAPConnectionManager.extraConnectionsPerServer) + extraConnectionPools[serverId] = pool + return pool + } + /// Lists all configured IMAP servers with their connection details @MCPTool public func listServers() async throws -> [ServerInfo] { @@ -315,12 +347,12 @@ public actor PostServer { throw PostServerError.invalidUIDSet(uids) } - return try await withServer(serverId: serverId) { server in - _ = try await server.selectMailbox(mailbox) + return try await withExtraConnection(serverId: serverId) { connection in + _ = try await connection.selectMailbox(mailbox) var results: [RawMessage] = [] for uid in set.toArray() { - let data = try await server.fetchRawMessage(identifier: uid) + let data = try await connection.fetchRawMessage(identifier: uid) results.append(RawMessage(uid: Int(uid.value), rawData: data)) } @@ -395,8 +427,8 @@ public actor PostServer { limit: Int = 100, afterUid: Int? = nil ) async throws -> SearchResult { - try await withServer(serverId: serverId) { server in - _ = try await server.selectMailbox(mailbox) + try await withExtraConnection(serverId: serverId) { connection in + _ = try await connection.selectMailbox(mailbox) var criteria: [SearchCriteria] = [] @@ -454,7 +486,7 @@ public actor PostServer { // Use extended search with PARTIAL for server-side paging. // RFC 9394 window is 1-based relative to the current scoped search result set. - let result = try await server.extendedSearch( + let result = try await connection.searchUIDs( identifierSet: identifierSet, criteria: criteria, partialRange: makeFirstPartialRange(limit: limit) @@ -483,9 +515,8 @@ public actor PostServer { ) } - // Fetch headers for the limited UIDs - let uidSet = MessageIdentifierSet(limitedUIDs) - let headers = try await collectHeaders(from: server.fetchMessages(using: uidSet)) + // Fetch only the headers for the limited UIDs, not the full messages + let headers = try await fetchHeaders(for: limitedUIDs, using: connection) let sortedHeaders = headers.sorted { $0.uid > $1.uid } // Calculate returned range @@ -697,8 +728,8 @@ public actor PostServer { seen: Bool? = nil, flagged: Bool? = nil ) async throws -> SearchCount { - try await withServer(serverId: serverId) { server in - _ = try await server.selectMailbox(mailbox) + try await withExtraConnection(serverId: serverId) { connection in + _ = try await connection.selectMailbox(mailbox) var criteria: [SearchCriteria] = [] if let fromAddress, !fromAddress.isEmpty { criteria.append(.from(fromAddress)) } @@ -711,7 +742,11 @@ public actor PostServer { if flagged == true { criteria.append(.flagged) } if criteria.isEmpty { criteria.append(.all) } - let result: ExtendedSearchResult = try await server.extendedSearch(criteria: criteria) + let result: ExtendedSearchResult = try await connection.searchUIDs( + identifierSet: nil, + criteria: criteria, + partialRange: nil + ) return SearchCount( count: result.count, minUID: result.min.map { Int($0.value) }, @@ -817,10 +852,10 @@ public actor PostServer { guard (1...Int(UInt32.max)).contains(uid) else { throw PostServerError.invalidUID(uid) } - return try await withServer(serverId: serverId) { server in - _ = try await server.selectMailbox(mailbox) + return try await withExtraConnection(serverId: serverId) { connection in + _ = try await connection.selectMailbox(mailbox) let uid = UID(UInt32(uid)) - return try await server.fetchRawMessage(identifier: uid) + return try await connection.fetchRawMessage(identifier: uid) } } @@ -1436,10 +1471,30 @@ public actor PostServer { return results.sorted() } + /// Runs `operation` on the account's shared primary connection, one operation at a time. + /// + /// IMAP keeps the selected mailbox per connection, and operations select a mailbox and + /// then run further commands. Holding the account's lock for the whole operation stops a + /// concurrent request's SELECT from landing in between, which made searches return + /// another mailbox's results and could make move, trash or expunge act on the same UIDs + /// in the wrong mailbox. internal func withServer(serverId: String, operation: (SwiftMail.IMAPServer) async throws -> T) async throws -> T { try await assertScopeAllowed("imap") try await assertServerAccessAllowed(serverId) + let lock = primaryLock(for: serverId) + await lock.lock() + do { + let result = try await runOnPrimaryConnection(serverId: serverId, operation: operation) + await lock.unlock() + return result + } catch { + await lock.unlock() + throw error + } + } + + private func runOnPrimaryConnection(serverId: String, operation: (SwiftMail.IMAPServer) async throws -> T) async throws -> T { do { let server = try await connectionManager.connection(for: serverId) return try await operation(server) @@ -1454,6 +1509,62 @@ public actor PostServer { } } + /// Runs `operation` on one of the account's extra connections, leased for the whole + /// operation, so searches and raw downloads neither wait for nor disturb the work on + /// the primary connection. If the server refuses an extra connection (for example + /// because of its connection limit), the operation runs on the primary connection. + internal func withExtraConnection(serverId: String, operation: (any MailboxCommandRunning) async throws -> T) async throws -> T { + try await assertScopeAllowed("imap") + try await assertServerAccessAllowed(serverId) + + let pool = extraConnectionPool(for: serverId) + let slot = await pool.acquire() + do { + let result = try await runOnExtraConnection(slot: slot, serverId: serverId, operation: operation) + await pool.release(slot) + return result + } catch { + await pool.release(slot) + throw error + } + } + + private func runOnExtraConnection(slot: Int, serverId: String, operation: (any MailboxCommandRunning) async throws -> T) async throws -> T { + let connection: IMAPNamedConnection + do { + connection = try await extraConnection(slot: slot, serverId: serverId) + } catch { + logger.warning("No extra connection for \(serverId), using the primary connection: \(String(describing: error))") + return try await withServer(serverId: serverId) { server in + try await operation(server) + } + } + + do { + return try await operation(connection) + } catch { + guard shouldReconnect(after: error) else { + throw error + } + + logger.warning("Extra connection \(slot) failed for \(serverId), reconnecting: \(String(describing: error))") + try? await connection.disconnect() + let fresh = try await extraConnection(slot: slot, serverId: serverId) + return try await operation(fresh) + } + } + + /// The account's extra connection for `slot`, connected and logged in. SwiftMail keeps + /// one connection per name, so the same slot reuses the same connection. + private func extraConnection(slot: Int, serverId: String) async throws -> IMAPNamedConnection { + let server = try await connectionManager.connection(for: serverId) + let connection = try await server.connection(named: "post-extra-\(slot)") + if await !connection.isConnected { + try await connection.connect() + } + return connection + } + private func shouldReconnect(after error: Error) -> Bool { if let imapError = error as? IMAPError { switch imapError { @@ -1559,18 +1670,37 @@ public actor PostServer { } internal func messageHeader(from message: Message) -> MessageHeader { - let sanitizedSubject = UnicodeAbuseSummary.sanitize(message.subject ?? "(No Subject)", field: "Subject") + messageHeader(from: message.header) + } + + internal func messageHeader(from info: MessageInfo) -> MessageHeader { + let sanitizedSubject = UnicodeAbuseSummary.sanitize(info.subject ?? "(No Subject)", field: "Subject") return MessageHeader( - uid: messageUID(from: message), - from: message.from ?? "Unknown", + uid: info.uid.map { Int($0.value) } ?? Int(info.sequenceNumber.value), + from: info.from ?? "Unknown", subject: sanitizedSubject.text, - date: formatDate(message.date), - flags: MessageFlags(message.flags), + date: formatDate(info.date), + flags: MessageFlags(info.flags), unicodeAbuse: sanitizedSubject.unicodeAbuse ) } + /// Fetches the header data for `uids` in chunks, so a large result set never becomes one + /// oversized response. + internal func fetchHeaders(for uids: [UID], using connection: any MailboxCommandRunning) async throws -> [MessageHeader] { + let chunkSize = 200 + var headers: [MessageHeader] = [] + var start = uids.startIndex + while start < uids.endIndex { + let end = uids.index(start, offsetBy: chunkSize, limitedBy: uids.endIndex) ?? uids.endIndex + let infos = try await connection.fetchHeaderInfos(using: MessageIdentifierSet(Array(uids[start.. MessageDetail { let attachments = message.attachments.map { AttachmentInfo( diff --git a/Tests/PostCLITests/AsyncSerialLockTests.swift b/Tests/PostCLITests/AsyncSerialLockTests.swift new file mode 100644 index 0000000..db6791d --- /dev/null +++ b/Tests/PostCLITests/AsyncSerialLockTests.swift @@ -0,0 +1,84 @@ +import XCTest +@testable import PostServer + +final class AsyncSerialLockTests: XCTestCase { + /// Counts how many holders are inside the critical section at once. + private actor Occupancy { + private(set) var current = 0 + private(set) var maximum = 0 + + func enter() { + current += 1 + maximum = max(maximum, current) + } + + func leave() { + current -= 1 + } + } + + private actor Recorder { + private(set) var values: [Int] = [] + + func append(_ value: Int) { + values.append(value) + } + } + + func testHoldersNeverOverlap() async { + let lock = AsyncSerialLock() + let occupancy = Occupancy() + + await withTaskGroup(of: Void.self) { group in + for _ in 0..<20 { + group.addTask { + await lock.lock() + await occupancy.enter() + try? await Task.sleep(nanoseconds: 1_000_000) + await occupancy.leave() + await lock.unlock() + } + } + } + + let maximum = await occupancy.maximum + XCTAssertEqual(maximum, 1) + } + + func testWaitersAreServedInArrivalOrder() async { + let lock = AsyncSerialLock() + let order = Recorder() + await lock.lock() + + var tasks: [Task] = [] + for index in 0..<5 { + tasks.append(Task { + await lock.lock() + await order.append(index) + await lock.unlock() + }) + // Queue the waiters one by one so their arrival order is known. + while await lock.waiterCount < index + 1 { + await Task.yield() + } + } + + await lock.unlock() + for task in tasks { + await task.value + } + + let values = await order.values + XCTAssertEqual(values, [0, 1, 2, 3, 4]) + } + + func testLockCanBeTakenAgainAfterUnlock() async { + let lock = AsyncSerialLock() + await lock.lock() + await lock.unlock() + await lock.lock() + let waiting = await lock.waiterCount + XCTAssertEqual(waiting, 0) + await lock.unlock() + } +} diff --git a/Tests/PostCLITests/ConcurrentMailboxAccessLiveTests.swift b/Tests/PostCLITests/ConcurrentMailboxAccessLiveTests.swift new file mode 100644 index 0000000..01aa4ae --- /dev/null +++ b/Tests/PostCLITests/ConcurrentMailboxAccessLiveTests.swift @@ -0,0 +1,83 @@ +import XCTest +@testable import PostServer + +/// Concurrent requests on one account must return exactly what the same requests return one +/// at a time. Before the per-account lock and the extra connections, a search in one mailbox +/// returned another mailbox's results whenever requests overlapped. +/// +/// Runs against a real account and only reads: set POST_LIVE_TEST_SERVER to a server ID from +/// ~/.post.json, optionally POST_LIVE_TEST_MAILBOXES (comma-separated, default INBOX,Sent,Archive). +final class ConcurrentMailboxAccessLiveTests: XCTestCase { + private struct Snapshot: Equatable, Sendable { + var searchUIDs: [String: [Int]] = [:] + var counts: [String: Int] = [:] + var rawMessages: [String: Data] = [:] + } + + func testConcurrentRequestsMatchSerialOnes() async throws { + let environment = ProcessInfo.processInfo.environment + guard let serverId = environment["POST_LIVE_TEST_SERVER"], !serverId.isEmpty else { + throw XCTSkip("Set POST_LIVE_TEST_SERVER to run against a real account") + } + let mailboxes = (environment["POST_LIVE_TEST_MAILBOXES"] ?? "INBOX,Sent,Archive") + .split(separator: ",") + .map { $0.trimmingCharacters(in: .whitespaces) } + + let server = PostServer(configuration: try PostConfiguration.load()) + // Counts return every matching UID; limit them to the last year so a large + // mailbox's reply stays within SwiftMail's response buffer. + let since = ISO8601DateFormatter().string(from: Date().addingTimeInterval(-365 * 24 * 3600)) + + // One request at a time. + var baseline = Snapshot() + for mailbox in mailboxes { + let uids = try await server.searchMessages(serverId: serverId, mailbox: mailbox, limit: 200).messages.map(\.uid) + baseline.searchUIDs[mailbox] = uids + baseline.counts[mailbox] = try await server.countMessages(serverId: serverId, mailbox: mailbox, since: since).count ?? -1 + if let uid = uids.min() { + baseline.rawMessages[mailbox] = try await server.downloadEml(serverId: serverId, uid: uid, mailbox: mailbox) + } + } + + // The same requests all at once, mixed with listings on the primary connection. + for round in 0..<3 { + let concurrent = try await withThrowingTaskGroup(of: Snapshot.self) { group in + for mailbox in mailboxes { + group.addTask { + let uids = try await server.searchMessages(serverId: serverId, mailbox: mailbox, limit: 200).messages.map(\.uid) + return Snapshot(searchUIDs: [mailbox: uids]) + } + group.addTask { + let count = try await server.countMessages(serverId: serverId, mailbox: mailbox, since: since).count ?? -1 + return Snapshot(counts: [mailbox: count]) + } + if let uid = baseline.searchUIDs[mailbox]?.min() { + group.addTask { + let data = try await server.downloadEml(serverId: serverId, uid: uid, mailbox: mailbox) + return Snapshot(rawMessages: [mailbox: data]) + } + } + group.addTask { + _ = try await server.listMessages(serverId: serverId, mailbox: mailbox, limit: 3) + return Snapshot() + } + } + + var merged = Snapshot() + for try await part in group { + merged.searchUIDs.merge(part.searchUIDs) { $1 } + merged.counts.merge(part.counts) { $1 } + merged.rawMessages.merge(part.rawMessages) { $1 } + } + return merged + } + + XCTAssertEqual(concurrent.searchUIDs, baseline.searchUIDs, "search results differ in round \(round)") + XCTAssertEqual(concurrent.counts, baseline.counts, "counts differ in round \(round)") + XCTAssertEqual(concurrent.rawMessages.mapValues { $0.count }, baseline.rawMessages.mapValues { $0.count }, "raw message sizes differ in round \(round)") + XCTAssertTrue(concurrent.rawMessages == baseline.rawMessages, "raw messages differ in round \(round)") + } + + await server.shutdown() + } +} diff --git a/Tests/PostCLITests/ConnectionManagerSwitchTests.swift b/Tests/PostCLITests/ConnectionManagerSwitchTests.swift new file mode 100644 index 0000000..f8703cd --- /dev/null +++ b/Tests/PostCLITests/ConnectionManagerSwitchTests.swift @@ -0,0 +1,44 @@ +import XCTest +@testable import PostServer + +/// A configuration reload replaces the connection manager. The per-account lock and the +/// extra-connection leases must survive that, or an operation that started before the reload +/// and one that starts after it could use the same new connection at the same time. +final class ConnectionManagerSwitchTests: XCTestCase { + private let configuration = PostConfiguration(servers: ["mail": .init()], httpPort: nil) + + func testLockAndLeasesSurviveAManagerSwitch() async { + let server = PostServer(configuration: configuration) + let lock = await server.primaryLock(for: "mail") + let pool = await server.extraConnectionPool(for: "mail") + + await server.replaceConnectionManager(configuration: configuration) + + let lockAfter = await server.primaryLock(for: "mail") + let poolAfter = await server.extraConnectionPool(for: "mail") + XCTAssertTrue(lock === lockAfter) + XCTAssertTrue(pool === poolAfter) + } + + func testHolderFromBeforeASwitchStillExcludesLaterCallers() async { + let server = PostServer(configuration: configuration) + let lock = await server.primaryLock(for: "mail") + await lock.lock() + + await server.replaceConnectionManager(configuration: configuration) + + let lockAfter = await server.primaryLock(for: "mail") + let later = Task { + await lockAfter.lock() + await lockAfter.unlock() + } + while await lock.waiterCount < 1 { + await Task.yield() + } + let waiting = await lock.waiterCount + XCTAssertEqual(waiting, 1, "a caller after the switch must wait for the holder from before it") + + await lock.unlock() + await later.value + } +} diff --git a/Tests/PostCLITests/ConnectionSlotPoolTests.swift b/Tests/PostCLITests/ConnectionSlotPoolTests.swift new file mode 100644 index 0000000..b9f6ddf --- /dev/null +++ b/Tests/PostCLITests/ConnectionSlotPoolTests.swift @@ -0,0 +1,69 @@ +import XCTest +@testable import PostServer + +final class ConnectionSlotPoolTests: XCTestCase { + /// Tracks which slots are leased and whether any slot was leased twice at once. + private actor Leases { + private var held: Set = [] + private(set) var doubleLeases = 0 + private(set) var maximum = 0 + + func take(_ slot: Int) { + if held.contains(slot) { + doubleLeases += 1 + } + held.insert(slot) + maximum = max(maximum, held.count) + } + + func give(_ slot: Int) { + held.remove(slot) + } + } + + func testNoSlotIsLeasedTwiceAtOnce() async { + let pool = ConnectionSlotPool(size: 2) + let leases = Leases() + + await withTaskGroup(of: Void.self) { group in + for _ in 0..<20 { + group.addTask { + let slot = await pool.acquire() + await leases.take(slot) + try? await Task.sleep(nanoseconds: 1_000_000) + await leases.give(slot) + await pool.release(slot) + } + } + } + + let doubleLeases = await leases.doubleLeases + let maximum = await leases.maximum + XCTAssertEqual(doubleLeases, 0) + XCTAssertLessThanOrEqual(maximum, 2) + } + + func testOneAtATimeUseKeepsTheFirstSlot() async { + let pool = ConnectionSlotPool(size: 2) + for _ in 0..<3 { + let slot = await pool.acquire() + XCTAssertEqual(slot, 0) + await pool.release(slot) + } + } + + func testWaiterReceivesTheReleasedSlot() async { + let pool = ConnectionSlotPool(size: 1) + let first = await pool.acquire() + + let waiter = Task { await pool.acquire() } + while await pool.waiterCount < 1 { + await Task.yield() + } + + await pool.release(first) + let second = await waiter.value + XCTAssertEqual(second, first) + await pool.release(second) + } +}