factor out the sync pipeline (#21027)

* factor out the sync pipeline (fetch, validate and DB-insert) into a generic pipeline and simpler functions for each step

* review comments

* review comment

* review comment

chia/_tests/util/test_task_pipeline.py

@@ -0,0 +1,380 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncIterator
+
+import pytest
+
+from chia.util.task_pipeline import TaskPipeline
+
+
+async def _count_up(n: int) -> AsyncIterator[int]:
+    for i in range(n):
+        yield i
+
+
+@pytest.mark.anyio
+async def test_basic_pipeline() -> None:
+    """Items flow through all stages in order."""
+    results: list[int] = []
+
+    async def double(x: int) -> int:
+        return x * 2
+
+    async def collect(x: int) -> None:
+        results.append(x)
+
+    pipeline = TaskPipeline(source=_count_up(5), stages=[double, collect], queue_size=2)
+    await pipeline.run()
+
+    assert results == [0, 2, 4, 6, 8]
+
+
+@pytest.mark.anyio
+async def test_single_consumer_stage() -> None:
+    """Pipeline works with just one consumer stage (no transforms)."""
+    results: list[int] = []
+
+    async def collect(x: int) -> None:
+        results.append(x)
+
+    pipeline = TaskPipeline(source=_count_up(3), stages=[collect], queue_size=4)
+    await pipeline.run()
+
+    assert results == [0, 1, 2]
+
+
+@pytest.mark.anyio
+async def test_filter_with_none() -> None:
+    """Returning None from a transform stage filters the item."""
+    results: list[int] = []
+
+    async def keep_even(x: int) -> int | None:
+        if x % 2 == 0:
+            return x
+        return None
+
+    async def collect(x: int) -> None:
+        results.append(x)
+
+    pipeline = TaskPipeline(source=_count_up(6), stages=[keep_even, collect], queue_size=4)
+    await pipeline.run()
+
+    assert results == [0, 2, 4]
+
+
+@pytest.mark.anyio
+async def test_source_exception_propagates() -> None:
+    """An exception in the source async iterator propagates from run()."""
+
+    async def bad_source() -> AsyncIterator[int]:
+        yield 1
+        raise RuntimeError("source failed")
+
+    results: list[int] = []
+
+    # Required stage argument. May process the one yielded item if the
+    # consumer runs before the source exception propagates; the assertion
+    # below accounts for both outcomes.
+    async def collect(x: int) -> None:
+        results.append(x)  # pragma: no cover
+
+    pipeline = TaskPipeline(source=bad_source(), stages=[collect], queue_size=4)
+    with pytest.raises(RuntimeError, match="source failed"):
+        await pipeline.run()
+
+    assert len(results) <= 1
+
+
+@pytest.mark.anyio
+async def test_stage_exception_propagates() -> None:
+    """An exception in a stage propagates from run()."""
+    processed: list[int] = []
+
+    async def failing_stage(x: int) -> int:
+        if x == 3:
+            raise ValueError("stage failed on 3")
+        processed.append(x)
+        return x
+
+    async def collect(x: int) -> None:
+        pass
+
+    pipeline = TaskPipeline(source=_count_up(10), stages=[failing_stage, collect], queue_size=2)
+    with pytest.raises(ValueError, match="stage failed on 3"):
+        await pipeline.run()
+
+    # Items 0, 1, 2 should have been processed before the failure
+    assert processed == [0, 1, 2]
+
+
+@pytest.mark.anyio
+async def test_drain_after_failure() -> None:
+    """drain() returns unconsumed items from a queue after a failure."""
+
+    async def identity(x: int) -> int:
+        return x
+
+    async def failing_consumer(x: int) -> None:
+        if x == 2:
+            raise ValueError("boom")
+
+    pipeline = TaskPipeline(source=_count_up(10), stages=[identity, failing_consumer], queue_size=5)
+    with pytest.raises(ValueError, match="boom"):
+        await pipeline.run()
+
+    # Some items may remain in the queue between identity and failing_consumer
+    remaining = pipeline.drain(1)
+    assert all(isinstance(item, int) for item in remaining)
+
+
+@pytest.mark.anyio
+async def test_backpressure() -> None:
+    """Slow consumer creates back-pressure without deadlock."""
+    results: list[int] = []
+
+    async def slow_collect(x: int) -> None:
+        await asyncio.sleep(0.01)
+        results.append(x)
+
+    pipeline = TaskPipeline(source=_count_up(20), stages=[slow_collect], queue_size=2)
+    await pipeline.run()
+
+    assert results == list(range(20))
+
+
+@pytest.mark.anyio
+async def test_empty_source() -> None:
+    """An empty source completes the pipeline cleanly."""
+    results: list[int] = []
+
+    # Required stage argument. Never called because the source is empty.
+    async def collect(x: int) -> None:
+        results.append(x)  # pragma: no cover
+
+    pipeline = TaskPipeline(source=_count_up(0), stages=[collect], queue_size=4)
+    await pipeline.run()
+
+    assert results == []
+
+
+@pytest.mark.anyio
+async def test_three_stage_pipeline() -> None:
+    """A 3-stage pipeline (transform, transform, consumer) works correctly."""
+    results: list[str] = []
+
+    async def add_one(x: int) -> int:
+        return x + 1
+
+    async def to_string(x: int) -> str:
+        return str(x)
+
+    async def collect(x: str) -> None:
+        results.append(x)
+
+    pipeline = TaskPipeline(source=_count_up(4), stages=[add_one, to_string, collect], queue_size=3)
+    await pipeline.run()
+
+    assert results == ["1", "2", "3", "4"]
+
+
+@pytest.mark.anyio
+async def test_no_stages_raises() -> None:
+    """Constructing a pipeline with no stages raises ValueError."""
+    with pytest.raises(ValueError, match="at least one stage"):
+        TaskPipeline(source=_count_up(1), stages=[], queue_size=4)
+
+
+@pytest.mark.anyio
+async def test_names_wrong_length_raises() -> None:
+    """names must have exactly len(stages) + 1 entries."""
+
+    async def collect(x: int) -> None:
+        pass  # pragma: no cover
+
+    with pytest.raises(ValueError, match="len\\(stages\\) \\+ 1"):
+        TaskPipeline(source=_count_up(1), stages=[collect], queue_size=4, names=["only_one"])
+
+    with pytest.raises(ValueError, match="len\\(stages\\) \\+ 1"):
+        TaskPipeline(source=_count_up(1), stages=[collect], queue_size=4, names=["a", "b", "c"])
+
+
+@pytest.mark.anyio
+async def test_queues_property() -> None:
+    """The queues property exposes the inter-stage queues created during run()."""
+    results: list[int] = []
+
+    async def double(x: int) -> int:
+        return x * 2
+
+    async def collect(x: int) -> None:
+        results.append(x)
+
+    pipeline = TaskPipeline(source=_count_up(3), stages=[double, collect], queue_size=4)
+    assert pipeline.queues == []
+    await pipeline.run()
+    assert len(pipeline.queues) == 2
+    assert results == [0, 2, 4]
+
+
+@pytest.mark.anyio
+async def test_drain_includes_dropped_results() -> None:
+    """drain() includes results that a stage produced but could not enqueue
+    because the pipeline was already shutting down."""
+    consumed: list[int] = []
+    produced: list[int] = []
+
+    async def slow_transform(x: int) -> int:
+        produced.append(x)
+        if x == 0:
+            # give the consumer time to fail before we return
+            await asyncio.sleep(0.1)
+        return x * 10
+
+    async def failing_consumer(x: int) -> None:
+        consumed.append(x)
+        raise ValueError("consumer failed")
+
+    pipeline = TaskPipeline(source=_count_up(5), stages=[slow_transform, failing_consumer], queue_size=2)
+    with pytest.raises(ValueError, match="consumer failed"):
+        await pipeline.run()
+
+    remaining = pipeline.drain(1)
+    # The transform produced item 0 (which the consumer saw and failed on),
+    # but it may also have produced further items that couldn't be enqueued.
+    # drain() should include both queued and dropped items.
+    all_seen = consumed + remaining
+    for item in all_seen:
+        assert isinstance(item, int)
+    # The dropped result from slow_transform(0) should appear somewhere —
+    # either consumed by the failing consumer or recovered via drain.
+    assert 0 in produced
+
+
+@pytest.mark.anyio
+async def test_failed_during_stage_saves_result_to_dropped() -> None:
+    """If the pipeline fails while a transform is awaiting and the transform
+    suppresses CancelledError and returns a result, that result is saved to
+    _dropped and recovered via drain()."""
+
+    async def resilient_transform(x: int) -> int:
+        try:
+            await asyncio.sleep(0.1)
+        except asyncio.CancelledError:
+            pass
+        return x * 10
+
+    async def failing_consumer(x: int) -> None:
+        raise ValueError("boom")
+
+    pipeline = TaskPipeline(source=_count_up(10), stages=[resilient_transform, failing_consumer], queue_size=2)
+    with pytest.raises(ValueError, match="boom"):
+        await pipeline.run()
+
+    remaining = pipeline.drain(1)
+    assert any(item % 10 == 0 for item in remaining)
+
+
+@pytest.mark.anyio
+async def test_stage_blocked_on_empty_queue_bails_on_failure() -> None:
+    """When a stage is blocked in _get_or_bail's slow path (queue empty) and
+    another stage fails, the blocked stage exits via the fail event."""
+
+    async def stalling_source() -> AsyncIterator[int]:
+        yield 0
+        await asyncio.Event().wait()  # stall forever; never yields again
+
+    async def transform(x: int) -> int:
+        return x
+
+    async def failing_consumer(x: int) -> None:
+        raise ValueError("consumer exploded")
+
+    pipeline = TaskPipeline(
+        source=stalling_source(),
+        stages=[transform, failing_consumer],
+        queue_size=2,
+    )
+    with pytest.raises(ValueError, match="consumer exploded"):
+        await pipeline.run()
+
+
+@pytest.mark.anyio
+async def test_get_or_bail_slow_path_item_arrives() -> None:
+    """Exercise _get_or_bail's slow path where the queue is empty when polled
+    but an item arrives before any failure (get_task wins the race)."""
+    results: list[int] = []
+
+    async def delayed_source() -> AsyncIterator[int]:
+        for i in range(5):
+            await asyncio.sleep(0.02)
+            yield i
+
+    async def collect(x: int) -> None:
+        results.append(x)
+
+    pipeline = TaskPipeline(source=delayed_source(), stages=[collect], queue_size=2)
+    await pipeline.run()
+
+    assert results == [0, 1, 2, 3, 4]
+
+
+@pytest.mark.anyio
+async def test_cancellation_calls_cleanup() -> None:
+    """Cancelling the run() task still invokes the cleanup callback."""
+    cleanup_called = asyncio.Event()
+
+    async def cleanup(p: TaskPipeline) -> None:
+        cleanup_called.set()
+
+    async def slow_source() -> AsyncIterator[int]:
+        for i in range(1000):
+            await asyncio.sleep(0.1)
+            yield i  # pragma: no cover
+
+    async def collect(x: int) -> None:
+        pass  # pragma: no cover
+
+    pipeline = TaskPipeline(source=slow_source(), stages=[collect], queue_size=2, cleanup=cleanup)
+    task = asyncio.ensure_future(pipeline.run())
+    await asyncio.sleep(0.05)
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    assert cleanup_called.is_set()
+
+
+@pytest.mark.anyio
+async def test_cancellation_does_not_leak_tasks() -> None:
+    """Cancelling run() awaits all internal tasks so nothing is left dangling."""
+
+    async def slow_source() -> AsyncIterator[int]:
+        for i in range(1000):
+            await asyncio.sleep(0.1)
+            yield i
+
+    stage_tasks_done: list[bool] = []
+
+    async def slow_stage(x: int) -> int:
+        try:
+            await asyncio.sleep(10)
+        except asyncio.CancelledError:
+            stage_tasks_done.append(True)
+            raise
+        return x  # pragma: no cover
+
+    async def collect(x: int) -> None:
+        pass  # pragma: no cover
+
+    pipeline = TaskPipeline(source=slow_source(), stages=[slow_stage, collect], queue_size=2)
+    task = asyncio.ensure_future(pipeline.run())
+    await asyncio.sleep(0.15)
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    # Give the event loop a tick to finalize everything
+    await asyncio.sleep(0)
+    # The stage task should have been cancelled and awaited, not left dangling
+    assert len(stage_tasks_done) > 0

chia/full_node/full_node.py

@@ -99,6 +99,7 @@ from chia.util.path import path_from_root
 from chia.util.priority_thread_pool_executor import Executor, PriorityThreadPoolExecutor
 from chia.util.profiler import enable_profiler, mem_profile_task, profile_task
 from chia.util.safe_cancel_task import cancel_task_safe
+from chia.util.task_pipeline import TaskPipeline
 from chia.util.task_referencer import create_referenced_task
 
 
@@ -1250,13 +1251,12 @@ class FullNode:
         blockchain = AugmentedBlockchain(self.blockchain)
         peers_with_peak: list[WSChiaConnection] = self.get_peers_with_peak(peak_hash)
 
-        async def fetch_blocks(output_queue: asyncio.Queue[tuple[WSChiaConnection, list[FullBlock]] | None]) -> None:
+        async def fetch_blocks() -> AsyncIterator[tuple[WSChiaConnection, list[FullBlock]]]:
             # the rate limit for respond_blocks is 100 messages / 60 seconds.
             # But the limit is scaled to 30% for outbound messages, so that's 30
             # messages per 60 seconds.
             # That's 2 seconds per request.
             seconds_per_request = 2
-            start_height, end_height = 0, 0
 
             # the timestamp of when the next request_block message is allowed to
             # be sent. It's initialized to the current time, and bumped by the
@@ -1273,243 +1273,192 @@ class FullNode:
             new_peers_with_peak: list[tuple[WSChiaConnection, float]] = [(c, now) for c in peers_with_peak[:]]
             self.log.info(f"peers with peak: {len(new_peers_with_peak)}")
             random.shuffle(new_peers_with_peak)
-            try:
-                # block request ranges are *inclusive*, this requires some
-                # gymnastics of this range (+1 to make it exclusive, like normal
-                # ranges) and then -1 when forming the request message
-                for start_height in range(fork_point_height, target_peak_sb_height + 1, batch_size):
-                    end_height = min(target_peak_sb_height, start_height + batch_size - 1)
-                    request = RequestBlocks(uint32(start_height), uint32(end_height), True)
-                    new_peers_with_peak.sort(key=lambda pair: pair[1])
-                    fetched = False
-                    for idx, (peer, timestamp) in enumerate(new_peers_with_peak):
-                        if peer.closed:
-                            continue
-
-                        start = time.monotonic()
-                        if start < timestamp:
-                            # rate limit ourselves, since we sent a message to
-                            # this peer too recently
-                            await asyncio.sleep(timestamp - start)
-                            start = time.monotonic()
-
-                        # update the timestamp, now that we're sending a request
-                        # it's OK for the timestamp to fall behind wall-clock
-                        # time. It just means we're allowed to send more
-                        # requests to catch up
-                        if is_localhost(peer.peer_info.host):
-                            # we don't apply rate limits to localhost, and our
-                            # tests depend on it
-                            bump = 0.1
-                        else:
-                            bump = seconds_per_request
-
-                        new_peers_with_peak[idx] = (
-                            new_peers_with_peak[idx][0],
-                            new_peers_with_peak[idx][1] + bump,
-                        )
-                        # the fewer peers we have, the more willing we should be
-                        # to wait for them.
-                        timeout = int(30 + 30 / len(new_peers_with_peak))
-                        response = await peer.call_api(FullNodeAPI.request_blocks, request, timeout=timeout)
-                        end = time.monotonic()
-                        if response is None:
-                            self.log.info(f"peer timed out after {end - start:.1f} s")
-                            await peer.close()
-                        elif isinstance(response, RespondBlocks):
-                            if end - start > 5:
-                                self.log.info(f"peer took {end - start:.1f} s to respond to request_blocks")
-                                # this isn't a great peer, reduce its priority
-                                # to prefer any peers that had to wait for it.
-                                # By setting the next allowed timestamp to now,
-                                # means that any other peer that has waited for
-                                # this will have its next allowed timestamp in
-                                # the passed, and be preferred multiple times
-                                # over this peer.
-                                new_peers_with_peak[idx] = (
-                                    new_peers_with_peak[idx][0],
-                                    end,
-                                )
-                            start = time.monotonic()
-                            await output_queue.put((peer, response.blocks))
-                            end = time.monotonic()
-                            if end - start > 1:
-                                self.log.info(
-                                    f"sync pipeline back-pressure. stalled {end - start:0.2f} "
-                                    "seconds on prevalidate block"
-                                )
-                            fetched = True
-                            break
-                    if fetched is False:
-                        self.log.error(f"failed fetching {start_height} to {end_height} from peers")
-                        return
-                    if self.sync_store.peers_changed.is_set():
-                        existing_peers = {id(c): timestamp for c, timestamp in new_peers_with_peak}
-                        peers = self.get_peers_with_peak(peak_hash)
-                        new_peers_with_peak = [(c, existing_peers.get(id(c), end)) for c in peers]
-                        random.shuffle(new_peers_with_peak)
-                        self.sync_store.peers_changed.clear()
-                        self.log.info(f"peers with peak: {len(new_peers_with_peak)}")
-            except Exception:
-                self.log.exception(f"Exception fetching {start_height} to {end_height} from peer")
-                raise
-            finally:
-                # finished signal with None
-                await output_queue.put(None)
-
-        async def validate_blocks(
-            input_queue: asyncio.Queue[tuple[WSChiaConnection, list[FullBlock]] | None],
-            output_queue: asyncio.Queue[
-                tuple[WSChiaConnection, ValidationState, list[Awaitable[PreValidationResult]], list[FullBlock]] | None
-            ],
-        ) -> None:
-            nonlocal blockchain
-            nonlocal fork_info
-            first_batch = True
-
-            vs = ValidationState(ssi, diff, prev_ses_block)
-
-            try:
-                while True:
-                    res: tuple[WSChiaConnection, list[FullBlock]] | None = await input_queue.get()
-                    if res is None:
-                        self.log.debug("done fetching blocks")
-                        return None
-                    peer, blocks = res
-
-                    # Keep the augmented chain's MMR snapshot aligned with the
-                    # underlying blockchain before validating each fetched
-                    # batch.
-                    # TODO: Consider per-batch AugmentedBlockchain instances.
-                    # The current sync pipeline shares one augmented overlay
-                    # across fetch/validate/ingest, so we refresh only the MMR
-                    # snapshot to avoid stale roots after ingest advances the
-                    # canonical chain.
-                    blockchain.mmr_manager = self.blockchain.mmr_manager.copy()
-
-                    # skip_blocks is only relevant at the start of the sync,
-                    # to skip blocks we already have in the database (and have
-                    # been validated). Once we start validating blocks, we
-                    # shouldn't be skipping any.
-                    blocks_to_validate = await self.skip_blocks(blockchain, blocks, fork_info, vs)
-                    assert first_batch or len(blocks_to_validate) == len(blocks)
-                    next_validation_state = copy.copy(vs)
-
-                    if len(blocks_to_validate) == 0:
+            # block request ranges are *inclusive*, this requires some
+            # gymnastics of this range (+1 to make it exclusive, like normal
+            # ranges) and then -1 when forming the request message
+            for start_height in range(fork_point_height, target_peak_sb_height + 1, batch_size):
+                end_height = min(target_peak_sb_height, start_height + batch_size - 1)
+                request = RequestBlocks(uint32(start_height), uint32(end_height), True)
+                new_peers_with_peak.sort(key=lambda pair: pair[1])
+                fetched = False
+                for idx, (peer, timestamp) in enumerate(new_peers_with_peak):
+                    if peer.closed:
                         continue
 
-                    first_batch = False
-
-                    futures: list[Awaitable[PreValidationResult]] = []
-                    for block in blocks_to_validate:
-                        futures.extend(
-                            await self.prevalidate_blocks(
-                                blockchain,
-                                [block],
-                                vs,
-                                summaries,
-                            )
-                        )
                     start = time.monotonic()
-                    await output_queue.put((peer, next_validation_state, list(futures), blocks_to_validate))
-                    end = time.monotonic()
-                    if end - start > 1:
-                        self.log.info(f"sync pipeline back-pressure. stalled {end - start:0.2f} seconds on add_block()")
-            except Exception:
-                self.log.exception("Exception validating")
-                raise
-            finally:
-                # finished signal with None
-                await output_queue.put(None)
+                    if start < timestamp:
+                        # rate limit ourselves, since we sent a message to
+                        # this peer too recently
+                        await asyncio.sleep(timestamp - start)
+                        start = time.monotonic()
 
-        async def ingest_blocks(
-            input_queue: asyncio.Queue[
-                tuple[WSChiaConnection, ValidationState, list[Awaitable[PreValidationResult]], list[FullBlock]] | None
-            ],
-        ) -> None:
-            nonlocal fork_info
-            block_rate = 0.0
-            block_rate_time = time.monotonic()
-            block_rate_height = -1
-            while True:
-                res = await input_queue.get()
-                if res is None:
-                    self.log.debug("done validating blocks")
-                    return None
-                peer, vs, futures, blocks = res
-                start_height = blocks[0].height
-                end_height = blocks[-1].height
+                    # update the timestamp, now that we're sending a request
+                    # it's OK for the timestamp to fall behind wall-clock
+                    # time. It just means we're allowed to send more
+                    # requests to catch up
+                    if is_localhost(peer.peer_info.host):
+                        # we don't apply rate limits to localhost, and our
+                        # tests depend on it
+                        bump = 0.1
+                    else:
+                        bump = seconds_per_request
 
-                if block_rate_height == -1:
-                    block_rate_height = start_height
-
-                pre_validation_results = list(await asyncio.gather(*futures))
-                # The ValidationState object (vs) is an in-out parameter. the add_block_batch()
-                # call will update it
-                state_change_summary, err = await self.add_prevalidated_blocks(
-                    blockchain,
-                    blocks,
-                    pre_validation_results,
-                    fork_info,
-                    peer.peer_info,
-                    vs,
-                )
-                if err is not None:
-                    await peer.close(CONSENSUS_ERROR_BAN_SECONDS)
-                    raise ValueError(f"Failed to validate block batch {start_height} to {end_height}: {err}")
-                if end_height - block_rate_height > 100:
-                    now = time.monotonic()
-                    block_rate = (end_height - block_rate_height) / (now - block_rate_time)
-                    block_rate_time = now
-                    block_rate_height = end_height
-
-                self.log.info(
-                    f"Added blocks {start_height} to {end_height} "
-                    f"({block_rate:.3g} blocks/s) (from: {peer.peer_info.ip})"
-                )
-                peak: BlockRecord | None = self.blockchain.get_peak()
-                if state_change_summary is not None:
-                    assert peak is not None
-                    # Hints must be added to the DB. The other post-processing tasks are not required when syncing
-                    hints_to_add, _ = get_hints_and_subscription_coin_ids(
-                        state_change_summary,
-                        self.subscriptions.has_coin_subscription,
-                        self.subscriptions.has_puzzle_subscription,
+                    new_peers_with_peak[idx] = (
+                        new_peers_with_peak[idx][0],
+                        new_peers_with_peak[idx][1] + bump,
                     )
-                    await self.hint_store.add_hints(hints_to_add)
-                # Note that end_height is not necessarily the peak at this
-                # point. In case of a re-org, it may even be significantly
-                # higher than _peak_height, and still not be the peak.
-                # clean_block_record() will not necessarily honor this cut-off
-                # height, in that case.
-                self.blockchain.clean_block_record(end_height - self.constants.BLOCKS_CACHE_SIZE)
+                    # the fewer peers we have, the more willing we should be
+                    # to wait for them.
+                    timeout = int(30 + 30 / len(new_peers_with_peak))
+                    response = await peer.call_api(FullNodeAPI.request_blocks, request, timeout=timeout)
+                    end = time.monotonic()
+                    if response is None:
+                        self.log.info(f"peer timed out after {end - start:.1f} s")
+                        await peer.close()
+                    elif isinstance(response, RespondBlocks):
+                        if end - start > 5:
+                            self.log.info(f"peer took {end - start:.1f} s to respond to request_blocks")
+                            # this isn't a great peer, reduce its priority
+                            # to prefer any peers that had to wait for it.
+                            # By setting the next allowed timestamp to now,
+                            # means that any other peer that has waited for
+                            # this will have its next allowed timestamp in
+                            # the passed, and be preferred multiple times
+                            # over this peer.
+                            new_peers_with_peak[idx] = (
+                                new_peers_with_peak[idx][0],
+                                end,
+                            )
+                        start = time.monotonic()
+                        yield (peer, response.blocks)
+                        end = time.monotonic()
+                        if end - start > 1:
+                            self.log.info(
+                                f"sync pipeline back-pressure. stalled {end - start:0.2f} seconds on prevalidate block"
+                            )
+                        fetched = True
+                        break
+                if fetched is False:
+                    self.log.error(f"failed fetching {start_height} to {end_height} from peers")
+                    return
+                if self.sync_store.peers_changed.is_set():
+                    existing_peers = {id(c): timestamp for c, timestamp in new_peers_with_peak}
+                    peers = self.get_peers_with_peak(peak_hash)
+                    new_peers_with_peak = [(c, existing_peers.get(id(c), end)) for c in peers]
+                    random.shuffle(new_peers_with_peak)
+                    self.sync_store.peers_changed.clear()
+                    self.log.info(f"peers with peak: {len(new_peers_with_peak)}")
 
-        block_queue: asyncio.Queue[tuple[WSChiaConnection, list[FullBlock]] | None] = asyncio.Queue(maxsize=10)
-        validation_queue: asyncio.Queue[
-            tuple[WSChiaConnection, ValidationState, list[Awaitable[PreValidationResult]], list[FullBlock]] | None
-        ] = asyncio.Queue(maxsize=10)
+        first_batch = True
+        vs = ValidationState(ssi, diff, prev_ses_block)
 
-        fetch_task = create_referenced_task(fetch_blocks(block_queue))
-        validate_task = create_referenced_task(validate_blocks(block_queue, validation_queue))
-        ingest_task = create_referenced_task(ingest_blocks(validation_queue))
+        async def validate_batch(
+            item: tuple[WSChiaConnection, list[FullBlock]],
+        ) -> tuple[WSChiaConnection, ValidationState, list[Awaitable[PreValidationResult]], list[FullBlock]] | None:
+            nonlocal first_batch, blockchain, fork_info
+            peer, blocks = item
+
+            # Keep the augmented chain's MMR snapshot aligned with the
+            # underlying blockchain before validating each fetched batch.
+            blockchain.mmr_manager = self.blockchain.mmr_manager.copy()
+
+            # skip_blocks is only relevant at the start of the sync,
+            # to skip blocks we already have in the database (and have
+            # been validated). Once we start validating blocks, we
+            # shouldn't be skipping any.
+            blocks_to_validate = await self.skip_blocks(blockchain, blocks, fork_info, vs)
+            assert first_batch or len(blocks_to_validate) == len(blocks)
+            next_validation_state = copy.copy(vs)
+
+            if len(blocks_to_validate) == 0:
+                return None
+
+            first_batch = False
+
+            futures: list[Awaitable[PreValidationResult]] = []
+            for block in blocks_to_validate:
+                futures.extend(
+                    await self.prevalidate_blocks(
+                        blockchain,
+                        [block],
+                        vs,
+                        summaries,
+                    )
+                )
+            return (peer, next_validation_state, list(futures), blocks_to_validate)
+
+        block_rate = 0.0
+        block_rate_time = time.monotonic()
+        block_rate_height = -1
+
+        async def ingest_batch(
+            item: tuple[WSChiaConnection, ValidationState, list[Awaitable[PreValidationResult]], list[FullBlock]],
+        ) -> None:
+            nonlocal fork_info, block_rate, block_rate_time, block_rate_height
+            peer, vs, futures, blocks = item
+            start_height = blocks[0].height
+            end_height = blocks[-1].height
+
+            if block_rate_height == -1:
+                block_rate_height = start_height
+
+            pre_validation_results = list(await asyncio.gather(*futures))
+            # The ValidationState object (vs) is an in-out parameter. the add_block_batch()
+            # call will update it
+            state_change_summary, err = await self.add_prevalidated_blocks(
+                blockchain,
+                blocks,
+                pre_validation_results,
+                fork_info,
+                peer.peer_info,
+                vs,
+            )
+            if err is not None:
+                await peer.close(CONSENSUS_ERROR_BAN_SECONDS)
+                raise ValueError(f"Failed to validate block batch {start_height} to {end_height}: {err}")
+            if end_height - block_rate_height > 100:
+                now = time.monotonic()
+                block_rate = (end_height - block_rate_height) / (now - block_rate_time)
+                block_rate_time = now
+                block_rate_height = end_height
+
+            self.log.info(
+                f"Added blocks {start_height} to {end_height} ({block_rate:.3g} blocks/s) (from: {peer.peer_info.ip})"
+            )
+            peak: BlockRecord | None = self.blockchain.get_peak()
+            if state_change_summary is not None:
+                assert peak is not None
+                # Hints must be added to the DB. The other post-processing tasks are not required when syncing
+                hints_to_add, _ = get_hints_and_subscription_coin_ids(
+                    state_change_summary,
+                    self.subscriptions.has_coin_subscription,
+                    self.subscriptions.has_puzzle_subscription,
+                )
+                await self.hint_store.add_hints(hints_to_add)
+            # Note that end_height is not necessarily the peak at this
+            # point. In case of a re-org, it may even be significantly
+            # higher than _peak_height, and still not be the peak.
+            # clean_block_record() will not necessarily honor this cut-off
+            # height, in that case.
+            self.blockchain.clean_block_record(end_height - self.constants.BLOCKS_CACHE_SIZE)
+
+        async def drain_pending_futures(p: TaskPipeline) -> None:
+            for item in p.drain(1):
+                _, _, futures, _ = item
+                await asyncio.gather(*futures)
+
+        pipeline = TaskPipeline(
+            source=fetch_blocks(),
+            stages=[validate_batch, ingest_batch],
+            queue_size=10,
+            names=["fetching", "validating", "ingesting"],
+            log=self.log,
+            cleanup=drain_pending_futures,
+        )
         try:
-            await asyncio.gather(fetch_task, validate_task, ingest_task)
+            await pipeline.run()
         except Exception:
             self.log.exception("sync from fork point failed")
-        finally:
-            cancel_task_safe(validate_task, self.log)
-            cancel_task_safe(fetch_task)
-            cancel_task_safe(ingest_task)
-
-            # we still need to await all the pending futures of the
-            # prevalidation steps posted to the thread pool
-            while not validation_queue.empty():
-                result = validation_queue.get_nowait()
-                if result is None:
-                    continue
-
-                _, _, futures, _ = result
-                await asyncio.gather(*futures)
 
     def get_peers_with_peak(self, peak_hash: bytes32) -> list[WSChiaConnection]:
         peer_ids: set[bytes32] = self.sync_store.get_peers_that_have_peak([peak_hash])

chia/util/task_pipeline.py

@@ -0,0 +1,217 @@
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import logging
+from collections.abc import AsyncIterator, Awaitable, Callable
+from typing import Any
+
+
+class TaskPipeline:
+    """
+    An N-stage async pipeline that connects an async-iterator source to a
+    chain of per-item async callables via bounded queues.
+
+    The source (an AsyncIterator) produces items that flow through each stage.
+    Each stage is called once per item and may return a value to forward
+    downstream, or None to skip/filter the item. The last stage is a pure
+    consumer (its return value is discarded).
+
+    The pipeline manages internal queues, back-pressure, and graceful shutdown
+    on failure (sentinel propagation, no task cancellation).
+    """
+
+    def __init__(
+        self,
+        source: AsyncIterator[Any],
+        stages: list[Callable[[Any], Awaitable[Any]]],
+        queue_size: int = 10,
+        names: list[str] | None = None,
+        log: logging.Logger | None = None,
+        cleanup: Callable[[TaskPipeline], Awaitable[None]] | None = None,
+    ) -> None:
+        if len(stages) == 0:
+            raise ValueError("pipeline requires at least one stage")
+        if names is not None and len(names) != len(stages) + 1:
+            raise ValueError("names must have len(stages) + 1 entries (source + each stage)")
+        self._source = source
+        self._stages = stages
+        self._queue_size = queue_size
+        self._names = names
+        self._log = log
+        self._cleanup = cleanup
+        self._queues: list[asyncio.Queue[Any]] = []
+        # Results produced by a stage after the pipeline began shutting down,
+        # which could not be enqueued. Recovered via drain().
+        self._dropped: list[list[Any]] = []
+        self._failed = asyncio.Event()
+        self._exception: BaseException | None = None
+
+    @property
+    def queues(self) -> list[asyncio.Queue[Any]]:
+        return self._queues
+
+    def drain(self, queue_index: int) -> list[Any]:
+        """Drain all non-sentinel items remaining in the queue at the given
+        index, plus any results that a stage produced but could not enqueue
+        because the pipeline was shutting down.
+
+        Only indices 1 .. len(stages)-1 are useful: _dropped[0] is never
+        written to (the feeder doesn't use it), and there is no queue after
+        the last stage."""
+        items: list[Any] = []
+        q = self._queues[queue_index]
+        while not q.empty():
+            item = q.get_nowait()
+            if item is not None:
+                items.append(item)
+        items.extend(self._dropped[queue_index])
+        self._dropped[queue_index].clear()
+        return items
+
+    async def run(self) -> None:
+        """Run the pipeline to completion. Re-raises the first stage exception."""
+        self._queues = [asyncio.Queue(maxsize=self._queue_size) for _ in range(len(self._stages))]
+        self._dropped = [[] for _ in range(len(self._stages))]
+        self._failed.clear()
+        self._exception = None
+
+        tasks = [asyncio.ensure_future(self._run_feeder())]
+        tasks.extend(asyncio.ensure_future(self._run_stage(i)) for i in range(len(self._stages)))
+
+        try:
+            all_done: asyncio.Task[Any] = asyncio.ensure_future(asyncio.gather(*tasks, return_exceptions=True))
+            fail_wait: asyncio.Task[Any] = asyncio.ensure_future(self._failed.wait())
+            try:
+                done, _ = await asyncio.wait({all_done, fail_wait}, return_when=asyncio.FIRST_COMPLETED)
+                if all_done not in done:
+                    for t in tasks:
+                        t.cancel()
+                    await asyncio.gather(*tasks, return_exceptions=True)
+            finally:
+                fail_wait.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await fail_wait
+                all_done.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await all_done
+        except asyncio.CancelledError:
+            self._failed.set()
+            for t in tasks:
+                t.cancel()
+            await asyncio.gather(*tasks, return_exceptions=True)
+            raise
+        finally:
+            if self._cleanup is not None:
+                await self._cleanup(self)
+
+        if self._exception is not None:
+            raise self._exception
+
+    async def _put_or_bail(self, queue: asyncio.Queue[Any], item: object) -> bool:
+        """Put item on queue. Returns False if the pipeline has failed (item discarded)."""
+        if self._failed.is_set():
+            return False
+        try:
+            queue.put_nowait(item)
+            return True
+        except asyncio.QueueFull:
+            pass
+        # Slow path: queue is full, race between put completing and pipeline failure.
+        put_task = asyncio.ensure_future(queue.put(item))
+        fail_task = asyncio.ensure_future(self._failed.wait())
+        try:
+            done, _ = await asyncio.wait({put_task, fail_task}, return_when=asyncio.FIRST_COMPLETED)
+            if put_task in done:
+                put_task.result()
+                return True
+            # Defensive: run() cancels tasks on failure before this
+            # await can observe fail_task completing, so this path
+            # is not reachable in practice.
+            else:  # pragma: no cover
+                put_task.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await put_task
+                return False
+        finally:
+            fail_task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await fail_task
+
+    async def _get_or_bail(self, queue: asyncio.Queue[Any]) -> tuple[bool, Any]:
+        """Get item from queue. Returns (False, None) if the pipeline has failed."""
+        if self._failed.is_set():
+            return False, None
+        try:
+            item = queue.get_nowait()
+            return True, item
+        except asyncio.QueueEmpty:
+            pass
+        # Slow path: queue is empty, race between get completing and pipeline failure.
+        get_task = asyncio.ensure_future(queue.get())
+        fail_task = asyncio.ensure_future(self._failed.wait())
+        try:
+            done, _ = await asyncio.wait({get_task, fail_task}, return_when=asyncio.FIRST_COMPLETED)
+            if get_task in done:
+                return True, get_task.result()
+            # Defensive: run() cancels tasks on failure before this
+            # await can observe fail_task completing, so this path
+            # is not reachable in practice.
+            else:  # pragma: no cover
+                get_task.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await get_task
+                return False, None
+        finally:
+            fail_task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await fail_task
+
+    async def _run_feeder(self) -> None:
+        output_queue = self._queues[0]
+        try:
+            async for item in self._source:
+                if not await self._put_or_bail(output_queue, item):
+                    return
+        except Exception as e:
+            if self._names is not None and self._log is not None:
+                self._log.exception("Exception %s", self._names[0])
+            if self._exception is None:
+                self._exception = e
+            self._failed.set()
+        finally:
+            # Propagate sentinel downstream. Uses _put_or_bail so that on
+            # normal completion it waits for space (downstream is alive), and
+            # on failure it bails immediately (downstream will see _failed).
+            await self._put_or_bail(output_queue, None)
+
+    async def _run_stage(self, idx: int) -> None:
+        input_queue = self._queues[idx]
+        output_queue = self._queues[idx + 1] if idx < len(self._stages) - 1 else None
+        try:
+            while True:
+                ok, item = await self._get_or_bail(input_queue)
+                if not ok or item is None:
+                    return
+                result = await self._stages[idx](item)
+                if self._failed.is_set():
+                    if output_queue is not None and result is not None:
+                        self._dropped[idx + 1].append(result)
+                    return
+                if output_queue is None or result is None:
+                    continue
+                # Defensive: run() cancels tasks on failure before
+                # _put_or_bail can return False from its slow path,
+                # so this branch is not reachable in practice.
+                if not await self._put_or_bail(output_queue, result):  # pragma: no cover
+                    self._dropped[idx + 1].append(result)
+                    return
+        except Exception as e:
+            if self._names is not None and self._log is not None:
+                self._log.exception("Exception %s", self._names[idx + 1])
+            if self._exception is None:
+                self._exception = e
+            self._failed.set()
+        finally:
+            if output_queue is not None:
+                await self._put_or_bail(output_queue, None)