[SS-348] implement Batches' for_each_diff_pair as an iterator (to support async operations)

[ublubu]

We want to write rows to S3 as soon as we read them from the batch
instead of queuing them in memory for the next operator (write_data_files).

That means the body of our for_each_diff_pair must support async operations.
This PR refactors for_each_diff_pair into an iterator so we can do that (in the stacked PR).

[linear-code]

SS-348

[martykulma]

lgtm, thanks @ublubu !

[def-]

QA LLM review found this: https://github.com/MaterializeInc/qa-llm-review/blob/master/commit-bugs/done/analysis-pr-37759.md
I added an extra test:

diff --git a/test/bounded-memory/mzcompose.py b/test/bounded-memory/mzcompose.py
index e725829d25..ea30f759c2 100644
--- a/test/bounded-memory/mzcompose.py
+++ b/test/bounded-memory/mzcompose.py
@@ -57,6 +57,16 @@ STRING_PAD = "x" * PAD_LEN
 REPEAT = 16 * 1024
 ITERATIONS = 128

+# Multiplicity of the single hot key in the `kafka-sink-hot-key-fanout` scenario.
+# A sink's arrangement walker (`mz_interchange::envelopes::for_each_diff_pair`)
+# fans one consolidated `(key, value, diff=N)` update out to N `DiffPair`s. The
+# regression under test (`iter_diff_pairs`) collects all N of them into a per-key
+# `Vec` before the wrapper yields the first one; its streaming predecessor
+# consumed the fan-out one pair at a time. So the buggy walk retains roughly
+# `N * sizeof((Timestamp, DiffPair<Row>))` extra on top of whatever the encoder
+# path holds, and N is the knob that sizes that regression. See the scenario.
+HOT_KEY_FANOUT = 10 * 1000 * 1000
+
 BOUNDED_MEMORY_FRAMEWORK_VERSION = "1.0.0"

 SERVICES = [
@@ -1336,6 +1346,91 @@ SCENARIOS = [
         materialized_memory="1.8Gb",
         clusterd_memory="0.5Gb",
     ),
+    # Regression test for the hot-key fan-out in a sink's arrangement walker.
+    #
+    # `hot` is a single row repeated `HOT_KEY_FANOUT` times. Differential
+    # consolidates it to one `(row, diff=N)` tuple, so the stored arrangement is
+    # O(1); nothing upstream of the sink ever holds N rows. When the Kafka sink
+    # walks that arrangement it hits a single `(key, timestamp)` group whose diff
+    # is N and fans it out to N `DiffPair`s.
+    #
+    # The streaming walker consumed that fan-out one pair at a time (peak O(1)
+    # rows; the snapshot streams out to Kafka incrementally). The eager-collect
+    # regression first materializes all N `DiffPair`s for the key into a `Vec`
+    # before the encoder sees any, so clusterd holds ~`N * sizeof((Timestamp,
+    # DiffPair<Row>))` extra at once. Sized past `clusterd_memory` that spike gets
+    # the replica OOM-killed (restart: no), the sink never commits, and the
+    # `messages_committed` step below times out -> the scenario fails. The
+    # streaming walker stays within budget and commits all N.
+    #
+    # NOTE: `clusterd_memory` is the load-bearing threshold and is a provisional
+    # estimate. It must sit above the streaming walker's peak and below the
+    # eager-collect spike, and where that spike lands depends on clusterd's worker
+    # scheduling and how fast Kafka drains, so it needs one calibration run:
+    # `bin/mzcompose --find bounded-memory run main kafka-sink-hot-key-fanout`
+    # on a build with the fix and on one without (watch clusterd RSS via
+    # `docker stats`), then set the limit between the two peaks. `HOT_KEY_FANOUT`
+    # widens the gap if more headroom is needed. NOTE: `run main <name>` filters
+    # to one scenario; `run default` ignores the name and runs the whole suite.
+    Scenario(
+        name="kafka-sink-hot-key-fanout",
+        pre_restart=dedent(f"""
+            # Lower the statistics scrape interval so the `messages_committed`
+            # check resolves quickly rather than on the 60s default. The whole
+            # snapshot commits in a single transaction that stays open while all N
+            # messages stream out, so `kafka_transaction_timeout` must cover that;
+            # the default is far too short for N this large.
+            $ postgres-execute connection=postgres://mz_system:materialize@${{testdrive.materialize-internal-sql-addr}}
+            ALTER SYSTEM SET storage_statistics_collection_interval = 1000;
+            ALTER SYSTEM SET storage_statistics_interval = 2000;
+            ALTER SYSTEM SET kafka_transaction_timeout = 600000;
+
+            > CREATE CONNECTION kafka_conn
+              FOR KAFKA BROKER '${{testdrive.kafka-addr}}', SECURITY PROTOCOL PLAINTEXT
+
+            # A table-provided bound keeps `generate_series` in the dataflow. A
+            # constant bound would be folded and expanded in environmentd instead.
+            > CREATE TABLE mult (n BIGINT)
+
+            > INSERT INTO mult VALUES ({HOT_KEY_FANOUT})
+
+            > CREATE MATERIALIZED VIEW hot IN CLUSTER clusterd AS
+              SELECT 0::int AS k, 0::int AS v FROM mult, generate_series(1, mult.n)
+
+            > SELECT count(*) FROM hot
+            {HOT_KEY_FANOUT}
+
+            > CREATE SINK hot_sink
+              IN CLUSTER clusterd
+              FROM hot
+              INTO KAFKA CONNECTION kafka_conn (TOPIC 'bounded-memory-hot-sink-${{testdrive.seed}}')
+              KEY (k) NOT ENFORCED
+              FORMAT JSON
+              ENVELOPE UPSERT
+
+            # Waiting for every message to commit both drives the fan-out and
+            # proves clusterd walked it within budget. A dead (OOM-killed) replica
+            # never advances this counter, so the step times out.
+            > SELECT SUM(u.messages_committed) >= {HOT_KEY_FANOUT}
+              FROM mz_sinks s
+              JOIN mz_internal.mz_sink_statistics_raw u ON s.id = u.id
+              WHERE s.name = 'hot_sink'
+            true
+            """),
+        # After the snapshot has committed the sink resumes past it, so
+        # re-hydration re-reads no data and does no fan-out.
+        post_restart=dedent("""
+            > SELECT status FROM mz_internal.mz_sink_statuses WHERE name = 'hot_sink'
+            running
+            """),
+        # Not the target of the test; generous so environmentd is never the
+        # binding constraint.
+        materialized_memory="4Gb",
+        clusterd_memory="3Gb",
+        # The fix streams ~N messages out to Kafka in a single transaction, which
+        # takes a while for N this large.
+        timeout="1800s",
+    ),
 ]

