/src/swift-nio/Sources/NIOPosix/NIOThreadPool.swift

Source
//===----------------------------------------------------------------------===//
//
// This source file is part of the SwiftNIO open source project
//
// Copyright (c) 2017-2024 Apple Inc. and the SwiftNIO project authors
// Licensed under Apache License v2.0
//
// See LICENSE.txt for license information
// See CONTRIBUTORS.txt for the list of SwiftNIO project authors
//
// SPDX-License-Identifier: Apache-2.0
//
//===----------------------------------------------------------------------===//

#if !os(WASI)

import Atomics
import DequeModule
import NIOConcurrencyHelpers
import NIOCore

#if canImport(Dispatch)
import Dispatch
#endif

/// Errors that may be thrown when executing work on a `NIOThreadPool`
public enum NIOThreadPoolError: Sendable {

    /// The `NIOThreadPool` was not active.
    public struct ThreadPoolInactive: Error {
        public init() {}
    }

    /// The `NIOThreadPool` operation is unsupported (e.g. shutdown of a perpetual pool).
    public struct UnsupportedOperation: Error {
        public init() {}
    }
}

/// A thread pool that should be used if some (kernel thread) blocking work
/// needs to be performed for which no non-blocking API exists.
///
/// When using NIO it is crucial not to block any of the `EventLoop`s as that
/// leads to slow downs or stalls of arbitrary other work. Unfortunately though
/// there are tasks that applications need to achieve for which no non-blocking
/// APIs exist. In those cases `NIOThreadPool` can be used but should be
/// treated as a last resort.
///
/// - Note: The prime example for missing non-blocking APIs is file IO on UNIX.
///   The OS does not provide a usable and truly non-blocking API but with
///   `NonBlockingFileIO` NIO provides a high-level API for file IO that should
///   be preferred to running blocking file IO system calls directly on
///   `NIOThreadPool`. Under the covers `NonBlockingFileIO` will use
///   `NIOThreadPool` on all currently supported platforms though.
public final class NIOThreadPool {

    /// The state of the `WorkItem`.
    public enum WorkItemState: Sendable {
        /// The `WorkItem` is active now and in process by the `NIOThreadPool`.
        case active
        /// The `WorkItem` was cancelled and will not be processed by the `NIOThreadPool`.
        case cancelled
    }

    /// The work that should be done by the `NIOThreadPool`.
    public typealias WorkItem = @Sendable (WorkItemState) -> Void

    @usableFromInline
    struct IdentifiableWorkItem: Sendable {
        @usableFromInline
        var workItem: WorkItem

        @usableFromInline
        var id: Int?
    }

    @usableFromInline
    internal enum State: Sendable {
        /// The `NIOThreadPool` is already stopped.
        case stopped
        /// The `NIOThreadPool` is shutting down, the array has one boolean entry for each thread indicating if it has shut down already.
        case shuttingDown([Bool])
        /// The `NIOThreadPool` is up and running, the `Deque` containing the yet unprocessed `IdentifiableWorkItem`s.
        case running(Deque<IdentifiableWorkItem>)
        /// Temporary state used when mutating the .running(items). Used to avoid CoW copies.
        /// It should never be "leaked" outside of the lock block.
        case modifying
    }

    @usableFromInline
    internal let _workAvailable: NIOThreadPoolWorkAvailable
    private var threads: [NIOThread]? = nil  // protected by `_workAvailable`
    @usableFromInline
    internal var _state: State = .stopped

    // WorkItems don't have a handle so they can't be cancelled directly. Instead an ID is assigned
    // to each cancellable work item and the IDs of each work item to cancel is stored in this set.
    // The set is checked when dequeuing work items prior to running them, the presence of an ID
    // indicates it should be cancelled. This approach makes cancellation cheap, but slow, as the
    // task isn't cancelled until it's dequeued.
    //
    // Possible alternatives:
    // - Removing items from the work queue on cancellation. This is linear and runs the risk of
    //   being expensive if a task tree with many enqueued work items is cancelled.
    // - Storing an atomic 'is cancelled' flag with each work item. This adds an allocation per
    //   work item.
    //
    // If a future version of this thread pool has work items which do have a handle this set should
    // be removed.
    //
    // Note: protected by 'lock'.
    @usableFromInline
    internal var _cancelledWorkIDs: Set<Int> = []
    private let nextWorkID = ManagedAtomic(0)

    public let numberOfThreads: Int
    private let canBeStopped: Bool
    private let pinnedCPUIDs: [Int]?

    /// Gracefully shutdown this `NIOThreadPool`. All tasks will be run before shutdown will take place.
    ///
    /// - Parameters:
    ///   - queue: The `DispatchQueue` used to executed the callback
    ///   - callback: The function to be executed once the shutdown is complete.
    @preconcurrency
    public func shutdownGracefully(queue: DispatchQueue, _ callback: @escaping @Sendable (Error?) -> Void) {
        self._shutdownGracefully(queue: queue, callback)
    }

