[SC] Manejador de timeout artesanal con TaskGroup

#swift #ios #concurrency

¿Cuál es la estructura del manejador de Timeout?

enum TaskTimeoutError: Swift.Error, Sendable {
  case timeout
  case invalidState
  case operationError(any Swift.Error & Sendable)
}

func withTaskTimeoutHandler<Success: Sendable>(timeout: Duration,
  operation: @escaping () async throws -> Success
) async throws -> Success {
  try await withThrowingTaskGroup(returning: Success.self) { group in
    _ = group.addTaskUnlessCancelled {
      do {
        return try await operation()
      } catch {
        throw TaskTimeoutError.operationError(error)
      }
    }
    _ = group.addTaskUnlessCancelled {
      try await Task.sleep(for: timeout, clock: .continuous)
      throw TaskTimeoutError.timeout
    }
    guard let result = try await group.next() else {
      throw TaskTimeoutError.invalidState
    }
    return result
  }
}

En este caso se crea una función global llamada withTaskTimeoutHandler(timeout:operation:) que recibe por parámetro la duración del timeout y la operación asíncrona a ejecutar, que también puede arrojar un error.

En su interior se crea un TaskGroup con withThrowingTaskGroup(of:returning:isolation:body:), que puede arrojar un error de tipo TaskTimeoutError. Dentro del TaskGroup se agregan dos tareas, en caso de que no se haya cancelado con addTaskUnlessCancelled(priority:operation:): la operación operation, envuelta con un do catch para poder envolver el error con un TaskTimeoutError.operationError; y Task.sleep que, al expirar, arroja un TaskTimeoutError.timeout.

Luego se espera a que la primera tarea responda con try await group.next() -next() retorna nil si no hay ninguna tarea pendiente en el TaskGroup, así que hay que arrojar un error en este caso porque sabemos a priori que tenemos dos tareas en el grupo.

Si primero termina la tarea del timeout, entonces se va a arrojar el error TaskTimeoutError.timeout. En caso contrario, habrá terminado la tarea operation y se retornará el valor generado por ella. Tener en cuenta que cuando la tarea del timeout arroja un error en el next(), esto automáticamente cancela todas las tareas.

Para probar la función global se usaron los siguientes casos de prueba:

struct TaskTimeoutHandlerTests {
  @Test
  func timeoutReached() async {
    do {
      _ = try await withTaskTimeoutHandler(timeout: .milliseconds(100)) {
        try? await Task.sleep(for: .milliseconds(110))
        return "Hola"
      }
      Issue.record("Se esperaba un error pero no se lanzó ninguno")
    } catch let error as TaskTimeoutError {
      #expect(error == .timeout)
    } catch {
      Issue.record("Error inesperado: \(error)")
    }
  }

  @Test
  func timeoutPassed() async throws {
    let result = try await withTaskTimeoutHandler(timeout: .milliseconds(110)) {
      try? await Task.sleep(for: .milliseconds(100))
      return "Hola"
    }
    #expect(result == "Hola")
  }
}

¿Cuál es la estructura del constructor que usa timeout?

extension Task where Failure == any Error {
  @discardableResult
  init<C>(
    name: String? = nil,
    priority: TaskPriority? = nil,
    timeout: C.Instant.Duration,
    clock: C = .continuous,
    operation: @Sendable @escaping @isolated(any) () async throws -> Success
  ) where C: Clock {
    self = Task(name: name, priority: priority, operation: {
      try await withThrowingTaskGroup { group in
        _ = group.addTaskUnlessCancelled {
          do {
            return try await operation()
          } catch {
            throw TaskTimeoutError.operationError(error)
          }
        }
        _ = group.addTaskUnlessCancelled { () -> Success in
          // OJO: No es obligatorio poner () -> Success in.
          try await Task<Never, Never>.sleep(for: timeout, clock: clock)
          throw TaskTimeoutError.timeout
        }
        guard let result = try await group.next() else {
          throw TaskTimeoutError.invalidState
        }
        return result
      }
    })
  }
}

El caso del constructor de Task modificado es muy parecido a la implementación de la función global. Las diferencias radican en la implementación de la tarea que hace Task.sleep:

Como Task.sleep es un método estático, cuando se llama dentro del contexto (el TaskGroup), el compilador lo resuelve como Task<Sucess, any Error>.sleep(...) -es decir, el Task del contexto exterior. Sin embargo, Task.sleep() requiere Failure == Never y Success == Never, mientras que el contexto exterior tiene valores distintos. Por esta razón, hay que marcar explícitamente que el tipo es Task<Never, Never>.sleep.
El compilador es capaz de inferir que la primera tarea responde Success gracias al tipo de dato de operation. Gracias a esto es capaz de inferir que el segundo grupo también debería retornar Success, incluso aunque no retorne nada. Sin embargo, haciendo pruebas manuales se detectó que algunas veces había un error de compilación. Por esta razón, en caso de ser necesario puede ser útil declarar el tipo del closure como () -> Success in.

¿Qué sucede con la tarea restante cuando una de las dos termina primero?

Se cancela:

Si termina primero operation, entonces el scope del Task termina y el Task.sleep recibe una cancelación automática.
Si termina primero Task.sleep entonces arroja un error, el try await group.next() lo detecta y al arrojarlo, cancela al resto de tareas del grupo.

¿Por qué `group.next()` devuelve opcional?

group.next() devuelve nil si no hay tareas pendientes en el grupo. Esto permite invocar group.next() para inspeccionar manualmente cada valor generado.

¿Qué patrones se pueden destacar en la implementación de la solución?

addTaskUnlessCancelled(name:priority:operation:) para agregar una tarea solo si no se ha cancelado el contexto.
group.next() para inspeccionar el resultado del primer valor que termina.