Add timeout to streaming shard snapshot writer (#9239)

* Add block level timeout to snapshot stream writer

* Add tooling to exercise streaming snapshot stalls

- tests/manual/slow_snapshot_download.py: slowly / partially download a
  streaming shard snapshot from a URL to exercise sender-side backpressure.
  Supports hold / rst / fin / blackhole termination to simulate a stalled,
  killed, or offline (network-partitioned) consumer against a remote node.
  Stdlib only; read-only against the target.
- tests/consensus_tests/test_streaming_snapshot_receiver_kill.py: throttled
  receiver killed mid-flight + a second receiver, to observe whether the
  sender releases the SegmentHolder lock and recovers.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: timvisee <tim@visee.me>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

lib/collection/src/common/mod.rs

@@ -12,4 +12,5 @@ pub mod snapshot_stream;
 pub mod snapshots_manager;
 pub mod stoppable_task;
 pub mod stoppable_task_async;
+pub mod timeout_writer;
 pub mod transpose_iterator;

lib/collection/src/common/timeout_writer.rs

@@ -0,0 +1,155 @@
+use std::future::Future;
+use std::io;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+use std::time::Duration;
+
+use tokio::io::AsyncWrite;
+use tokio::time::{Sleep, sleep};
+
+/// An [`AsyncWrite`] adapter that aborts a write when it makes no progress for a
+/// given idle duration.
+///
+/// This guards the writing end of a streamed snapshot. The writer is driven by
+/// a remote consumer through a small in-memory pipe ([`tokio::io::duplex`]); if
+/// that consumer stops reading, the inner write stays [`Poll::Pending`]
+/// indefinitely. Because the snapshot holds a read lock on the shard's segment
+/// holder (and occupies a blocking thread) while streaming, a stalled consumer
+/// would otherwise wedge the whole shard and leak the thread forever.
+///
+/// When no bytes can be handed to the inner writer for `timeout`, the next poll
+/// resolves with an [`io::ErrorKind::TimedOut`] error. That error unwinds the
+/// snapshot, releases the segment holder lock, and frees the blocking thread.
+///
+/// The timer measures *idle* time only: it is armed while the inner writer is
+/// not ready and reset every time the inner writer accepts data or flushes. So
+/// the timeout fires only after `timeout` of continuous back-pressure,
+/// regardless of how long the overall transfer takes.
+pub struct TimeoutWriter<W> {
+    inner: W,
+    timeout: Duration,
+    /// Idle timer, armed lazily while the inner writer is not ready and cleared
+    /// as soon as it makes progress.
+    idle: Option<Pin<Box<Sleep>>>,
+}
+
+impl<W> TimeoutWriter<W> {
+    pub fn new(inner: W, timeout: Duration) -> Self {
+        Self {
+            inner,
+            timeout,
+            idle: None,
+        }
+    }
+
+    /// Run `inner_poll`, arming/resetting the idle timer around it.
+    ///
+    /// Resolves with [`io::ErrorKind::TimedOut`] if the inner poll stays
+    /// [`Poll::Pending`] for longer than `timeout` without making progress.
+    fn poll_with_timeout<T>(
+        &mut self,
+        cx: &mut Context<'_>,
+        inner_poll: impl FnOnce(Pin<&mut W>, &mut Context<'_>) -> Poll<io::Result<T>>,
+    ) -> Poll<io::Result<T>>
+    where
+        W: Unpin,
+    {
+        match inner_poll(Pin::new(&mut self.inner), cx) {
+            // Made progress (or failed outright): disarm the idle timer.
+            Poll::Ready(result) => {
+                self.idle = None;
+                Poll::Ready(result)
+            }
+            // No progress: arm the idle timer (keeping any existing deadline) and
+            // abort if it has elapsed.
+            Poll::Pending => {
+                let timeout = self.timeout;
+                let idle = self.idle.get_or_insert_with(|| Box::pin(sleep(timeout)));
+                match idle.as_mut().poll(cx) {
+                    Poll::Ready(()) => {
+                        self.idle = None;
+                        Poll::Ready(Err(io::Error::new(
+                            io::ErrorKind::TimedOut,
+                            format!("write stalled: no data was read for {}s", timeout.as_secs()),
+                        )))
+                    }
+                    Poll::Pending => Poll::Pending,
+                }
+            }
+        }
+    }
+}
+
+impl<W: AsyncWrite + Unpin> AsyncWrite for TimeoutWriter<W> {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<io::Result<usize>> {
+        self.get_mut()
+            .poll_with_timeout(cx, |inner, cx| inner.poll_write(cx, buf))
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        self.get_mut()
+            .poll_with_timeout(cx, |inner, cx| inner.poll_flush(cx))
+    }
+
+    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        self.get_mut()
+            .poll_with_timeout(cx, |inner, cx| inner.poll_shutdown(cx))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
+
+    use super::*;
+
+    /// A stalled consumer must abort the write after the idle timeout. Our
+    /// `sleep` is the only timer in play, so a short real timeout is enough to
+    /// keep this deterministic.
+    #[tokio::test]
+    async fn aborts_write_when_consumer_stalls() {
+        // Keep `read_half` alive but never read from it, so the pipe fills up and
+        // stays full (dropping it would yield a `BrokenPipe` error instead).
+        let (_read_half, write_half) = tokio::io::duplex(8);
+        let mut writer = TimeoutWriter::new(write_half, Duration::from_millis(100));
+
+        // First chunk fits in the buffer.
+        writer.write_all(&[0u8; 8]).await.unwrap();
+
+        // The buffer is now full and nobody drains it: the next write stalls and
+        // must abort once the idle timeout elapses.
+        let err = writer.write_all(&[0u8; 8]).await.unwrap_err();
+        assert_eq!(err.kind(), io::ErrorKind::TimedOut);
+    }
+
+    /// A consumer that keeps reading must never trip the idle timeout, no matter
+    /// how small the pipe buffer is relative to the payload. The timeout is set
+    /// generously so it cannot fire during the near-instant transfer.
+    #[tokio::test]
+    async fn does_not_abort_while_consumer_reads() {
+        let (mut read_half, write_half) = tokio::io::duplex(8);
+        let mut writer = TimeoutWriter::new(write_half, Duration::from_secs(30));
+
+        let reader = tokio::spawn(async move {
+            let mut buf = [0u8; 16];
+            let mut total = 0;
+            while let Ok(n) = read_half.read(&mut buf).await {
+                if n == 0 {
+                    break;
+                }
+                total += n;
+            }
+            total
+        });
+
+        writer.write_all(&[0u8; 64]).await.unwrap();
+        writer.shutdown().await.unwrap();
+        drop(writer);
+
+        assert_eq!(reader.await.unwrap(), 64);
+    }
+}

lib/collection/src/shards/shard_holder/mod.rs

@@ -7,6 +7,7 @@ use std::collections::{BTreeMap, HashMap, HashSet};
 use std::ops::Deref as _;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
+use std::time::Duration;
 
 use ahash::AHashMap;
 use api::rest::ShardKeyWithFallback;
@@ -39,6 +40,7 @@ use crate::collection::payload_index_schema::PayloadIndexSchema;
 use crate::common::adaptive_handle::AdaptiveSearchHandle;
 use crate::common::collection_size_stats::CollectionSizeStats;
 use crate::common::snapshot_stream::SnapshotStream;
+use crate::common::timeout_writer::TimeoutWriter;
 use crate::config::{CollectionConfigInternal, ShardingMethod};
 use crate::hash_ring::HashRingRouter;
 use crate::operations::cluster_ops::ReshardingDirection;
@@ -65,6 +67,13 @@ const SHARD_TRANSFERS_FILE: &str = "shard_transfers";
 const RESHARDING_STATE_FILE: &str = "resharding_state.json";
 pub const SHARD_KEY_MAPPING_FILE: &str = "shard_key_mapping.json";
 
+/// Abort a streamed snapshot if its consumer reads no data for this long.
+///
+/// The streaming write holds a read lock on the shard's segment holder and
+/// occupies a blocking thread, so a stalled consumer must not be allowed to
+/// block it forever. See [`TimeoutWriter`].
+const SNAPSHOT_STREAM_WRITE_IDLE_TIMEOUT: Duration = Duration::from_secs(5 * 60);
+
 pub struct ShardHolder {
     /// `BTreeMap` for deterministic iteration order by `ShardId` — iteration is externally
     /// observable (fan-out, telemetry, consensus state apply).
@@ -1183,6 +1192,13 @@ impl ShardHolder {
 
         let (read_half, write_half) = tokio::io::duplex(4096);
 
+        // Abort the snapshot if the consumer stops draining the stream. This
+        // write holds a read lock on the shard's segment holder and occupies a
+        // blocking thread; without a timeout, a stalled consumer (e.g. a slow or
+        // hung client) would block the segment holder and leak the thread
+        // indefinitely, wedging the whole shard.
+        let write_half = TimeoutWriter::new(write_half, SNAPSHOT_STREAM_WRITE_IDLE_TIMEOUT);
+
         let tar = BuilderExt::new_streaming_owned(SyncIoBridge::new(write_half));
 
         let snapshot_creator = shard

tests/consensus_tests/test_streaming_snapshot_receiver_kill.py

@@ -0,0 +1,213 @@
+"""
+Observe how the *sender* of a streaming shard snapshot reacts when a throttled
+receiver is killed mid-flight and a second receiver then starts.
+
+Scenario (single node = the sender):
+  1. Create a collection with enough data that the snapshot tar is many MB,
+     so a slow reader cannot drain it and the sender ends up parked mid-write.
+  2. Receiver #1 pulls `GET /collections/{c}/shards/0/snapshot` at a throttled
+     speed (read a chunk, sleep). This fills the sender's 4 KB duplex + socket
+     buffers and parks the snapshot writer **while it holds the SegmentHolder
+     upgradable-read lock** (see `proxy_all_segments_and_apply`).
+  3. Kill receiver #1 mid-flight (like a receiver peer dying).
+  4. Start receiver #2 and watch what the sender does.
+
+What we observe ("how the sender reacts"):
+  - Does the node stay responsive (`/`, `/cluster`)?  -> hard assert
+  - Does receiver #2 complete its download?           -> the key signal:
+      it can only finish if the sender RELEASED the SegmentHolder lock that
+      receiver #1 was pinning. If the sender is wedged, #2 streams 200 OK but
+      then never receives bytes (its snapshot blocks on `upgradable_read`).
+  - When/whether the sender logged "broken pipe" / "Failed to stream shard
+    snapshot" for receiver #1.
+
+IMPORTANT — localhost caveat:
+  Killing a local process makes the kernel send a prompt RST, so the sender's
+  actix usually notices the disconnect quickly and recovers. The 18h
+  production hang needed the RST to NOT be delivered (half-open connection via
+  a cloud LB / NAT). To reproduce *that*, route the download through a
+  throttling TCP proxy and, on "kill", drop the receiver<->proxy side while
+  keeping the proxy<->sender socket open. This test deliberately uses the
+  simple process-kill so it documents the baseline behavior; the proxy variant
+  is noted at the bottom.
+
+Usage:
+    pytest tests/consensus_tests/test_streaming_snapshot_receiver_kill.py -s -v
+"""
+
+import multiprocessing
+import pathlib
+import time
+
+import requests
+
+from .fixtures import create_collection, upsert_random_points
+from .utils import *
+
+COLLECTION = "test_collection"
+NUM_POINTS = 10_000
+SHARD_ID = 0
+
+# Throttle: ~20 KB/s, far slower than the sender can produce -> backpressure.
+THROTTLE_CHUNK = 4096
+THROTTLE_SLEEP = 0.2
+
+# How long to let backpressure build before killing receiver #1.
+BACKPRESSURE_BUILD_SECS = 2
+
+# How long we give the sender to recover (receiver #2 to finish) after the kill.
+RECOVER_TIMEOUT_SECS = 120
+
+RESPONSIVE_TIMEOUT = 10
+
+
+def _snapshot_url(peer_url: str) -> str:
+    return f"{peer_url}/collections/{COLLECTION}/shards/{SHARD_ID}/snapshot"
+
+
+def throttled_receiver(peer_url: str, ready_path: str):
+    """Receiver #1: pull the snapshot slowly, forever. Run as a separate OS
+    process so the test can `.kill()` it and drop its socket like a dead peer."""
+    try:
+        r = requests.get(_snapshot_url(peer_url), stream=True, timeout=300)
+        if r.status_code != 200:
+            print(f"[recv1] unexpected status {r.status_code}")
+            return
+        # Signal that the stream actually started (sender is now producing/parked).
+        pathlib.Path(ready_path).write_text("streaming")
+        total = 0
+        for chunk in r.iter_content(chunk_size=THROTTLE_CHUNK):
+            total += len(chunk)
+            time.sleep(THROTTLE_SLEEP)
+    except Exception as e:  # noqa: BLE001 - diagnostic process
+        print(f"[recv1] ended: {e}")
+
+
+def full_receiver(peer_url: str, ready_path: str, done_path: str):
+    """Receiver #2: download to completion as fast as possible and record the
+    result. Completing proves the sender released the lock receiver #1 held."""
+    try:
+        start = time.time()
+        r = requests.get(_snapshot_url(peer_url), stream=True, timeout=RECOVER_TIMEOUT_SECS)
+        if r.status_code != 200:
+            pathlib.Path(done_path).write_text(f"status={r.status_code}")
+            return
+        pathlib.Path(ready_path).write_text("streaming")
+        total = 0
+        for chunk in r.iter_content(chunk_size=1 << 20):
+            total += len(chunk)
+        elapsed = time.time() - start
+        pathlib.Path(done_path).write_text(f"ok bytes={total} secs={elapsed:.1f}")
+    except Exception as e:  # noqa: BLE001 - diagnostic process
+        pathlib.Path(done_path).write_text(f"error={e}")
+
+
+def _responsive(peer_url: str) -> bool:
+    try:
+        requests.get(f"{peer_url}/", timeout=RESPONSIVE_TIMEOUT)
+        requests.get(f"{peer_url}/collections/{COLLECTION}/cluster", timeout=RESPONSIVE_TIMEOUT)
+        return True
+    except requests.exceptions.RequestException:
+        return False
+
+
+def _sender_log_text() -> str:
+    log_path = pathlib.Path(init_pytest_log_folder()) / "peer_0_0.log"
+    try:
+        return log_path.read_text(errors="replace")
+    except OSError:
+        return ""
+
+
+def _wait_for_file(path: pathlib.Path, timeout: float) -> bool:
+    deadline = time.time() + timeout
+    while time.time() < deadline:
+        if path.exists():
+            return True
+        time.sleep(0.1)
+    return False
+
+
+def test_streaming_snapshot_receiver_kill(tmp_path: pathlib.Path):
+    assert_project_root()
+
+    peer_urls, _, _ = start_cluster(tmp_path, 1)
+    peer_url = peer_urls[0]
+
+    create_collection(peer_url, collection=COLLECTION)
+    upsert_random_points(peer_url, NUM_POINTS, collection_name=COLLECTION)
+    print(f"[test] inserted {NUM_POINTS} points")
+
+    ready1 = tmp_path / "recv1_ready"
+    ready2 = tmp_path / "recv2_ready"
+    done2 = tmp_path / "recv2_done"
+
+    # --- Receiver #1: throttled, killed mid-flight ---------------------------
+    recv1 = multiprocessing.Process(
+        target=throttled_receiver, args=(peer_url, str(ready1)), daemon=True
+    )
+    recv1.start()
+
+    assert _wait_for_file(ready1, timeout=30), "receiver #1 never started streaming"
+    print("[test] receiver #1 is streaming (sender now parked holding the lock)")
+
+    # Let the sender fill its buffers and park mid-write.
+    time.sleep(BACKPRESSURE_BUILD_SECS)
+    assert _responsive(peer_url), "node already unresponsive before the kill"
+
+    recv1.kill()
+    recv1.join(timeout=10)
+    kill_time = time.time()
+    print("[test] receiver #1 KILLED mid-flight")
+
+    # --- Receiver #2: the "another receiver" --------------------------------
+    recv2 = multiprocessing.Process(
+        target=full_receiver, args=(peer_url, str(ready2), str(done2)), daemon=True
+    )
+    recv2.start()
+
+    # Observe the sender's reaction.
+    recovered = _wait_for_file(done2, timeout=RECOVER_TIMEOUT_SECS)
+    recv2.join(timeout=5)
+
+    # The node must stay responsive the whole time (blast-radius containment).
+    assert _responsive(peer_url), "node became unresponsive after receiver kill"
+
+    log = _sender_log_text()
+    noticed_broken_pipe = "broken pipe" in log or "Failed to stream shard snapshot" in log
+
+    print("\n==================== SENDER REACTION ====================")
+    print(f"  node responsive after kill : True")
+    print(f"  sender logged disconnect   : {noticed_broken_pipe}")
+    if recovered:
+        print(f"  receiver #2 result         : {done2.read_text()}")
+        print(f"  time kill -> #2 done       : {time.time() - kill_time:.1f}s")
+    else:
+        print(f"  receiver #2 result         : DID NOT FINISH within {RECOVER_TIMEOUT_SECS}s")
+        print(f"  -> sender appears WEDGED: the lock from receiver #1 was not released")
+    print("=========================================================\n")
+
+    # Key assertion: a second receiver must be able to snapshot, which is only
+    # possible if the sender released the SegmentHolder lock held by receiver #1.
+    assert recovered, (
+        f"sender did not recover within {RECOVER_TIMEOUT_SECS}s after receiver #1 "
+        f"was killed: receiver #2 could not obtain the snapshot stream"
+    )
+    assert done2.read_text().startswith("ok"), f"receiver #2 failed: {done2.read_text()}"
+
+
+# -----------------------------------------------------------------------------
+# Reproducing the TRUE half-open hang (optional, manual):
+#
+# Replace the direct download URL with a tiny throttling TCP proxy in front of
+# the sender's REST port:
+#
+#   downloader  ──>  proxy (throttles, you control)  ──>  sender:REST
+#
+# "Kill the receiver" then means: close the downloader<->proxy socket but KEEP
+# the proxy<->sender socket open and stop forwarding. The sender never receives
+# an RST, its write stays parked, the SegmentHolder lock is held indefinitely,
+# and receiver #2 hangs -> `recovered` stays False. That is the production
+# deadlock, and it is exactly what a sender-side inactivity write-timeout
+# (TimeoutWriter) would bound.
+# -----------------------------------------------------------------------------

tests/manual/slow_snapshot_download.py

@@ -0,0 +1,346 @@
+#!/usr/bin/env python3
+"""
+Slowly download a streaming shard snapshot from a given URL, to manually
+reproduce sender-side backpressure against a *remote* node (where, unlike
+localhost, a killed/stalled consumer does not produce a prompt RST).
+
+The sender serves the shard snapshot at:
+    GET /collections/{collection}/shards/{shard}/snapshot
+
+While this client throttles (or stops) reading, the sender's snapshot writer
+parks on the 4 KiB duplex while holding the SegmentHolder upgradable-read lock.
+That lets you watch, on the remote node, whether/when it logs a broken pipe,
+releases the lock, and whether concurrent ops / other transfers wedge.
+
+Examples
+--------
+# Drip-read at ~20 KiB/s forever (Ctrl-C to stop):
+  ./slow_snapshot_download.py http://NODE:6333/collections/c/shards/1/snapshot
+
+# Read the first 1 MiB, then HOLD the connection open without reading
+# (simulates a stalled/half-open receiver — connection stays ESTABLISHED):
+  ./slow_snapshot_download.py URL --read-bytes 1MiB --then hold
+
+# Read the first 1 MiB, then abruptly RST the socket
+# (simulates the client process being kill -9'd / crashing):
+  ./slow_snapshot_download.py URL --read-bytes 1MiB --then rst
+
+# Read the first 1 MiB, then BLACKHOLE the connection: drop all packets so the
+# sender gets no FIN/RST at all (simulates the machine suddenly going offline /
+# a network partition). Needs root or passwordless sudo for iptables; if it
+# can't, it prints the exact commands to run by hand. THIS is the 18h-hang case.
+  sudo ./slow_snapshot_download.py URL --read-bytes 1MiB --then blackhole
+
+# 8 KiB every 0.5s, with an API key, ignore TLS verification:
+  ./slow_snapshot_download.py URL --chunk 8KiB --delay 0.5 --api-key KEY --insecure
+
+# Open N concurrent slow downloads to pile up streams on the sender:
+  ./slow_snapshot_download.py URL --concurrency 10 --then hold
+
+Uses only the Python standard library (no requests), so it runs anywhere.
+"""
+
+import argparse
+import socket
+import ssl
+import sys
+import threading
+import time
+from urllib.parse import urlsplit
+
+
+def parse_size(s: str) -> int:
+    """Parse sizes like '4096', '8KiB', '1MiB', '2MB'."""
+    s = s.strip()
+    units = {
+        "": 1,
+        "b": 1,
+        "kib": 1024,
+        "kb": 1000,
+        "k": 1024,
+        "mib": 1024**2,
+        "mb": 1000**2,
+        "m": 1024**2,
+        "gib": 1024**3,
+        "gb": 1000**3,
+        "g": 1024**3,
+    }
+    num = s
+    unit = ""
+    for i, ch in enumerate(s):
+        if not (ch.isdigit() or ch == "."):
+            num, unit = s[:i], s[i:].strip().lower()
+            break
+    if unit not in units:
+        raise argparse.ArgumentTypeError(f"unknown size unit in {s!r}")
+    return int(float(num) * units[unit])
+
+
+def human(n: int) -> str:
+    for unit in ("B", "KiB", "MiB", "GiB"):
+        if n < 1024 or unit == "GiB":
+            return f"{n:.1f}{unit}" if unit != "B" else f"{n}B"
+        n /= 1024
+    return f"{n}"
+
+
+def open_socket(url: str, api_key: str | None, insecure: bool, timeout: float) -> socket.socket:
+    """Send the GET request and return a connected socket positioned at the body."""
+    parts = urlsplit(url)
+    host = parts.hostname
+    if host is None:
+        raise ValueError(f"no host in URL {url!r}")
+    is_tls = parts.scheme == "https"
+    port = parts.port or (443 if is_tls else 80)
+    path = parts.path or "/"
+    if parts.query:
+        path += "?" + parts.query
+
+    raw = socket.create_connection((host, port), timeout=timeout)
+    if is_tls:
+        ctx = ssl.create_default_context()
+        if insecure:
+            ctx.check_hostname = False
+            ctx.verify_mode = ssl.CERT_NONE
+        sock = ctx.wrap_socket(raw, server_hostname=host)
+    else:
+        sock = raw
+
+    req = [
+        f"GET {path} HTTP/1.1",
+        f"Host: {parts.netloc}",
+        "User-Agent: slow-snapshot-download",
+        "Accept: */*",
+        "Connection: close",
+    ]
+    if api_key:
+        req.append(f"Api-Key: {api_key}")
+    req.append("")
+    req.append("")
+    sock.sendall("\r\n".join(req).encode())
+    return sock
+
+
+def drain_headers(sock: socket.socket, tag: str) -> bytes:
+    """Read up to the end of HTTP headers; return any leftover body bytes."""
+    buf = b""
+    while b"\r\n\r\n" not in buf:
+        chunk = sock.recv(4096)
+        if not chunk:
+            raise ConnectionError("connection closed before headers completed")
+        buf += chunk
+    head, _, rest = buf.partition(b"\r\n\r\n")
+    status_line = head.split(b"\r\n", 1)[0].decode(errors="replace")
+    print(f"[{tag}] {status_line}")
+    if " 200 " not in status_line:
+        # Print a few header lines for context on non-200.
+        for line in head.split(b"\r\n")[1:8]:
+            print(f"[{tag}]   {line.decode(errors='replace')}")
+    return rest
+
+
+def download(idx: int, args) -> None:
+    tag = f"dl{idx}"
+    started = time.time()
+    total = 0
+    try:
+        sock = open_socket(args.url, args.api_key, args.insecure, args.connect_timeout)
+    except Exception as e:  # noqa: BLE001
+        print(f"[{tag}] connect failed: {e}")
+        return
+
+    try:
+        leftover = drain_headers(sock, tag)
+        total += len(leftover)
+
+        target = args.read_bytes  # None => read forever (throttled)
+        last_report = started
+        # Use a recv timeout so we can periodically report even when the
+        # server stops sending (backpressure -> our recv blocks).
+        sock.settimeout(args.delay if args.delay > 0 else 5.0)
+
+        while target is None or total < target:
+            if args.delay > 0:
+                time.sleep(args.delay)
+            try:
+                chunk = sock.recv(args.chunk)
+            except socket.timeout:
+                # No data within delay window — server is backpressured/stalled.
+                now = time.time()
+                if now - last_report >= 1.0:
+                    print(f"[{tag}] (no data) total={human(total)} elapsed={now-started:.1f}s")
+                    last_report = now
+                continue
+            if not chunk:
+                print(f"[{tag}] server closed stream: total={human(total)} "
+                      f"elapsed={time.time()-started:.1f}s")
+                return
+            total += len(chunk)
+            now = time.time()
+            if now - last_report >= 1.0:
+                rate = total / max(now - started, 1e-9)
+                print(f"[{tag}] read={human(total)} rate={human(int(rate))}/s "
+                      f"elapsed={now-started:.1f}s")
+                last_report = now
+
+        # Reached the read target: decide what to do next.
+        print(f"[{tag}] reached read target {human(total)} after {time.time()-started:.1f}s "
+              f"-> action={args.then}")
+        if args.then in ("rst", "close"):
+            # Abortive close: kernel sends RST, no graceful FIN.
+            # This is what a kill -9 / process crash looks like to the sender.
+            try:
+                sock.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, _linger_zero())
+            except OSError:
+                pass
+            sock.close()
+            print(f"[{tag}] socket RST (abortive close) — like kill -9 / crash")
+        elif args.then == "fin":
+            sock.close()
+            print(f"[{tag}] socket closed (graceful FIN)")
+        elif args.then == "blackhole":
+            # Drop all packets on this connection so the sender receives no
+            # FIN/RST — as if the machine vanished / the network partitioned.
+            cleanup = _blackhole(sock, tag)
+            print(f"[{tag}] connection BLACKHOLED, sender gets silence. "
+                  f"Watch the remote node's logs. Ctrl-C to restore.")
+            try:
+                while True:
+                    time.sleep(3600)
+            finally:
+                if cleanup:
+                    cleanup()
+        else:  # hold
+            print(f"[{tag}] HOLDING connection open (still ESTABLISHED), not reading. "
+                  f"Watch the remote node's logs. Ctrl-C to release.")
+            while True:
+                time.sleep(3600)
+    except KeyboardInterrupt:
+        print(f"[{tag}] interrupted, closing")
+    finally:
+        try:
+            sock.close()
+        except OSError:
+            pass
+
+
+def _linger_zero() -> bytes:
+    import struct
+    return struct.pack("ii", 1, 0)
+
+
+def _blackhole(sock: socket.socket, tag: str):
+    """Drop every packet on this TCP connection so the remote sender receives
+    no FIN/RST — simulating the machine suddenly going offline / a network
+    partition. Returns a cleanup callable (or None if rules couldn't be set).
+
+    Linux only; needs root or passwordless sudo. Falls back to printing the
+    exact iptables commands so they can be run by hand.
+    """
+    import os
+    import shutil
+    import subprocess
+
+    try:
+        laddr = sock.getsockname()
+        raddr = sock.getpeername()
+    except OSError as e:
+        print(f"[{tag}] cannot read socket addresses for blackhole: {e}")
+        return None
+
+    lport = laddr[1]
+    rip, rport = raddr[0], raddr[1]
+    ipt = "ip6tables" if sock.family == socket.AF_INET6 else "iptables"
+
+    out_match = ["-p", "tcp", "-d", rip, "--dport", str(rport), "--sport", str(lport), "-j", "DROP"]
+    in_match = ["-p", "tcp", "-s", rip, "--sport", str(rport), "--dport", str(lport), "-j", "DROP"]
+    adds = [[ipt, "-I", "OUTPUT", *out_match], [ipt, "-I", "INPUT", *in_match]]
+    dels = [[ipt, "-D", "OUTPUT", *out_match], [ipt, "-D", "INPUT", *in_match]]
+
+    prefix: list[str] = []
+    if hasattr(os, "geteuid") and os.geteuid() != 0:
+        if shutil.which("sudo"):
+            prefix = ["sudo", "-n"]
+
+    def show_manual():
+        print(f"[{tag}] run these to blackhole, and the matching -D to restore:")
+        for cmd in adds:
+            print("    sudo " + " ".join(cmd))
+
+    if shutil.which(ipt) is None:
+        print(f"[{tag}] {ipt} not found (Linux only); cannot blackhole automatically.")
+        show_manual()
+        return None
+
+    for cmd in adds:
+        res = subprocess.run(prefix + cmd, capture_output=True, text=True)
+        if res.returncode != 0:
+            print(f"[{tag}] failed to install rule ({res.stderr.strip() or 'permission denied'}).")
+            show_manual()
+            # Best-effort: try to undo whatever we managed to add.
+            for d in dels:
+                subprocess.run(prefix + d, capture_output=True, text=True)
+            return None
+
+    print(f"[{tag}] BLACKHOLE active: all packets to/from {rip}:{rport} (our port {lport}) dropped.")
+
+    def cleanup():
+        for d in dels:
+            subprocess.run(prefix + d, capture_output=True, text=True)
+        print(f"[{tag}] blackhole rules removed")
+
+    return cleanup
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(
+        description="Slowly download a streaming shard snapshot to test sender backpressure.",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=__doc__,
+    )
+    p.add_argument("url", help="snapshot URL, e.g. http://NODE:6333/collections/c/shards/1/snapshot")
+    p.add_argument("--chunk", type=parse_size, default=4096,
+                   help="bytes per recv() (default 4KiB)")
+    p.add_argument("--delay", type=float, default=0.2,
+                   help="seconds to sleep between reads (default 0.2 => ~20KiB/s)")
+    p.add_argument("--read-bytes", type=parse_size, default=None,
+                   help="stop reading after this many bytes, then run --then "
+                        "(default: read forever, throttled)")
+    p.add_argument("--then", choices=("hold", "rst", "close", "fin", "blackhole"), default="hold",
+                   help="after --read-bytes: "
+                        "'hold' keep connection ESTABLISHED w/o reading (stalled/half-open); "
+                        "'rst' (alias 'close') abortive RST, like the client kill -9'd / crashed; "
+                        "'fin' graceful close; "
+                        "'blackhole' drop all packets so the sender gets no FIN/RST, like the "
+                        "machine went offline (Linux, needs root/sudo for iptables). "
+                        "(default hold)")
+    p.add_argument("--concurrency", type=int, default=1,
+                   help="number of parallel downloads to open (default 1)")
+    p.add_argument("--api-key", default=None, help="value for the Api-Key header")
+    p.add_argument("--insecure", action="store_true", help="skip TLS certificate verification")
+    p.add_argument("--connect-timeout", type=float, default=10.0,
+                   help="TCP connect timeout seconds (default 10)")
+    args = p.parse_args()
+
+    print(f"target: {args.url}")
+    print(f"chunk={human(args.chunk)} delay={args.delay}s "
+          f"read_bytes={human(args.read_bytes) if args.read_bytes else 'unlimited'} "
+          f"then={args.then} concurrency={args.concurrency}")
+
+    if args.concurrency == 1:
+        download(0, args)
+        return
+
+    threads = [threading.Thread(target=download, args=(i, args), daemon=True)
+               for i in range(args.concurrency)]
+    for t in threads:
+        t.start()
+    try:
+        while any(t.is_alive() for t in threads):
+            time.sleep(0.5)
+    except KeyboardInterrupt:
+        print("interrupted, exiting (daemon downloads will be torn down)")
+
+
+if __name__ == "__main__":
+    main()