    private func _shutdownGracefully(queue: DispatchQueue, _ callback: @escaping @Sendable (Error?) -> Void) {
        guard self.canBeStopped else {
            queue.async {
                callback(NIOThreadPoolError.UnsupportedOperation())
            }
            return
        }

        let threadsToJoin = self._workAvailable.withLockSettingWorkAvailable {
            switch self._state {
            case .running(let items):
                self._state = .modifying
                queue.async {
                    for item in items {
                        item.workItem(.cancelled)
                    }
                }
                self._state = .shuttingDown(Array(repeating: true, count: self.numberOfThreads))

                let threads = self.threads!
                self.threads = nil

                // Wake all threads so they can see .shuttingDown and exit.
                return (workAvailable: self.numberOfThreads, signal: .broadcastAll, result: threads)

            case .shuttingDown, .stopped:
                // Don't modify workAvailable - threads may still be waking up to see .shuttingDown
                return (workAvailable: nil, signal: .none, result: [])

            case .modifying:
                fatalError(".modifying state misuse")
            }
        }

        DispatchQueue(label: "io.swiftnio.NIOThreadPool.shutdownGracefully").async {
            for thread in threadsToJoin {
                thread.join()
            }
            queue.async {
                callback(nil)
            }
        }
    }

    /// Submit a `WorkItem` to process.
    ///
    /// - Note: This is a low-level method, in most cases the `runIfActive` method should be used.
    ///
    /// - Parameters:
    ///   - body: The `WorkItem` to process by the `NIOThreadPool`.
    @preconcurrency
    public func submit(_ body: @escaping WorkItem) {
        self._submit(id: nil, body)
    }

    private func _submit(id: Int?, _ body: @escaping WorkItem) {
        let submitted = self._workAvailable.withLock {
            let submitted: Bool

            switch self._state {
            case .running(var items):
                self._state = .modifying
                items.append(.init(workItem: body, id: id))
                self._state = .running(items)
                submitted = true

            case .shuttingDown, .stopped:
                submitted = false

            case .modifying:
                fatalError(".modifying state misuse")
            }

            return (workDelta: submitted ? 1 : 0, signal: submitted ? .signalOne : .none, result: submitted)
        }

        // if item couldn't be added run it immediately indicating that it couldn't be run
        if !submitted {
            body(.cancelled)
        }
    }

    /// Initialize a `NIOThreadPool` thread pool with `numberOfThreads` threads.
    ///
    /// - Parameters:
    ///   - numberOfThreads: The number of threads to use for the thread pool.
    public convenience init(numberOfThreads: Int) {
        self.init(
            numberOfThreads: numberOfThreads,
            pinnedCPUIDs: nil,
            canBeStopped: true
        )
    }

    #if !os(WASI)
    #if os(Linux) || os(Android)
    /// Initialize a `NIOThreadPool` thread pool with threads pinned to specific CPUs.
    ///
    /// - Parameters:
    ///   - pinnedCPUIDs: The IDs of the CPUs to use for the thread pool. One thread will be
    ///      created per entry in this array.
    public convenience init(pinnedCPUIDs: [Int]) {
        self.init(
            numberOfThreads: pinnedCPUIDs.count,
            pinnedCPUIDs: pinnedCPUIDs,
            canBeStopped: true
        )
    }

    #endif
    #endif
    /// Create a ``NIOThreadPool`` that is already started, cannot be shut down and must not be `deinit`ed.
    ///
    /// This is only useful for global singletons.
    public static func _makePerpetualStartedPool(numberOfThreads: Int, threadNamePrefix: String) -> NIOThreadPool {
        let pool = self.init(numberOfThreads: numberOfThreads, pinnedCPUIDs: nil, canBeStopped: false)
        pool._start(threadNamePrefix: threadNamePrefix)
        return pool
    }

    private init(numberOfThreads: Int, pinnedCPUIDs: [Int]?, canBeStopped: Bool) {
        self.numberOfThreads = numberOfThreads
        self.pinnedCPUIDs = pinnedCPUIDs
        self.canBeStopped = canBeStopped
        self._workAvailable = NIOThreadPoolWorkAvailable()
    }

