apple · simonjbeaumont · Jun 28, 2024 · Jun 28, 2024 · Jun 28, 2024 · Jun 28, 2024
diff --git a/Benchmarks/Benchmarks/NIOPosixBenchmarks/Benchmarks.swift b/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.setTimer(for:_:)",
  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: NIOTimerHandler {
  func timerFired(loop: any EventLoop) {}
  }
  let timer = Timer()
 
  benchmark.startMeasurement()
  for _ in benchmark.scaledIterations {
  let handle = loop.setTimer(for: .hours(1), timer)
  }
  benchmark.stopMeasurement()
  }
 }
diff --git a/Sources/NIOCore/EventLoop.swift b/Sources/NIOCore/EventLoop.swift
@@ -356,6 356,14 @@ 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)?
 
  /// Set a timer that will call a handler at the given time.
  @discardableResult
  func setTimer(for deadline: NIODeadline, _ handler: any NIOTimerHandler) -> NIOTimer
 
  /// Set a timer that will call a handler after a given amount of time.
  @discardableResult
  func setTimer(for duration: TimeAmount, _ handler: any NIOTimerHandler) -> NIOTimer
 }
 
 extension EventLoop {

diff --git a/Sources/NIOCore/NIOTimer.swift b/Sources/NIOCore/NIOTimer.swift
@@ -0,0 1,94 @@
 //===----------------------------------------------------------------------===//
 //
 // 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 timer callbacks scheduled with ``EventLoop/setTimer(for:_:)-5e37g``.
 ///
 /// - Seealso: ``EventLoop/setTimer(for:_:)-5e37g``.
 public protocol NIOTimerHandler {
  func timerFired(loop: any EventLoop)
 }
 
 /// An opaque handle that can be used to cancel a timer.
 ///
 /// Users cannot create an instance of this type; it is returned by ``EventLoop/setTimer(for:_:)-5e37g``.
 ///
 /// - Seealso: ``EventLoop/setTimer(for:_:)-5e37g``.
 public struct NIOTimer {
  @usableFromInline
  enum Backing {
  /// A task created using `EventLoop.scheduleTask(deadline:_:)`, used by default for `EventLoop` implementations.
  case scheduledTask(Scheduled<Void>)
  /// An identifier for a timer, used by `EventLoop` implementations that conform to `CustomTimerImplementation`.
  case custom(eventLoop: any NIOCustomTimerImplementation, id: UInt64)
  }
 
  @usableFromInline
  var backing: Backing
 
  fileprivate init(_ scheduled: Scheduled<Void>) {
  self.backing = .scheduledTask(scheduled)
  }
 
  @inlinable
  init(_ eventLoop: any NIOCustomTimerImplementation, id: UInt64) {
  self.backing = .custom(eventLoop: eventLoop, id: id)
  }
 
  /// Cancel the timer associated with this handle.
  @inlinable
  public func cancel() {
  switch self.backing {
  case .scheduledTask(let scheduled):
  scheduled.cancel()
  case .custom(let eventLoop, let id):
  eventLoop.cancelTimer(id)
  }
  }
 }
 
 /// Default implementation of `setSimpleTimer(for deadline:_:)`, backed by `EventLoop.scheduleTask`.
 extension EventLoop {
  @discardableResult
  public func setTimer(for deadline: NIODeadline, _ handler: any NIOTimerHandler) -> NIOTimer {
  NIOTimer(self.scheduleTask(deadline: deadline) { handler.timerFired(loop: self) })
  }
 }
 
 /// Default implementation of `setSimpleTimer(for duration:_:)`, delegating to `setSimpleTimer(for deadline:_:)`.
 extension EventLoop {
  @discardableResult
  @inlinable
  public func setTimer(for duration: TimeAmount, _ handler: any NIOTimerHandler) -> NIOTimer {
  self.setTimer(for: .now() duration, handler)
  }
 }
 
 /// Extension point for `EventLoop` implementations implement a custom timer.
 public protocol NIOCustomTimerImplementation: EventLoop {
  /// Set a timer that calls handler at the given time.
  ///
  /// Implementations must return an integer identifier that uniquely identifies the timer.
  func setTimer(for deadline: NIODeadline, _ handler: any NIOTimerHandler) -> UInt64
 
  /// Cancel a timer with a given timer identifier.
  func cancelTimer(_ id: UInt64)
 }
 
 /// Default implementation of `setSimpleTimer(for deadline:_:)` for `EventLoop` types that opted in to `CustomeTimerImplementation`.
 extension EventLoop where Self: NIOCustomTimerImplementation {
  @inlinable
  public func setTimer(for deadline: NIODeadline, _ handler: any NIOTimerHandler) -> NIOTimer {
  NIOTimer(self, id: self.setTimer(for: deadline, handler))
  }
 }
diff --git a/Sources/NIOEmbedded/AsyncTestingEventLoop.swift b/Sources/NIOEmbedded/AsyncTestingEventLoop.swift
@@ -171,6 171,14 @@ public final class NIOAsyncTestingEventLoop: EventLoop, @unchecked Sendable {
  return self.scheduleTask(deadline: self.now `in`, task)
  }
 
  @discardableResult
  public func setTimer(for duration: TimeAmount, _ handler: any NIOTimerHandler) -> NIOTimer {
  /// 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.setTimer(for: self.now duration, 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) {

diff --git a/Sources/NIOEmbedded/Embedded.swift b/Sources/NIOEmbedded/Embedded.swift
@@ -127,6 127,14 @@ public final class EmbeddedEventLoop: EventLoop {
  return scheduleTask(deadline: self._now `in`, task)
  }
 
  @discardableResult
  public func setTimer(for duration: TimeAmount, _ handler: any NIOTimerHandler) -> NIOTimer {
  /// 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.setTimer(for: self._now duration, 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) {

diff --git a/Sources/NIOPosix/MultiThreadedEventLoopGroup.swift b/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 timer(any NIOTimerHandler)
  }
 
  @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 NIOTimerHandler, _ time: NIODeadline) {
  self.id = id
  self.readyTime = time
  self.kind = .timer(handler)
  }
 }
 

diff --git a/Sources/NIOPosix/SelectableEventLoop.swift b/Sources/NIOPosix/SelectableEventLoop.swift
@@ -474,6 474,8 @@ Further information:
  fatalError("Tried to run an UnownedJob without runtime support")
  }
  #endif
  case .timer(let handler):
  handler.timerFired(loop: 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 .timer(let handler): tasksCopy.append(.timer(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 (timers 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 .timer:
  break
  }
  }
 
  iterations = 1
@@ -869,6 879,7 @@ enum UnderlyingTask {
  #if compiler(>=5.9)
  case unownedJob(ErasedUnownedJob)
  #endif
  case timer(any NIOTimerHandler)
 }
 
 @usableFromInline
@@ -883,3 894,20 @@ internal func assertExpression(_ body: () -> Bool) {
  return body()
  }())
 }
 
 extension SelectableEventLoop: NIOCustomTimerImplementation {
  @inlinable
  internal func setTimer(for deadline: NIODeadline, _ handler: any NIOTimerHandler) -> UInt64 {
  let taskId = self.scheduledTaskCounter.loadThenWrappingIncrement(ordering: .relaxed)
  let task = ScheduledTask(id: taskId, handler, deadline)
  try! self._schedule0(.scheduled(task))
  return taskId
  }
 
  @inlinable
  internal func cancelTimer(_ id: UInt64) {
  self._tasksLock.withLock {
  self._scheduledTasks.removeFirst(where: { $0.id == id })
  }
  }
 }