engine(runtime): broadcast Shutdown to extensions on error path with bounded drain

Address review feedback on PR #2860:

- runtime_pipeline.rs: refactor the run loop into an inner async
  block whose Result is captured, so cleanup runs on every exit
  path ("async finally"). Previously, on any error the loop did
  `return Err(e)` immediately and dropped already-started
  extensions without sending Shutdown — extensions owning sockets,
  files, or background work missed their cleanup signal. Now the
  outer block always calls broadcast_shutdown() (idempotent on the
  normal path) and drain_until_deadline().await before propagating
  the loop's result. Behavior on the normal path is unchanged.
  Refs: #2860 (comment)

- extension_lifecycle.rs: add drain_until_deadline() that waits for
  remaining active+background extension tasks to finish but never
  past the broadcast deadline (plus a small slack to absorb the
  ControlChannel adapter's deferred Shutdown delivery and the
  task-return / JoinHandle-observed latency). A misbehaving
  extension that ignores Shutdown can no longer hang the pipeline
  indefinitely — it is dropped after EXTENSION_SHUTDOWN_GRACE +
  EXTENSION_SHUTDOWN_DRAIN_SLACK with a warning. broadcast_shutdown
  now also stashes the deadline so the drain stays consistent with
  what the wire-level message advertised.
  Refs: #2860 (comment)

- extension_e2e.rs: regression test
  test_other_extensions_receive_shutdown_when_pipeline_errors —
  configures one failing extension and one shutdown-recording
  extension (each bound by its own probe receiver to survive
  defined-but-unbound pruning). When the failing extension's
  start() errors, the pipeline aborts but the recording extension
  must still observe Shutdown before run_forever returns.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: gouslu

rust/otap-dataflow/crates/engine/src/extension_lifecycle.rs

@@ -31,13 +31,28 @@ use otap_df_telemetry::otel_warn;
 use otap_df_telemetry::reporter::MetricsReporter;
 use std::time::{Duration, Instant};
 use tokio::task::{JoinError, JoinHandle, LocalSet};
+use tokio::time::Instant as TokioInstant;
 
 /// Cleanup window granted to extensions after the data path has
 /// drained. Extensions that don't terminate within this window will
 /// be left to the runtime's natural drop semantics when
 /// `run_forever` returns.
 pub(crate) const EXTENSION_SHUTDOWN_GRACE: Duration = Duration::from_secs(5);
 
+/// Slack added past `EXTENSION_SHUTDOWN_GRACE` before the runtime
+/// hard-stops draining extension tasks. The cooperative deadline that
+/// `broadcast_shutdown` puts on the wire and the deadline the runtime
+/// waits on are intentionally different: extensions get exactly
+/// `EXTENSION_SHUTDOWN_GRACE` to wind down (matching the contract in
+/// the `Shutdown` message), and the runtime waits a small additional
+/// window to absorb the time it takes the
+/// [`crate::extension::ControlChannel`] adapter to deliver the
+/// `Shutdown` message through to the extension's `recv()` and for the
+/// extension's task to return. Without this slack a well-behaved
+/// extension that wakes up exactly at the deadline races against the
+/// runtime's drain timeout.
+pub(crate) const EXTENSION_SHUTDOWN_DRAIN_SLACK: Duration = Duration::from_millis(500);
+
 const SHUTDOWN_REASON: &str = "pipeline data-path drained";
 
 /// Holds the spawned extension tasks, control senders, and passive
@@ -57,6 +72,11 @@ pub(crate) struct ExtensionLifecycle {
     /// One-shot latch: `true` after `Shutdown` has been broadcast.
     /// Prevents re-firing on subsequent loop iterations.
     shutdown_broadcast_fired: bool,
+    /// Deadline established when [`Self::broadcast_shutdown`] fires.
+    /// Used by [`Self::drain_until_deadline`] to bound how long the
+    /// runtime will wait for extensions to honour `Shutdown` so a
+    /// misbehaving extension can't hang the pipeline indefinitely.
+    shutdown_deadline: Option<Instant>,
 }
 
 impl ExtensionLifecycle {
@@ -105,6 +125,7 @@ impl ExtensionLifecycle {
             shutdown_senders,
             _passive: passive,
             shutdown_broadcast_fired: false,
+            shutdown_deadline: None,
         }
     }
 
