Add EventLoop APIs for simpler scheduling of callbacks

Benchmarks/Benchmarks/NIOPosixBenchmarks/Benchmarks.swift

@@ -13,6 +13,7 @@
 //===----------------------------------------------------------------------===//
 
 import Benchmark
+import NIOCore
 import NIOPosix
 
 private let eventLoop = MultiThreadedEventLoopGroup.singleton.next()
@@ -64,4 +65,45 @@ let benchmarks = {
         )
     }
     #endif
+
+    Benchmark(
+        "MTELG.scheduleTask(in:_:)",
+        configuration: Benchmark.Configuration(
+            metrics: [.mallocCountTotal, .cpuTotal],
+            scalingFactor: .kilo
+        )
+    ) { benchmark in
+        let group = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        defer { try! group.syncShutdownGracefully() }
+        let loop = group.next()
+
+        benchmark.startMeasurement()
+        for _ in benchmark.scaledIterations {
+            loop.scheduleTask(in: .hours(1), {})
+        }
+        benchmark.stopMeasurement()
+    }
+
+    Benchmark(
+        "MTELG.scheduleCallback(in:_:)",
+        configuration: Benchmark.Configuration(
+            metrics: [.mallocCountTotal, .cpuTotal],
+            scalingFactor: .kilo
+        )
+    ) { benchmark in
+        let group = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        defer { try! group.syncShutdownGracefully() }
+        let loop = group.next()
+
+        final class Timer: NIOScheduledCallbackHandler {
+            func onSchedule(eventLoop: some EventLoop) {}
+        }
+        let timer = Timer()
+
+        benchmark.startMeasurement()
+        for _ in benchmark.scaledIterations {
+            let handle = loop.scheduleCallback(in: .hours(1), handler: timer)
+        }
+        benchmark.stopMeasurement()
+    }
 }

Sources/NIOCore/EventLoop.swift

