【The Composable Architecture (TCA) 源码解析】04 - Effect 扩展之 Cancellation

此源码解析适用于 0.41.0 之前的版本。最新的源码已经把 Reducer 重构成了一个协议 ReducerProtocol，但核心的实现方法还是离不开这个版本的源码。

这个扩展的作用是将普通 Effect 变成可取消的。无论 Effect 的 operation 是 .publisher 还是 .run，它们的核心逻辑都是利用 AnyCancellable 来完成。

extension Effect {
  public func cancellable(id: AnyHashable, cancelInFlight: Bool = false) -> Self {
    switch self.operation {
    case .none:
      return .none
    case let .publisher(publisher):
      return Self(
        operation: .publisher(
			// `Deferred` 可以让闭包在接收到订阅者时才执行
          Deferred {
            ()
              -> Publishers.HandleEvents<
                Publishers.PrefixUntilOutput<
                  AnyPublisher<Action, Failure>, PassthroughSubject<Void, Never>
                >
              > in
			  // 加锁，保证相关数据的准确性
            cancellablesLock.lock()
            defer { cancellablesLock.unlock() }

            let id = CancelToken(id: id)
            if cancelInFlight {
			    // 取消正在运行中的 `Effect`
              cancellationCancellables[id]?.forEach { $0.cancel() }
            }

			  // 用于取消正在运行中的 `Effect`
            let cancellationSubject = PassthroughSubject<Void, Never>()

			  // 声明用于取消的 `AnyCancellable`
            var cancellationCancellable: AnyCancellable!

			  // 创建用于取消的 `AnyCancellable`
            cancellationCancellable = AnyCancellable {
			    // 这个代码块将在调用 `cancel()` 的时候调用
              cancellablesLock.sync {
				  // `cancellationSubject` 发送一个值，并完成
                cancellationSubject.send(())
                cancellationSubject.send(completion: .finished)

				  // 将 `cancellationCancellable` 从 `id` 对应的集合中移除
                cancellationCancellables[id]?.remove(cancellationCancellable)

				  // 如果 `id` 对应的集合为空，则把 `id` 对应的值设置为 `nil`
                if cancellationCancellables[id]?.isEmpty == .some(true) {
                  cancellationCancellables[id] = nil
                }
              }
            }

			  // 最终实现取消功能的核心在这里：当 `cancellationSubject` 发出一个值后，`publisher` 就不会再发出值。
			  // 所以上面的 `AnyCancellable` 调用 `cancel()` 后，就可以达到取消的目的。
            return publisher.prefix(untilOutputFrom: cancellationSubject)
              .handleEvents(
                receiveSubscription: { _ in
                  _ = cancellablesLock.sync {
					   // 接收到订阅，把 `cancellationCancellable` 缓存起来
                    cancellationCancellables[id, default: []].insert(
                      cancellationCancellable
                    )
                  }
                },
				   // 接收到完成或取消，都调用 `cancellationCancellable.cancel()`，完成数据清理
                receiveCompletion: { _ in cancellationCancellable.cancel() },
                receiveCancel: cancellationCancellable.cancel
              )
          }
          .eraseToAnyPublisher()
        )
      )
    case let .run(priority, operation):
      return Self(
        operation: .run(priority) { send in
			// 调用了一个具体实现的方法，查看此方法解析
          await withTaskCancellation(id: id, cancelInFlight: cancelInFlight) {
            await operation(send)
          }
        }
      )
    }
  }

  // 上面那个方法的重载方法，这里的 `id` 是一个 `Any.Type`
  public func cancellable(id: Any.Type, cancelInFlight: Bool = false) -> Self {
    self.cancellable(id: ObjectIdentifier(id), cancelInFlight: cancelInFlight)
  }

	// 取消对应 `id` 的所有正在进行中的 Effect
  public static func cancel(id: AnyHashable) -> Self {
    .fireAndForget {
      cancellablesLock.sync {
        cancellationCancellables[.init(id: id)]?.forEach { $0.cancel() }
      }
    }
  }

  // 上面那个方法的重载方法，这里的 `id` 是一个 `Any.Type`
  public static func cancel(id: Any.Type) -> Self {
    .cancel(id: ObjectIdentifier(id))
  }

	// 取消 `ids` 数组元素对应的所有正在进行中的 Effect
  public static func cancel(ids: [AnyHashable]) -> Self {
    // 每个 `id` 转换为一个 Effect，然后再合并
    .merge(ids.map(Effect.cancel(id:)))
  }

