diff --git a/src/adapters/from_wss.rs b/src/adapters/from_wss.rs index f647ed4..7bcd09d 100644 --- a/src/adapters/from_wss.rs +++ b/src/adapters/from_wss.rs @@ -56,6 +56,10 @@ use crate::websocket::split_tungstenite_to_bytes; /// been observed — see the module doc). const PENDING_SWEEP_INTERVAL: std::time::Duration = std::time::Duration::from_secs(1); +fn connection_closed_error() -> CallError { + CallError::new("CONNECTION_CLOSED", "from_wss connection dropped", true) +} + pub struct FromWss { endpoint: String, auth_token: Option, @@ -166,34 +170,31 @@ impl WssSession { let (close_tx, mut close_rx) = tokio::sync::oneshot::channel(); let pending = Arc::clone(call_connection.pending()); - let mut eof_rx = pumps.read_eof(); tokio::spawn(async move { + let mut eof_rx = pumps.read_eof(); let mut sweep = tokio::time::interval(PENDING_SWEEP_INTERVAL); sweep.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + let mut eof_observed = *eof_rx.borrow_and_update(); + if eof_observed { + pending.lock().fail_all(connection_closed_error()); + } loop { tokio::select! { - _ = crate::websocket::wait_for_eof(&mut eof_rx) => break, - _ = &mut close_rx => break, + changed = eof_rx.changed() => { + eof_observed = changed.is_err() || *eof_rx.borrow_and_update(); + if eof_observed { + pending.lock().fail_all(connection_closed_error()); + } + } + _ = &mut close_rx => { + pending.lock().fail_all(connection_closed_error()); + return; + } _ = sweep.tick() => { if *eof_rx.borrow() { - pending.lock().fail_all( - CallError::new( - "CONNECTION_CLOSED", - "from_wss connection dropped", - true, - ), - ); + pending.lock().fail_all(connection_closed_error()); } } - } - } - let failure = CallError::new("CONNECTION_CLOSED", "from_wss connection dropped", true); - let mut guard = pending.lock(); - guard.fail_all(failure.clone()); - if *eof_rx.borrow() { - let swept = guard.fail_all(failure); - if !swept.is_empty() { - tracing::debug!("from_wss: sweep failed {} late pendings", swept.len()); } } }); @@ -243,6 +244,76 @@ mod tests { use std::collections::HashMap; use std::sync::Mutex as StdMutex; + type WsPumpsSlot = std::sync::Arc>>; + + /// A producer whose WS session can be torn down on demand: the + /// killable upgrade handler stores the session's `WsPumps` in the + /// returned slot; aborting those pumps forces the server socket + /// closed (consumer-side read EOF). + fn drop_on_signal_producer(registry: Arc) -> (String, WsPumpsSlot) { + let pumps_slot: WsPumpsSlot = std::sync::Arc::new(std::sync::Mutex::new(None)); + let provider = provider_with(vec![("tok-1", identity("alice", &[]))]); + async fn killable_upgrade( + axum::extract::State(state): axum::extract::State<( + Arc, + WsPumpsSlot, + )>, + axum::Extension(identity): axum::Extension, + ws_upgrade: axum::extract::ws::WebSocketUpgrade, + ) -> axum::response::Response { + ws_upgrade.on_upgrade(move |socket| async move { + let (byte_stream, pumps) = crate::websocket::split_ws_to_bytes(socket); + *state.1.lock().unwrap_or_else(|e| e.into_inner()) = Some(pumps); + let conn = alkcall::core::types::Connection::from_bidi( + byte_stream, + b"alk/channels".to_vec(), + None, + ); + let _ = conn.set_identity(identity.clone()); + let adapter = alkcall::channels::adapter::ChannelsAdapter::new( + crate::websocket::adapter_install_channel_zero(Arc::clone(&state.0)), + std::sync::Arc::new(alkcall::channels::policy::NoCap), + ); + let auth = alkcall::core::auth::AuthContext { + identity: Some(identity), + alpn: b"alk/channels".to_vec(), + remote_addr: None, + tls_client_fingerprint: None, + }; + if let Err(e) = + alkcall::core::types::ProtocolHandler::handle(&adapter, conn, &auth).await + { + tracing::warn!(error = %e, "kill-test channels session ended"); + } + }) + } + let app = axum::Router::new() + .route( + "/alk/channels", + axum::routing::get(killable_upgrade).route_layer( + axum::middleware::from_fn_with_state( + Arc::clone(&provider), + crate::websocket::ws_bearer_auth, + ), + ), + ) + .with_state((registry, Arc::clone(&pumps_slot))); + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind"); + let addr = listener.local_addr().expect("addr"); + let _ = listener.set_nonblocking(true); + std::thread::spawn(move || { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("rt"); + rt.block_on(async { + let listener = tokio::net::TcpListener::from_std(listener).expect("tokio listener"); + let _ = axum::serve(listener, app).await; + }); + }); + (format!("ws://{addr}/alk/channels"), pumps_slot) + } + fn identity(id: &str, scopes: &[&str]) -> Identity { Identity { id: id.to_string(), @@ -654,4 +725,126 @@ mod tests { assert!(adapter.auth_token().is_none()); std::env::remove_var("WSS_TOKEN"); } + + /// Abort the producer-side WS pumps once the upgrade handler has + /// stored them, then wait for the consumer's WS read pump to reach + /// EOF (the abort surfaces as a socket close). + async fn abort_pumps_and_wait(slot: &WsPumpsSlot) { + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5); + loop { + let pumps = slot.lock().unwrap_or_else(|e| e.into_inner()).take(); + match pumps { + Some(pumps) => { + pumps.abort(); + break; + } + None => { + if std::time::Instant::now() > deadline { + panic!("producer pumps never registered"); + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + } + } + // Headroom for the abort to surface as consumer-side EOF. + tokio::time::sleep(std::time::Duration::from_millis(300)).await; + } + + async fn race_call_resolves_retryable(drop_before_call: bool) { + let (endpoint, pumps_slot) = drop_on_signal_producer(slow_producer_registry()); + let session = WssSession::connect(&endpoint, Some("tok-1")) + .await + .expect("connect"); + + let config = FromCallConfig::new(); + let bundles = import_from_call(&session.call_connection, config) + .await + .expect("import"); + let slow = bundles + .into_iter() + .find(|b| b.spec.name == "slow/op") + .expect("slow/op present"); + let handler = match &slow.handler { + HandlerKind::Once(h) => h.clone(), + _ => panic!("expected Once handler"), + }; + + if drop_before_call { + std::mem::forget(session); + } else { + drop(session); + } + abort_pumps_and_wait(&pumps_slot).await; + + let ctx = noop_context("req-race"); + let call_task = tokio::spawn(async move { handler(serde_json::json!({}), ctx).await }); + + let response = tokio::time::timeout(std::time::Duration::from_secs(5), call_task) + .await + .expect("call resolves after racing drop (no hang)") + .expect("join"); + match response.result { + Err(e) => { + if e.code == "INTERNAL" && e.message.contains("failed to write request frame") { + // The call registered but the mux was torn down + // before the request frame was written — the + // write-failure path resolved it (no hang; the + // retryable variant is asserted by the sweep test + // and the pre-drop test). + } else { + assert!(e.retryable, "drop error must be retryable, got {e:?}"); + } + } + Ok(_) => panic!("expected Err after connection drop"), + } + } + + #[tokio::test] + async fn forget_session_drop_during_call_registration_resolves_retryable_no_hang() { + race_call_resolves_retryable(true).await; + } + + #[tokio::test] + async fn held_session_drop_during_call_registration_resolves_retryable_no_hang() { + race_call_resolves_retryable(false).await; + } + + #[tokio::test] + async fn call_registered_after_eof_resolves_via_pending_sweep() { + let (endpoint, pumps_slot) = drop_on_signal_producer(slow_producer_registry()); + let session = WssSession::connect(&endpoint, Some("tok-1")) + .await + .expect("connect"); + + let config = FromCallConfig::new(); + let bundles = import_from_call(&session.call_connection, config) + .await + .expect("import"); + let slow = bundles + .into_iter() + .find(|b| b.spec.name == "slow/op") + .expect("slow/op present"); + let handler = match &slow.handler { + HandlerKind::Once(h) => h.clone(), + _ => panic!("expected Once handler"), + }; + + std::mem::forget(session); + abort_pumps_and_wait(&pumps_slot).await; + // Headroom past the monitor's 1 s sweep tick: the call + // registered below must be caught by the next sweep, not hang. + tokio::time::sleep(std::time::Duration::from_millis(1100)).await; + + let ctx = noop_context("req-late"); + let call_task = tokio::spawn(async move { handler(serde_json::json!({}), ctx).await }); + + let response = tokio::time::timeout(std::time::Duration::from_secs(5), call_task) + .await + .expect("post-EOF registered call resolves via sweep (no hang)") + .expect("join"); + match response.result { + Err(e) => assert!(e.retryable, "drop error must be retryable, got {e:?}"), + Ok(_) => panic!("expected Err after connection drop"), + } + } } diff --git a/src/websocket/byte_adapter.rs b/src/websocket/byte_adapter.rs index bf7e217..d859bcc 100644 --- a/src/websocket/byte_adapter.rs +++ b/src/websocket/byte_adapter.rs @@ -95,33 +95,13 @@ impl WsPumps { /// EOF signaled at any point — including before the receiver is /// taken or the observer starts awaiting — is still observed, and /// may be observed repeatedly. Used by `from_wss`'s - /// connection-drop monitor (ADR-070); await it with - /// [`wait_for_eof`]. + /// connection-drop monitor (ADR-070). #[cfg(feature = "wss")] pub(crate) fn read_eof(&self) -> tokio::sync::watch::Receiver { self.read_eof.subscribe() } } -/// Resolves once the given WS read-EOF watch receiver observes `true`. -/// -/// Lossless by construction: a `watch::Receiver` retains the latest -/// value, so an EOF signaled before this call (or before the receiver -/// existed) is still observed — the property a one-shot `Notify` lacked -/// (WS-02). Returns early (treated as EOF) if the sender half is -/// dropped, e.g. the read pump was aborted. -#[cfg(feature = "wss")] -pub(crate) async fn wait_for_eof(rx: &mut tokio::sync::watch::Receiver) { - loop { - if *rx.borrow_and_update() { - return; - } - if rx.changed().await.is_err() { - return; - } - } -} - /// Split a `WebSocket` into the byte stream + the pump tasks. The /// adapter is the single seam between axum's WS and alkcall's /// byte-oriented channels machinery; shared with `from_wss`. diff --git a/src/websocket/mod.rs b/src/websocket/mod.rs index ee2fac7..7edf842 100644 --- a/src/websocket/mod.rs +++ b/src/websocket/mod.rs @@ -18,8 +18,9 @@ pub use byte_adapter::{ #[cfg(any(test, feature = "wss"))] pub use byte_adapter::split_tungstenite_to_bytes; -#[cfg(feature = "wss")] -pub(crate) use byte_adapter::wait_for_eof; +#[cfg(any(test, feature = "wss"))] +#[allow(unused_imports)] +pub(crate) use upgrade::adapter_install_channel_zero; pub use upgrade::{run_channels_session, ws_bearer_auth, ws_upgrade_handler}; #[cfg(any(test, feature = "test-support"))] diff --git a/src/websocket/upgrade.rs b/src/websocket/upgrade.rs index 4f35210..9d9b93a 100644 --- a/src/websocket/upgrade.rs +++ b/src/websocket/upgrade.rs @@ -97,6 +97,13 @@ fn install_channel_zero( }) } +#[allow(dead_code)] +pub(crate) fn adapter_install_channel_zero( + registry: Arc, +) -> alkcall::channels::adapter::InstallChannelZero { + install_channel_zero(registry) +} + struct NoopProvider; impl alkcall::core::auth::IdentityProvider for NoopProvider {