A Task that defers its start until it is awaited

LazyTask.swift

import Foundation

/// A Task that starts when its value is awaited for the first time.
final class LazyTask<Success, Failure>: @unchecked Sendable where Success: Sendable, Failure: Error {
    private let priority: TaskPriority
    private let operation: @Sendable () async throws(Failure) -> Success
    private let lock = NSLock()
    private var _protectedTask: Task<Success, Failure>?
    private var _task: Task<Success, Failure>? {
        get { lock.withLock { _protectedTask } }
        set { lock.withLock { _protectedTask = newValue } }
    }
    
    /// A task that starts when its value is awaited for the first time.
    /// - Parameters:
    ///   - priority: The priority of the task. Pass nil to use the priority from Task.currentPriority.
    ///   - operation: The operation to perform.
    init(
        priority: TaskPriority = .medium,
        operation: @Sendable @escaping @isolated(any) () async throws(Failure) -> Success
    ) {
        self.priority = priority
        self.operation = operation
    }
}

extension LazyTask where Failure == any Error {
    /// The result from a throwing task, after it completes.
    /// - Returns: The task’s result.
    var value: Success {
        get async throws {
            try await (_task ?? createTask()).value
        }
    }

    /// The result from a throwing task, after it completes.
    /// - Returns: The result or error from a throwing task, after it completes.
    var result: Result<Success, Failure> {
        get async {
            await (_task ?? createTask()).result
        }
    }

    @discardableResult
    private func createTask() -> Task<Success, Failure> {
        let task = Task<Success, Failure>(priority: priority) {
            try await self.operation()
        }
        _task = task
        return task
    }

    /// Cancels this task.
    func cancel() where Failure == any Error {
        // if the task hasnt been started yet, we need to start it before cancelling, otherwise it will await forever
        (_task ?? createTask()).cancel()
    }
}

extension LazyTask where Failure == Never {
    /// The result from a nonthrowing task, after it completes.
    /// - Returns: The task’s result.
    var value: Success {
        get async {
            await (_task ?? createTask()).value
        }
    }

    @discardableResult
    private func createTask() -> Task<Success, Failure> {
        let task = Task<Success, Failure>(priority: priority) {
            await self.operation()
        }
        _task = task
        return task
    }
    
    /// Cancels this task.
    func cancel() where Failure == Never {
        // if the task hasnt been started yet, we need to start it before cancelling, otherwise it will await forever
        (_task ?? createTask()).cancel()
    }
}

extension Task where Failure == any Error {

    /// A task that starts when its value is awaited for the first time.
    /// - Parameters:
    ///   - priority: The priority of the task. Pass nil to use the priority from Task.currentPriority.
    ///   - operation: The operation to perform.
    /// - Returns: A reference to the task.
    static func lazy(
        priority: TaskPriority = .medium,
        operation: @Sendable @escaping () async throws(Failure) -> Success
    ) -> LazyTask<Success, Failure> {
        LazyTask(priority: priority, operation: operation)
    }
}

extension Task where Failure == Never {
    /// A task that starts when its value is awaited for the first time.
    /// - Parameters:
    ///   - priority: The priority of the task. Pass nil to use the priority from Task.currentPriority.
    ///   - operation: The operation to perform.
    /// - Returns: A reference to the task.
    static func lazy(
        priority: TaskPriority = .medium,
        operation: @Sendable @escaping () async throws(Failure) -> Success
    ) -> LazyTask<Success, Failure> {
        LazyTask(priority: priority, operation: operation)
    }
}

extension LazyTask: Hashable {
    func hash(into hasher: inout Hasher) {
        hasher.combine(_task)
    }

    static func == (lhs: LazyTask<Success, Failure>, rhs: LazyTask<Success, Failure>) -> Bool {
        return lhs._task == rhs._task
    }
}