A generic task storage and management system

TaskStore.swift

import struct Foundation.UUID
import Observation

extension Task where Failure == Never {
  var cancellableValue: Success {
    get async {
      await withTaskCancellationHandler {
        await self.value
      } onCancel: {
        self.cancel()
      }
    }
  }
}

extension Task where Failure == any Error {
  var cancellableValue: Success {
    get async throws {
      try await withTaskCancellationHandler {
        try await self.value
      } onCancel: {
        self.cancel()
      }
    }
  }
}

struct TaskStoreDuplicateKeyBehavior: Hashable, Sendable {
  fileprivate struct PreferNewOptions: Hashable, Sendable {
    var cancelPrevious: Bool
    var waitForPrevious: Bool
  }

  fileprivate var preferNewOptions: PreferNewOptions?

  var preferPrevious: Bool {
    preferNewOptions == nil
  }

  var cancelPrevious: Bool {
    preferNewOptions?.cancelPrevious ?? false
  }

  var waitForPrevious: Bool {
    preferNewOptions?.waitForPrevious ?? true
  }

  var runNewTask: Bool {
    preferNewOptions != nil
  }

  static var wait: Self {
    Self(preferNewOptions: .init(cancelPrevious: false, waitForPrevious: true))
  }

  static func cancelPrevious(wait: Bool) -> Self {
    Self(preferNewOptions: .init(cancelPrevious: true, waitForPrevious: wait))
  }

  static var runConcurrently: Self {
    Self(preferNewOptions: .init(cancelPrevious: false, waitForPrevious: false))
  }

  static var preferPrevious: Self { Self() }
}

@Observable
final class TaskStore<Key: Hashable> {
  private struct TaskData {
    var id: UUID
    var task: Task<Void, Never>
  }

  typealias DuplicateKeyBehavior = TaskStoreDuplicateKeyBehavior

  private var currentTasks: [Key: TaskData] = [:]

  @discardableResult
  func addTask(
    forKey key: Key,
    duplicateKeyBehavior: TaskStoreDuplicateKeyBehavior = .cancelPrevious(wait: false),
    priority: TaskPriority? = nil,
    isolation: isolated (any Actor)? = #isolation,
    operation: @escaping @Sendable () async -> Void
  ) -> Task<Void, Never> {
    let preferNewOptions = duplicateKeyBehavior.preferNewOptions
    let previousTask = self.currentTasks[key]?.task
    if let previousTask {
      guard let preferNewOptions else { return previousTask }
      if preferNewOptions.cancelPrevious {
        previousTask.cancel()
      }
    }
    let newTaskID = UUID()
    let newTask = Task(
      priority: priority,
      operation: {
        await withTaskCancellationHandler {
          if let previousTask, let preferNewOptions, preferNewOptions.waitForPrevious {
            await previousTask.value
          }
          await operation()
          self.taskFinished(key: key, id: newTaskID, isolation: isolation)
        } onCancel: {
          previousTask?.cancel()
        }
      }
    )
    self.currentTasks[key] = TaskData(id: newTaskID, task: newTask)
    return newTask
  }

  private func taskFinished(key: Key, id: UUID, isolation: (any Actor)?)  {
    if currentTasks[key]?.id == id {
      currentTasks.removeValue(forKey: key)
    }
  }

  func cancelTask(forKey key: Key) {
    currentTasks[key]?.task.cancel()
  }

  func taskIsRunning(forKey key: Key) -> Bool {
    currentTasks.keys.contains(key)
  }
}