Timeout operation Swift Concurrency

Timeout operation Swift Concurrency

Muukii

-

Oct 6, 2023

Depends on what type of timeout you need but

I made one of the timeout approaches.

This approach raises cancel of the current task by timeout occurring and moving forward discarding awaiting the task.

Canceling tasks is not actually stopping current operations it’s just fact that telling them it’s canceled.

Below code calling cancel to task but “done” print will be done after longOperation has been finished except the operation has handling for cancel.

let task = Task {
  await longOperation()
  print("done")
}

...

task.cancel()

So let’s imagine the following case

await fetchFlag() // want this to make timeout in 5sec.

applyFlag()

This case needs a timeout for fetching the flag in 5sec in any case.

Whether raising canceling makes actually stop the current operation and pop out their step is up to the implementation.

To make timeout anyway at all costs of discarding handling them.

(So actually fetch request is still ongoing but moving forward anyway because it’s timed out.)

await withTimeout(5) {

  await fetchFlag() // want this to make timeout in 5sec.

}

applyFlag() // apply new flag or current flag if timeout occurred.

To achieve this behavior needs to harness unstructured concurrency API inside.

Source code

public enum TimeoutHandlerError: Error {

  case timeoutOccured

}

@_unsafeInheritExecutor

public func withTimeout<Return: Sendable>(

  nanoseconds: UInt64,

  @_inheritActorContext _ operation: @escaping @Sendable () async throws -> Return

) async throws -> Return {

  let task = Ref<Task<(), Never>>(value: nil)

  let timeoutTask = Ref<Task<(), any Error>>(value: nil)

  let flag = Flag()

  return try await withTaskCancellationHandler {

    try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Return, Error>) in

      do {

        try Task.checkCancellation()

      } catch {

        continuation.resume(throwing: error)

        return

      }

      let _task = Task {

        do {

          let taskResult = try await operation()

          await flag.performIf(expected: false) {

            continuation.resume(returning: taskResult)

            return true

          }

        } catch {

          await flag.performIf(expected: false) {

            continuation.resume(throwing: error)

            return true

          }

        }

      }

      task.value = _task

      let _timeoutTask = Task {

        try await Task.sleep(nanoseconds: nanoseconds)

        _task.cancel()

        await flag.performIf(expected: false) {

          continuation.resume(throwing: TimeoutHandlerError.timeoutOccured)

          return true

        }

      }

      timeoutTask.value = _timeoutTask

    }

  } onCancel: {

    task.value?.cancel()

    timeoutTask.value?.cancel()

  }

}

private final class Ref<T>: @unchecked Sendable {

  var value: T?

  init(value: T?) {

    self.value = value

  }

}

private actor Flag {

  var value: Bool = false

  func set(value: Bool) {

    self.value = value

  }

  func performIf(expected: Bool, perform: @Sendable () -> Bool) {

    if value == expected {

      value = perform()

    }

  }

}