    // Do not rename or remove this function.
    //
    // When doing on-/off-CPU analysis, for example with continuous profiling, it's
    // important to recognise certain functions that are purely there to wait.
    //
    // This function is one of those and giving it a consistent name makes it much easier to remove from the profiles
    // when only interested in on-CPU work.
    @inlinable
    internal func _blockingWaitForWork(identifier: Int) -> (item: WorkItem, state: WorkItemState)? {
        self._workAvailable.withLockWaitingForWork {
            let result: (WorkItem, WorkItemState)?

            switch self._state {
            case .running(var items):
                self._state = .modifying
                let itemAndID = items.removeFirst()

                let state: WorkItemState
                if let id = itemAndID.id, !self._cancelledWorkIDs.isEmpty {
                    state = self._cancelledWorkIDs.remove(id) == nil ? .active : .cancelled
                } else {
                    state = .active
                }

                self._state = .running(items)

                result = (itemAndID.workItem, state)

            case .shuttingDown(var aliveStates):
                self._state = .modifying
                assert(aliveStates[identifier])
                aliveStates[identifier] = false
                self._state = .shuttingDown(aliveStates)

                result = nil

            case .stopped:
                // Unreachable: 'stopped' is the initial state which is left when starting the
                // thread pool, and before any thread calls this function.
                fatalError("Invalid state")

            case .modifying:
                fatalError(".modifying state misuse")
            }

            return (workDelta: -1, signal: .none, result: result)
        }
    }

    private func process(identifier: Int) {
        repeat {
            let itemAndState = self._blockingWaitForWork(identifier: identifier)

            if let (item, state) = itemAndState {
                // if there was a work item popped, run it
                item(state)
            } else {
                break  // Otherwise, we're done
            }
        } while true
    }

    /// Start the `NIOThreadPool` if not already started.
    public func start() {
        self._start(threadNamePrefix: "TP-#")
    }

    public func _start(threadNamePrefix: String) {
        let alreadyRunning = self._workAvailable.withLock {
            switch self._state {
            case .running:
                // Already running, this has no effect on whether there is more work for the
                // threads to run.
                return (workDelta: 0, signal: .none, result: true)

            case .shuttingDown:
                // This should never happen
                fatalError("start() called while in shuttingDown")

            case .stopped:
                self._state = .running(Deque(minimumCapacity: 16))
                assert(self.threads == nil)
                self.threads = []
                self.threads!.reserveCapacity(self.numberOfThreads)
                return (workDelta: 0, signal: .none, result: false)

            case .modifying:
                fatalError(".modifying state misuse")
            }
        }

        if alreadyRunning {
            return
        }

        // We use this condition lock as a tricky kind of semaphore.
        // This is done to sidestep the thread performance checker warning
        // that would otherwise be emitted.
        let readyThreads = ConditionLock(value: 0)
        for id in 0..<self.numberOfThreads {
            let pinnedCPUID = self.pinnedCPUIDs?[id]

            // We should keep thread names under 16 characters because Linux doesn't allow more.
            NIOThread.spawnAndRun(name: "\(threadNamePrefix)\(id)") { thread in
                #if !os(WASI)
                #if os(Linux) || os(Android)
                if let pinnedCPUID {
                    precondition(thread.isCurrentSlow)
                    NIOThread.currentAffinity = LinuxCPUSet(pinnedCPUID)
                }
                #endif
                #endif
                readyThreads.withLock {
                    let threadCount = self._workAvailable.withLock {
                        self.threads!.append(thread)
                        let threadCount = self.threads!.count
                        return (workDelta: 0, signal: .none, result: threadCount)
                    }

                    return (unlockWith: threadCount, result: ())
                }

                self.process(identifier: id)
                return ()
            }
        }

        readyThreads.lock(whenValue: self.numberOfThreads)
        readyThreads.unlock()

        func threadCount() -> Int {
            self._workAvailable.withLock {
                (workDelta: 0, signal: .none, result: self.threads?.count ?? -1)
            }
        }
        assert(threadCount() == self.numberOfThreads)
    }

    deinit {
        assert(
            self.canBeStopped,
            "Perpetual NIOThreadPool has been deinited, you must make sure that perpetual pools don't deinit"
        )
        switch self._state {
        case .stopped, .shuttingDown:
            ()
        default:
            assertionFailure("wrong state \(self._state)")
        }
    }
}

extension NIOThreadPool: @unchecked Sendable {}

extension NIOThreadPool {

    /// Runs the submitted closure if the thread pool is still active, otherwise fails the promise.
    /// The closure will be run on the thread pool so can do blocking work.
    ///
    /// - Parameters:
    ///   - eventLoop: The `EventLoop` the returned `EventLoopFuture` will fire on.
    ///   - body: The closure which performs some blocking work to be done on the thread pool.
    /// - Returns: The `EventLoopFuture` of `promise` fulfilled with the result (or error) of the passed closure.
    @preconcurrency
    public func runIfActive<T: Sendable>(
        eventLoop: EventLoop,
        _ body: @escaping @Sendable () throws -> T
    ) -> EventLoopFuture<T> {
        self._runIfActive(eventLoop: eventLoop, body)
    }

    private func _runIfActive<T: Sendable>(
        eventLoop: EventLoop,
        _ body: @escaping @Sendable () throws -> T
    ) -> EventLoopFuture<T> {
        let promise = eventLoop.makePromise(of: T.self)
        self.submit { shouldRun in
            guard case shouldRun = NIOThreadPool.WorkItemState.active else {
                promise.fail(NIOThreadPoolError.ThreadPoolInactive())
                return
            }
            do {
                try promise.succeed(body())
            } catch {
                promise.fail(error)
            }
        }
        return promise.futureResult
    }