@@ -135,6 +156,7 @@ impl ExtensionLifecycle {
         self.shutdown_broadcast_fired = true;
 
         let deadline = Instant::now() + EXTENSION_SHUTDOWN_GRACE;
+        self.shutdown_deadline = Some(deadline);
         for sender in &self.shutdown_senders {
             // `try_send` is intentional: the extension's control
             // channel is a small mpsc and we don't want shutdown
@@ -148,4 +170,63 @@ impl ExtensionLifecycle {
             });
         }
     }
+
+    /// Drain remaining active+background extension tasks, but never
+    /// past the shutdown deadline.
+    ///
+    /// `Shutdown` is cooperative — extensions may ignore it or take
+    /// longer than the grace window to exit. Without this bound, an
+    /// extension that never returns from `start()` would hang the
+    /// pipeline forever. After the deadline elapses, any still-running
+    /// futures are dropped with a warning; the runtime's natural drop
+    /// semantics take over once the lifecycle holder itself is dropped.
+    ///
+    /// No-op if there are no remaining futures or if shutdown has not
+    /// been broadcast (in which case there is no deadline yet).
+    pub async fn drain_until_deadline(&mut self) {
+        if self.futures.is_empty() {
+            return;
+        }
+        // If the caller invokes drain without a prior broadcast there
+        // is no deadline yet — synthesize one from the same grace
+        // window so we still bound the wait.
+        let deadline = self
+            .shutdown_deadline
+            .get_or_insert_with(|| Instant::now() + EXTENSION_SHUTDOWN_GRACE);
+        // Wait slightly past the cooperative deadline to absorb the
+        // time it takes the ControlChannel adapter to deliver the
+        // queued `Shutdown` to the extension and for the extension's
+        // task to return.
+        let drain_deadline = TokioInstant::from_std(*deadline + EXTENSION_SHUTDOWN_DRAIN_SLACK);
+
+        let drain = async {
+            while let Some(result) = self.futures.next().await {
+                match result {
+                    Ok(Ok(())) => {}
+                    Ok(Err(e)) => {
+                        otel_warn!("extension.shutdown.task.error", error = format!("{e}"));
+                    }
+                    Err(e) => {
+                        otel_warn!(
+                            "extension.shutdown.task.join_error",
+                            is_canceled = e.is_cancelled(),
+                            is_panic = e.is_panic(),
+                            error = e.to_string()
+                        );
+                    }
+                }
+            }
+        };
+
+        if tokio::time::timeout_at(drain_deadline, drain)
+            .await
+            .is_err()
+        {
+            otel_warn!(
+                "extension.shutdown.timeout",
+                grace_secs = EXTENSION_SHUTDOWN_GRACE.as_secs(),
+                remaining = self.futures.len()
+            );
+        }
+    }
 }

rust/otap-dataflow/crates/engine/src/runtime_pipeline.rs

@@ -539,56 +539,90 @@ impl<PData: 'static + Debug + Clone + ReceivedAtNode + Unwindable + FlowMetricHo
         // `extension_lifecycle`) run concurrently. When the data-path
         // drains, broadcast `Shutdown` to extensions so they can
         // terminate gracefully, then continue draining extension
