153 lines · 3.9 KB · Swift
RawDownload
import CodexBarCoreimport Foundation final class ProviderRefreshTaskState: @unchecked Sendable {    let generation: UInt64     private let lock = NSLock()    private var task: Task<Void, Never>?    private var waiterIDs: Set<UInt64> = []    private var completed = false    private var retryRequired = false     init(generation: UInt64) {        self.generation = generation    }     func install(task: Task<Void, Never>) {        self.lock.withLock {            self.task = task        }    }     func addWaiter(_ waiterID: UInt64) -> Task<Void, Never>? {        self.lock.withLock {            self.waiterIDs.insert(waiterID)            return self.task        }    }     func cancelWaiter(_ waiterID: UInt64) {        let taskToCancel = self.lock.withLock {            guard self.waiterIDs.remove(waiterID) != nil else { return nil as Task<Void, Never>? }            return self.waiterIDs.isEmpty && !self.completed ? self.task : nil        }        taskToCancel?.cancel()    }     func finishWaiter(_ waiterID: UInt64) {        _ = self.lock.withLock {            self.waiterIDs.remove(waiterID)        }    }     func markCompleted(retryRequired: Bool) {        self.lock.withLock {            self.completed = true            self.retryRequired = retryRequired        }    }     func cancelTask() {        let task = self.lock.withLock {            self.completed ? nil : self.task        }        task?.cancel()    }     func waitForTaskCompletion() async {        let task = self.lock.withLock { self.task }        await task?.value    }     var isCompleted: Bool {        self.lock.withLock { self.completed }    }     var shouldRetry: Bool {        self.lock.withLock { self.retryRequired }    }     var canRemove: Bool {        self.lock.withLock { self.completed && self.waiterIDs.isEmpty }    }} enum ProviderStatusIndicator: String {    case none    case minor    case major    case critical    case maintenance    case unknown     var hasIssue: Bool {        switch self {        case .none: false        default: true        }    }     var label: String {        switch self {        case .none: L("status_operational")        case .minor: L("status_partial_outage")        case .major: L("status_major_outage")        case .critical: L("status_critical_issue")        case .maintenance: L("status_maintenance")        case .unknown: L("status_unknown")        }    }} struct ProviderStatus {    let indicator: ProviderStatusIndicator    let description: String?    let updatedAt: Date?} /// Tracks consecutive failures so we can ignore a single flake when we previously had fresh data.struct ConsecutiveFailureGate {    private(set) var streak: Int = 0     mutating func recordSuccess() {        self.streak = 0    }     mutating func reset() {        self.streak = 0    }     /// Returns true when the caller should surface the error to the UI.    mutating func shouldSurfaceError(onFailureWithPriorData hadPriorData: Bool) -> Bool {        self.streak += 1        if hadPriorData, self.streak == 1 { return false }        return true    }} #if DEBUGextension UsageStore {    func _setSnapshotForTesting(_ snapshot: UsageSnapshot?, provider: UsageProvider) {        self.snapshots[provider] = snapshot?.scoped(to: provider)    }     func _setTokenSnapshotForTesting(_ snapshot: CostUsageTokenSnapshot?, provider: UsageProvider) {        self.tokenSnapshots[provider] = snapshot    }     func _setTokenErrorForTesting(_ error: String?, provider: UsageProvider) {        self.tokenErrors[provider] = error    }     func _setErrorForTesting(_ error: String?, provider: UsageProvider) {        self.errors[provider] = error    }     func _setCodexHistoricalDatasetForTesting(_ dataset: CodexHistoricalDataset?, accountKey: String? = nil) {        self.codexHistoricalDataset = dataset        self.codexHistoricalDatasetAccountKey = accountKey        self.historicalPaceRevision += 1    }}#endif