zbus/connection/
pending_method_calls.rs

1use std::{
2    collections::HashMap,
3    io::{self, ErrorKind},
4    num::NonZeroU32,
5    pin::Pin,
6    sync::{Arc, Mutex as SyncMutex},
7    task::{Context, Poll},
8};
9
10use async_broadcast::{Receiver, Sender, broadcast};
11use futures_core::Future;
12use futures_lite::Stream;
13use ordered_stream::OrderedFuture;
14
15use crate::{Error, Message, Result, message::Sequence};
16
17#[derive(Clone, Debug)]
18pub struct PendingMethodCalls {
19    inner: Arc<SyncMutex<PendingMethodCallsState>>,
20}
21
22impl PendingMethodCalls {
23    pub fn register_call(
24        &self,
25        serial: NonZeroU32,
26    ) -> impl Future<Output = Result<Message>>
27    + OrderedFuture<Output = Result<Message>, Ordering = Sequence> {
28        use std::collections::hash_map::Entry;
29
30        let (reply_sender, reply_receiver) = broadcast(1);
31        let closed_error = {
32            let mut state = self.inner.lock().unwrap();
33            if let Some(error) = &state.closed_error {
34                Some(error.clone())
35            } else {
36                match state.calls.entry(serial) {
37                    Entry::Vacant(entry) => {
38                        entry.insert(reply_sender.clone());
39                    }
40                    Entry::Occupied(_) => {
41                        unreachable!(
42                            "Serial number `{serial}` reused while a method call is still pending"
43                        );
44                    }
45                }
46
47                None
48            }
49        };
50
51        if let Some(error) = closed_error {
52            send_reply(reply_sender, Sequence::LAST, Err(error));
53        }
54
55        PendingMethodCall {
56            serial,
57            reply_receiver,
58            pending_method_calls: self.clone(),
59        }
60    }
61
62    pub fn complete_call(&self, serial: NonZeroU32, ordering: Sequence, reply: Result<Message>) {
63        let reply_sender = self.inner.lock().unwrap().calls.remove(&serial);
64        let Some(reply_sender) = reply_sender else {
65            return;
66        };
67
68        send_reply(reply_sender, ordering, reply);
69    }
70
71    pub fn fail_all(&self, error: Error) {
72        let reply_senders: Vec<_> = {
73            let mut state = self.inner.lock().unwrap();
74            state.closed_error.get_or_insert_with(|| error.clone());
75            state
76                .calls
77                .drain()
78                .map(|(_, reply_sender)| reply_sender)
79                .collect()
80        };
81
82        for reply_sender in reply_senders {
83            send_reply(reply_sender, Sequence::LAST, Err(error.clone()));
84        }
85    }
86
87    fn remove_call(&self, serial: NonZeroU32) {
88        self.inner.lock().unwrap().calls.remove(&serial);
89    }
90}
91
92impl Default for PendingMethodCalls {
93    fn default() -> Self {
94        Self {
95            inner: Arc::new(SyncMutex::new(PendingMethodCallsState::default())),
96        }
97    }
98}
99
100/// A method call whose completion can be awaited or joined with other streams.
101///
102/// This is useful for cache population method calls, where joining the call with an update signal
103/// stream can be used to ensure that cache updates are not overwritten by a cache population whose
104/// task is scheduled later.
105#[derive(Debug)]
106struct PendingMethodCall {
107    serial: NonZeroU32,
108    reply_receiver: Receiver<PendingMethodReply>,
109    pending_method_calls: PendingMethodCalls,
110}
111
112impl Drop for PendingMethodCall {
113    fn drop(&mut self) {
114        self.pending_method_calls.remove_call(self.serial);
115    }
116}
117
118impl Future for PendingMethodCall {
119    type Output = Result<Message>;
120
121    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
122        self.poll_before(cx, None).map(|ret| {
123            ret.map(|(_, r)| r).unwrap_or_else(|| {
124                Err(Error::InputOutput(
125                    io::Error::new(ErrorKind::BrokenPipe, "socket closed").into(),
126                ))
127            })
128        })
129    }
130}
131
132impl OrderedFuture for PendingMethodCall {
133    type Output = Result<Message>;
134    type Ordering = Sequence;
135
136    fn poll_before(
137        self: Pin<&mut Self>,
138        cx: &mut Context<'_>,
139        before: Option<&Self::Ordering>,
140    ) -> Poll<Option<(Self::Ordering, Self::Output)>> {
141        let this = self.get_mut();
142
143        match Pin::new(&mut this.reply_receiver).poll_next(cx) {
144            Poll::Ready(Some(reply)) => Poll::Ready(Some(reply)),
145            Poll::Ready(None) => Poll::Ready(None),
146            // `before` is only provided after another stream has produced that sequence. Since the
147            // socket reader dispatches replies synchronously as it reads messages, any earlier
148            // matching reply would already be queued above. For `OrderedFuture`, `Ready(None)` is
149            // the `NoneBefore` equivalent; return it only after polling the receiver so the current
150            // task is still woken when the reply arrives.
151            Poll::Pending if before.is_some() => Poll::Ready(None),
152            Poll::Pending => Poll::Pending,
153        }
154    }
155}
156
157#[derive(Debug, Default)]
158struct PendingMethodCallsState {
159    calls: HashMap<NonZeroU32, PendingMethodReplySender>,
160    closed_error: Option<Error>,
161}
162
163type PendingMethodReply = (Sequence, Result<Message>);
164type PendingMethodReplySender = Sender<PendingMethodReply>;
165
166fn send_reply(reply_sender: PendingMethodReplySender, ordering: Sequence, reply: Result<Message>) {
167    let _ = reply_sender.try_broadcast((ordering, reply));
168}
zbus/connection/pending_method_calls.rs

zbus/connection/
pending_method_calls.rs
