AsyncStreamを`for await`する`task`をラップしようとしてハマった話

はじめに

RxSwift.PublishRelay や Combine.PassthroughSubject に代わるものを作ろうとするとAsyncStreamを使用することになります。

例えば以下のコードを見てみましょう。

private var event: ((any Sendable) -> Void)?

var stream: AsyncStream<any Sendable> {
    AsyncStream { cont in
        cont.onTermination = { print($0) }
        event = { cont.yield($0) }
    }
}

let msgs = ["foo", "bar", "baz"]

struct MyView1: View {
    @State private var count = 0
    @Binding var navigationPath: NavigationPath

    var body: some View {
        VStack(spacing: 16) {
            Button("Hello World!") {
                defer { count += 1 }
                event?(msgs[count % msgs.count])
            }
            Text("count: \(count)")
            Button("Submit") {
                navigationPath.append("submit")
            }
            Button("Reset") {
                count = 0
            }
        }
        .padding(20)
        .task {
            for await event in stream {
                print(event) // 流れて来たイベント
            }
        }
        .navigationDestination(for: String.self) {
            switch $0 {
            case "submit":
                MyView2(
                    count: count,
                    navigationPath: $navigationPath
                )
            default: EmptyView()
            }
        }
    }
}

struct MyView2: View {
    var count: Int
    @Binding var navigationPath: NavigationPath

    var body: some View {
        Text("count: \(count)")
    }
}

struct RootView: View {
    @State private var navigationPath = NavigationPath()

    var body: some View {
        NavigationStack(path: $navigationPath) {
            MyView1(navigationPath: $navigationPath)
        }
    }
}

このコードをPlaygroundで動かした結果がこちら

streamに流した値がtaskの中で処理されていることが分かります。
また、画面遷移によってtaskがキャンセルされるため、stream内の cont.onTermination が発火していることも確認できます。

問題のコード

PassthroughSubjectを使っていた時は onReceive を使って

.onReceive(subject) {
    // イベントを処理する
}

と書いていました。
このview側のコードをなるべく変えないよう以下のextensionを切ってみました。

extension View {
    func onReceive<E: Sendable>(
        _ stream: AsyncStream<E>,
        action: @escaping (E) -> Void
    ) -> some View {
        task {
            for await event in stream {
                action(event)
            }
        }
    }
}

そしてMyView1のbodyを以下のように書き換えます

VStack(spacing: 16) {
    Button("Hello World!") {
        defer { count += 1 }
        event?(msgs[count % msgs.count])
    }
    Text("count: \(count)")
    Button("Submit") {
        navigationPath.append("submit")
    }
    Button("Reset") {
        count = 0
    }
}
.padding(20)
- .task {
-     for await event in stream {
-         print(event) // 流れて来たイベント
-     }
- }
+ .onReceive(stream) { print($0) }
.navigationDestination(for: String.self) {
    switch $0 {
    case "submit":
        MyView2(
            count: count,
            navigationPath: $navigationPath
        )
    default: EmptyView()
    }
}

一見問題無いように見えますが、実際に動かすと想定しない動きをします。

原因

taskの中はclosureなので、画面表示されてから初めて実行されそのタイミングでstreamが発行されます。
また、taskは画面表示後の1度しか実行されない(@Stateの変更に左右されない)ため想定通りの動きになります。
一方、今回作成した onReceive は初期表示時の他@Stateが変更される度に呼ばれてしまうため、その度にstreamが発行されてしまいますが task の中は実行されないためそのまま捨てられてしまいます。

解決策

この現象を回避する手段として、onReceiveに渡すstreamを遅延評価させる方法があります。
具体的には値をそのまま渡すのではなくClosureに閉じ込めて渡します。

extension View {
    func onReceive<E: Sendable>(
-       _ stream: AsyncStream<E>,
+       _ stream: @escaping () -> AsyncStream<E>,
        action: @escaping (E) -> Void
    ) -> some View {
        task {
-           for await event in stream {
+           for await event in stream() {
                action(event)
            }
        }
    }
}

呼び出し方は

- .onReceive(stream) { print($0) }
+ .onReceive({stream}) { print($0) }

または

- .onReceive(stream) { print($0) }
+ .onReceive { stream } action: { print($0) }

のようになります。

...ちょっと見栄えが悪いですね😅

シン・解決策

Swiftには @autoclosure という独特の機能があります。
この機能を使うことで、呼び出し元に手を加えることなく遅延評価を実現することができます。

extension View {
    func onReceive<E: Sendable>(
-       _ stream: AsyncStream<E>,
+       _ stream: @autoclosure @escaping () -> AsyncStream<E>,
        action: @escaping (E) -> Void
    ) -> some View {
        task {
-           for await event in stream {
+           for await event in stream() {
                action(event)
            }
        }
    }
}

そして最終的にできたコードで動かした結果がこちら

想定通り動くようになりました

まとめ

task(_:)をラップして関数を作る場合、遅延評価をしないと意図しない動きをする場合があることに注意しましょう。

参考