    /// Runs the submitted closure if the thread pool is still active, otherwise throw an error.
    /// The closure will be run on the thread pool, such that we can do blocking work.
    ///
    /// - Parameters:
    ///   - body: The closure which performs some blocking work to be done on the thread pool.
    /// - Returns: Result of the passed closure.
    @available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
    public func runIfActive<T: Sendable>(_ body: @escaping @Sendable () throws -> T) async throws -> T {
        let workID = self.nextWorkID.loadThenWrappingIncrement(ordering: .relaxed)

        return try await withTaskCancellationHandler {
            try await withCheckedThrowingContinuation { (cont: CheckedContinuation<T, Error>) in
                self._submit(id: workID) { shouldRun in
                    switch shouldRun {
                    case .active:
                        let result = Result(catching: body)
                        cont.resume(with: result)
                    case .cancelled:
                        cont.resume(throwing: CancellationError())
                    }
                }
            }
        } onCancel: {
            self._workAvailable.withLock {
                self._cancelledWorkIDs.insert(workID)
                return (workDelta: 0, signal: .none, result: ())
            }
        }
    }
}

extension NIOThreadPool {
    @preconcurrency
    public func shutdownGracefully(_ callback: @escaping @Sendable (Error?) -> Void) {
        self.shutdownGracefully(queue: .global(), callback)
    }

    /// Shuts down the thread pool gracefully.
    @available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
    @inlinable
    public func shutdownGracefully() async throws {
        try await withCheckedThrowingContinuation { (cont: CheckedContinuation<Void, Error>) in
            self.shutdownGracefully { error in
                if let error = error {
                    cont.resume(throwing: error)
                } else {
                    cont.resume()
                }
            }
        }
    }

    @available(*, noasync, message: "this can end up blocking the calling thread", renamed: "shutdownGracefully()")
    public func syncShutdownGracefully() throws {
        try self._syncShutdownGracefully()
    }

    private func _syncShutdownGracefully() throws {
        let errorStorageLock = NIOLockedValueBox<Swift.Error?>(nil)
        let continuation = ConditionLock(value: 0)
        self.shutdownGracefully { error in
            if let error = error {
                errorStorageLock.withLockedValue {
                    $0 = error
                }
            }
            continuation.lock(whenValue: 0)
            continuation.unlock(withValue: 1)
        }
        continuation.lock(whenValue: 1)
        continuation.unlock()
        try errorStorageLock.withLockedValue {
            if let error = $0 {
                throw error
            }
        }
    }
}

extension ConditionLock {
    @inlinable
    func _lock(when value: T?) {
        if let value = value {
            self.lock(whenValue: value)
        } else {
            self.lock()
        }
    }

    @inlinable
    func _unlock(with value: T?) {
        if let value = value {
            self.unlock(withValue: value)
        } else {
            self.unlock()
        }
    }

    @inlinable
    func withLock<Result>(when value: T? = nil, _ body: () -> (unlockWith: T?, result: Result)) -> Result {
        self._lock(when: value)
        let (unlockValue, result) = body()
        self._unlock(with: unlockValue)
        return result
    }
}
#endif  // !os(WASI)

Coverage Report

Created: 2026-10-13 06:43