@@ -356,6 +356,29 @@ public protocol EventLoop: EventLoopGroup {
     /// It is valid for an `EventLoop` not to implement any of the two `_promise` functions. If either of them are implemented,
     /// however, both of them should be implemented.
     func _promiseCompleted(futureIdentifier: _NIOEventLoopFutureIdentifier) -> (file: StaticString, line: UInt)?
+
+    /// Schedule a callback at a given time.
+    ///
+    /// - NOTE: Event loops that provide a custom scheduled callback implementation **must** implement _both_
+    ///         `sheduleCallback(at deadline:handler:)` _and_ `cancelScheduledCallback(_:)`. Failure to do so will
+    ///         result in a runtime error.
+    @discardableResult
+    func scheduleCallback(at deadline: NIODeadline, handler: some NIOScheduledCallbackHandler) -> NIOScheduledCallback
+
+    /// Schedule a callback after given time.
+    ///
+    /// - NOTE: Event loops that provide a custom scheduled callback implementation **must** implement _both_
+    ///         `sheduleCallback(at deadline:handler:)` _and_ `cancelScheduledCallback(_:)`. Failure to do so will
+    ///         result in a runtime error.
+    @discardableResult
+    func scheduleCallback(in amount: TimeAmount, handler: some NIOScheduledCallbackHandler) -> NIOScheduledCallback
+
+    /// Cancel a scheduled callback.
+    ///
+    /// - NOTE: Event loops that provide a custom scheduled callback implementation **must** implement _both_
+    ///         `sheduleCallback(at deadline:handler:)` _and_ `cancelScheduledCallback(_:)`. Failure to do so will
+    ///         result in a runtime error.
+    func cancelScheduledCallback(_ scheduledCallback: NIOScheduledCallback)
 }
 
 extension EventLoop {

Sources/NIOCore/NIOScheduledCallback.swift

@@ -0,0 +1,107 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the SwiftNIO open source project
+//
+// Copyright (c) 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
+//
+//===----------------------------------------------------------------------===//
+
+/// A type that handles callbacks scheduled with `EventLoop.scheduleCallback(at:handler:)`.
+///
+/// - Seealso: `EventLoop.scheduleCallback(at:handler:)`.
+public protocol NIOScheduledCallbackHandler {
+    /// This function is called at the scheduled time, unless the scheduled callback is cancelled.
+    ///
+    /// - Parameter eventLoop: The event loop on which the callback was scheduled.
+    func onSchedule(eventLoop: some EventLoop)
+}
+
+/// An opaque handle that can be used to cancel a scheduled callback.
+///
+/// Users should not create an instance of this type; it is returned by `EventLoop.scheduleCallback(at:handler:)`.
+///
+/// - Seealso: `EventLoop.scheduleCallback(at:handler:)`.
+public struct NIOScheduledCallback: Sendable {
+    @usableFromInline
+    enum Backing: Sendable {
+        /// A task created using `EventLoop.scheduleTask(deadline:_:)` by the default implementation.
+        case `default`(_ task: Scheduled<Void>)
+        /// A custom callback identifier, used by event loops that provide a custom implementation.
+        case custom(id: UInt64)
+    }
+
+    @usableFromInline
+    var eventLoop: any EventLoop
+
+    @usableFromInline
+    var backing: Backing
+
+    /// This initializer is only for the default implementation and is fileprivate to avoid use in EL implementations.
+    fileprivate init(_ eventLoop: any EventLoop, _ task: Scheduled<Void>) {
+        self.eventLoop = eventLoop
+        self.backing = .default(task)
+    }
+
+    /// Create a handle for the scheduled callback with an opaque identifier managed by the event loop.
+    ///
+    /// - NOTE: This initializer is for event loop implementors only, end users should use `EventLoop.scheduleCallback`.
+    ///
+    /// - Seealso: `EventLoop.scheduleCallback(at:handler:)`.
+    @inlinable
+    public init(_ eventLoop: any EventLoop, id: UInt64) {
+        self.eventLoop = eventLoop
+        self.backing = .custom(id: id)
+    }
+
+    /// Cancel the scheduled callback associated with this handle.
+    @inlinable
+    public func cancel() {
+        self.eventLoop.cancelScheduledCallback(self)
+    }
+
+    /// The callback identifier, if the event loop uses a custom scheduled callback implementation; nil otherwise.
+    ///
+    /// - NOTE: This property is for event loop implementors only.
+    @inlinable
+    public var customCallbackID: UInt64? {
+        guard case .custom(let id) = backing else { return nil }
+        return id
+    }
+}
+
+extension EventLoop {
+    /// Default implementation of `scheduleCallback(at deadline:handler:)`: backed by `EventLoop.scheduleTask`.
+    @discardableResult
+    public func scheduleCallback(at deadline: NIODeadline, handler: some NIOScheduledCallbackHandler) -> NIOScheduledCallback {
+        let task = self.scheduleTask(deadline: deadline) { handler.onSchedule(eventLoop: self) }
+        return NIOScheduledCallback(self, task)
+    }
+
+    /// Default implementation of `scheduleCallback(in amount:handler:)`: calls `scheduleCallback(at deadline:handler:)`.
+    @discardableResult
+    @inlinable
+    public func scheduleCallback(in amount: TimeAmount, handler: some NIOScheduledCallbackHandler) -> NIOScheduledCallback {
+        self.scheduleCallback(at: .now() + amount, handler: handler)
+    }
+
+    /// Default implementation of `cancelScheduledCallback(_:)`: only cancels callbacks scheduled by the default implementation of `scheduleCallback`.
+    ///
+    /// - NOTE: Event loops that provide a custom scheduled callback implementation **must** implement _both_
+    ///         `sheduleCallback(at deadline:handler:)` _and_ `cancelScheduledCallback(_:)`. Failure to do so will
+    ///         result in a runtime error.
+    @inlinable
+    public func cancelScheduledCallback(_ scheduledCallback: NIOScheduledCallback) {
+        switch scheduledCallback.backing {
+        case .default(let task):
+            task.cancel()
+        case .custom:
+            preconditionFailure("EventLoop missing custom implementation of cancelTimer(_:)")
+        }
+    }
+}

Sources/NIOEmbedded/AsyncTestingEventLoop.swift

@@ -171,6 +171,15 @@ public final class NIOAsyncTestingEventLoop: EventLoop, @unchecked Sendable {
         return self.scheduleTask(deadline: self.now + `in`, task)
     }
 
+    @discardableResult
+    public func scheduleCallback(in amount: TimeAmount, handler: some NIOScheduledCallbackHandler) -> NIOScheduledCallback {
+        // TODO: docs
+        /// Even though this type does not conform to `CustomTimerImplemenation`, it has a manual clock so we cannot
+        /// rely on the default implemntation of `setTimer(for duration:_:)`, which computes the deadline for
+        /// `setTimer(for deadline:_:)` naively using `NIODeadline.now`, but we must use `self.now`.
+        self.scheduleCallback(at: self.now + amount, handler: handler)
+    }
+
     /// On an `NIOAsyncTestingEventLoop`, `execute` will simply use `scheduleTask` with a deadline of _now_. Unlike with the other operations, this will
     /// immediately execute, to eliminate a common class of bugs.
     public func execute(_ task: @escaping () -> Void) {

Sources/NIOEmbedded/Embedded.swift

@@ -127,6 +127,15 @@ public final class EmbeddedEventLoop: EventLoop {
         return scheduleTask(deadline: self._now + `in`, task)
     }
 
+    @discardableResult
+    public func scheduleCallback(in amount: TimeAmount, handler: some NIOScheduledCallbackHandler) -> NIOScheduledCallback {
+        // TODO: docs
+        /// Even though this type does not conform to `CustomTimerImplemenation`, it has a manual clock so we cannot
+        /// rely on the default implemntation of `setTimer(for duration:_:)`, which computes the deadline for
+        /// `setTimer(for deadline:_:)` naively using `NIODeadline.now`, but we must use `self._now`.
+        self.scheduleCallback(at: self._now + amount, handler: handler)
+    }
+
     /// On an `EmbeddedEventLoop`, `execute` will simply use `scheduleTask` with a deadline of _now_. This means that
     /// `task` will be run the next time you call `EmbeddedEventLoop.run`.
     public func execute(_ task: @escaping () -> Void) {

Sources/NIOPosix/MultiThreadedEventLoopGroup.swift

@@ -460,21 +460,32 @@ internal struct ScheduledTask {
     ///     This means, the ids need to be unique for a given ``SelectableEventLoop`` and they need to be in ascending order.
     @usableFromInline
     let id: UInt64
-    let task: () -> Void
-    private let failFn: (Error) -> Void
+
     @usableFromInline
     internal let readyTime: NIODeadline
 
+    @usableFromInline
+    enum Kind {
+        case task(task: () -> Void, failFn: (Error) -> Void)
+        case callback(any NIOScheduledCallbackHandler)
+    }
+
+    @usableFromInline
+    let kind: Kind
+
+    // TODO: Should these be .init() or should they be static functions?
     @usableFromInline
     init(id: UInt64, _ task: @escaping () -> Void, _ failFn: @escaping (Error) -> Void, _ time: NIODeadline) {
         self.id = id
-        self.task = task
-        self.failFn = failFn
         self.readyTime = time
+        self.kind = .task(task: task, failFn: failFn)
     }
 
-    func fail(_ error: Error) {
-        failFn(error)
+    @usableFromInline
+    init(id: UInt64, _ handler: any NIOScheduledCallbackHandler, _ time: NIODeadline) {
+        self.id = id
+        self.readyTime = time
+        self.kind = .callback(handler)
     }
 }
 

Sources/NIOPosix/SelectableEventLoop.swift

@@ -474,6 +474,8 @@ Further information:
                     fatalError("Tried to run an UnownedJob without runtime support")
                 }
             #endif
+            case .callback(let handler):
+                handler.onSchedule(eventLoop: self)
             }
         }
     }