  // 上面那个方法的重载方法，这里的 `ids` 是一个 `Any.Type` 数组
  public static func cancel(ids: [Any.Type]) -> Self {
    .merge(ids.map(Effect.cancel(id:)))
  }
}

public func withTaskCancellation<T: Sendable>(
  id: AnyHashable,
  cancelInFlight: Bool = false,
  operation: @Sendable @escaping () async throws -> T
) async rethrows -> T {
  // 创建任务
  let task = { () -> Task<T, Error> in
	  // 加锁，保证相关数据的准确性
    cancellablesLock.lock()

    let id = CancelToken(id: id)
    if cancelInFlight {
	    // 取消正在运行中的 `Effect`
      cancellationCancellables[id]?.forEach { $0.cancel() }
    }

    // 创建任务
    let task = Task { try await operation() }

	  // 声明用于取消的 `AnyCancellable`
    var cancellable: AnyCancellable!

	  // 创建用于取消的 `AnyCancellable`
    cancellable = AnyCancellable {
		// 这个代码块将在调用 `AnyCancellable.cancel()` 的时候调用

		// 取消任务
      task.cancel()
      cancellablesLock.sync {
		  // 将 `cancellable ` 从 `id` 对应的集合中移除
        cancellationCancellables[id]?.remove(cancellable)

		  // 如果 `id` 对应的集合为空，则把 `id` 对应的值设置为 `nil`
        if cancellationCancellables[id]?.isEmpty == .some(true) {
          cancellationCancellables[id] = nil
        }
      }
    }

    // 把 `cancellable` 缓存起来
    cancellationCancellables[id, default: []].insert(cancellable)
    cancellablesLock.unlock()
    return task
  }()
  do {
    // 执行异步任务
    return try await task.cancellableValue
  } catch {
	  // 如果执行出错，则重新抛出错误
    return try Result<T, Error>.failure(error)._rethrowGet()
  }
}

// 上面那个方法的重载方法，这里的 `id` 是一个 `Any.Type`
public func withTaskCancellation<T: Sendable>(
  id: Any.Type,
  cancelInFlight: Bool = false,
  operation: @Sendable @escaping () async throws -> T
) async rethrows -> T {
  try await withTaskCancellation(
    id: ObjectIdentifier(id),
    cancelInFlight: cancelInFlight,
    operation: operation
  )
}

extension Task where Success == Never, Failure == Never {
  // `Task` 的扩展方法，取消对应 `id` 的所有正在进行中副作用
  public static func cancel<ID: Hashable & Sendable>(id: ID) async {
	  // 在使用 `Task.cancel(id:)` 的地方，通常是伴随着 `async/await` 的使用，
	  // 所以把这个方法定义为一个 `async` 方法。
    // 为了把 cancel 的操作放在 `async/await` 的环境下执行，就使用了 `MainActor.run`。
    await MainActor.run {
      cancellablesLock.sync { cancellationCancellables[.init(id: id)]?.forEach { $0.cancel() } }
    }
  }

  // 上面那个方法的重载方法，这里的 `id` 是一个 `Any.Type`
  public static func cancel(id: Any.Type) async {
    await self.cancel(id: ObjectIdentifier(id))
  }
}

// 作为 `cancellationCancellables` 字典的 Key
struct CancelToken: Hashable {
  let id: AnyHashable

  // 不知道 `discriminator` 存在的意义是什么？`id` 已经足够保证唯一性。
  let discriminator: ObjectIdentifier

  init(id: AnyHashable) {
    self.id = id
    self.discriminator = ObjectIdentifier(type(of: id.base))
  }
}

// 存储 cancellables
var cancellationCancellables: [CancelToken: Set<AnyCancellable>] = [:]

// 递归锁，保证 cancellables 相关数据的准确性。
// 使用递归锁的原因是在 `Store.send(_:originatingFrom)` 方法中有递归操作。
let cancellablesLock = NSRecursiveLock()

// 定义 `_ErrorMechanism` 是为了让 `Result` 拥有 `_rethrowGet` 方法，方便重抛错误。
// 在 `withTaskCancellation` 的具体实现中用到。
@rethrows
private protocol _ErrorMechanism {
  associatedtype Output
  func get() throws -> Output
}

extension _ErrorMechanism {
  internal func _rethrowError() rethrows -> Never {
    _ = try _rethrowGet()
    fatalError()
  }

  internal func _rethrowGet() rethrows -> Output {
    return try get()
  }
}

extension Result: _ErrorMechanism {}