-        // futures. Errors from either side short-circuit and abort.
+        // futures.
+        //
+        // Errors from either side short-circuit out of the inner
+        // async block. The outer code then unconditionally
+        // broadcasts `Shutdown` and bounded-drains extensions before
+        // propagating the error. This guarantees that extensions
+        // owning sockets, files, or background work get the same
+        // cleanup signal on the error path that they would on the
+        // normal path, and that a misbehaving extension that ignores
+        // `Shutdown` cannot hang the pipeline beyond
+        // `EXTENSION_SHUTDOWN_GRACE`.
         rt.block_on(async {
             local_tasks
                 .run_until(async {
-                    let mut task_results = Vec::new();
-
-                    loop {
-                        // `biased;`: prefer data-path completions when both
-                        // arms are simultaneously ready. Functionally
-                        // optional — kept to make intent explicit.
-                        tokio::select! {
-                            biased;
-                            Some(result) = futures.next(), if !futures.is_empty() => {
-                                match result {
-                                    Ok(Ok(res)) => task_results.push(res),
-                                    Ok(Err(e)) => return Err(e),
-                                    Err(e) => return Err(Error::JoinTaskError {
-                                        is_canceled: e.is_cancelled(),
-                                        is_panic: e.is_panic(),
-                                        error: e.to_string(),
-                                    }),
+                    // Inner async block isolates the loop's
+                    // `return Err(...)` short-circuits from the
+                    // outer cleanup, giving us an "async finally"
+                    // shape: cleanup always runs, then the loop's
+                    // result is propagated.
+                    let loop_result: Result<Vec<_>, Error> = async {
+                        let mut task_results = Vec::new();
+                        loop {
+                            // `biased;`: prefer data-path completions when both
+                            // arms are simultaneously ready. Functionally
+                            // optional — kept to make intent explicit.
+                            tokio::select! {
+                                biased;
+                                Some(result) = futures.next(), if !futures.is_empty() => {
+                                    match result {
+                                        Ok(Ok(res)) => task_results.push(res),
+                                        Ok(Err(e)) => return Err(e),
+                                        Err(e) => return Err(Error::JoinTaskError {
+                                            is_canceled: e.is_cancelled(),
+                                            is_panic: e.is_panic(),
+                                            error: e.to_string(),
+                                        }),
+                                    }
                                 }
-                            }
-                            Some(result) = extension_lifecycle.next_completion(), if !extension_lifecycle.is_empty() => {
-                                match result {
-                                    Ok(Ok(())) => {}
-                                    Ok(Err(e)) => return Err(e),
-                                    Err(e) => return Err(Error::JoinTaskError {
-                                        is_canceled: e.is_cancelled(),
-                                        is_panic: e.is_panic(),
-                                        error: e.to_string(),
-                                    }),
+                                Some(result) = extension_lifecycle.next_completion(), if !extension_lifecycle.is_empty() => {
+                                    match result {
+                                        Ok(Ok(())) => {}
+                                        Ok(Err(e)) => return Err(e),
+                                        Err(e) => return Err(Error::JoinTaskError {
+                                            is_canceled: e.is_cancelled(),
+                                            is_panic: e.is_panic(),
+                                            error: e.to_string(),
+                                        }),
+                                    }
                                 }
+                                else => break,
                             }
-                            else => break,
-                        }
 
-                        // Once data-path is drained, fire the shutdown
-                        // broadcast exactly once. The lifecycle holder
-                        // gates on its own one-shot latch and uses a
-                        // local `now() + EXTENSION_SHUTDOWN_GRACE`
-                        // deadline — extensions shut down after every
-                        // data-path task has terminated, so this is
-                        // the start of a fresh extension cleanup
-                        // window, not a continuation of the
-                        // pipeline-wide deadline.
-                        if futures.is_empty() {
-                            extension_lifecycle.broadcast_shutdown();
+                            // Once data-path is drained, fire the shutdown
+                            // broadcast exactly once. The lifecycle holder
+                            // gates on its own one-shot latch and uses a
+                            // local `now() + EXTENSION_SHUTDOWN_GRACE`
+                            // deadline — extensions shut down after every
+                            // data-path task has terminated, so this is
+                            // the start of a fresh extension cleanup
+                            // window, not a continuation of the
+                            // pipeline-wide deadline.
+                            if futures.is_empty() {
+                                extension_lifecycle.broadcast_shutdown();
+                            }
                         }
+                        Ok(task_results)
                     }
+                    .await;
+
+                    // "Async finally": unconditional cleanup that
+                    // runs on both the normal and error paths.
+                    // `broadcast_shutdown` is idempotent (latched by
+                    // `shutdown_broadcast_fired`), so a no-op on the
+                    // normal path; on the error path it ensures
+                    // already-started extensions still receive
+                    // `Shutdown` before the pipeline exits, so they
+                    // can release sockets, files, or background work
+                    // cleanly. `drain_until_deadline` then bounds the
+                    // wait so a misbehaving extension can't hang the
+                    // runtime.
+                    extension_lifecycle.broadcast_shutdown();
+                    extension_lifecycle.drain_until_deadline().await;
+
+                    let task_results = loop_result?;
                     let mut final_metrics_reporter = final_metrics_reporter.clone();
                     if let Err(err) = report_node_metrics_with_handles(
                         &final_node_metric_handles,

rust/otap-dataflow/crates/engine/tests/extension_e2e.rs

@@ -2239,6 +2239,101 @@ connections:
     );
 }
 
+// ─────────────────────────────────────────────────────────────────────
+// Error-path shutdown hygiene — when one active extension errors out,
+// any *other* already-started extensions must still receive `Shutdown`
+// before the pipeline returns, so they can release sockets, files, or
+// background work cleanly. Regression test for review feedback on
+// PR #2860 (discussion_r3228775534).
+// ─────────────────────────────────────────────────────────────────────
+
+#[test]
+fn test_other_extensions_receive_shutdown_when_pipeline_errors() {
+    let receiver_key = "err-shutdown-recv";
+    let ext_probe_key = "err-shutdown-ext-probe";
+    let _probe = make_probe(receiver_key, CallSequence::Local);
+
+    let ext_shutdown_at: Arc<parking_lot::Mutex<Option<Instant>>> =
+        Arc::new(parking_lot::Mutex::new(None));
+    register_shutdown_recording_probe(
+        ext_probe_key,
+        ShutdownRecordingProbe {
+            shutdown_at: Arc::clone(&ext_shutdown_at),
+        },
+    );
+
+    // Both extensions must be bound by some node so neither gets
+    // pruned as "defined-but-unbound" before the run starts. We use
+    // two probe receivers, each binding a distinct extension, fanning
+    // into the same exporter. The failing extension errors at start
+    // → pipeline aborts → the recording extension must still receive
+    // `Shutdown` on the way out.
+    let receiver_b_key = "err-shutdown-recv-b";
+    let _probe_b = make_probe(receiver_b_key, CallSequence::Local);
+    let yaml = format!(
+        r#"
+nodes:
+  receiver-a:
+    type: "{PROBE_RECEIVER_URN}"
+    config:
+      probe_key: "{receiver_key}"
+    capabilities:
+      no_op_stateless: "shutdown-rec"
+  receiver-b:
+    type: "{PROBE_RECEIVER_URN}"
+    config:
+      probe_key: "{receiver_b_key}"
+    capabilities:
+      no_op_stateless: "failing"
+  exporter:
+    type: "{NOOP_EXPORTER_URN}"
+
+extensions:
+  shutdown-rec:
+    type: "{SHUTDOWN_RECORDING_EXTENSION_URN}"
+    config:
+      probe_key: "{ext_probe_key}"
+  failing:
+    type: "{FAILING_EXTENSION_URN}"
+
+connections:
+  - from: receiver-a
+    to: exporter
+  - from: receiver-b
+    to: exporter
+"#
+    );
+    let (runtime_pipeline, ctx, entity_key, ts) = build_test_runtime_pipeline(&yaml);
+
+    // Generous outer shutdown grace — the test must return well before
+    // it because the pipeline aborts on the failing extension's error
+    // and then bounded-drains the recording extension within
+    // `EXTENSION_SHUTDOWN_GRACE` (5s).
+    let started_at = Instant::now();
+    let result = run_pipeline_with_shutdown_after(
+        runtime_pipeline,
+        ctx,
+        entity_key,
+        ts,
+        Duration::from_secs(60),
+    );
+    let elapsed = started_at.elapsed();
+
+    assert!(
+        result.is_err(),
+        "pipeline should propagate the failing extension's error: {result:?}"
+    );
+    assert!(
+        elapsed < Duration::from_secs(10),
+        "pipeline should not hang after the error path; took {elapsed:?}"
+    );
+    let shutdown = ext_shutdown_at.lock();
+    assert!(
+        shutdown.is_some(),
+        "recording extension must receive Shutdown on the error path"
+    );
+}
+
 // ─────────────────────────────────────────────────────────────────────
 // Test 5 — shutdown ordering: extension records Shutdown timestamp
 //          inside the pipeline run window