Running bin/mzcompose --find bounded-memory down && bin/mzcompose --find bounded-memory run main kafka-sink-hot-key-fanout hangs with OoM kills.

[martykulma]

oof - I missed the collect! Thanks @def- , @ublubu is out, I'll take a look.

[martykulma]

Running bin/mzcompose --find bounded-memory down && bin/mzcompose --find bounded-memory run main kafka-sink-hot-key-fanout hangs with OoM kills.

I ran the test locally, but test passed. I'll try adjusting to increase the message size.

==> mzcompose: test case Scenario 'kafka-sink-hot-key-fanout' succeeded
==> Sanity Restart: Restart Mz and verify source/sink/replica health
$ docker compose ps materialized --quiet
$ docker compose kill -sSIGKILL materialized
[+] kill 1/1
 ✔ Container bounded-memory-materialized-1 Killed                                                                                                                                                             0.1s
$ docker compose logs --no-color --timestamps --since 1784749041.9365747
$ docker compose up --detach --wait materialized
[+] up 1/1
 ✔ Container bounded-memory-materialized-1 Healthy                                                                                                                                                            5.6s
> SELECT 1
==> mzcompose: test case workflow-main succeeded

[martykulma]

The test seems to fail with and without my change from Vec to Iterator, and this is likely because this is synchronous code, and operator generates all pairs before yielding. A unit test can cover it

[martykulma]

fix on the way: #37817