@@ -565,7 +567,10 @@ Further information:
             if moreScheduledTasksToConsider, tasksCopy.count < tasksCopyBatchSize, let task = scheduledTasks.peek() {
                 if task.readyTime.readyIn(now) <= .nanoseconds(0) {
                     scheduledTasks.pop()
-                    tasksCopy.append(.function(task.task))
+                    switch task.kind {
+                    case .task(let task, _): tasksCopy.append(.function(task))
+                    case .callback(let handler): tasksCopy.append(.callback(handler))
+                    }
                 } else {
                     nextScheduledTaskDeadline = task.readyTime
                     moreScheduledTasksToConsider = false
@@ -658,9 +663,14 @@ Further information:
                 for task in immediateTasksCopy {
                     self.run(task)
                 }
-                // Fail all the scheduled tasks.
+                // Fail all the scheduled tasks (callbacks have no failFn and can just be dropped).
                 for task in scheduledTasksCopy {
-                    task.fail(EventLoopError.shutdown)
+                    switch task.kind {
+                    case .task(_, let failFn):
+                        failFn(EventLoopError.shutdown)
+                    case .callback:
+                        break
+                    }
                 }
 
                 iterations += 1
@@ -869,6 +879,7 @@ enum UnderlyingTask {
     #if compiler(>=5.9)
     case unownedJob(ErasedUnownedJob)
     #endif
+    case callback(any NIOScheduledCallbackHandler)
 }
 
 @usableFromInline
@@ -883,3 +894,23 @@ internal func assertExpression(_ body: () -> Bool) {
         return body()
     }())
 }
+
+extension SelectableEventLoop {
+    @inlinable
+    func scheduleCallback(at deadline: NIODeadline, handler: some NIOScheduledCallbackHandler) -> NIOScheduledCallback {
+        let taskId = self.scheduledTaskCounter.loadThenWrappingIncrement(ordering: .relaxed)
+        let task = ScheduledTask(id: taskId, handler, deadline)
+        try! self._schedule0(.scheduled(task))
+        return NIOScheduledCallback(self, id: taskId)
+    }
+
+    @inlinable
+    func cancelScheduledCallback(_ scheduledCallback: NIOScheduledCallback) {
+        guard let id = scheduledCallback.customCallbackID else {
+            preconditionFailure("No custom ID for callback")
+        }
+        self._tasksLock.withLock {
+            self._scheduledTasks.removeFirst(where: { $0.id == id })
+        }
+    }
+}