diff --git a/Sources/RequestDL/Image/Core/RDLImageLoader.swift b/Sources/RequestDL/Image/Core/RDLImageLoader.swift index f6c528d5..92035a2e 100644 --- a/Sources/RequestDL/Image/Core/RDLImageLoader.swift +++ b/Sources/RequestDL/Image/Core/RDLImageLoader.swift @@ -20,9 +20,10 @@ import struct Foundation.Data /// starting a second request. This is separate from, and on top of, RequestDL's own response /// cache, which is what makes a *later*, non-concurrent request for the same image cheap. /// -/// The shared download is cancelled once every call waiting for it has been cancelled. While at -/// least one is still waiting it keeps running, so cancelling one call never takes the image -/// away from another. +/// A call that is cancelled stops waiting at once, with a `CancellationError`, whether it was +/// alone or one of several. The shared download is cancelled once every call waiting for it has +/// been cancelled. While at least one is still waiting it keeps running, so cancelling one call +/// never takes the image away from another. public actor RDLImageLoader { /// The shared loader instance, used by RequestDL's SwiftUI, UIKit, AppKit and WatchKit @@ -48,14 +49,20 @@ public actor RDLImageLoader { private var nextGeneration = 0 - /// A download in flight and how many calls are waiting for it. + private var nextWaiter = 0 + + /// A download in flight and the calls waiting for it, each with its own continuation. + /// + /// A call has its own continuation, rather than awaiting the download's task, so that it can + /// stop waiting on its own: awaiting a `Task` is not interrupted by the cancellation of the + /// caller, and the download has to go on for the others. /// /// `generation` tells one download for an `id` from the next one, so a late cancellation of /// a call that waited for an earlier download never touches a newer one. private struct InFlight { - let task: Task + let task: Task let generation: Int - var waiters: Int + var waiters: [Int: CheckedContinuation] } // MARK: - Inits @@ -107,33 +114,38 @@ public actor RDLImageLoader { id: String, task: Content ) async throws -> PlatformImage { - if var existing = tasks[id] { - existing.waiters += 1 - tasks[id] = existing - - return try await waitFor(existing.task, id: id, generation: existing.generation) + if let existing = tasks[id] { + return try await waitFor(id: id, generation: existing.generation).image } + let generation = nextGeneration + nextGeneration += 1 + // `.detached` rather than a plain `Task { ... }`: a plain `Task` would inherit this actor's // executor and run the CPU-bound `PlatformImage(data:)` decode on it. The decode touches no // actor state, and running it there would serialize every decode behind one another and // block unrelated `load()` calls from reaching their `tasks[id]` lookup. let newTask = Task.detached { - let result = try await task.result() + let result: Result - guard let image = PlatformImage(data: result.payload) else { - throw RDLImageDecodingError() + do { + let response = try await task.result() + + guard let image = PlatformImage(data: response.payload) else { + throw RDLImageDecodingError() + } + + result = .success(SendableImage(image)) + } catch { + result = .failure(error) } - return SendableImage(image) + await self.complete(id: id, generation: generation, with: result) } - let generation = nextGeneration - nextGeneration += 1 - - tasks[id] = InFlight(task: newTask, generation: generation, waiters: 1) + tasks[id] = InFlight(task: newTask, generation: generation, waiters: [:]) - return try await waitFor(newTask, id: id, generation: generation) + return try await waitFor(id: id, generation: generation).image } /// @@ -181,34 +193,46 @@ public actor RDLImageLoader { // MARK: - Private methods - /// Waits for `task`, which the caller already counted as one of its waiters. + /// Waits for the download of `id`, as one of the calls it is shared by. /// - /// A plain `await task.value` would ignore this caller's cancellation, since the download - /// runs in a detached task of its own. The cancellation handler gives the caller's place - /// back instead, and the last one to leave cancels the download. - private func waitFor( - _ task: Task, - id: String, - generation: Int - ) async throws -> PlatformImage { - defer { finish(id: id, generation: generation) } + /// The call's place is taken in the same turn of the actor that found or made the download, + /// so the download cannot finish before there is a continuation to hand its result to. When + /// the call is cancelled it leaves at once: its continuation is resumed with a + /// `CancellationError` and the download goes on for the others, or is cancelled if it was + /// the last. + private func waitFor(id: String, generation: Int) async throws -> SendableImage { + let waiter = nextWaiter + nextWaiter += 1 return try await withTaskCancellationHandler { - try await task.value.image + try await withCheckedThrowingContinuation { continuation in + guard var inFlight = tasks[id], inFlight.generation == generation else { + continuation.resume(throwing: CancellationError()) + return + } + + inFlight.waiters[waiter] = continuation + tasks[id] = inFlight + } } onCancel: { - Task { await self.leave(id: id, generation: generation) } + Task { await self.leave(id: id, generation: generation, waiter: waiter) } } } - /// Gives up one waiter's place, cancelling the download when it was the last. - private func leave(id: String, generation: Int) { - guard var inFlight = tasks[id], inFlight.generation == generation else { + /// Gives up one call's place, resuming it with a `CancellationError`, and cancels the + /// download when it was the last. + private func leave(id: String, generation: Int, waiter: Int) { + guard + var inFlight = tasks[id], + inFlight.generation == generation, + let continuation = inFlight.waiters.removeValue(forKey: waiter) + else { return } - inFlight.waiters -= 1 + continuation.resume(throwing: CancellationError()) - guard inFlight.waiters <= 0 else { + guard inFlight.waiters.isEmpty else { tasks[id] = inFlight return } @@ -217,12 +241,24 @@ public actor RDLImageLoader { tasks[id] = nil } - /// Forgets the download once a call has its outcome, unless a newer one took its place. - private func finish(id: String, generation: Int) { - if tasks[id]?.generation == generation { - tasks[id] = nil + /// Hands the outcome of the download to every call still waiting for it, and forgets it. + /// Nothing is left to do when the download was cancelled with its last call. + private func complete(id: String, generation: Int, with result: Result) { + guard let inFlight = tasks[id], inFlight.generation == generation else { + return + } + + tasks[id] = nil + + for continuation in inFlight.waiters.values { + continuation.resume(with: result) } } + + /// How many calls are waiting for the download of `id`. + func waiterCountForTesting(id: String) -> Int { + tasks[id]?.waiters.count ?? 0 + } } #endif diff --git a/Tests/RequestDLTests/Image/Core/RDLImageLoaderTests.swift b/Tests/RequestDLTests/Image/Core/RDLImageLoaderTests.swift index 7938b357..0b56a871 100644 --- a/Tests/RequestDLTests/Image/Core/RDLImageLoaderTests.swift +++ b/Tests/RequestDLTests/Image/Core/RDLImageLoaderTests.swift @@ -203,6 +203,130 @@ struct RDLImageLoaderTests { #expect(await probe.startCount == 1) } + /// The cancelled caller does not wait for the shared request: it stops at once, with a + /// `CancellationError`, while the request goes on for the caller that is still waiting. + @Test + func cancellingOneOfTwoWaiters_returnsTheCancelledOneAtOnce_whileTheRequestStillRuns() async throws { + // Given: two callers sharing one request that stays in flight until the gate is opened. + let probe = RequestProbe() + let gate = Gate() + let task = StubImageTask { + await probe.started() + await gate.wait() + + return .init(head: .stub, payload: .onePixelPNG) + } + + let loader = RDLImageLoader() + let first = Task { try await loader.load(id: "shared", task: task) } + try await eventually { await probe.didStart } + let second = Task { try await loader.load(id: "shared", task: task) } + try await eventually { await loader.waiterCountForTesting(id: "shared") == 2 } + + // When + first.cancel() + + // Then: the first stops without the request having finished. + let outcome = try await completing(within: 30) { () -> String in + do { + _ = try await first.value + return "returned an image" + } catch is CancellationError { + return "cancelled" + } catch { + return "failed with \(error)" + } + } + + #expect(outcome == "cancelled") + #expect(await gate.isOpen == false) + #expect(await loader.waiterCountForTesting(id: "shared") == 1) + + // And the other still gets its image from the same request, which was never cancelled. + await gate.open() + + let image = try await second.value + #expect(image.size.width > 0) + #expect(await probe.wasCancelled == false) + #expect(await probe.startCount == 1) + } + + @Test + func cancellingEveryWaiterOfASharedRequest_cancelsIt() async throws { + // Given: two callers sharing one request that would run for a long time. + let probe = RequestProbe() + let task = StubImageTask { + await probe.started() + + do { + try await Task.sleep(nanoseconds: 30_000_000_000) + } catch { + await probe.cancelled() + throw error + } + + return .init(head: .stub, payload: .onePixelPNG) + } + + let loader = RDLImageLoader() + let first = Task { try await loader.load(id: "shared", task: task) } + try await eventually { await probe.didStart } + let second = Task { try await loader.load(id: "shared", task: task) } + try await eventually { await loader.waiterCountForTesting(id: "shared") == 2 } + + // When: the first leaves, and then the second does. + first.cancel() + try await eventually { await loader.waiterCountForTesting(id: "shared") == 1 } + #expect(await probe.wasCancelled == false) + + second.cancel() + + // Then: the request is cancelled once nobody is left, and both callers stopped. + try await eventually { await probe.wasCancelled } + try await eventually { await loader.waiterCountForTesting(id: "shared") == 0 } + + for load in [first, second] { + do { + _ = try await load.value + Issue.record("Expected the cancelled load to throw") + } catch is CancellationError { + // expected + } + } + } + + /// A caller that arrives after the others left, but while the request still runs, joins it. + @Test + func aCallerThatJoinsAfterAnotherLeft_getsTheSameRequest() async throws { + // Given + let probe = RequestProbe() + let gate = Gate() + let task = StubImageTask { + await probe.started() + await gate.wait() + + return .init(head: .stub, payload: .onePixelPNG) + } + + let loader = RDLImageLoader() + let staying = Task { try await loader.load(id: "shared", task: task) } + try await eventually { await probe.didStart } + let leaving = Task { try await loader.load(id: "shared", task: task) } + try await eventually { await loader.waiterCountForTesting(id: "shared") == 2 } + leaving.cancel() + try await eventually { await loader.waiterCountForTesting(id: "shared") == 1 } + + // When + let joining = Task { try await loader.load(id: "shared", task: task) } + try await eventually { await loader.waiterCountForTesting(id: "shared") == 2 } + await gate.open() + + // Then: both get the image, from one request. + #expect(try await staying.value.size.width > 0) + #expect(try await joining.value.size.width > 0) + #expect(await probe.startCount == 1) + } + @Test func loadingAgain_afterTheOnlyWaiterWasCancelled_startsANewRequest() async throws { // Given: a request that was cancelled together with its only waiter. @@ -264,6 +388,31 @@ private actor RequestProbe { } } +/// Holds a request until the test lets it finish. +private actor Gate { + + private(set) var isOpen = false + private var waiting: [CheckedContinuation] = [] + + func wait() async { + guard !isOpen else { + return + } + + await withCheckedContinuation { waiting.append($0) } + } + + func open() { + isOpen = true + + for continuation in waiting { + continuation.resume() + } + + waiting = [] + } +} + private actor Counter { private(set) var value = 0