fix: Unit 9 — abort-cancels-Pub (P-06)

- Add in_flight_sink_aborts map to Dispatcher for cross-stream abort
  signaling via oneshot channels shared across handle_stream tasks
- handle_abort signals in-flight sink aborts before cascading to
  PendingRequestMap
- pump_sink registers an abort receiver, selects on it alongside the
  handler and reader; on abort (same-stream or cross-stream), drops
  the handler future instead of awaiting it — no response written
- InFlightSink gains handler_handle: JoinHandle<()>; single-stream
  EVENT_ABORTED aborts the handler task via JoinHandle::abort()
- 4 new tests: same-stream abort drops handler, cross-stream abort
  drops handler, cleanup of in_flight_sink_aborts map, handler
  resource release via DropGuard

Verification:
- cargo test: 483 passed, 0 failed
- cargo clippy --all-targets -- -D warnings: clean
- cargo fmt --check: clean
- cargo doc --no-deps: clean

src/protocol/dispatch.rs

@@ -22,7 +22,9 @@ use crate::core::types::StreamError;
 use futures::channel::mpsc;
 use futures::stream::StreamExt;
 use futures::SinkExt;
+use parking_lot::Mutex as ParkingLotMutex;
 use serde_json::Value;
+use tokio::sync::oneshot;
 use tokio::task::JoinHandle;
 use tracing::{debug, warn};
 
