xrpc/
client.rs

1use parking_lot::Mutex;
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4use std::fmt::Debug;
5use std::sync::Arc;
6use std::time::Duration;
7use tokio::sync::{Notify, RwLock, oneshot, watch};
8use tokio::task::AbortHandle;
9
10use crate::channel::message::MessageChannel;
11use crate::codec::{BincodeCodec, Codec};
12use crate::error::{Result, RpcError, TransportError};
13use crate::message::Message;
14use crate::message::types::{MessageId, MessageType};
15use crate::streaming::{StreamId, StreamManager, StreamReceiver, next_stream_id};
16
17type PendingRequests<C> = Arc<Mutex<HashMap<MessageId, oneshot::Sender<Result<Message<C>>>>>>;
18type ReceiveTaskAbort = Arc<Mutex<Option<AbortHandle>>>;
19
20/// Application-level response deadline policy for an RPC call.
21///
22/// This policy does not change transport send/read timeouts or terminal connection handling.
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24#[non_exhaustive]
25pub enum CallTimeout {
26    /// Use the default response timeout configured on the [`RpcClient`].
27    ClientDefault,
28    /// Apply the supplied application-level response timeout after the request send commits.
29    After(Duration),
30    /// Wait without an application-level response deadline.
31    ///
32    /// Explicit client close, terminal connection failure, and dropping the call future still
33    /// cancel the local wait. They do not cancel work already running on the server.
34    Disabled,
35}
36
37/// Per-call RPC options.
38#[derive(Debug, Clone, Copy, PartialEq, Eq)]
39pub struct CallOptions {
40    timeout: CallTimeout,
41}
42
43impl Default for CallOptions {
44    fn default() -> Self {
45        Self {
46            timeout: CallTimeout::ClientDefault,
47        }
48    }
49}
50
51impl CallOptions {
52    /// Use a specific application-level response timeout for this call.
53    pub const fn with_timeout(timeout: Duration) -> Self {
54        Self {
55            timeout: CallTimeout::After(timeout),
56        }
57    }
58
59    /// Disable only the application-level response deadline for this call.
60    ///
61    /// The local wait remains interruptible by explicit client close, terminal connection
62    /// failure, and future cancellation. No cancellation is propagated to the server.
63    pub const fn without_timeout() -> Self {
64        Self {
65            timeout: CallTimeout::Disabled,
66        }
67    }
68}
69
70#[derive(Debug, Default)]
71struct ClientState {
72    closed: bool,
73    started: bool,
74    running: bool,
75    terminal_error: Option<RpcError>,
76}
77
78#[derive(Debug, Clone, Default)]
79enum TransportCloseState {
80    #[default]
81    Open,
82    Closing,
83    Closed(Result<()>),
84}
85
86struct ClientTaskFinalizer<C: Codec> {
87    state: Arc<Mutex<ClientState>>,
88    pending: PendingRequests<C>,
89    stream_manager: Arc<StreamManager<C>>,
90    send_cancel_tx: watch::Sender<Option<RpcError>>,
91    receive_abort: ReceiveTaskAbort,
92    cleanup_work: bool,
93}
94
95impl<C: Codec> Drop for ClientTaskFinalizer<C> {
96    fn drop(&mut self) {
97        let (error, _) = mark_client_terminal(&self.state, RpcError::ConnectionClosed);
98        self.receive_abort.lock().take();
99        if self.cleanup_work {
100            self.send_cancel_tx.send_replace(Some(error.clone()));
101            close_client_work(&self.pending, &self.stream_manager, error);
102        }
103    }
104}
105
106struct PendingRequestRegistration<C: Codec> {
107    pending: PendingRequests<C>,
108    message_id: MessageId,
109}
110
111impl<C: Codec> PendingRequestRegistration<C> {
112    fn new(
113        pending: &PendingRequests<C>,
114        message_id: MessageId,
115        sender: oneshot::Sender<Result<Message<C>>>,
116    ) -> Self {
117        pending.lock().insert(message_id, sender);
118        Self {
119            pending: Arc::clone(pending),
120            message_id,
121        }
122    }
123}
124
125impl<C: Codec> Drop for PendingRequestRegistration<C> {
126    fn drop(&mut self) {
127        self.pending.lock().remove(&self.message_id);
128    }
129}
130
131fn close_client_work<C: Codec>(
132    pending: &PendingRequests<C>,
133    stream_manager: &StreamManager<C>,
134    error: RpcError,
135) {
136    let pending = {
137        let mut pending = pending.lock();
138        std::mem::take(&mut *pending)
139    };
140    for sender in pending.into_values() {
141        let _ = sender.send(Err(error.clone()));
142    }
143    stream_manager.close_all(error);
144}
145
146fn normalize_transport_error(error: TransportError) -> RpcError {
147    match error {
148        TransportError::ConnectionClosed => RpcError::ConnectionClosed,
149        error => RpcError::Transport(error),
150    }
151}
152
153fn mark_client_terminal(state: &Mutex<ClientState>, error: RpcError) -> (RpcError, bool) {
154    let mut state = state.lock();
155    state.closed = true;
156    state.running = false;
157    if let Some(error) = &state.terminal_error {
158        return (error.clone(), false);
159    }
160    state.terminal_error = Some(error.clone());
161    (error, true)
162}
163
164fn is_terminal_outbound_error(error: &TransportError, is_connected: bool) -> bool {
165    matches!(error, TransportError::ConnectionClosed) || !is_connected
166}
167
168struct TransportCloseFinalizer {
169    state: Arc<Mutex<TransportCloseState>>,
170    notify: Arc<Notify>,
171    completed: bool,
172}
173
174struct ClientCloseFinalizer {
175    state: Arc<Mutex<TransportCloseState>>,
176    notify: Arc<Notify>,
177    completed: bool,
178}
179
180impl ClientCloseFinalizer {
181    fn complete(&mut self, result: Result<()>) {
182        *self.state.lock() = TransportCloseState::Closed(result);
183        self.completed = true;
184        self.notify.notify_waiters();
185    }
186}
187
188impl Drop for ClientCloseFinalizer {
189    fn drop(&mut self) {
190        if self.completed {
191            return;
192        }
193        *self.state.lock() = TransportCloseState::Closed(Err(RpcError::ClientError(
194            "Graceful client close task was cancelled or panicked".to_string(),
195        )));
196        self.notify.notify_waiters();
197    }
198}
199
200impl TransportCloseFinalizer {
201    fn complete(&mut self, result: Result<()>) {
202        *self.state.lock() = TransportCloseState::Closed(result);
203        self.completed = true;
204        self.notify.notify_waiters();
205    }
206}
207
208impl Drop for TransportCloseFinalizer {
209    fn drop(&mut self) {
210        if self.completed {
211            return;
212        }
213        *self.state.lock() = TransportCloseState::Closed(Err(RpcError::ClientError(
214            "Transport close task was cancelled or panicked".to_string(),
215        )));
216        self.notify.notify_waiters();
217    }
218}
219
220async fn close_transport_once<T, C>(
221    transport: &Arc<T>,
222    close_state: &Arc<Mutex<TransportCloseState>>,
223    close_notify: &Arc<Notify>,
224) -> Result<()>
225where
226    T: MessageChannel<C> + 'static,
227    C: Codec + 'static,
228{
229    let should_start = {
230        let mut state = close_state.lock();
231        match &*state {
232            TransportCloseState::Open => {
233                *state = TransportCloseState::Closing;
234                true
235            }
236            TransportCloseState::Closing => false,
237            TransportCloseState::Closed(result) => return result.clone(),
238        }
239    };
240
241    if should_start {
242        let transport = Arc::clone(transport);
243        let close_state = Arc::clone(close_state);
244        let close_notify = Arc::clone(close_notify);
245        tokio::spawn(async move {
246            let mut finalizer = TransportCloseFinalizer {
247                state: close_state,
248                notify: close_notify,
249                completed: false,
250            };
251            let result = transport.close().await.map_err(normalize_transport_error);
252            finalizer.complete(result);
253        });
254    }
255
256    loop {
257        let notified = close_notify.notified();
258        if let TransportCloseState::Closed(result) = &*close_state.lock() {
259            return result.clone();
260        }
261        notified.await;
262    }
263}
264
265#[allow(clippy::too_many_arguments)]
266async fn close_client_gracefully_once<T, C>(
267    transport: &Arc<T>,
268    pending: &PendingRequests<C>,
269    stream_manager: &Arc<StreamManager<C>>,
270    admission: &Arc<RwLock<()>>,
271    stop_tx: &watch::Sender<bool>,
272    transport_close_state: &Arc<Mutex<TransportCloseState>>,
273    transport_close_notify: &Arc<Notify>,
274    client_close_state: &Arc<Mutex<TransportCloseState>>,
275    client_close_notify: &Arc<Notify>,
276    error: RpcError,
277) -> Result<()>
278where
279    T: MessageChannel<C> + 'static,
280    C: Codec + 'static,
281{
282    let should_start = {
283        let mut state = client_close_state.lock();
284        match &*state {
285            TransportCloseState::Open => {
286                *state = TransportCloseState::Closing;
287                true
288            }
289            TransportCloseState::Closing => false,
290            TransportCloseState::Closed(result) => return result.clone(),
291        }
292    };
293
294    if should_start {
295        let transport = Arc::clone(transport);
296        let pending = Arc::clone(pending);
297        let stream_manager = Arc::clone(stream_manager);
298        let admission = Arc::clone(admission);
299        let stop_tx = stop_tx.clone();
300        let transport_close_state = Arc::clone(transport_close_state);
301        let transport_close_notify = Arc::clone(transport_close_notify);
302        let client_close_state = Arc::clone(client_close_state);
303        let client_close_notify = Arc::clone(client_close_notify);
304        tokio::spawn(async move {
305            let mut finalizer = ClientCloseFinalizer {
306                state: client_close_state,
307                notify: client_close_notify,
308                completed: false,
309            };
310            let admission = admission.write().await;
311            stop_tx.send_replace(true);
312            close_client_work(&pending, &stream_manager, error);
313            drop(admission);
314
315            let result =
316                close_transport_once(&transport, &transport_close_state, &transport_close_notify)
317                    .await;
318            finalizer.complete(result);
319        });
320    }
321
322    loop {
323        let notified = client_close_notify.notified();
324        if let TransportCloseState::Closed(result) = &*client_close_state.lock() {
325            return result.clone();
326        }
327        notified.await;
328    }
329}
330
331/// RPC client for making remote procedure calls.
332pub struct RpcClient<T, C: Codec = BincodeCodec>
333where
334    T: MessageChannel<C>,
335{
336    transport: Arc<T>,
337    pending: PendingRequests<C>,
338    stream_manager: Arc<StreamManager<C>>,
339    codec: C,
340    state: Arc<Mutex<ClientState>>,
341    admission: Arc<RwLock<()>>,
342    send_cancel_tx: watch::Sender<Option<RpcError>>,
343    stop_tx: watch::Sender<bool>,
344    receive_abort: ReceiveTaskAbort,
345    transport_close_state: Arc<Mutex<TransportCloseState>>,
346    transport_close_notify: Arc<Notify>,
347    client_close_state: Arc<Mutex<TransportCloseState>>,
348    client_close_notify: Arc<Notify>,
349    default_timeout: Duration,
350}
351
352impl<T, C> Drop for RpcClient<T, C>
353where
354    T: MessageChannel<C>,
355    C: Codec,
356{
357    fn drop(&mut self) {
358        let (error, _) = mark_client_terminal(&self.state, RpcError::ConnectionClosed);
359        self.send_cancel_tx.send_replace(Some(error.clone()));
360        self.stop_tx.send_replace(true);
361        if let Some(abort_handle) = self.receive_abort.lock().take() {
362            abort_handle.abort();
363        }
364        close_client_work(&self.pending, &self.stream_manager, error);
365    }
366}
367
368impl<T: MessageChannel<BincodeCodec> + 'static> RpcClient<T, BincodeCodec> {
369    pub fn new(transport: T) -> Self {
370        Self::with_timeout(transport, Duration::from_secs(30))
371    }
372
373    pub fn with_timeout(transport: T, default_timeout: Duration) -> Self {
374        let (send_cancel_tx, _) = watch::channel(None::<RpcError>);
375        let (stop_tx, _) = watch::channel(false);
376        Self {
377            transport: Arc::new(transport),
378            pending: Arc::new(Mutex::new(HashMap::new())),
379            stream_manager: Arc::new(StreamManager::new()),
380            codec: BincodeCodec,
381            state: Arc::new(Mutex::new(ClientState::default())),
382            admission: Arc::new(RwLock::new(())),
383            send_cancel_tx,
384            stop_tx,
385            receive_abort: Arc::new(Mutex::new(None)),
386            transport_close_state: Arc::new(Mutex::new(TransportCloseState::default())),
387            transport_close_notify: Arc::new(Notify::new()),
388            client_close_state: Arc::new(Mutex::new(TransportCloseState::default())),
389            client_close_notify: Arc::new(Notify::new()),
390            default_timeout,
391        }
392    }
393}
394
395impl<T, C> RpcClient<T, C>
396where
397    T: MessageChannel<C> + 'static,
398    C: Codec + Clone + Default + 'static,
399{
400    pub fn with_codec(transport: T, codec: C) -> Self {
401        let (send_cancel_tx, _) = watch::channel(None::<RpcError>);
402        let (stop_tx, _) = watch::channel(false);
403        Self {
404            transport: Arc::new(transport),
405            pending: Arc::new(Mutex::new(HashMap::new())),
406            stream_manager: Arc::new(StreamManager::with_codec(codec.clone())),
407            codec,
408            state: Arc::new(Mutex::new(ClientState::default())),
409            admission: Arc::new(RwLock::new(())),
410            send_cancel_tx,
411            stop_tx,
412            receive_abort: Arc::new(Mutex::new(None)),
413            transport_close_state: Arc::new(Mutex::new(TransportCloseState::default())),
414            transport_close_notify: Arc::new(Notify::new()),
415            client_close_state: Arc::new(Mutex::new(TransportCloseState::default())),
416            client_close_notify: Arc::new(Notify::new()),
417            default_timeout: Duration::from_secs(30),
418        }
419    }
420
421    pub fn with_codec_and_timeout(transport: T, codec: C, default_timeout: Duration) -> Self {
422        let (send_cancel_tx, _) = watch::channel(None::<RpcError>);
423        let (stop_tx, _) = watch::channel(false);
424        Self {
425            transport: Arc::new(transport),
426            pending: Arc::new(Mutex::new(HashMap::new())),
427            stream_manager: Arc::new(StreamManager::with_codec(codec.clone())),
428            codec,
429            state: Arc::new(Mutex::new(ClientState::default())),
430            admission: Arc::new(RwLock::new(())),
431            send_cancel_tx,
432            stop_tx,
433            receive_abort: Arc::new(Mutex::new(None)),
434            transport_close_state: Arc::new(Mutex::new(TransportCloseState::default())),
435            transport_close_notify: Arc::new(Notify::new()),
436            client_close_state: Arc::new(Mutex::new(TransportCloseState::default())),
437            client_close_notify: Arc::new(Notify::new()),
438            default_timeout,
439        }
440    }
441
442    /// Start the client's background receive loop.
443    pub fn start(&self) -> RpcClientHandle {
444        self.try_start().expect("RPC client cannot be started")
445    }
446
447    /// Start the background receive loop and reject duplicate or post-close startup.
448    pub fn try_start(&self) -> Result<RpcClientHandle> {
449        {
450            let mut state = self.state.lock();
451            if state.closed {
452                return Err(RpcError::ConnectionClosed);
453            }
454            if state.started {
455                return Err(RpcError::ClientError(
456                    "RPC client has already been started".to_string(),
457                ));
458            }
459            state.started = true;
460            state.running = true;
461        }
462
463        let transport = self.transport.clone();
464        let pending = self.pending.clone();
465        let stream_manager = self.stream_manager.clone();
466        let state = self.state.clone();
467        let admission = self.admission.clone();
468        let codec = self.codec.clone();
469        let send_cancel_tx = self.send_cancel_tx.clone();
470        let mut stop_rx = self.stop_tx.subscribe();
471        let receive_abort = self.receive_abort.clone();
472        let transport_close_state = self.transport_close_state.clone();
473        let transport_close_notify = self.transport_close_notify.clone();
474        let (task_ready_tx, task_ready_rx) = oneshot::channel();
475        let finalizer = ClientTaskFinalizer {
476            state: state.clone(),
477            pending: pending.clone(),
478            stream_manager: stream_manager.clone(),
479            send_cancel_tx: send_cancel_tx.clone(),
480            receive_abort: receive_abort.clone(),
481            cleanup_work: true,
482        };
483
484        let handle = tokio::spawn(async move {
485            let mut finalizer = finalizer;
486            if task_ready_rx.await.is_err() {
487                return;
488            }
489            loop {
490                tokio::select! {
491                    stop_result = stop_rx.changed() => {
492                        if stop_result.is_err() || *stop_rx.borrow() {
493                            finalizer.cleanup_work = false;
494                            break;
495                        }
496                    }
497                    receive_result = transport.recv() => match receive_result {
498                        Ok(message) => match message.msg_type {
499                            MessageType::Reply => {
500                                if let Some(tx) = pending.lock().remove(&message.id) {
501                                    let _ = tx.send(Ok(message));
502                                }
503                            }
504                            MessageType::Error => {
505                                if let Some(stream_id) = message.metadata.stream_id {
506                                    let error_msg: String = codec
507                                        .decode(&message.payload)
508                                        .unwrap_or_else(|_| "Unknown error".to_string());
509                                    stream_manager.send_error(stream_id, error_msg);
510                                } else if let Some(tx) = pending.lock().remove(&message.id) {
511                                    let _ = tx.send(Ok(message));
512                                }
513                            }
514                            MessageType::StreamChunk | MessageType::StreamEnd => {
515                                stream_manager.handle_message(&message);
516                            }
517                            _ => {}
518                        },
519                        Err(error) => {
520                            if matches!(&error, TransportError::Timeout { .. })
521                                && transport.is_connected()
522                            {
523                                continue;
524                            }
525
526                            let error = normalize_transport_error(error);
527                            let (error, _) = mark_client_terminal(&state, error);
528                            send_cancel_tx.send_replace(Some(error.clone()));
529                            close_client_work(&pending, &stream_manager, error);
530                            finalizer.cleanup_work = false;
531                            {
532                                let _admission = admission.write().await;
533                            }
534                            let _ = close_transport_once(
535                                &transport,
536                                &transport_close_state,
537                                &transport_close_notify,
538                            ).await;
539                            break;
540                        },
541                    }
542                }
543            }
544        });
545
546        *receive_abort.lock() = Some(handle.abort_handle());
547        let _ = task_ready_tx.send(());
548
549        Ok(RpcClientHandle { handle })
550    }
551
552    fn ensure_running(state: &ClientState) -> Result<()> {
553        if !state.started && !state.closed {
554            return Err(RpcError::ClientError(
555                "RPC client has not been started".to_string(),
556            ));
557        }
558        if state.closed || !state.running {
559            return Err(state
560                .terminal_error
561                .clone()
562                .unwrap_or(RpcError::ConnectionClosed));
563        }
564        Ok(())
565    }
566
567    async fn send_until_committed(&self, message: &Message<C>) -> Result<()> {
568        let mut cancel_rx = self.send_cancel_tx.subscribe();
569        if let Some(error) = cancel_rx.borrow().clone() {
570            return Err(error);
571        }
572
573        tokio::select! {
574            biased;
575            result = self.transport.send(message) => match result {
576                Ok(()) => Ok(()),
577                Err(error) => {
578                    let terminal = is_terminal_outbound_error(
579                        &error,
580                        self.transport.is_connected(),
581                    );
582                    let error = normalize_transport_error(error);
583                    if terminal {
584                        Err(self.handle_terminal_outbound_error(error))
585                    } else {
586                        Err(error)
587                    }
588                }
589            },
590            cancel_result = cancel_rx.changed() => {
591                let _ = cancel_result;
592                Err(cancel_rx
593                    .borrow()
594                    .clone()
595                    .unwrap_or(RpcError::ConnectionClosed))
596            }
597        }
598    }
599
600    fn handle_terminal_outbound_error(&self, error: RpcError) -> RpcError {
601        let (error, claimed) = mark_client_terminal(&self.state, error);
602        if !claimed {
603            return error;
604        }
605
606        self.send_cancel_tx.send_replace(Some(error.clone()));
607        close_client_work(&self.pending, &self.stream_manager, error.clone());
608        self.stop_tx.send_replace(true);
609
610        let admission = Arc::clone(&self.admission);
611        let transport = Arc::clone(&self.transport);
612        let close_state = Arc::clone(&self.transport_close_state);
613        let close_notify = Arc::clone(&self.transport_close_notify);
614        tokio::spawn(async move {
615            {
616                let _admission = admission.write().await;
617            }
618            let _ = close_transport_once(&transport, &close_state, &close_notify).await;
619        });
620
621        error
622    }
623
624    pub fn transport(&self) -> Arc<T> {
625        self.transport.clone()
626    }
627
628    pub fn stream_manager(&self) -> Arc<StreamManager<C>> {
629        self.stream_manager.clone()
630    }
631
632    async fn wait_for_response(
633        &self,
634        msg_id: MessageId,
635        rx: oneshot::Receiver<Result<Message<C>>>,
636        timeout: CallTimeout,
637    ) -> Result<Message<C>> {
638        let timeout = match timeout {
639            CallTimeout::ClientDefault => Some(self.default_timeout),
640            CallTimeout::After(timeout) => Some(timeout),
641            CallTimeout::Disabled => None,
642        };
643
644        match timeout {
645            Some(timeout) => tokio::time::timeout(timeout, rx)
646                .await
647                .map_err(|_| {
648                    RpcError::Timeout(format!("Request {} timed out after {:?}", msg_id, timeout))
649                })?
650                .map_err(|_| RpcError::ConnectionClosed)?,
651            None => {
652                let mut cancel_rx = self.send_cancel_tx.subscribe();
653                if let Some(error) = cancel_rx.borrow().clone() {
654                    return Err(error);
655                }
656
657                tokio::select! {
658                    biased;
659                    response = rx => response.map_err(|_| RpcError::ConnectionClosed)?,
660                    changed = cancel_rx.changed() => {
661                        let _ = changed;
662                        Err(cancel_rx
663                            .borrow()
664                            .clone()
665                            .unwrap_or(RpcError::ConnectionClosed))
666                    }
667                }
668            }
669        }
670    }
671
672    /// Make a typed RPC call.
673    pub async fn call<Req, Resp>(&self, method: &str, request: &Req) -> Result<Resp>
674    where
675        Req: Serialize,
676        Resp: for<'de> Deserialize<'de>,
677    {
678        self.call_with_options(method, request, CallOptions::default())
679            .await
680    }
681
682    /// Make a typed RPC call with custom timeout.
683    pub async fn call_with_timeout<Req, Resp>(
684        &self,
685        method: &str,
686        request: &Req,
687        timeout: Duration,
688    ) -> Result<Resp>
689    where
690        Req: Serialize,
691        Resp: for<'de> Deserialize<'de>,
692    {
693        self.call_with_options(method, request, CallOptions::with_timeout(timeout))
694            .await
695    }
696
697    /// Make a typed RPC call with explicit per-call options.
698    ///
699    /// [`CallOptions::without_timeout`] disables only the application response deadline. The
700    /// local wait still ends on explicit client close or terminal connection failure. Dropping
701    /// the call future removes its local pending registration but does not cancel server work.
702    pub async fn call_with_options<Req, Resp>(
703        &self,
704        method: &str,
705        request: &Req,
706        options: CallOptions,
707    ) -> Result<Resp>
708    where
709        Req: Serialize,
710        Resp: for<'de> Deserialize<'de>,
711    {
712        let _admission = self.admission.read().await;
713        {
714            let state = self.state.lock();
715            Self::ensure_running(&state)?;
716        }
717
718        let message = Message::call_with_codec(method, request, self.codec.clone())?;
719        let msg_id = message.id;
720        let (tx, rx) = oneshot::channel();
721        let _registration = PendingRequestRegistration::new(&self.pending, msg_id, tx);
722
723        self.send_until_committed(&message).await?;
724
725        let response = self.wait_for_response(msg_id, rx, options.timeout).await?;
726
727        match response.msg_type {
728            MessageType::Reply => self.codec.decode(&response.payload),
729            MessageType::Error => {
730                let error_msg: String = self
731                    .codec
732                    .decode(&response.payload)
733                    .unwrap_or_else(|_| "Unknown error".to_string());
734                Err(RpcError::ServerError(error_msg))
735            }
736            _ => Err(RpcError::InvalidMessage(format!(
737                "Unexpected message type: {:?}",
738                response.msg_type
739            ))),
740        }
741    }
742
743    /// Initiate a server streaming RPC call.
744    pub async fn call_server_stream<Req, Resp>(
745        &self,
746        method: &str,
747        request: &Req,
748    ) -> Result<StreamReceiver<Resp, C>>
749    where
750        Req: Serialize,
751        Resp: for<'de> Deserialize<'de>,
752    {
753        self.call_server_stream_with_id(method, request, next_stream_id())
754            .await
755    }
756
757    pub(crate) async fn call_server_stream_with_id<Req, Resp>(
758        &self,
759        method: &str,
760        request: &Req,
761        stream_id: StreamId,
762    ) -> Result<StreamReceiver<Resp, C>>
763    where
764        Req: Serialize,
765        Resp: for<'de> Deserialize<'de>,
766    {
767        let _admission = self.admission.read().await;
768        {
769            let state = self.state.lock();
770            Self::ensure_running(&state)?;
771        }
772
773        let mut message = Message::call_with_codec(method, request, self.codec.clone())?;
774        message.metadata = message.metadata.with_stream(stream_id, 0);
775        let receiver = self.stream_manager.create_receiver::<Resp>(stream_id);
776
777        self.send_until_committed(&message).await?;
778
779        Ok(receiver)
780    }
781
782    /// Send a one-way notification (no response expected).
783    pub async fn notify<Req: Serialize>(&self, method: &str, request: &Req) -> Result<()> {
784        let _admission = self.admission.read().await;
785        {
786            let state = self.state.lock();
787            Self::ensure_running(&state)?;
788        }
789        let message = Message::notification_with_codec(method, request, self.codec.clone())?;
790
791        self.send_until_committed(&message).await
792    }
793
794    /// Make a raw bytes RPC call.
795    pub async fn call_raw(&self, method: &str, payload: Vec<u8>) -> Result<Vec<u8>> {
796        self.call_raw_with_options(method, payload, CallOptions::default())
797            .await
798    }
799
800    /// Make a raw bytes RPC call with custom timeout.
801    pub async fn call_raw_with_timeout(
802        &self,
803        method: &str,
804        payload: Vec<u8>,
805        timeout: Duration,
806    ) -> Result<Vec<u8>> {
807        self.call_raw_with_options(method, payload, CallOptions::with_timeout(timeout))
808            .await
809    }
810
811    /// Make a raw bytes RPC call with explicit per-call options.
812    ///
813    /// [`CallOptions::without_timeout`] disables only the application response deadline. The
814    /// local wait remains interruptible by explicit client close, terminal connection failure,
815    /// and future cancellation. These local events do not cancel server work.
816    pub async fn call_raw_with_options(
817        &self,
818        method: &str,
819        payload: Vec<u8>,
820        options: CallOptions,
821    ) -> Result<Vec<u8>> {
822        let _admission = self.admission.read().await;
823        {
824            let state = self.state.lock();
825            Self::ensure_running(&state)?;
826        }
827
828        let message = Message::new_with_codec(
829            MessageId::new(),
830            MessageType::Call,
831            method,
832            payload.into(),
833            Default::default(),
834            self.codec.clone(),
835        );
836        let msg_id = message.id;
837        let (tx, rx) = oneshot::channel();
838        let _registration = PendingRequestRegistration::new(&self.pending, msg_id, tx);
839
840        self.send_until_committed(&message).await?;
841
842        let response = self.wait_for_response(msg_id, rx, options.timeout).await?;
843
844        match response.msg_type {
845            MessageType::Reply => Ok(response.payload.to_vec()),
846            MessageType::Error => {
847                let error_msg: String = self
848                    .codec
849                    .decode(&response.payload)
850                    .unwrap_or_else(|_| "Unknown error".to_string());
851                Err(RpcError::ServerError(error_msg))
852            }
853            _ => Err(RpcError::InvalidMessage(format!(
854                "Unexpected message type: {:?}",
855                response.msg_type
856            ))),
857        }
858    }
859
860    pub fn is_connected(&self) -> bool {
861        let state = self.state.lock();
862        state.running && !state.closed && self.transport.is_connected()
863    }
864
865    pub fn is_started(&self) -> bool {
866        self.state.lock().started
867    }
868
869    pub fn is_running(&self) -> bool {
870        self.state.lock().running
871    }
872
873    pub fn is_closed(&self) -> bool {
874        self.state.lock().closed
875    }
876
877    pub fn active_streams(&self) -> usize {
878        self.stream_manager.active_stream_count()
879    }
880
881    pub async fn close(&self) -> Result<()> {
882        let (error, _) = mark_client_terminal(&self.state, RpcError::ConnectionClosed);
883        self.send_cancel_tx.send_replace(Some(error.clone()));
884
885        close_client_gracefully_once(
886            &self.transport,
887            &self.pending,
888            &self.stream_manager,
889            &self.admission,
890            &self.stop_tx,
891            &self.transport_close_state,
892            &self.transport_close_notify,
893            &self.client_close_state,
894            &self.client_close_notify,
895            error,
896        )
897        .await
898    }
899}
900
901impl<T, C> Debug for RpcClient<T, C>
902where
903    T: MessageChannel<C>,
904    C: Codec + Clone,
905{
906    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
907        let state = self.state.lock();
908        f.debug_struct("RpcClient")
909            .field("closed", &state.closed)
910            .field("started", &state.started)
911            .field("running", &state.running)
912            .field("pending_requests", &self.pending.lock().len())
913            .field("active_streams", &self.stream_manager.active_stream_count())
914            .finish()
915    }
916}
917
918/// Handle for the client's background task.
919pub struct RpcClientHandle {
920    handle: tokio::task::JoinHandle<()>,
921}
922
923impl RpcClientHandle {
924    /// Wait for the receive task to finish without aborting it.
925    pub async fn join(self) -> std::result::Result<(), tokio::task::JoinError> {
926        self.handle.await
927    }
928
929    /// Abort the receive task when graceful client close is not possible.
930    pub async fn shutdown(self) {
931        self.handle.abort();
932        let _ = self.handle.await;
933    }
934}
935
936#[cfg(test)]
937mod tests {
938    use super::*;
939    use crate::channel::message::MessageChannelAdapter;
940    use crate::codec::JsonCodec;
941    use crate::error::TransportResult;
942    use crate::transport::channel::{ChannelConfig, ChannelFrameTransport};
943    use std::collections::VecDeque;
944    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
945    use tokio::sync::Notify;
946
947    #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
948    struct AddRequest {
949        a: i32,
950        b: i32,
951    }
952
953    #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
954    struct AddResponse {
955        result: i32,
956    }
957
958    #[derive(Debug)]
959    struct LifecycleTestChannel {
960        block_close: bool,
961        block_send: bool,
962        close_count: AtomicUsize,
963        close_release: Notify,
964        close_started: Notify,
965        fail_send: bool,
966        fail_close: bool,
967        panic_close: bool,
968        receive_error: TransportError,
969        receive_failed: Notify,
970        send_release: Notify,
971        sent: Notify,
972    }
973
974    impl LifecycleTestChannel {
975        fn new(fail_send: bool, fail_close: bool) -> Self {
976            Self {
977                block_close: false,
978                block_send: false,
979                close_count: AtomicUsize::new(0),
980                close_release: Notify::new(),
981                close_started: Notify::new(),
982                fail_send,
983                fail_close,
984                panic_close: false,
985                receive_error: TransportError::ConnectionClosed,
986                receive_failed: Notify::new(),
987                send_release: Notify::new(),
988                sent: Notify::new(),
989            }
990        }
991
992        fn with_blocked_send() -> Self {
993            Self {
994                block_send: true,
995                ..Self::new(false, false)
996            }
997        }
998
999        fn with_blocked_close() -> Self {
1000            Self {
1001                block_close: true,
1002                ..Self::new(false, false)
1003            }
1004        }
1005
1006        fn with_blocked_failing_close() -> Self {
1007            Self {
1008                block_close: true,
1009                ..Self::new(false, true)
1010            }
1011        }
1012
1013        fn with_receive_error(receive_error: TransportError) -> Self {
1014            Self {
1015                receive_error,
1016                ..Self::new(false, false)
1017            }
1018        }
1019
1020        fn with_panicking_close() -> Self {
1021            Self {
1022                panic_close: true,
1023                ..Self::new(false, false)
1024            }
1025        }
1026    }
1027
1028    #[async_trait::async_trait]
1029    impl MessageChannel for LifecycleTestChannel {
1030        async fn send(&self, _message: &Message) -> TransportResult<()> {
1031            self.sent.notify_one();
1032            if self.block_send {
1033                self.send_release.notified().await;
1034                return Err(TransportError::ConnectionClosed);
1035            }
1036            if self.fail_send {
1037                return Err(TransportError::ConnectionClosed);
1038            }
1039            Ok(())
1040        }
1041
1042        async fn recv(&self) -> TransportResult<Message> {
1043            self.receive_failed.notified().await;
1044            Err(self.receive_error.clone())
1045        }
1046
1047        fn is_connected(&self) -> bool {
1048            true
1049        }
1050
1051        async fn close(&self) -> TransportResult<()> {
1052            self.close_count.fetch_add(1, Ordering::SeqCst);
1053            self.close_started.notify_one();
1054            if self.panic_close {
1055                panic!("test transport close panic");
1056            }
1057            if self.block_close {
1058                self.close_release.notified().await;
1059            }
1060            self.send_release.notify_one();
1061            if self.fail_close {
1062                return Err(TransportError::ConnectionClosed);
1063            }
1064            Ok(())
1065        }
1066    }
1067
1068    #[derive(Debug, Default)]
1069    struct TimeoutThenReplyChannel {
1070        message_id: Mutex<Option<MessageId>>,
1071        message_sent: Notify,
1072        receive_count: AtomicUsize,
1073    }
1074
1075    #[async_trait::async_trait]
1076    impl MessageChannel for TimeoutThenReplyChannel {
1077        async fn send(&self, message: &Message) -> TransportResult<()> {
1078            *self.message_id.lock() = Some(message.id);
1079            self.message_sent.notify_one();
1080            Ok(())
1081        }
1082
1083        async fn recv(&self) -> TransportResult<Message> {
1084            let receive_count = self.receive_count.fetch_add(1, Ordering::SeqCst);
1085            if receive_count == 0 {
1086                return Err(TransportError::Timeout {
1087                    duration_ms: 10,
1088                    operation: "receive".to_string(),
1089                });
1090            }
1091            if receive_count > 1 {
1092                return std::future::pending().await;
1093            }
1094
1095            loop {
1096                let notified = self.message_sent.notified();
1097                if let Some(message_id) = *self.message_id.lock() {
1098                    return Message::reply(message_id, AddResponse { result: 42 })
1099                        .map_err(|error| TransportError::Protocol(error.to_string()));
1100                }
1101                notified.await;
1102            }
1103        }
1104
1105        fn is_connected(&self) -> bool {
1106            true
1107        }
1108
1109        async fn close(&self) -> TransportResult<()> {
1110            Ok(())
1111        }
1112    }
1113
1114    #[derive(Debug, Default)]
1115    struct CommitTestChannel {
1116        close_count: AtomicUsize,
1117        committed: Notify,
1118        message_id: Mutex<Option<MessageId>>,
1119        reply_release: Notify,
1120        send_release: Notify,
1121        send_started: Notify,
1122    }
1123
1124    #[async_trait::async_trait]
1125    impl MessageChannel for CommitTestChannel {
1126        async fn send(&self, message: &Message) -> TransportResult<()> {
1127            self.send_started.notify_one();
1128            self.send_release.notified().await;
1129            *self.message_id.lock() = Some(message.id);
1130            self.committed.notify_one();
1131            Ok(())
1132        }
1133
1134        async fn recv(&self) -> TransportResult<Message> {
1135            self.reply_release.notified().await;
1136            let message_id = self
1137                .message_id
1138                .lock()
1139                .expect("send must commit before reply release");
1140            Message::reply(message_id, AddResponse { result: 3 })
1141                .map_err(|error| TransportError::Protocol(error.to_string()))
1142        }
1143
1144        fn is_connected(&self) -> bool {
1145            true
1146        }
1147
1148        async fn close(&self) -> TransportResult<()> {
1149            self.close_count.fetch_add(1, Ordering::SeqCst);
1150            Ok(())
1151        }
1152    }
1153
1154    #[derive(Debug)]
1155    struct TerminalSendChannel {
1156        close_count: AtomicUsize,
1157        close_started: Notify,
1158        connected: AtomicBool,
1159        fail_send: AtomicBool,
1160        sent: Notify,
1161    }
1162
1163    impl Default for TerminalSendChannel {
1164        fn default() -> Self {
1165            Self {
1166                close_count: AtomicUsize::new(0),
1167                close_started: Notify::new(),
1168                connected: AtomicBool::new(true),
1169                fail_send: AtomicBool::new(false),
1170                sent: Notify::new(),
1171            }
1172        }
1173    }
1174
1175    #[async_trait::async_trait]
1176    impl MessageChannel for TerminalSendChannel {
1177        async fn send(&self, _message: &Message) -> TransportResult<()> {
1178            if self.fail_send.load(Ordering::SeqCst) {
1179                return Err(TransportError::NotConnected);
1180            }
1181            self.sent.notify_one();
1182            Ok(())
1183        }
1184
1185        async fn recv(&self) -> TransportResult<Message> {
1186            std::future::pending().await
1187        }
1188
1189        fn is_connected(&self) -> bool {
1190            self.connected.load(Ordering::SeqCst)
1191        }
1192
1193        async fn close(&self) -> TransportResult<()> {
1194            self.close_count.fetch_add(1, Ordering::SeqCst);
1195            self.close_started.notify_one();
1196            Ok(())
1197        }
1198    }
1199
1200    #[derive(Debug)]
1201    struct ConnectedTimeoutChannel {
1202        close_count: AtomicUsize,
1203        fail_next: AtomicBool,
1204    }
1205
1206    impl Default for ConnectedTimeoutChannel {
1207        fn default() -> Self {
1208            Self {
1209                close_count: AtomicUsize::new(0),
1210                fail_next: AtomicBool::new(true),
1211            }
1212        }
1213    }
1214
1215    #[async_trait::async_trait]
1216    impl MessageChannel for ConnectedTimeoutChannel {
1217        async fn send(&self, _message: &Message) -> TransportResult<()> {
1218            if self.fail_next.swap(false, Ordering::SeqCst) {
1219                return Err(TransportError::Timeout {
1220                    duration_ms: 10,
1221                    operation: "send".to_string(),
1222                });
1223            }
1224            Ok(())
1225        }
1226
1227        async fn recv(&self) -> TransportResult<Message> {
1228            std::future::pending().await
1229        }
1230
1231        fn is_connected(&self) -> bool {
1232            true
1233        }
1234
1235        async fn close(&self) -> TransportResult<()> {
1236            self.close_count.fetch_add(1, Ordering::SeqCst);
1237            Ok(())
1238        }
1239    }
1240
1241    #[derive(Debug)]
1242    struct DelayedReplyChannel {
1243        connected: AtomicBool,
1244        delay: Duration,
1245        raw: bool,
1246        requests: Mutex<VecDeque<(MessageId, Vec<u8>)>>,
1247        request_notify: Notify,
1248    }
1249
1250    impl DelayedReplyChannel {
1251        fn typed(delay: Duration) -> Self {
1252            Self::new(delay, false)
1253        }
1254
1255        fn raw(delay: Duration) -> Self {
1256            Self::new(delay, true)
1257        }
1258
1259        fn new(delay: Duration, raw: bool) -> Self {
1260            Self {
1261                connected: AtomicBool::new(true),
1262                delay,
1263                raw,
1264                requests: Mutex::new(VecDeque::new()),
1265                request_notify: Notify::new(),
1266            }
1267        }
1268
1269        async fn next_request(&self) -> (MessageId, Vec<u8>) {
1270            loop {
1271                let notified = self.request_notify.notified();
1272                if let Some(request) = self.requests.lock().pop_front() {
1273                    return request;
1274                }
1275                notified.await;
1276            }
1277        }
1278    }
1279
1280    #[async_trait::async_trait]
1281    impl MessageChannel for DelayedReplyChannel {
1282        async fn send(&self, message: &Message) -> TransportResult<()> {
1283            self.requests
1284                .lock()
1285                .push_back((message.id, message.payload.to_vec()));
1286            self.request_notify.notify_one();
1287            Ok(())
1288        }
1289
1290        async fn recv(&self) -> TransportResult<Message> {
1291            let (message_id, payload) = self.next_request().await;
1292            tokio::time::sleep(self.delay).await;
1293            if self.raw {
1294                Ok(Message::new(
1295                    message_id,
1296                    MessageType::Reply,
1297                    "",
1298                    payload.into(),
1299                    Default::default(),
1300                ))
1301            } else {
1302                Message::reply(message_id, AddResponse { result: 42 })
1303                    .map_err(|error| TransportError::Protocol(error.to_string()))
1304            }
1305        }
1306
1307        fn is_connected(&self) -> bool {
1308            self.connected.load(Ordering::Acquire)
1309        }
1310
1311        async fn close(&self) -> TransportResult<()> {
1312            self.connected.store(false, Ordering::Release);
1313            self.request_notify.notify_waiters();
1314            Ok(())
1315        }
1316    }
1317
1318    #[derive(Debug, Clone)]
1319    struct StatefulCodec {
1320        tag: Arc<str>,
1321    }
1322
1323    impl StatefulCodec {
1324        fn configured() -> Self {
1325            Self {
1326                tag: Arc::from("configured"),
1327            }
1328        }
1329    }
1330
1331    impl Default for StatefulCodec {
1332        fn default() -> Self {
1333            Self {
1334                tag: Arc::from("default"),
1335            }
1336        }
1337    }
1338
1339    impl Codec for StatefulCodec {
1340        fn encode<T: Serialize>(&self, data: &T) -> Result<Vec<u8>> {
1341            let mut payload = self.tag.as_bytes().to_vec();
1342            payload.push(b':');
1343            payload.extend(serde_json::to_vec(data)?);
1344            Ok(payload)
1345        }
1346
1347        fn decode<T: for<'de> Deserialize<'de>>(&self, data: &[u8]) -> Result<T> {
1348            let mut prefix = self.tag.as_bytes().to_vec();
1349            prefix.push(b':');
1350            let payload = data.strip_prefix(prefix.as_slice()).ok_or_else(|| {
1351                RpcError::Serialization(format!("missing {} codec prefix", self.tag))
1352            })?;
1353            Ok(serde_json::from_slice(payload)?)
1354        }
1355    }
1356
1357    #[tokio::test]
1358    async fn test_call_options_default_and_explicit_timeout_resolution() {
1359        let default_client = RpcClient::with_timeout(
1360            DelayedReplyChannel::typed(Duration::from_millis(75)),
1361            Duration::from_millis(20),
1362        );
1363        let default_handle = default_client.start();
1364        let result: Result<AddResponse> = default_client
1365            .call_with_options("add", &AddRequest { a: 20, b: 22 }, CallOptions::default())
1366            .await;
1367        assert!(matches!(result, Err(RpcError::Timeout(_))));
1368        default_client.close().await.unwrap();
1369        default_handle.join().await.unwrap();
1370
1371        let override_client = RpcClient::with_timeout(
1372            DelayedReplyChannel::typed(Duration::from_millis(75)),
1373            Duration::from_millis(20),
1374        );
1375        let override_handle = override_client.start();
1376        let response: AddResponse = override_client
1377            .call_with_options(
1378                "add",
1379                &AddRequest { a: 20, b: 22 },
1380                CallOptions::with_timeout(Duration::from_millis(250)),
1381            )
1382            .await
1383            .unwrap();
1384        assert_eq!(response.result, 42);
1385        override_client.close().await.unwrap();
1386        override_handle.join().await.unwrap();
1387    }
1388
1389    #[tokio::test]
1390    async fn test_deadline_free_typed_call_exceeds_client_default_timeout() {
1391        let client = RpcClient::with_timeout(
1392            DelayedReplyChannel::typed(Duration::from_millis(75)),
1393            Duration::from_millis(20),
1394        );
1395        let handle = client.start();
1396
1397        let response: AddResponse = client
1398            .call_with_options(
1399                "add",
1400                &AddRequest { a: 20, b: 22 },
1401                CallOptions::without_timeout(),
1402            )
1403            .await
1404            .unwrap();
1405        assert_eq!(response.result, 42);
1406
1407        client.close().await.unwrap();
1408        handle.join().await.unwrap();
1409    }
1410
1411    #[tokio::test]
1412    async fn test_deadline_free_raw_call_exceeds_client_default_timeout() {
1413        let client = RpcClient::with_timeout(
1414            DelayedReplyChannel::raw(Duration::from_millis(75)),
1415            Duration::from_millis(20),
1416        );
1417        let handle = client.start();
1418        let payload = vec![1, 2, 3, 4];
1419
1420        let response = client
1421            .call_raw_with_options("raw", payload.clone(), CallOptions::without_timeout())
1422            .await
1423            .unwrap();
1424        assert_eq!(response, payload);
1425
1426        client.close().await.unwrap();
1427        handle.join().await.unwrap();
1428    }
1429
1430    #[tokio::test]
1431    async fn test_client_close_unblocks_deadline_free_typed_call() {
1432        let client = Arc::new(RpcClient::new(LifecycleTestChannel::new(false, false)));
1433        let transport = client.transport();
1434        let handle = client.start();
1435        let call_client = Arc::clone(&client);
1436        let call = tokio::spawn(async move {
1437            call_client
1438                .call_with_options::<_, AddResponse>(
1439                    "add",
1440                    &AddRequest { a: 1, b: 2 },
1441                    CallOptions::without_timeout(),
1442                )
1443                .await
1444        });
1445        tokio::time::timeout(Duration::from_secs(1), transport.sent.notified())
1446            .await
1447            .unwrap();
1448
1449        tokio::time::timeout(Duration::from_secs(1), client.close())
1450            .await
1451            .expect("client close was blocked by a deadline-free typed call")
1452            .unwrap();
1453        assert!(matches!(
1454            call.await.unwrap(),
1455            Err(RpcError::ConnectionClosed)
1456        ));
1457        assert!(client.pending.lock().is_empty());
1458        handle.join().await.unwrap();
1459    }
1460
1461    #[tokio::test]
1462    async fn test_client_close_unblocks_deadline_free_raw_call() {
1463        let client = Arc::new(RpcClient::new(LifecycleTestChannel::new(false, false)));
1464        let transport = client.transport();
1465        let handle = client.start();
1466        let call_client = Arc::clone(&client);
1467        let call = tokio::spawn(async move {
1468            call_client
1469                .call_raw_with_options("raw", vec![1, 2, 3], CallOptions::without_timeout())
1470                .await
1471        });
1472        tokio::time::timeout(Duration::from_secs(1), transport.sent.notified())
1473            .await
1474            .unwrap();
1475
1476        tokio::time::timeout(Duration::from_secs(1), client.close())
1477            .await
1478            .expect("client close was blocked by a deadline-free raw call")
1479            .unwrap();
1480        assert!(matches!(
1481            call.await.unwrap(),
1482            Err(RpcError::ConnectionClosed)
1483        ));
1484        assert!(client.pending.lock().is_empty());
1485        handle.join().await.unwrap();
1486    }
1487
1488    #[tokio::test]
1489    async fn test_terminal_failure_unblocks_deadline_free_typed_call() {
1490        let client = Arc::new(RpcClient::new(LifecycleTestChannel::with_receive_error(
1491            TransportError::Protocol("terminal typed".to_string()),
1492        )));
1493        let transport = client.transport();
1494        let handle = client.start();
1495        let call_client = Arc::clone(&client);
1496        let call = tokio::spawn(async move {
1497            call_client
1498                .call_with_options::<_, AddResponse>(
1499                    "add",
1500                    &AddRequest { a: 1, b: 2 },
1501                    CallOptions::without_timeout(),
1502                )
1503                .await
1504        });
1505        tokio::time::timeout(Duration::from_secs(1), transport.sent.notified())
1506            .await
1507            .unwrap();
1508        transport.receive_failed.notify_one();
1509
1510        assert!(matches!(
1511            call.await.unwrap(),
1512            Err(RpcError::Transport(TransportError::Protocol(message)))
1513                if message == "terminal typed"
1514        ));
1515        assert!(client.pending.lock().is_empty());
1516        handle.join().await.unwrap();
1517        client.close().await.unwrap();
1518    }
1519
1520    #[tokio::test]
1521    async fn test_terminal_failure_unblocks_deadline_free_raw_call() {
1522        let client = Arc::new(RpcClient::new(LifecycleTestChannel::with_receive_error(
1523            TransportError::Protocol("terminal raw".to_string()),
1524        )));
1525        let transport = client.transport();
1526        let handle = client.start();
1527        let call_client = Arc::clone(&client);
1528        let call = tokio::spawn(async move {
1529            call_client
1530                .call_raw_with_options("raw", vec![1, 2, 3], CallOptions::without_timeout())
1531                .await
1532        });
1533        tokio::time::timeout(Duration::from_secs(1), transport.sent.notified())
1534            .await
1535            .unwrap();
1536        transport.receive_failed.notify_one();
1537
1538        assert!(matches!(
1539            call.await.unwrap(),
1540            Err(RpcError::Transport(TransportError::Protocol(message)))
1541                if message == "terminal raw"
1542        ));
1543        assert!(client.pending.lock().is_empty());
1544        handle.join().await.unwrap();
1545        client.close().await.unwrap();
1546    }
1547
1548    #[tokio::test]
1549    async fn test_dropped_deadline_free_typed_call_removes_pending_registration() {
1550        let client = Arc::new(RpcClient::new(LifecycleTestChannel::new(false, false)));
1551        let transport = client.transport();
1552        let handle = client.start();
1553        let call_client = Arc::clone(&client);
1554        let call = tokio::spawn(async move {
1555            call_client
1556                .call_with_options::<_, AddResponse>(
1557                    "add",
1558                    &AddRequest { a: 1, b: 2 },
1559                    CallOptions::without_timeout(),
1560                )
1561                .await
1562        });
1563        tokio::time::timeout(Duration::from_secs(1), transport.sent.notified())
1564            .await
1565            .unwrap();
1566        call.abort();
1567        assert!(call.await.unwrap_err().is_cancelled());
1568        assert!(client.pending.lock().is_empty());
1569
1570        client.close().await.unwrap();
1571        handle.join().await.unwrap();
1572    }
1573
1574    #[tokio::test]
1575    async fn test_dropped_deadline_free_raw_call_removes_pending_registration() {
1576        let client = Arc::new(RpcClient::new(LifecycleTestChannel::new(false, false)));
1577        let transport = client.transport();
1578        let handle = client.start();
1579        let call_client = Arc::clone(&client);
1580        let call = tokio::spawn(async move {
1581            call_client
1582                .call_raw_with_options("raw", vec![1, 2, 3], CallOptions::without_timeout())
1583                .await
1584        });
1585        tokio::time::timeout(Duration::from_secs(1), transport.sent.notified())
1586            .await
1587            .unwrap();
1588        call.abort();
1589        assert!(call.await.unwrap_err().is_cancelled());
1590        assert!(client.pending.lock().is_empty());
1591
1592        client.close().await.unwrap();
1593        handle.join().await.unwrap();
1594    }
1595
1596    #[derive(Debug)]
1597    struct CapturedMessage {
1598        codec_tag: Arc<str>,
1599        method: String,
1600        payload: Vec<u8>,
1601        stream_id: Option<StreamId>,
1602    }
1603
1604    #[derive(Debug)]
1605    struct ConfiguredCodecChannel {
1606        captured: Mutex<Vec<CapturedMessage>>,
1607        codec: StatefulCodec,
1608        replies: Mutex<VecDeque<(MessageId, String)>>,
1609        reply_ready: Notify,
1610    }
1611
1612    impl ConfiguredCodecChannel {
1613        fn new(codec: StatefulCodec) -> Self {
1614            Self {
1615                captured: Mutex::new(Vec::new()),
1616                codec,
1617                replies: Mutex::new(VecDeque::new()),
1618                reply_ready: Notify::new(),
1619            }
1620        }
1621    }
1622
1623    #[async_trait::async_trait]
1624    impl MessageChannel<StatefulCodec> for ConfiguredCodecChannel {
1625        async fn send(&self, message: &Message<StatefulCodec>) -> TransportResult<()> {
1626            self.captured.lock().push(CapturedMessage {
1627                codec_tag: Arc::clone(&message.codec.tag),
1628                method: message.method.clone(),
1629                payload: message.payload.to_vec(),
1630                stream_id: message.metadata.stream_id,
1631            });
1632            if message.metadata.stream_id.is_none()
1633                && matches!(message.method.as_str(), "typed" | "raw")
1634            {
1635                self.replies
1636                    .lock()
1637                    .push_back((message.id, message.method.clone()));
1638                self.reply_ready.notify_one();
1639            }
1640            Ok(())
1641        }
1642
1643        async fn recv(&self) -> TransportResult<Message<StatefulCodec>> {
1644            loop {
1645                let notified = self.reply_ready.notified();
1646                let reply = self.replies.lock().pop_front();
1647                if let Some((id, method)) = reply {
1648                    let payload = match method.as_str() {
1649                        "typed" => self
1650                            .codec
1651                            .encode(&AddResponse { result: 3 })
1652                            .map_err(|error| TransportError::Protocol(error.to_string()))?,
1653                        "raw" => vec![9, 8, 7],
1654                        _ => unreachable!(),
1655                    };
1656                    return Ok(Message::new_with_codec(
1657                        id,
1658                        MessageType::Reply,
1659                        "",
1660                        payload.into(),
1661                        Default::default(),
1662                        self.codec.clone(),
1663                    ));
1664                }
1665                notified.await;
1666            }
1667        }
1668
1669        fn is_connected(&self) -> bool {
1670            true
1671        }
1672
1673        async fn close(&self) -> TransportResult<()> {
1674            Ok(())
1675        }
1676    }
1677
1678    #[tokio::test]
1679    async fn test_client_call_reply() {
1680        let config = ChannelConfig::default();
1681        let (t1, t2) = ChannelFrameTransport::create_pair("test", config).unwrap();
1682
1683        let client_channel = MessageChannelAdapter::new(t1);
1684        let server_channel = MessageChannelAdapter::new(t2);
1685
1686        let client = RpcClient::new(client_channel);
1687        let handle = client.start();
1688
1689        let server_handle = tokio::spawn(async move {
1690            let msg = server_channel.recv().await.unwrap();
1691            assert_eq!(msg.method, "add");
1692
1693            let req: AddRequest = msg.deserialize_payload().unwrap();
1694            let resp = AddResponse {
1695                result: req.a + req.b,
1696            };
1697
1698            let reply: Message = Message::reply(msg.id, resp).unwrap();
1699            server_channel.send(&reply).await.unwrap();
1700        });
1701
1702        let response: AddResponse = client
1703            .call("add", &AddRequest { a: 10, b: 32 })
1704            .await
1705            .unwrap();
1706        assert_eq!(response.result, 42);
1707
1708        server_handle.await.unwrap();
1709        client.close().await.unwrap();
1710        handle.join().await.unwrap();
1711    }
1712
1713    #[tokio::test]
1714    async fn test_configured_codec_is_retained_for_all_outbound_messages() {
1715        let codec = StatefulCodec::configured();
1716        let client = RpcClient::with_codec(ConfiguredCodecChannel::new(codec.clone()), codec);
1717        let transport = client.transport();
1718        let handle = client.start();
1719
1720        let response: AddResponse = client
1721            .call("typed", &AddRequest { a: 1, b: 2 })
1722            .await
1723            .unwrap();
1724        assert_eq!(response.result, 3);
1725        let stream: StreamReceiver<i32, StatefulCodec> =
1726            client.call_server_stream("stream", &17).await.unwrap();
1727        drop(stream);
1728        client.notify("notification", &23).await.unwrap();
1729        assert_eq!(
1730            client.call_raw("raw", vec![1, 2, 3]).await.unwrap(),
1731            vec![9, 8, 7]
1732        );
1733
1734        {
1735            let captured = transport.captured.lock();
1736            assert_eq!(captured.len(), 4);
1737            assert_eq!(
1738                captured
1739                    .iter()
1740                    .map(|message| message.method.as_str())
1741                    .collect::<Vec<_>>(),
1742                vec!["typed", "stream", "notification", "raw"]
1743            );
1744            assert!(
1745                captured
1746                    .iter()
1747                    .all(|message| message.codec_tag.as_ref() == "configured")
1748            );
1749            assert!(
1750                captured[..3]
1751                    .iter()
1752                    .all(|message| message.payload.starts_with(b"configured:"))
1753            );
1754            assert!(captured[1].stream_id.is_some());
1755            assert_eq!(captured[3].payload, vec![1, 2, 3]);
1756        }
1757
1758        client.close().await.unwrap();
1759        handle.join().await.unwrap();
1760    }
1761
1762    #[tokio::test]
1763    async fn test_connected_receive_timeout_keeps_client_running() {
1764        let client = RpcClient::new(TimeoutThenReplyChannel::default());
1765        let handle = client.start();
1766
1767        let response: AddResponse = client
1768            .call("add", &AddRequest { a: 10, b: 32 })
1769            .await
1770            .unwrap();
1771
1772        assert_eq!(response.result, 42);
1773        assert!(client.is_running());
1774        client.close().await.unwrap();
1775        handle.join().await.unwrap();
1776    }
1777
1778    #[tokio::test]
1779    async fn test_committed_blocked_send_completes_before_close() {
1780        let client = Arc::new(RpcClient::new(CommitTestChannel::default()));
1781        let transport = client.transport();
1782        let handle = client.start();
1783        let call_client = Arc::clone(&client);
1784        let call = tokio::spawn(async move {
1785            call_client
1786                .call::<_, AddResponse>("add", &AddRequest { a: 1, b: 2 })
1787                .await
1788        });
1789
1790        tokio::time::timeout(Duration::from_secs(1), transport.send_started.notified())
1791            .await
1792            .unwrap();
1793        transport.send_release.notify_one();
1794        tokio::time::timeout(Duration::from_secs(1), transport.committed.notified())
1795            .await
1796            .unwrap();
1797
1798        let close_client = Arc::clone(&client);
1799        let mut close = tokio::spawn(async move { close_client.close().await });
1800        assert!(
1801            tokio::time::timeout(Duration::from_millis(50), &mut close)
1802                .await
1803                .is_err(),
1804            "close completed before the committed call received its reply"
1805        );
1806
1807        transport.reply_release.notify_one();
1808        let response = call.await.unwrap().unwrap();
1809        assert_eq!(response.result, 3);
1810        close.await.unwrap().unwrap();
1811        assert_eq!(transport.close_count.load(Ordering::SeqCst), 1);
1812        handle.join().await.unwrap();
1813    }
1814
1815    #[tokio::test]
1816    async fn test_client_server_stream() {
1817        let config = ChannelConfig::default();
1818        let (t1, t2) = ChannelFrameTransport::create_pair("test", config).unwrap();
1819
1820        let client_channel = MessageChannelAdapter::new(t1);
1821        let server_channel = Arc::new(MessageChannelAdapter::new(t2));
1822
1823        let client = RpcClient::new(client_channel);
1824        let handle = client.start();
1825
1826        let server_channel_clone = server_channel.clone();
1827        let server_handle = tokio::spawn(async move {
1828            let msg = server_channel_clone.recv().await.unwrap();
1829            let stream_id = msg.metadata.stream_id.unwrap();
1830
1831            for i in 1..=3 {
1832                let chunk: Message = Message::stream_chunk(stream_id, i - 1, i as i32).unwrap();
1833                server_channel_clone.send(&chunk).await.unwrap();
1834            }
1835
1836            let end: Message = Message::stream_end(stream_id);
1837            server_channel_clone.send(&end).await.unwrap();
1838        });
1839
1840        let mut stream: StreamReceiver<i32> =
1841            client.call_server_stream("get_numbers", &()).await.unwrap();
1842
1843        let mut items = Vec::new();
1844        while let Some(result) = stream.recv().await {
1845            items.push(result.unwrap());
1846        }
1847
1848        assert_eq!(items, vec![1, 2, 3]);
1849        server_handle.await.unwrap();
1850        client.close().await.unwrap();
1851        handle.join().await.unwrap();
1852    }
1853
1854    #[tokio::test]
1855    async fn test_server_stream_uses_caller_provided_id() {
1856        let config = ChannelConfig::default();
1857        let (t1, t2) = ChannelFrameTransport::create_pair("test", config).unwrap();
1858
1859        let client = RpcClient::new(MessageChannelAdapter::new(t1));
1860        let server_channel = MessageChannelAdapter::new(t2);
1861        let handle = client.start();
1862        let stream_id = 42;
1863
1864        let server_handle = tokio::spawn(async move {
1865            server_channel
1866                .recv()
1867                .await
1868                .unwrap()
1869                .metadata
1870                .stream_id
1871                .unwrap()
1872        });
1873        let stream: StreamReceiver<i32> = client
1874            .call_server_stream_with_id("numbers", &(), stream_id)
1875            .await
1876            .unwrap();
1877
1878        assert_eq!(stream.stream_id(), stream_id);
1879        assert_eq!(server_handle.await.unwrap(), stream_id);
1880        client.close().await.unwrap();
1881        handle.join().await.unwrap();
1882    }
1883
1884    #[tokio::test]
1885    async fn test_client_stream_error_uses_configured_codec() {
1886        let config = ChannelConfig::default();
1887        let (t1, t2) = ChannelFrameTransport::create_pair("test", config).unwrap();
1888
1889        let client_channel = MessageChannelAdapter::<_, JsonCodec>::with_codec(t1);
1890        let server_channel = MessageChannelAdapter::<_, JsonCodec>::with_codec(t2);
1891
1892        let client = RpcClient::with_codec(client_channel, JsonCodec);
1893        let handle = client.start();
1894
1895        let server_handle = tokio::spawn(async move {
1896            let msg = server_channel.recv().await.unwrap();
1897            let stream_id = msg.metadata.stream_id.unwrap();
1898
1899            let error: Message<JsonCodec> =
1900                Message::stream_error(msg.id, stream_id, "method not found");
1901            server_channel.send(&error).await.unwrap();
1902        });
1903
1904        let mut stream: StreamReceiver<i32, JsonCodec> = client
1905            .call_server_stream("unknown_method", &())
1906            .await
1907            .unwrap();
1908
1909        assert!(matches!(
1910            stream.recv().await,
1911            Some(Err(RpcError::ServerError(error))) if error == "method not found"
1912        ));
1913
1914        server_handle.await.unwrap();
1915        client.close().await.unwrap();
1916        handle.join().await.unwrap();
1917    }
1918
1919    #[tokio::test]
1920    async fn test_call_before_start_fails_promptly() {
1921        let client = RpcClient::new(LifecycleTestChannel::new(false, false));
1922
1923        let result: Result<AddResponse> = client.call("add", &AddRequest { a: 1, b: 2 }).await;
1924
1925        assert!(matches!(result, Err(RpcError::ClientError(_))));
1926        assert!(client.pending.lock().is_empty());
1927        assert!(!client.is_started());
1928    }
1929
1930    #[tokio::test]
1931    async fn test_duplicate_and_post_close_start_fail() {
1932        let client = RpcClient::new(LifecycleTestChannel::new(false, false));
1933        let handle = client.try_start().unwrap();
1934
1935        assert!(client.is_started());
1936        assert!(client.is_running());
1937        assert!(matches!(client.try_start(), Err(RpcError::ClientError(_))));
1938
1939        client.close().await.unwrap();
1940
1941        assert!(client.is_closed());
1942        assert!(!client.is_running());
1943        assert!(matches!(
1944            client.try_start(),
1945            Err(RpcError::ConnectionClosed)
1946        ));
1947        handle.join().await.unwrap();
1948    }
1949
1950    #[tokio::test]
1951    async fn test_close_waits_for_committed_call_timeout() {
1952        let client = Arc::new(RpcClient::new(LifecycleTestChannel::new(false, false)));
1953        let transport = client.transport();
1954        let handle = client.start();
1955        let call_client = client.clone();
1956        let call_handle = tokio::spawn(async move {
1957            call_client
1958                .call_with_timeout::<_, AddResponse>(
1959                    "add",
1960                    &AddRequest { a: 1, b: 2 },
1961                    Duration::from_millis(50),
1962                )
1963                .await
1964        });
1965
1966        tokio::time::timeout(Duration::from_secs(1), transport.sent.notified())
1967            .await
1968            .unwrap();
1969        client.close().await.unwrap();
1970
1971        let result = call_handle.await.unwrap();
1972        assert!(matches!(result, Err(RpcError::Timeout(_))));
1973        assert!(client.pending.lock().is_empty());
1974        handle.join().await.unwrap();
1975    }
1976
1977    #[tokio::test]
1978    async fn test_terminal_receive_error_is_preserved_for_pending_work() {
1979        let client = Arc::new(RpcClient::new(LifecycleTestChannel::with_receive_error(
1980            TransportError::Protocol("invalid frame".to_string()),
1981        )));
1982        let transport = client.transport();
1983        let handle = client.start();
1984        let call_client = Arc::clone(&client);
1985        let call = tokio::spawn(async move {
1986            call_client
1987                .call::<_, AddResponse>("add", &AddRequest { a: 1, b: 2 })
1988                .await
1989        });
1990
1991        tokio::time::timeout(Duration::from_secs(1), transport.sent.notified())
1992            .await
1993            .unwrap();
1994        let mut stream: StreamReceiver<i32> =
1995            client.call_server_stream("numbers", &()).await.unwrap();
1996        transport.receive_failed.notify_one();
1997
1998        assert!(matches!(
1999            call.await.unwrap(),
2000            Err(RpcError::Transport(TransportError::Protocol(error)))
2001                if error == "invalid frame"
2002        ));
2003        assert!(matches!(
2004            stream.recv().await,
2005            Some(Err(RpcError::Transport(TransportError::Protocol(error))))
2006                if error == "invalid frame"
2007        ));
2008        assert!(client.is_closed());
2009        handle.join().await.unwrap();
2010        client.close().await.unwrap();
2011    }
2012
2013    #[tokio::test]
2014    async fn test_active_stream_fails_when_client_closes() {
2015        let client = RpcClient::new(LifecycleTestChannel::new(false, false));
2016        let handle = client.start();
2017        let mut stream: StreamReceiver<i32> =
2018            client.call_server_stream("numbers", &()).await.unwrap();
2019
2020        assert_eq!(client.active_streams(), 1);
2021        client.close().await.unwrap();
2022
2023        assert!(matches!(
2024            stream.recv().await,
2025            Some(Err(RpcError::ConnectionClosed))
2026        ));
2027        assert_eq!(client.active_streams(), 0);
2028        handle.join().await.unwrap();
2029    }
2030
2031    #[tokio::test]
2032    async fn test_terminal_outbound_send_drains_work_and_closes_once() {
2033        let client = Arc::new(RpcClient::new(TerminalSendChannel::default()));
2034        let transport = client.transport();
2035        let handle = client.start();
2036        let call_client = Arc::clone(&client);
2037        let call = tokio::spawn(async move {
2038            call_client
2039                .call::<_, AddResponse>("pending", &AddRequest { a: 1, b: 2 })
2040                .await
2041        });
2042
2043        tokio::time::timeout(Duration::from_secs(1), transport.sent.notified())
2044            .await
2045            .unwrap();
2046        let mut stream: StreamReceiver<i32> =
2047            client.call_server_stream("stream", &()).await.unwrap();
2048        transport.connected.store(false, Ordering::SeqCst);
2049        transport.fail_send.store(true, Ordering::SeqCst);
2050
2051        assert!(matches!(
2052            client.notify("fail", &()).await,
2053            Err(RpcError::Transport(TransportError::NotConnected))
2054        ));
2055        assert!(matches!(
2056            call.await.unwrap(),
2057            Err(RpcError::Transport(TransportError::NotConnected))
2058        ));
2059        assert!(matches!(
2060            stream.recv().await,
2061            Some(Err(RpcError::Transport(TransportError::NotConnected)))
2062        ));
2063        tokio::time::timeout(Duration::from_secs(1), transport.close_started.notified())
2064            .await
2065            .unwrap();
2066        tokio::time::timeout(Duration::from_secs(1), handle.join())
2067            .await
2068            .unwrap()
2069            .unwrap();
2070
2071        assert!(client.is_closed());
2072        assert!(!client.is_running());
2073        client.close().await.unwrap();
2074        assert_eq!(transport.close_count.load(Ordering::SeqCst), 1);
2075    }
2076
2077    #[tokio::test]
2078    async fn test_connected_outbound_timeout_is_not_terminal() {
2079        let client = RpcClient::new(ConnectedTimeoutChannel::default());
2080        let transport = client.transport();
2081        let handle = client.start();
2082
2083        assert!(matches!(
2084            client.notify("first", &()).await,
2085            Err(RpcError::Transport(TransportError::Timeout { .. }))
2086        ));
2087        assert!(client.is_running());
2088        assert!(!client.is_closed());
2089        assert_eq!(transport.close_count.load(Ordering::SeqCst), 0);
2090        client.notify("second", &()).await.unwrap();
2091
2092        client.close().await.unwrap();
2093        handle.join().await.unwrap();
2094        assert_eq!(transport.close_count.load(Ordering::SeqCst), 1);
2095    }
2096
2097    #[tokio::test]
2098    async fn test_stream_registration_rolls_back_when_send_fails() {
2099        let client = RpcClient::new(LifecycleTestChannel::new(true, false));
2100        let handle = client.start();
2101
2102        let result: Result<StreamReceiver<i32>> = client.call_server_stream("numbers", &()).await;
2103
2104        assert!(matches!(result, Err(RpcError::ConnectionClosed)));
2105        assert_eq!(client.active_streams(), 0);
2106        client.close().await.unwrap();
2107        handle.join().await.unwrap();
2108    }
2109
2110    #[tokio::test]
2111    async fn test_dropped_typed_call_cleans_pending_registration_during_send() {
2112        let client = Arc::new(RpcClient::new(LifecycleTestChannel::with_blocked_send()));
2113        let transport = client.transport();
2114        let handle = client.start();
2115        let call_client = Arc::clone(&client);
2116        let call = tokio::spawn(async move {
2117            call_client
2118                .call::<_, AddResponse>("add", &AddRequest { a: 1, b: 2 })
2119                .await
2120        });
2121
2122        tokio::time::timeout(Duration::from_secs(1), transport.sent.notified())
2123            .await
2124            .unwrap();
2125        call.abort();
2126        assert!(call.await.unwrap_err().is_cancelled());
2127        assert!(client.pending.lock().is_empty());
2128
2129        client.close().await.unwrap();
2130        handle.join().await.unwrap();
2131    }
2132
2133    #[tokio::test]
2134    async fn test_dropped_raw_call_cleans_pending_registration_while_waiting() {
2135        let client = Arc::new(RpcClient::new(LifecycleTestChannel::new(false, false)));
2136        let transport = client.transport();
2137        let handle = client.start();
2138        let call_client = Arc::clone(&client);
2139        let call = tokio::spawn(async move { call_client.call_raw("raw", vec![1, 2, 3]).await });
2140
2141        tokio::time::timeout(Duration::from_secs(1), transport.sent.notified())
2142            .await
2143            .unwrap();
2144        call.abort();
2145        assert!(call.await.unwrap_err().is_cancelled());
2146        assert!(client.pending.lock().is_empty());
2147
2148        client.close().await.unwrap();
2149        handle.join().await.unwrap();
2150    }
2151
2152    #[tokio::test]
2153    async fn test_dropped_stream_setup_cleans_registration_during_send() {
2154        let client = Arc::new(RpcClient::new(LifecycleTestChannel::with_blocked_send()));
2155        let transport = client.transport();
2156        let handle = client.start();
2157        let stream_client = Arc::clone(&client);
2158        let stream = tokio::spawn(async move {
2159            stream_client
2160                .call_server_stream::<_, i32>("numbers", &())
2161                .await
2162        });
2163
2164        tokio::time::timeout(Duration::from_secs(1), transport.sent.notified())
2165            .await
2166            .unwrap();
2167        stream.abort();
2168        assert!(matches!(stream.await, Err(error) if error.is_cancelled()));
2169        assert_eq!(client.active_streams(), 0);
2170
2171        client.close().await.unwrap();
2172        handle.join().await.unwrap();
2173    }
2174
2175    #[tokio::test]
2176    async fn test_close_is_idempotent_after_transport_error() {
2177        let client = RpcClient::new(LifecycleTestChannel::new(false, true));
2178        let transport = client.transport();
2179        let handle = client.start();
2180
2181        assert!(matches!(
2182            client.close().await,
2183            Err(RpcError::ConnectionClosed)
2184        ));
2185        assert!(matches!(
2186            client.close().await,
2187            Err(RpcError::ConnectionClosed)
2188        ));
2189        assert_eq!(transport.close_count.load(Ordering::SeqCst), 1);
2190        assert!(matches!(
2191            client.notify("event", &()).await,
2192            Err(RpcError::ConnectionClosed)
2193        ));
2194        let stream_result: Result<StreamReceiver<i32>> =
2195            client.call_server_stream("numbers", &()).await;
2196        assert!(matches!(stream_result, Err(RpcError::ConnectionClosed)));
2197        handle.join().await.unwrap();
2198    }
2199
2200    #[tokio::test]
2201    async fn test_close_proceeds_while_outbound_send_is_blocked() {
2202        let client = Arc::new(RpcClient::new(LifecycleTestChannel::with_blocked_send()));
2203        let transport = client.transport();
2204        let handle = client.start();
2205        let call_client = client.clone();
2206        let call_handle = tokio::spawn(async move {
2207            call_client
2208                .call::<_, AddResponse>("add", &AddRequest { a: 1, b: 2 })
2209                .await
2210        });
2211
2212        tokio::time::timeout(Duration::from_secs(1), transport.sent.notified())
2213            .await
2214            .unwrap();
2215        tokio::time::timeout(Duration::from_secs(1), client.close())
2216            .await
2217            .unwrap()
2218            .unwrap();
2219
2220        assert!(matches!(
2221            call_handle.await.unwrap(),
2222            Err(RpcError::ConnectionClosed)
2223        ));
2224        assert_eq!(transport.close_count.load(Ordering::SeqCst), 1);
2225        handle.join().await.unwrap();
2226    }
2227
2228    #[tokio::test]
2229    async fn test_receive_failure_claims_close_before_transport_await() {
2230        let client = RpcClient::new(LifecycleTestChannel::with_blocked_close());
2231        let transport = client.transport();
2232        let handle = client.start();
2233
2234        transport.receive_failed.notify_one();
2235        tokio::time::timeout(Duration::from_secs(1), transport.close_started.notified())
2236            .await
2237            .unwrap();
2238
2239        assert!(client.is_closed());
2240        assert!(!client.is_running());
2241        let result: Result<AddResponse> = client.call("add", &AddRequest { a: 1, b: 2 }).await;
2242        assert!(matches!(result, Err(RpcError::ConnectionClosed)));
2243        let close = client.close();
2244        tokio::pin!(close);
2245        assert!(
2246            tokio::time::timeout(Duration::from_millis(50), &mut close)
2247                .await
2248                .is_err(),
2249            "concurrent close returned before the in-flight transport close completed"
2250        );
2251        assert_eq!(transport.close_count.load(Ordering::SeqCst), 1);
2252
2253        transport.close_release.notify_one();
2254        close.await.unwrap();
2255        client.close().await.unwrap();
2256        tokio::time::timeout(Duration::from_secs(1), handle.join())
2257            .await
2258            .unwrap()
2259            .unwrap();
2260
2261        assert!(client.receive_abort.lock().is_none());
2262        assert_eq!(transport.close_count.load(Ordering::SeqCst), 1);
2263    }
2264
2265    #[tokio::test]
2266    async fn test_concurrent_close_waiters_share_transport_error() {
2267        let client = Arc::new(RpcClient::new(
2268            LifecycleTestChannel::with_blocked_failing_close(),
2269        ));
2270        let transport = client.transport();
2271        let handle = client.start();
2272
2273        let first_client = Arc::clone(&client);
2274        let first = tokio::spawn(async move { first_client.close().await });
2275        tokio::time::timeout(Duration::from_secs(1), transport.close_started.notified())
2276            .await
2277            .unwrap();
2278        let second_client = Arc::clone(&client);
2279        let second = tokio::spawn(async move { second_client.close().await });
2280        tokio::task::yield_now().await;
2281        assert!(!first.is_finished());
2282        assert!(!second.is_finished());
2283
2284        transport.close_release.notify_one();
2285        assert!(matches!(
2286            first.await.unwrap(),
2287            Err(RpcError::ConnectionClosed)
2288        ));
2289        assert!(matches!(
2290            second.await.unwrap(),
2291            Err(RpcError::ConnectionClosed)
2292        ));
2293        assert_eq!(transport.close_count.load(Ordering::SeqCst), 1);
2294        handle.join().await.unwrap();
2295    }
2296
2297    #[tokio::test]
2298    async fn test_cancelled_close_owner_does_not_strand_waiters() {
2299        let client = Arc::new(RpcClient::new(LifecycleTestChannel::with_blocked_close()));
2300        let transport = client.transport();
2301        let handle = client.start();
2302
2303        let first_client = Arc::clone(&client);
2304        let first = tokio::spawn(async move { first_client.close().await });
2305        tokio::time::timeout(Duration::from_secs(1), transport.close_started.notified())
2306            .await
2307            .unwrap();
2308        first.abort();
2309        assert!(first.await.unwrap_err().is_cancelled());
2310
2311        let second_client = Arc::clone(&client);
2312        let second = tokio::spawn(async move { second_client.close().await });
2313        tokio::task::yield_now().await;
2314        assert!(!second.is_finished());
2315
2316        transport.close_release.notify_one();
2317        second.await.unwrap().unwrap();
2318        assert_eq!(transport.close_count.load(Ordering::SeqCst), 1);
2319        handle.join().await.unwrap();
2320    }
2321
2322    #[tokio::test]
2323    async fn test_cancelled_close_before_admission_continues_in_background() {
2324        let client = Arc::new(RpcClient::new(CommitTestChannel::default()));
2325        let transport = client.transport();
2326        let handle = client.start();
2327        let call_client = Arc::clone(&client);
2328        let call = tokio::spawn(async move {
2329            call_client
2330                .call::<_, AddResponse>("add", &AddRequest { a: 1, b: 2 })
2331                .await
2332        });
2333
2334        tokio::time::timeout(Duration::from_secs(1), transport.send_started.notified())
2335            .await
2336            .unwrap();
2337        transport.send_release.notify_one();
2338        tokio::time::timeout(Duration::from_secs(1), transport.committed.notified())
2339            .await
2340            .unwrap();
2341
2342        let close_client = Arc::clone(&client);
2343        let first_close = tokio::spawn(async move { close_client.close().await });
2344        tokio::time::timeout(Duration::from_secs(1), async {
2345            while !matches!(
2346                &*client.client_close_state.lock(),
2347                TransportCloseState::Closing
2348            ) {
2349                tokio::task::yield_now().await;
2350            }
2351        })
2352        .await
2353        .unwrap();
2354        assert_eq!(transport.close_count.load(Ordering::SeqCst), 0);
2355
2356        first_close.abort();
2357        assert!(first_close.await.unwrap_err().is_cancelled());
2358        transport.reply_release.notify_one();
2359        assert_eq!(call.await.unwrap().unwrap().result, 3);
2360
2361        tokio::time::timeout(Duration::from_secs(1), async {
2362            while transport.close_count.load(Ordering::SeqCst) == 0 {
2363                tokio::task::yield_now().await;
2364            }
2365        })
2366        .await
2367        .unwrap();
2368        client.close().await.unwrap();
2369        assert_eq!(transport.close_count.load(Ordering::SeqCst), 1);
2370        handle.join().await.unwrap();
2371    }
2372
2373    #[tokio::test]
2374    async fn test_panicking_transport_close_releases_all_waiters() {
2375        let client = RpcClient::new(LifecycleTestChannel::with_panicking_close());
2376        let transport = client.transport();
2377        let handle = client.start();
2378
2379        let first = tokio::time::timeout(Duration::from_secs(1), client.close())
2380            .await
2381            .unwrap();
2382        let second = tokio::time::timeout(Duration::from_secs(1), client.close())
2383            .await
2384            .unwrap();
2385
2386        assert!(matches!(first, Err(RpcError::ClientError(error)) if error.contains("panicked")));
2387        assert!(matches!(second, Err(RpcError::ClientError(error)) if error.contains("panicked")));
2388        assert_eq!(transport.close_count.load(Ordering::SeqCst), 1);
2389        handle.join().await.unwrap();
2390    }
2391
2392    #[test]
2393    fn test_transport_error_normalization() {
2394        assert!(matches!(
2395            normalize_transport_error(TransportError::ConnectionClosed),
2396            RpcError::ConnectionClosed
2397        ));
2398        assert!(matches!(
2399            normalize_transport_error(TransportError::NotConnected),
2400            RpcError::Transport(TransportError::NotConnected)
2401        ));
2402    }
2403
2404    #[tokio::test]
2405    async fn test_drop_aborts_detached_task_and_drains_work() {
2406        let client = RpcClient::new(LifecycleTestChannel::new(false, false));
2407        let state = client.state.clone();
2408        let receive_abort = client.receive_abort.clone();
2409        let stream_manager = client.stream_manager();
2410        let (pending_tx, pending_rx) = oneshot::channel::<Result<Message>>();
2411        client.pending.lock().insert(MessageId::new(), pending_tx);
2412        let mut stream: StreamReceiver<i32> = stream_manager.create_receiver(next_stream_id());
2413        let handle = client.start();
2414
2415        assert!(receive_abort.lock().is_some());
2416        drop(handle);
2417        drop(client);
2418
2419        assert!(state.lock().closed);
2420        assert!(!state.lock().running);
2421        assert!(receive_abort.lock().is_none());
2422        assert!(matches!(
2423            pending_rx.await.unwrap(),
2424            Err(RpcError::ConnectionClosed)
2425        ));
2426        assert!(matches!(
2427            stream.recv().await,
2428            Some(Err(RpcError::ConnectionClosed))
2429        ));
2430        assert_eq!(stream_manager.active_stream_count(), 0);
2431    }
2432
2433    #[tokio::test]
2434    async fn test_handle_join_after_graceful_close() {
2435        let client = RpcClient::new(LifecycleTestChannel::new(false, false));
2436        let handle = client.start();
2437
2438        client.close().await.unwrap();
2439
2440        tokio::time::timeout(Duration::from_secs(1), handle.join())
2441            .await
2442            .unwrap()
2443            .unwrap();
2444    }
2445
2446    #[tokio::test]
2447    async fn test_forced_shutdown_cancels_blocked_send() {
2448        let client = Arc::new(RpcClient::new(LifecycleTestChannel::with_blocked_send()));
2449        let transport = client.transport();
2450        let handle = client.start();
2451        let call_client = Arc::clone(&client);
2452        let call = tokio::spawn(async move {
2453            call_client
2454                .call::<_, AddResponse>("add", &AddRequest { a: 1, b: 2 })
2455                .await
2456        });
2457
2458        tokio::time::timeout(Duration::from_secs(1), transport.sent.notified())
2459            .await
2460            .unwrap();
2461        handle.shutdown().await;
2462
2463        assert!(matches!(
2464            tokio::time::timeout(Duration::from_secs(1), call)
2465                .await
2466                .unwrap()
2467                .unwrap(),
2468            Err(RpcError::ConnectionClosed)
2469        ));
2470        tokio::time::timeout(Duration::from_secs(1), client.close())
2471            .await
2472            .unwrap()
2473            .unwrap();
2474        assert_eq!(transport.close_count.load(Ordering::SeqCst), 1);
2475    }
2476
2477    #[tokio::test]
2478    async fn test_forced_shutdown_drains_pending_call() {
2479        let client = Arc::new(RpcClient::new(LifecycleTestChannel::new(false, false)));
2480        let transport = client.transport();
2481        let handle = client.start();
2482        let call_client = client.clone();
2483        let call_handle = tokio::spawn(async move {
2484            call_client
2485                .call::<_, AddResponse>("add", &AddRequest { a: 1, b: 2 })
2486                .await
2487        });
2488
2489        tokio::time::timeout(Duration::from_secs(1), transport.sent.notified())
2490            .await
2491            .unwrap();
2492        handle.shutdown().await;
2493
2494        assert!(matches!(
2495            call_handle.await.unwrap(),
2496            Err(RpcError::ConnectionClosed)
2497        ));
2498        assert!(client.is_closed());
2499        assert!(!client.is_running());
2500        assert!(client.receive_abort.lock().is_none());
2501        client.close().await.unwrap();
2502    }
2503}
xrpc/client.rs

xrpc/
client.rs