Line	Count	Source
1		//===----------------------------------------------------------------------===//
2		//
3		// This source file is part of the SwiftNIO open source project
4		//
5		// Copyright (c) 2017-2024 Apple Inc. and the SwiftNIO project authors
6		// Licensed under Apache License v2.0
7		//
8		// See LICENSE.txt for license information
9		// See CONTRIBUTORS.txt for the list of SwiftNIO project authors
10		//
11		// SPDX-License-Identifier: Apache-2.0
12		//
13		//===----------------------------------------------------------------------===//
14
15		#if !os(WASI)
16
17		import Atomics
18		import DequeModule
19		import NIOConcurrencyHelpers
20		import NIOCore
21
22		#if canImport(Dispatch)
23		import Dispatch
24		#endif
25
26		/// Errors that may be thrown when executing work on a `NIOThreadPool`
27		public enum NIOThreadPoolError: Sendable {
28
29		/// The `NIOThreadPool` was not active.
30		public struct ThreadPoolInactive: Error {
31	0	public init() {}
32		}
33
34		/// The `NIOThreadPool` operation is unsupported (e.g. shutdown of a perpetual pool).
35		public struct UnsupportedOperation: Error {
36	0	public init() {}
37		}
38		}
39
40		/// A thread pool that should be used if some (kernel thread) blocking work
41		/// needs to be performed for which no non-blocking API exists.
42		///
43		/// When using NIO it is crucial not to block any of the `EventLoop`s as that
44		/// leads to slow downs or stalls of arbitrary other work. Unfortunately though
45		/// there are tasks that applications need to achieve for which no non-blocking
46		/// APIs exist. In those cases `NIOThreadPool` can be used but should be
47		/// treated as a last resort.
48		///
49		/// - Note: The prime example for missing non-blocking APIs is file IO on UNIX.
50		/// The OS does not provide a usable and truly non-blocking API but with
51		/// `NonBlockingFileIO` NIO provides a high-level API for file IO that should
52		/// be preferred to running blocking file IO system calls directly on
53		/// `NIOThreadPool`. Under the covers `NonBlockingFileIO` will use
54		/// `NIOThreadPool` on all currently supported platforms though.
55		public final class NIOThreadPool {
56
57		/// The state of the `WorkItem`.
58		public enum WorkItemState: Sendable {
59		/// The `WorkItem` is active now and in process by the `NIOThreadPool`.
60		case active
61		/// The `WorkItem` was cancelled and will not be processed by the `NIOThreadPool`.
62		case cancelled
63		}
64
65		/// The work that should be done by the `NIOThreadPool`.
66		public typealias WorkItem = @Sendable (WorkItemState) -> Void
67
68		@usableFromInline
69		struct IdentifiableWorkItem: Sendable {
70		@usableFromInline
71		var workItem: WorkItem
72
73		@usableFromInline
74		var id: Int?
75		}
76
77		@usableFromInline
78		internal enum State: Sendable {
79		/// The `NIOThreadPool` is already stopped.
80		case stopped
81		/// The `NIOThreadPool` is shutting down, the array has one boolean entry for each thread indicating if it has shut down already.
82		case shuttingDown([Bool])
83		/// The `NIOThreadPool` is up and running, the `Deque` containing the yet unprocessed `IdentifiableWorkItem`s.
84		case running(Deque<IdentifiableWorkItem>)
85		/// Temporary state used when mutating the .running(items). Used to avoid CoW copies.
86		/// It should never be "leaked" outside of the lock block.
87		case modifying
88		}
89
90		@usableFromInline
91		internal let _workAvailable: NIOThreadPoolWorkAvailable
92	0	private var threads: [NIOThread]? = nil // protected by `_workAvailable`
93		@usableFromInline
94	0	internal var _state: State = .stopped
95
96		// WorkItems don't have a handle so they can't be cancelled directly. Instead an ID is assigned
97		// to each cancellable work item and the IDs of each work item to cancel is stored in this set.
98		// The set is checked when dequeuing work items prior to running them, the presence of an ID
99		// indicates it should be cancelled. This approach makes cancellation cheap, but slow, as the
100		// task isn't cancelled until it's dequeued.
101		//
102		// Possible alternatives:
103		// - Removing items from the work queue on cancellation. This is linear and runs the risk of
104		// being expensive if a task tree with many enqueued work items is cancelled.
105		// - Storing an atomic 'is cancelled' flag with each work item. This adds an allocation per
106		// work item.
107		//
108		// If a future version of this thread pool has work items which do have a handle this set should
109		// be removed.
110		//
111		// Note: protected by 'lock'.
112		@usableFromInline
113	0	internal var _cancelledWorkIDs: Set<Int> = []
114	0	private let nextWorkID = ManagedAtomic(0)
115
116		public let numberOfThreads: Int
117		private let canBeStopped: Bool
118		private let pinnedCPUIDs: [Int]?
119
120		/// Gracefully shutdown this `NIOThreadPool`. All tasks will be run before shutdown will take place.
121		///
122		/// - Parameters:
123		/// - queue: The `DispatchQueue` used to executed the callback
124		/// - callback: The function to be executed once the shutdown is complete.
125		@preconcurrency
126	0	public func shutdownGracefully(queue: DispatchQueue, _ callback: @escaping @Sendable (Error?) -> Void) {
127	0	self._shutdownGracefully(queue: queue, callback)
128	0	}
129
130	0	private func _shutdownGracefully(queue: DispatchQueue, _ callback: @escaping @Sendable (Error?) -> Void) {
131	0	guard self.canBeStopped else {
132	0	queue.async {
133	0	callback(NIOThreadPoolError.UnsupportedOperation())
134	0	}
135	0	return
136	0	}
137	0
138	0	let threadsToJoin = self._workAvailable.withLockSettingWorkAvailable {
139	0	switch self._state {
140	0	case .running(let items):
141	0	self._state = .modifying
142	0	queue.async {
143	0	for item in items {
144	0	item.workItem(.cancelled)
145	0	}
146	0	}
147	0	self._state = .shuttingDown(Array(repeating: true, count: self.numberOfThreads))
148	0
149	0	let threads = self.threads!
150	0	self.threads = nil
151	0
152	0	// Wake all threads so they can see .shuttingDown and exit.
153	0	return (workAvailable: self.numberOfThreads, signal: .broadcastAll, result: threads)
154	0
155	0	case .shuttingDown, .stopped:
156	0	// Don't modify workAvailable - threads may still be waking up to see .shuttingDown
157	0	return (workAvailable: nil, signal: .none, result: [])
158	0
159	0	case .modifying:
160	0	fatalError(".modifying state misuse")
161	0	}
162	0	}
163	0
164	0	DispatchQueue(label: "io.swiftnio.NIOThreadPool.shutdownGracefully").async {
165	0	for thread in threadsToJoin {
166	0	thread.join()
167	0	}
168	0	queue.async {
169	0	callback(nil)
170	0	}
171	0	}
172	0	}
173
174		/// Submit a `WorkItem` to process.
175		///
176		/// - Note: This is a low-level method, in most cases the `runIfActive` method should be used.
177		///
178		/// - Parameters:
179		/// - body: The `WorkItem` to process by the `NIOThreadPool`.
180		@preconcurrency
181	0	public func submit(_ body: @escaping WorkItem) {
182	0	self._submit(id: nil, body)
183	0	}
184
185	0	private func _submit(id: Int?, _ body: @escaping WorkItem) {
186	0	let submitted = self._workAvailable.withLock {
187	0	let submitted: Bool
188	0
189	0	switch self._state {
190	0	case .running(var items):
191	0	self._state = .modifying
192	0	items.append(.init(workItem: body, id: id))
193	0	self._state = .running(items)
194	0	submitted = true
195	0
196	0	case .shuttingDown, .stopped:
197	0	submitted = false
198	0
199	0	case .modifying:
200	0	fatalError(".modifying state misuse")
201	0	}
202	0
203	0	return (workDelta: submitted ? 1 : 0, signal: submitted ? .signalOne : .none, result: submitted)
204	0	}
205	0
206	0	// if item couldn't be added run it immediately indicating that it couldn't be run
207	0	if !submitted {
208	0	body(.cancelled)
209	0	}
210	0	}
211
212		/// Initialize a `NIOThreadPool` thread pool with `numberOfThreads` threads.
213		///
214		/// - Parameters:
215		/// - numberOfThreads: The number of threads to use for the thread pool.
216	0	public convenience init(numberOfThreads: Int) {
217	0	self.init(
218	0	numberOfThreads: numberOfThreads,
219	0	pinnedCPUIDs: nil,
220	0	canBeStopped: true
221	0	)
222	0	}
223
224		#if !os(WASI)
225		#if os(Linux) \|\| os(Android)
226		/// Initialize a `NIOThreadPool` thread pool with threads pinned to specific CPUs.
227		///
228		/// - Parameters:
229		/// - pinnedCPUIDs: The IDs of the CPUs to use for the thread pool. One thread will be
230		/// created per entry in this array.
231	0	public convenience init(pinnedCPUIDs: [Int]) {
232	0	self.init(
233	0	numberOfThreads: pinnedCPUIDs.count,
234	0	pinnedCPUIDs: pinnedCPUIDs,
235	0	canBeStopped: true
236	0	)
237	0	}
238
239		#endif
240		#endif
241		/// Create a ``NIOThreadPool`` that is already started, cannot be shut down and must not be `deinit`ed.
242		///
243		/// This is only useful for global singletons.
244	0	public static func _makePerpetualStartedPool(numberOfThreads: Int, threadNamePrefix: String) -> NIOThreadPool {
245	0	let pool = self.init(numberOfThreads: numberOfThreads, pinnedCPUIDs: nil, canBeStopped: false)
246	0	pool._start(threadNamePrefix: threadNamePrefix)
247	0	return pool
248	0	}
249
250	0	private init(numberOfThreads: Int, pinnedCPUIDs: [Int]?, canBeStopped: Bool) {
251	0	self.numberOfThreads = numberOfThreads
252	0	self.pinnedCPUIDs = pinnedCPUIDs
253	0	self.canBeStopped = canBeStopped
254	0	self._workAvailable = NIOThreadPoolWorkAvailable()
255	0	}
256
257		// Do not rename or remove this function.
258		//
259		// When doing on-/off-CPU analysis, for example with continuous profiling, it's
260		// important to recognise certain functions that are purely there to wait.
261		//
262		// This function is one of those and giving it a consistent name makes it much easier to remove from the profiles
263		// when only interested in on-CPU work.
264		@inlinable
265	0	internal func _blockingWaitForWork(identifier: Int) -> (item: WorkItem, state: WorkItemState)? {
266	0	self._workAvailable.withLockWaitingForWork {
267	0	let result: (WorkItem, WorkItemState)?
268	0
269	0	switch self._state {
270	0	case .running(var items):
271	0	self._state = .modifying
272	0	let itemAndID = items.removeFirst()
273	0
274	0	let state: WorkItemState
275	0	if let id = itemAndID.id, !self._cancelledWorkIDs.isEmpty {
276	0	state = self._cancelledWorkIDs.remove(id) == nil ? .active : .cancelled
277	0	} else {
278	0	state = .active
279	0	}
280	0
281	0	self._state = .running(items)
282	0
283	0	result = (itemAndID.workItem, state)
284	0
285	0	case .shuttingDown(var aliveStates):
286	0	self._state = .modifying
287	0	assert(aliveStates[identifier])
288	0	aliveStates[identifier] = false
289	0	self._state = .shuttingDown(aliveStates)
290	0
291	0	result = nil
292	0
293	0	case .stopped:
294	0	// Unreachable: 'stopped' is the initial state which is left when starting the
295	0	// thread pool, and before any thread calls this function.
296	0	fatalError("Invalid state")
297	0
298	0	case .modifying:
299	0	fatalError(".modifying state misuse")
300	0	}
301	0
302	0	return (workDelta: -1, signal: .none, result: result)
303	0	}
304	0	}
305
306	0	private func process(identifier: Int) {
307	0	repeat {
308	0	let itemAndState = self._blockingWaitForWork(identifier: identifier)
309	0
310	0	if let (item, state) = itemAndState {
311	0	// if there was a work item popped, run it
312	0	item(state)
313	0	} else {
314	0	break // Otherwise, we're done
315	0	}
316	0	} while true
317	0	}
318
319		/// Start the `NIOThreadPool` if not already started.
320	0	public func start() {
321	0	self._start(threadNamePrefix: "TP-#")
322	0	}
323
324	0	public func _start(threadNamePrefix: String) {
325	0	let alreadyRunning = self._workAvailable.withLock {
326	0	switch self._state {
327	0	case .running:
328	0	// Already running, this has no effect on whether there is more work for the
329	0	// threads to run.
330	0	return (workDelta: 0, signal: .none, result: true)
331	0
332	0	case .shuttingDown:
333	0	// This should never happen
334	0	fatalError("start() called while in shuttingDown")
335	0
336	0	case .stopped:
337	0	self._state = .running(Deque(minimumCapacity: 16))
338	0	assert(self.threads == nil)
339	0	self.threads = []
340	0	self.threads!.reserveCapacity(self.numberOfThreads)
341	0	return (workDelta: 0, signal: .none, result: false)
342	0
343	0	case .modifying:
344	0	fatalError(".modifying state misuse")
345	0	}
346	0	}
347	0
348	0	if alreadyRunning {
349	0	return
350	0	}
351	0
352	0	// We use this condition lock as a tricky kind of semaphore.
353	0	// This is done to sidestep the thread performance checker warning
354	0	// that would otherwise be emitted.
355	0	let readyThreads = ConditionLock(value: 0)
356	0	for id in 0..<self.numberOfThreads {
357	0	let pinnedCPUID = self.pinnedCPUIDs?[id]
358	0
359	0	// We should keep thread names under 16 characters because Linux doesn't allow more.
360	0	NIOThread.spawnAndRun(name: "\(threadNamePrefix)\(id)") { thread in
361		#if !os(WASI)
362		#if os(Linux) \|\| os(Android)
363	0	if let pinnedCPUID {
364	0	precondition(thread.isCurrentSlow)
365	0	NIOThread.currentAffinity = LinuxCPUSet(pinnedCPUID)
366	0	}
367		#endif
368		#endif
369	0	readyThreads.withLock {
370	0	let threadCount = self._workAvailable.withLock {
371	0	self.threads!.append(thread)
372	0	let threadCount = self.threads!.count
373	0	return (workDelta: 0, signal: .none, result: threadCount)
374	0	}
375	0
376	0	return (unlockWith: threadCount, result: ())
377	0	}
378	0
379	0	self.process(identifier: id)
380	0	return ()
381	0	}
382	0	}
383	0
384	0	readyThreads.lock(whenValue: self.numberOfThreads)
385	0	readyThreads.unlock()
386	0
387	0	func threadCount() -> Int {
388	0	self._workAvailable.withLock {
389	0	(workDelta: 0, signal: .none, result: self.threads?.count ?? -1)
390	0	}
391	0	}
392	0	assert(threadCount() == self.numberOfThreads)
393	0	}
394
395	0	deinit {
396	0	assert(
397	0	self.canBeStopped,
398	0	"Perpetual NIOThreadPool has been deinited, you must make sure that perpetual pools don't deinit"
399	0	)
400	0	switch self._state {
401	0	case .stopped, .shuttingDown:
402	0	()
403	0	default:
404	0	assertionFailure("wrong state \(self._state)")
405	0	}
406	0	}
407		}
408
409		extension NIOThreadPool: @unchecked Sendable {}
410
411		extension NIOThreadPool {
412
413		/// Runs the submitted closure if the thread pool is still active, otherwise fails the promise.
414		/// The closure will be run on the thread pool so can do blocking work.
415		///
416		/// - Parameters:
417		/// - eventLoop: The `EventLoop` the returned `EventLoopFuture` will fire on.
418		/// - body: The closure which performs some blocking work to be done on the thread pool.
419		/// - Returns: The `EventLoopFuture` of `promise` fulfilled with the result (or error) of the passed closure.
420		@preconcurrency
421		public func runIfActive<T: Sendable>(
422		eventLoop: EventLoop,
423		_ body: @escaping @Sendable () throws -> T
424	0	) -> EventLoopFuture<T> {
425	0	self._runIfActive(eventLoop: eventLoop, body)
426	0	}
427
428		private func _runIfActive<T: Sendable>(
429		eventLoop: EventLoop,
430		_ body: @escaping @Sendable () throws -> T
431	0	) -> EventLoopFuture<T> {
432	0	let promise = eventLoop.makePromise(of: T.self)
433	0	self.submit { shouldRun in
434	0	guard case shouldRun = NIOThreadPool.WorkItemState.active else {
435	0	promise.fail(NIOThreadPoolError.ThreadPoolInactive())
436	0	return
437	0	}
438	0	do {
439	0	try promise.succeed(body())
440	0	} catch {
441	0	promise.fail(error)
442	0	}
443	0	}
444	0	return promise.futureResult
445	0	}
446
447		/// Runs the submitted closure if the thread pool is still active, otherwise throw an error.
448		/// The closure will be run on the thread pool, such that we can do blocking work.
449		///
450		/// - Parameters:
451		/// - body: The closure which performs some blocking work to be done on the thread pool.
452		/// - Returns: Result of the passed closure.
453		@available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
454	0	public func runIfActive<T: Sendable>(_ body: @escaping @Sendable () throws -> T) async throws -> T {
455	0	let workID = self.nextWorkID.loadThenWrappingIncrement(ordering: .relaxed)
456	0
457	0	return try await withTaskCancellationHandler {
458	0	try await withCheckedThrowingContinuation { (cont: CheckedContinuation<T, Error>) in
459	0	self._submit(id: workID) { shouldRun in
460	0	switch shouldRun {
461	0	case .active:
462	0	let result = Result(catching: body)
463	0	cont.resume(with: result)
464	0	case .cancelled:
465	0	cont.resume(throwing: CancellationError())
466	0	}
467	0	}
468	0	}
469	0	} onCancel: {
470	0	self._workAvailable.withLock {
471	0	self._cancelledWorkIDs.insert(workID)
472	0	return (workDelta: 0, signal: .none, result: ())
473	0	}
474	0	}
475	0	}
476		}
477
478		extension NIOThreadPool {
479		@preconcurrency
480	0	public func shutdownGracefully(_ callback: @escaping @Sendable (Error?) -> Void) {
481	0	self.shutdownGracefully(queue: .global(), callback)
482	0	}
483
484		/// Shuts down the thread pool gracefully.
485		@available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
486		@inlinable
487	0	public func shutdownGracefully() async throws {
488	0	try await withCheckedThrowingContinuation { (cont: CheckedContinuation<Void, Error>) in
489	0	self.shutdownGracefully { error in
490	0	if let error = error {
491	0	cont.resume(throwing: error)
492	0	} else {
493	0	cont.resume()
494	0	}
495	0	}
496	0	}
497	0	}
498
499		@available(*, noasync, message: "this can end up blocking the calling thread", renamed: "shutdownGracefully()")
500	0	public func syncShutdownGracefully() throws {
501	0	try self._syncShutdownGracefully()
502	0	}
503
504	0	private func _syncShutdownGracefully() throws {
505	0	let errorStorageLock = NIOLockedValueBox<Swift.Error?>(nil)
506	0	let continuation = ConditionLock(value: 0)
507	0	self.shutdownGracefully { error in
508	0	if let error = error {
509	0	errorStorageLock.withLockedValue {
510	0	$0 = error
511	0	}
512	0	}
513	0	continuation.lock(whenValue: 0)
514	0	continuation.unlock(withValue: 1)
515	0	}
516	0	continuation.lock(whenValue: 1)
517	0	continuation.unlock()
518	0	try errorStorageLock.withLockedValue {
519	0	if let error = $0 {
520	0	throw error
521	0	}
522	0	}
523	0	}
524		}
525
526		extension ConditionLock {
527		@inlinable
528	0	func _lock(when value: T?) {
529	0	if let value = value {
530	0	self.lock(whenValue: value)
531	0	} else {
532	0	self.lock()
533	0	}
534	0	}
535
536		@inlinable
537	0	func _unlock(with value: T?) {
538	0	if let value = value {
539	0	self.unlock(withValue: value)
540	0	} else {
541	0	self.unlock()
542	0	}
543	0	}
544
545		@inlinable
546	0	func withLock<Result>(when value: T? = nil, _ body: () -> (unlockWith: T?, result: Result)) -> Result {
547	0	self._lock(when: value)
548	0	let (unlockValue, result) = body()
549	0	self._unlock(with: unlockValue)
550	0	return result
551	0	}
552		}
553		#endif // !os(WASI)