@@ -88,13 +90,16 @@ pub struct SinkDispatch {
 
 /// An in-flight Pub (Sink) on the single-stream dispatch loop
 /// (`run_loop_single_stream`), tracked by `request_id`. Carries the
-/// feed's `chunk_tx` and the per-chunk validator (built once at
-/// dispatch time from the op's `publish_schema`). `call.published`
-/// frames arriving after the `call.requested` are routed here;
-/// `call.completed` / `call.aborted` / `call.error` remove the entry.
+/// feed's `chunk_tx`, the per-chunk validator (built once at dispatch
+/// time from the op's `publish_schema`), and the spawned handler task's
+/// `JoinHandle` so that `call.aborted` can cancel the handler future
+/// (ADR-046 §5, review 001 Unit 9). `call.published` frames arriving
+/// after the `call.requested` are routed here; `call.completed` /
+/// `call.aborted` / `call.error` remove the entry.
 struct InFlightSink {
     chunk_tx: mpsc::Sender<Result<Value, CallError>>,
     publish_validator: Option<jsonschema::Validator>,
+    handler_handle: JoinHandle<()>,
 }
 
 impl std::fmt::Debug for DispatchResult {
@@ -113,12 +118,21 @@ impl std::fmt::Debug for DispatchResult {
 /// both `CallAdapter` (accept path) and `CallClient` (connect path) and used
 /// to run the dispatch loop. Holds no per-connection state; the
 /// `CallConnection` is passed into `run_loop`.
+///
+/// `in_flight_sink_aborts` maps in-flight Pub `request_id`s to
+/// `oneshot::Sender`s so that a cross-stream `call.aborted` (arriving on a
+/// different `BiStream` than the one the Pub is pumping) can signal the
+/// running `pump_sink` to cancel the handler future (ADR-046 §5, review 001
+/// Unit 9). The map is shared across all `handle_stream` tasks spawned by
+/// `run_loop`; each `pump_sink` inserts its abort sender on entry and
+/// removes it on exit.
 pub struct Dispatcher {
     pub registry: Arc<OperationRegistry>,
     pub identity_provider: Arc<dyn IdentityProvider>,
     pub session_source: Option<Arc<dyn SessionOverlaySource + Send + Sync>>,
     pub ownership_provider: Option<Arc<dyn OwnershipProvider>>,
     pub default_timeout: Duration,
+    in_flight_sink_aborts: Arc<ParkingLotMutex<HashMap<String, oneshot::Sender<()>>>>,
 }
 
 impl Dispatcher {
@@ -132,6 +146,7 @@ impl Dispatcher {
             session_source: None,
             ownership_provider: None,
             default_timeout: DEFAULT_TIMEOUT,
+            in_flight_sink_aborts: Arc::new(ParkingLotMutex::new(HashMap::new())),
         }
     }
 
@@ -363,6 +378,10 @@ impl Dispatcher {
     }
 
     pub async fn handle_abort(&self, connection: &Arc<CallConnection>, request_id: &str) {
+        if let Some(tx) = self.in_flight_sink_aborts.lock().remove(request_id) {
+            let _ = tx.send(());
+            debug!(request_id, "signalled in-flight sink abort");
+        }
         let mut pending = connection.pending().lock();
         let mut cascade = AbortCascade::new(&mut pending);
         let aborted = cascade.cascade_abort(request_id, AbortPolicy::AbortDependents);
@@ -498,10 +517,16 @@ impl Dispatcher {
     /// `Err` item, and terminates the feed — the handler sees the
     /// initiator's error, not a synthetic "aborted" message.
     ///
-    /// On `call.aborted` for this request ID, the feed injects an `Err`
-    /// into `chunk_tx` and ends; the handler future is awaited to
-    /// completion (it is not yet cancelled — the abort-cancels-Pub
-    /// mechanism is review 001 Unit 9).
+    /// On `call.aborted` for this request ID (same-stream or
+    /// cross-stream), the handler future is dropped — the abort
+    /// cancels the Pub (ADR-046 §5, review 001 Unit 9). A cross-stream
+    /// abort arrives via `handle_abort` on a different `BiStream` and
+    /// signals through the shared `in_flight_sink_aborts` map; a
+    /// same-stream abort arrives as a `call.aborted` frame on this
+    /// stream. In both cases the handler future is dropped (not
+    /// awaited), `chunk_tx` is dropped (the `PublishStream` sees EOF),
+    /// and no response is written to the wire — the aborting initiator
+    /// already knows the request is cancelled.
     pub(crate) async fn pump_sink<R, W>(
         &self,
         reader: &mut super::wire::FrameFramedReader<R>,
@@ -518,9 +543,15 @@ impl Dispatcher {
             publish_validator,
         } = sink;
 
+        let (abort_tx, mut abort_rx) = oneshot::channel::<()>();
+        self.in_flight_sink_aborts
+            .lock()
+            .insert(request_id.clone(), abort_tx);
+
         use futures::FutureExt;
         let mut handler = handler.fuse();
         let mut response: Option<ResponseEnvelope> = None;
+        let mut aborted = false;
 
         loop {
             tokio::select! {
@@ -529,6 +560,10 @@ impl Dispatcher {
                     response = Some(resp);
                     break;
                 }
+                _ = &mut abort_rx => {
+                    aborted = true;
+                    break;
+                }
                 read = reader.read_frame() => {
                     match read {
                         Err(super::wire::FrameError::ConnectionClosed) => break,
@@ -571,11 +606,7 @@ impl Dispatcher {
                                 }
                                 EVENT_COMPLETED => break,
                                 EVENT_ABORTED => {
-                                    let _ = chunk_tx
-                                        .send(Err(CallError::internal(
-                                            "publish aborted by initiator",
-                                        )))
-                                        .await;
+                                    aborted = true;
                                     break;
                                 }
                                 EVENT_ERROR => {
@@ -603,8 +634,14 @@ impl Dispatcher {
             }
         }
 
+        self.in_flight_sink_aborts.lock().remove(&request_id);
         drop(chunk_tx);
 
+        if aborted {
+            drop(handler);
+            return;
+        }
+
         let response = match response {
             Some(r) => r,
             None => handler.await,
@@ -767,16 +804,9 @@ impl Dispatcher {
                                 chunk_tx,
                                 publish_validator,
                             } = sink;
-                            in_flight_sinks.insert(
-                                request_id.clone(),
-                                InFlightSink {
-                                    chunk_tx,
-                                    publish_validator,
-                                },
-                            );
                             let writer_clone = Arc::clone(&writer);
                             let request_id_for_handler = request_id.clone();
-                            tokio::spawn(async move {
+                            let handle = tokio::spawn(async move {
                                 let response = handler.await;
                                 let event: EventEnvelope = response.into();
                                 if let Err(err) = writer_clone.write_frame(&event).await {
@@ -787,12 +817,21 @@ impl Dispatcher {
                                     );
                                 }
                             });
+                            in_flight_sinks.insert(
+                                request_id.clone(),
+                                InFlightSink {
+                                    chunk_tx,
+                                    publish_validator,
+                                    handler_handle: handle,
+                                },
+                            );
                         }
                     }
                 }
                 EVENT_ABORTED => {
                     let request_id = envelope.id.clone();
                     if let Some(mut entry) = in_flight_sinks.remove(&request_id) {
+                        entry.handler_handle.abort();
                         let _ = entry
                             .chunk_tx
                             .send(Err(CallError::internal("publish aborted by initiator")))
@@ -917,6 +956,7 @@ impl Clone for Dispatcher {
             session_source: self.session_source.clone(),
             ownership_provider: self.ownership_provider.clone(),
             default_timeout: self.default_timeout,
+            in_flight_sink_aborts: Arc::clone(&self.in_flight_sink_aborts),
         }
     }
 }
@@ -2376,4 +2416,181 @@ mod tests {
             "handler consumed exactly one chunk before returning"
         );
     }
+
+    // --- Unit 9: abort-cancels-Pub (P-06) ----------------------------------
+
+    #[tokio::test]
+    async fn pump_sink_same_stream_abort_drops_handler_future_no_response_written() {
+        use std::sync::atomic::{AtomicBool, Ordering};
+        use std::sync::Arc as StdArc;
+
+        let dropped = StdArc::new(AtomicBool::new(false));
+        let dropped_clone = StdArc::clone(&dropped);
+        let handler: crate::registry::registration::SinkHandler = {
+            use crate::registry::registration::make_sink_handler;
+            use futures::stream::StreamExt;
+            make_sink_handler(move |_input, ctx, mut stream| {
+                let flag = StdArc::clone(&dropped_clone);
+                struct DropGuard(StdArc<AtomicBool>);
+                impl Drop for DropGuard {
+                    fn drop(&mut self) {
+                        self.0.store(true, Ordering::SeqCst);
+                    }
+                }
+                async move {
+                    let _guard = DropGuard(StdArc::clone(&flag));
+                    while let Some(item) = stream.next().await {
+                        if item.is_err() {
+                            break;
+                        }
+                    }
+                    ResponseEnvelope::ok(ctx.request_id, serde_json::json!({"done": true}))
+                }
+            })
+        };
+        let registry = registry_with_pub("fs/upload", handler);
+        let provider: Arc<dyn IdentityProvider> = Arc::new(StaticIdentityProvider::new());
+        let dp = Dispatcher::new(registry, provider);
+        let conn = Arc::new(CallConnection::new(stub_connection()));
+
+        let request = EventEnvelope::requested(
+            "pub-abort-ss-1",
+            serde_json::json!({
+                "operationId": "/fs/upload",
+                "input": {},
+            }),
+        );
+        let c1 = EventEnvelope::published("pub-abort-ss-1", serde_json::json!({"chunk": 1}));
+        let abort = EventEnvelope::aborted("pub-abort-ss-1");
+
+        let mut frame_buf = Vec::new();
+        frame_buf.extend_from_slice(&encode_frame(&request));
+        frame_buf.extend_from_slice(&encode_frame(&c1));
+        frame_buf.extend_from_slice(&encode_frame(&abort));
+
+        let recv = tokio::io::BufReader::new(std::io::Cursor::new(frame_buf));
+        let (send, mut sink) = tokio::io::duplex(8 * 1024);
+        let stream = crate::core::types::BiStream::from_joined(recv, send);
+
+        dp.handle_stream(conn, stream).await;
+
+        let frames = read_all_frames(&mut sink).await;
+        assert!(
+            frames.is_empty(),
+            "no response written after same-stream abort (handler was dropped, not awaited)"
+        );
+        assert!(
+            dropped.load(Ordering::SeqCst),
+            "handler future was dropped (DropGuard fired)"
+        );
+    }
+
+    #[tokio::test]
+    async fn pump_sink_cross_stream_abort_drops_handler_future_no_response_written() {
+        use std::sync::atomic::{AtomicBool, Ordering};
+        use std::sync::Arc as StdArc;
+
+        let dropped = StdArc::new(AtomicBool::new(false));
+        let dropped_clone = StdArc::clone(&dropped);
+        let handler: crate::registry::registration::SinkHandler = {
+            use crate::registry::registration::make_sink_handler;
+            use futures::stream::StreamExt;
+            make_sink_handler(move |_input, ctx, mut stream| {
+                let flag = StdArc::clone(&dropped_clone);
+                struct DropGuard(StdArc<AtomicBool>);
+                impl Drop for DropGuard {
+                    fn drop(&mut self) {
+                        self.0.store(true, Ordering::SeqCst);
+                    }
+                }
+                async move {
+                    let _guard = DropGuard(StdArc::clone(&flag));
+                    while let Some(item) = stream.next().await {
+                        if item.is_err() {
+                            break;
+                        }
+                    }
+                    ResponseEnvelope::ok(ctx.request_id, serde_json::json!({"done": true}))
+                }
+            })
+        };
+        let registry = registry_with_pub("fs/upload", handler);
+        let provider: Arc<dyn IdentityProvider> = Arc::new(StaticIdentityProvider::new());
+        let dp = Dispatcher::new(registry, provider);
+        let conn = Arc::new(CallConnection::new(stub_connection()));
+
+        let request = EventEnvelope::requested(
+            "pub-abort-cs-1",
+            serde_json::json!({
+                "operationId": "/fs/upload",
+                "input": {},
+            }),
+        );
+        let c1 = EventEnvelope::published("pub-abort-cs-1", serde_json::json!({"chunk": 1}));
+
+        let (mut read_half, write_half) = tokio::io::duplex(8 * 1024);
+        {
+            use tokio::io::AsyncWriteExt;
+            read_half.write_all(&encode_frame(&request)).await.unwrap();
+            read_half.write_all(&encode_frame(&c1)).await.unwrap();
+        }
+
+        let recv = tokio::io::BufReader::new(write_half);
+        let (send, mut sink) = tokio::io::duplex(8 * 1024);
+        let stream = crate::core::types::BiStream::from_joined(recv, send);
+
+        let conn_clone = Arc::clone(&conn);
+        let dp_clone = dp.clone();
+        let handle = tokio::spawn(async move {
+            dp_clone.handle_stream(conn_clone, stream).await;
+        });
+
+        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+        dp.handle_abort(&conn, "pub-abort-cs-1").await;
+        drop(read_half);
+
+        let frames = read_all_frames(&mut sink).await;
+        handle.await.unwrap();
+
+        assert!(
+            frames.is_empty(),
+            "no response written after cross-stream abort"
+        );
+        assert!(
+            dropped.load(Ordering::SeqCst),
+            "handler future was dropped (DropGuard fired) via cross-stream abort signal"
+        );
+    }
+
+    #[tokio::test]
+    async fn pump_sink_abort_cleans_up_in_flight_sink_aborts_map() {
+        let registry = registry_with_pub("fs/upload", counting_sink_handler());
+        let provider: Arc<dyn IdentityProvider> = Arc::new(StaticIdentityProvider::new());
+        let dp = Dispatcher::new(registry, provider);
+        let conn = Arc::new(CallConnection::new(stub_connection()));
+
+        let request = EventEnvelope::requested(
+            "pub-abort-cleanup-1",
+            serde_json::json!({
+                "operationId": "/fs/upload",
+                "input": {},
+            }),
+        );
+        let abort = EventEnvelope::aborted("pub-abort-cleanup-1");
+
+        let mut frame_buf = Vec::new();
+        frame_buf.extend_from_slice(&encode_frame(&request));
+        frame_buf.extend_from_slice(&encode_frame(&abort));
+
+        let recv = tokio::io::BufReader::new(std::io::Cursor::new(frame_buf));
+        let (send, _sink) = tokio::io::duplex(8 * 1024);
+        let stream = crate::core::types::BiStream::from_joined(recv, send);
+
+        dp.handle_stream(conn, stream).await;
+
+        assert!(
+            dp.in_flight_sink_aborts.lock().is_empty(),
+            "in_flight_sink_aborts map is cleaned up after pump_sink exits"
+        );
+    }
 }