A concurrent and thread safe resource cache that spawns an unstructured Task to fetch a resource when the resource is first requested. All Tasks requesting the resource will wait on the completion of the spawned unstructured resource fetching Task. If all resource requesting Tasks are canceled, the resource cache will also cancel the spawned uns…

TaskSpawningCancelingResourceCache.swift

//
//  TaskSpawningCancelingResourceCache.swift
//

import Foundation
import Atomics

/// A thread/concurrency context safe value using [Swift Atomics](https://github.com/apple/swift-atomics).
/// Heavily inspired by this WWDC 2023 video: [Beyond the basics of structured concurrency](https://developer.apple.com/videos/play/wwdc2023/10170/).
final class AtomicCount: Sendable {
    private let protectedValue = ManagedAtomic<Int>(0)

    var value: Int {
        protectedValue.load(ordering: .acquiring)
    }

    func increment() {
        protectedValue.wrappingIncrement(ordering: .relaxed)
    }

    func decrement() {
        protectedValue.wrappingDecrement(ordering: .relaxed)
    }

    func reset() {
        protectedValue.store(0, ordering: .relaxed)
    }
}

/// A thread/concurrency context safe value using [NSLock](https://developer.apple.com/documentation/foundation/nslock).
/// Heavily inspired by this WWDC 2023 video: [Beyond the basics of structured concurrency](https://developer.apple.com/videos/play/wwdc2023/10170/).
final class LockedCount: @unchecked Sendable {
    private let lock = NSLock()
    private var protectedValue = 0

    var value: Int {
        lock.withLock { protectedValue }
    }

    func increment() {
        lock.withLock { protectedValue += 1 }
    }

    func decrement() {
        lock.withLock { protectedValue -= 1 }
    }

    func reset() {
        lock.withLock { protectedValue = 0 }
    }
}

/// A concurrent and thread safe resource cache that spawns an unstructured Task to fetch a resource when the resource is first requested.
/// All Tasks requesting the resource will wait on the completion of the spawned unstructured resource fetching Task.
/// If all resource requesting Tasks are canceled, the resource cache will also cancel the spawned unstructured resource fetching Task.
final actor TaskSpawningCancelingResourceCache {
    private let urlSession: URLSession

    private var cachedResource: Data?
    private var resourceFetchingTask: Task<Data?, Never>?

    private let waitingTasksCount = AtomicCount() // Alternative 1
    // private let waitingTasksCount = LockedCount() // Alternative 2

    init(urlSession: URLSession = .shared) {
        self.urlSession = urlSession
    }

    // MARK: - Private

    private func fetchResource() async -> Data? {
        let url = URL(string: "https://raw.githubusercontent.com/PokeAPI/sprites/master/sprites/pokemon/132.png")!
        let request = URLRequest(url: url)

        do {
            let (data, _) = try await urlSession.data(for: request)

            return data
        } catch {
            // Ignoring proper Error handling for simplicity.
            print("Failed to fetch resource:", error)
            return nil
        }
    }

    // MARK: - Public

    var resource: Data? {
        get async {
            if let cachedResource {
                return cachedResource
            }

            if let resourceFetchingTask {
                waitingTasksCount.increment()

                return await withTaskCancellationHandler {
                    await resourceFetchingTask.value
                } onCancel: {
                    waitingTasksCount.decrement()

                    if waitingTasksCount.value > 0 {
                        return
                    }

                    // Cancel the resource fetching Task,
                    // since there are no longer any other Tasks waiting for the result.
                    resourceFetchingTask.cancel()
                }
            } else {
                waitingTasksCount.reset()

                let task = Task<Data?, Never> { [weak self] in
                    guard let self else {
                        return nil
                    }

                    // One could imagine this Task having to fetch and combine multiple resources,
                    // then it might be desirable to check for cancelation before and after each chunk of async work.

                    if Task.isCancelled {
                        return nil
                    }

                    let resource = await self.fetchResource()

                    if Task.isCancelled {
                        return nil
                    }

                    return resource
                }

                resourceFetchingTask = task

                waitingTasksCount.increment()

                return await withTaskCancellationHandler {
                    let resource = await task.value

                    cachedResource = resource
                    resourceFetchingTask = nil

                    return resource
                } onCancel: {
                    waitingTasksCount.decrement()

                    if waitingTasksCount.value > 0 {
                        return
                    }

                    // Cancel the resource fetching Task,
                    // since there are no longer any other Tasks waiting for the result.
                    task.cancel()
                }
            }
        }
    }

    var freshResource: Data? {
        get async {
            cachedResource = nil

            return await resource
        }
    }
}