Bound queuing thread with hard limit instead of soft

Author: harshil-goel

src/bin/stream.rs

@@ -200,9 +200,6 @@ struct Args {
     /// `restore_command` reads from here
     #[arg(long)]
     out_dir: PathBuf,
-    /// Optional permanent physical slot name on source PG
-    #[arg(long)]
-    slot: Option<String>,
     /// Start LSN in `X/Y` hex form. Defaults to source's current
     /// `pg_current_wal_lsn` (per `IDENTIFY_SYSTEM`), aligned down to a
     /// segment boundary.
@@ -518,6 +515,8 @@ async fn run(args: Args) -> Result<()> {
     } else {
         None
     };
+    // Physical replication slot from config (`[source] slot`); None = slotless.
+    let source_slot: Option<String> = ch_config.as_ref().and_then(|c| c.source_slot.clone());
     let bootstrap_end_lsn: Option<u64> = if matches!(args.bootstrap_mode, BootstrapMode::Off) {
         None
     } else {
@@ -689,6 +688,14 @@ async fn run(args: Args) -> Result<()> {
         .await
         .set_invalidation_epoch(invalidation_epoch.clone());
 
+    // Create the configured slot before preflight, which requires it to exist.
+    if let Some(slot) = source_slot.as_deref() {
+        feed.ensure_physical_slot(slot)
+            .await
+            .with_context(|| format!("ensure physical replication slot {slot}"))?;
+        tracing::info!(target: "walshadow", slot, "physical replication slot ready");
+    }
+
     // Pre-flight validators run after both source + shadow SQL clients
     // are up so every check has its connection.
     if !args.skip_preflight {
@@ -708,7 +715,7 @@ async fn run(args: Args) -> Result<()> {
             source_version_num,
             source_sql,
             shadow_sql: &shadow_sql,
-            slot: args.slot.as_deref(),
+            slot: source_slot.as_deref(),
             ch_config: ch_config.as_ref(),
         })
         .await
@@ -766,7 +773,7 @@ async fn run(args: Args) -> Result<()> {
         }
     };
 
-    feed.start_physical_replication(args.slot.as_deref(), aligned, ident.timeline)
+    feed.start_physical_replication(source_slot.as_deref(), aligned, ident.timeline)
         .await
         .context("START_REPLICATION")?;
     // Spill dir wiped every startup: cursor file commits drains
@@ -1194,9 +1201,12 @@ async fn run(args: Args) -> Result<()> {
                 .context("write resume cursor")?;
             last_cursor_write = Some(Instant::now());
         }
+        // flush drives a physical slot's restart_lsn (what the source retains),
+        // so cap it at apply_ceiling: the source keeps WAL until CH has durably
+        // applied it (and shadow replayed), never recycling un-consumed WAL.
         let status = StandbyStatus {
             write_lsn: received,
-            flush_lsn: durable,
+            flush_lsn: durable.min(apply_ceiling),
             apply_lsn: apply_ceiling,
         };
         let dispatched_before = stream.dispatched_lsn();
@@ -1210,10 +1220,34 @@ async fn run(args: Args) -> Result<()> {
             // Idle tick so metrics/cursor keep tracking the draining pipeline
             // when no new WAL arrives.
             _ = tokio::time::sleep(metrics_tick) => None,
-            res = feed.next_chunk(status, &mut chunk_buf) => Some(match res? {
-                Some(c) => c,
-                None => break "CopyDone",
-            }),
+            res = feed.next_chunk(status, &mut chunk_buf) => match res {
+                Ok(Some(c)) => Some(c),
+                Ok(None) => break "CopyDone",
+                Err(e) => {
+                    let resume = stream.next_lsn();
+                    tracing::warn!(
+                        target: "walshadow",
+                        error = %e,
+                        resume_lsn = format_pg_lsn(resume).to_string(),
+                        "source stream error — recovering",
+                    );
+                    feed = reconnect_or_fatal(
+                        e,
+                        &cfg,
+                        source_slot.as_deref(),
+                        resume,
+                        ident.timeline,
+                        Duration::from_secs(args.status_interval),
+                    )
+                    .await?;
+                    tracing::info!(
+                        target: "walshadow",
+                        resume_lsn = format_pg_lsn(resume).to_string(),
+                        "source reconnected — resuming replication",
+                    );
+                    None
+                }
+            },
         };
         let server_end = chunk.as_ref().map(|c| c.server_wal_end).unwrap_or(received);
         if let Some(chunk) = chunk {
@@ -1963,6 +1997,35 @@ async fn populate_metrics(
     registry.set(snap).await;
 }
 
+/// Recover from a source stream error. A transient drop reconnects and resumes
+/// at `resume_lsn`; a recycled segment (58P01) is fatal — the resume point is
+/// gone, so the daemon exits and recovery is a re-seed via config
+/// `initial_load` on restart, not an in-place reconnect.
+async fn reconnect_or_fatal(
+    e: anyhow::Error,
+    cfg: &PgConfig,
+    slot: Option<&str>,
+    resume_lsn: u64,
+    timeline: u32,
+    status_interval: Duration,
+) -> Result<SourceFeed> {
+    if walshadow::source_feed::is_wal_segment_removed(&e) {
+        return Err(e).context(
+            "source WAL segment recycled past the resume point; \
+             re-seed the affected tables via config initial_load \
+             (base_backup/object_store), then restart",
+        );
+    }
+    match SourceFeed::reconnect(cfg, slot, resume_lsn, timeline, status_interval).await {
+        Ok(feed) => Ok(feed),
+        Err(e2) if walshadow::source_feed::is_wal_segment_removed(&e2) => Err(e2).context(
+            "source WAL segment recycled (seen on reconnect); \
+             re-seed via config initial_load, then restart",
+        ),
+        Err(e2) => Err(e2),
+    }
+}
+
 /// Orchestrate BASE_BACKUP into a fresh shadow data dir; returns the
 /// backup's `end_lsn` so the caller rebinds the WAL pump past it.
 ///

src/ch_emitter.rs

@@ -161,6 +161,10 @@ pub struct EmitterConfig {
     /// overlay tables. `None` (field empty or omitted) disables the whole
     /// overlay subsystem — no boot seed, no config_decoder, pure TOML+CLI.
     pub runtime_config_schema: Option<String>,
+    /// `[source] slot`: physical replication slot to create + stream from on
+    /// the source. `Some` reserves WAL so a stalled/disconnected consumer
+    /// resumes without recycling; `None` runs slotless. Boot-only.
+    pub source_slot: Option<String>,
 }
 
 pub const DEFAULT_INSERT_TIMEOUT_SECS: u64 = 30;
@@ -211,6 +215,7 @@ impl Default for EmitterConfig {
             toast: crate::toast::ToastConfig::default(),
             decode_chunk_rows: crate::pipeline::decode::DECODE_CHUNK_ROWS,
             runtime_config_schema: None,
+            source_slot: None,
         }
     }
 }
@@ -448,6 +453,13 @@ impl EmitterConfig {
             // Empty string == omitted == overlay disabled.
             out.runtime_config_schema = Some(schema.into());
         }
+        if let Some(src) = root.get("source").and_then(Value::as_table)
+            && let Some(slot) = src.get("slot").and_then(Value::as_str)
+            && !slot.is_empty()
+        {
+            // Empty string == omitted == slotless.
+            out.source_slot = Some(slot.into());
+        }
         if let Some(nss) = root.get("namespace").and_then(Value::as_table) {
             for (k, v) in nss {
                 let t = v.as_table().ok_or_else(|| {

src/queueing_record_sink.rs

@@ -40,10 +40,11 @@ use crate::xact_buffer::TxnSpanRegistry;
 /// over the clone-into-owned baseline.
 pub const DEFAULT_QUEUEING_BATCH_SIZE: usize = 64;
 
-/// Soft in-flight cap (channel batches + pump buffer). Past it the
-/// pump yields so the worker drains. No hard cap: a permanently
-/// stalled worker surfaces via the `wait_for_replay` timeout on the
-/// shared err slot.
+/// Hard in-flight cap (channel batches + pump buffer). The channel is bounded
+/// at `max_records / batch_size`, so past it the pump's `send` blocks — real
+/// backpressure to the source instead of unbounded RAM growth. Deadlock-safe:
+/// shadow is fed by an independent walsender task (+keepalive), so a parked
+/// pump can't starve `wait_for_replay`.
 pub const DEFAULT_QUEUEING_RECORD_SINK_CAPACITY: usize = 16_384;
 
 /// Worker `on_idle` cadence. Lets CH emitter's hold-INSERT-open
@@ -54,7 +55,7 @@ pub const DEFAULT_QUEUEING_IDLE_INTERVAL: Duration = Duration::from_millis(50);
 /// Construct via [`QueueingRecordSink::spawn`].
 pub struct QueueingRecordSink {
     /// Each batch carries its ship `Instant` for the worker's `queued_ms`.
-    tx: Option<mpsc::UnboundedSender<(Instant, Vec<Record<'static>>)>>,
+    tx: Option<mpsc::Sender<(Instant, Vec<Record<'static>>)>>,
     /// Pump-side accumulator; shipped as one message at `batch_size` (or `close`).
     buf: Vec<Record<'static>>,
     batch_size: usize,
@@ -63,7 +64,6 @@ pub struct QueueingRecordSink {
     /// Records the worker has dispatched; with `in_flight`, tells a draining
     /// queue from a stalled one.
     processed: Arc<AtomicU64>,
-    soft_cap: u64,
     worker: Option<JoinHandle<()>>,
     /// Per-txn span map; `Some` only with OTLP on. `flush_buf` stamps each
     /// shipped record's ship instant (`note_shipped`).
@@ -75,7 +75,7 @@ impl QueueingRecordSink {
     pub fn spawn<S>(
         inner: S,
         batch_size: usize,
-        soft_cap: usize,
+        max_records: usize,
         span_registry: Option<TxnSpanRegistry>,
     ) -> Self
     where
@@ -84,43 +84,46 @@ impl QueueingRecordSink {
         Self::spawn_with_idle(
             inner,
             batch_size,
-            soft_cap,
+            max_records,
             DEFAULT_QUEUEING_IDLE_INTERVAL,
             span_registry,
         )
     }
 
     /// Worker owns `inner`, drains batches, dispatches each record.
-    /// `soft_cap` triggers `yield_now` once in-flight (channel + pump
-    /// buffer) exceeds it. `idle_interval` paces `inner.on_idle()` on a
-    /// quiescent channel so time-based observer work (CH emitter's
-    /// hold-INSERT-open deadline) fires without fresh records.
+    /// `max_records` bounds records in flight (channel + pump buffer):
+    /// the channel holds `max_records / batch_size` batches, so a slow
+    /// worker blocks the pump's send instead of growing without limit.
+    /// `idle_interval` paces `inner.on_idle()` on a quiescent channel so
+    /// time-based observer work (CH emitter's hold-INSERT-open deadline)
+    /// fires without fresh records.
     pub fn spawn_with_idle<S>(
         mut inner: S,
         batch_size: usize,
-        soft_cap: usize,
+        max_records: usize,
         idle_interval: Duration,
         span_registry: Option<TxnSpanRegistry>,
     ) -> Self
     where
         S: RecordSink + Send + 'static,
     {
-        let (tx, mut rx) = mpsc::unbounded_channel::<(Instant, Vec<Record<'static>>)>();
+        let batch_size = batch_size.max(1);
+        let channel_cap = (max_records / batch_size).max(1);
+        let (tx, mut rx) = mpsc::channel::<(Instant, Vec<Record<'static>>)>(channel_cap);
         let err = Arc::new(StdMutex::new(None));
         let in_flight = Arc::new(AtomicU64::new(0));
         let processed = Arc::new(AtomicU64::new(0));
         let err_w = err.clone();
         let in_flight_w = in_flight.clone();
         let processed_w = processed.clone();
         let reg_w = span_registry.clone();
-        let batch_size = batch_size.max(1);
         let idle_interval = idle_interval.max(Duration::from_millis(1));
         let worker = tokio::spawn(async move {
             // Park error, drop in-flight (`n`, or 0 on idle path),
             // close+drain so `in_flight` settles. Caller breaks after.
             let park_err_and_drain = async |e: SinkError,
                                             n: u64,
-                                            rx: &mut mpsc::UnboundedReceiver<(
+                                            rx: &mut mpsc::Receiver<(
                 Instant,
                 Vec<Record<'static>>,
             )>| {
@@ -216,7 +219,6 @@ impl QueueingRecordSink {
             err,
             in_flight,
             processed,
-            soft_cap: soft_cap.max(1) as u64,
             worker: Some(worker),
             span_registry,
         }
@@ -239,10 +241,10 @@ impl QueueingRecordSink {
         if let Some(e) = self.take_pending_error() {
             return Err(e);
         }
-        self.flush_buf()
+        self.flush_buf().await
     }
 
-    fn flush_buf(&mut self) -> Result<(), SinkError> {
+    async fn flush_buf(&mut self) -> Result<(), SinkError> {
         if self.buf.is_empty() {
             return Ok(());
         }
@@ -260,7 +262,9 @@ impl QueueingRecordSink {
             .tx
             .as_ref()
             .ok_or_else(|| SinkError::Other("queueing record sink already closed".into()))?;
-        if tx.send((Instant::now(), batch)).is_err() {
+        // Blocking the pump here is deadlock-safe: shadow is fed by an independent
+        // walsender task (+keepalive), so a parked pump can't starve wait_for_replay.
+        if tx.send((Instant::now(), batch)).await.is_err() {
             self.in_flight.fetch_sub(n, Ordering::Relaxed);
             if let Some(e) = self.take_pending_error() {
                 return Err(e);
@@ -276,7 +280,7 @@ impl QueueingRecordSink {
     /// Call after the pump stops feeding records.
     pub async fn close(mut self) -> Result<(), SinkError> {
         // Flush tail before dropping sender so worker sees final batch.
-        self.flush_buf()?;
+        self.flush_buf().await?;
         self.tx.take();
         if let Some(handle) = self.worker.take() {
             // Treat a worker panic as a sink error so daemon shutdown
@@ -336,12 +340,7 @@ impl RecordSink for QueueingRecordSink {
                 catalog_signal: record.catalog_signal,
             });
             if self.buf.len() >= self.batch_size {
-                self.flush_buf()?;
-                // Soft backpressure: yield only when actually behind,
-                // checked at flush time to keep per-record cost low.
-                if self.in_flight.load(Ordering::Relaxed) > self.soft_cap {
-                    tokio::task::yield_now().await;
-                }
+                self.flush_buf().await?;
             }
             Ok(())
         })