Skip to main content

zbus/connection/
socket_reader.rs

1use std::{collections::HashMap, sync::Arc};
2
3use event_listener::Event;
4use tracing::{debug, instrument, trace};
5
6use crate::{
7    Executor, Message, OwnedMatchRule, Task,
8    async_lock::Mutex,
9    connection::{MsgBroadcaster, PendingMethodCalls},
10    message::Type,
11};
12
13use super::socket::ReadHalf;
14
15#[derive(Debug)]
16pub(crate) struct SocketReader {
17    socket: Box<dyn ReadHalf>,
18    senders: Arc<Mutex<HashMap<Option<OwnedMatchRule>, MsgBroadcaster>>>,
19    pending_method_calls: PendingMethodCalls,
20    already_received_bytes: Vec<u8>,
21    #[cfg(unix)]
22    already_received_fds: Vec<std::os::fd::OwnedFd>,
23    prev_seq: u64,
24    activity_event: Arc<Event>,
25}
26
27impl SocketReader {
28    pub fn new(
29        socket: Box<dyn ReadHalf>,
30        senders: Arc<Mutex<HashMap<Option<OwnedMatchRule>, MsgBroadcaster>>>,
31        pending_method_calls: PendingMethodCalls,
32        already_received_bytes: Vec<u8>,
33        #[cfg(unix)] already_received_fds: Vec<std::os::fd::OwnedFd>,
34        activity_event: Arc<Event>,
35    ) -> Self {
36        Self {
37            socket,
38            senders,
39            pending_method_calls,
40            already_received_bytes,
41            #[cfg(unix)]
42            already_received_fds,
43            prev_seq: 0,
44            activity_event,
45        }
46    }
47
48    pub fn spawn(self, executor: &Executor<'_>) -> Task<()> {
49        executor.spawn(self.receive_msg(), "socket reader")
50    }
51
52    // Keep receiving messages and put them on the queue.
53    #[instrument(name = "socket reader", skip(self), level = "trace")]
54    async fn receive_msg(mut self) {
55        loop {
56            trace!("Waiting for message on the socket..");
57            let msg = self.read_socket().await;
58            match &msg {
59                Ok(msg) => {
60                    trace!("Message received on the socket: {:?}", msg);
61                    if matches!(msg.message_type(), Type::MethodReturn | Type::Error) {
62                        self.dispatch_pending_reply(msg);
63                    }
64                }
65                Err(e) => {
66                    trace!("Error reading from the socket: {:?}", e);
67                    self.fail_pending_method_calls(e.clone());
68                }
69            };
70
71            let mut senders = self.senders.lock().await;
72            for (rule, sender) in &*senders {
73                if let Ok(msg) = &msg {
74                    if let Some(rule) = rule.as_ref() {
75                        match rule.matches(msg) {
76                            Ok(true) => (),
77                            Ok(false) => continue,
78                            Err(e) => {
79                                debug!("Error matching message against rule: {:?}", e);
80
81                                continue;
82                            }
83                        }
84                    }
85                }
86
87                if let Err(e) = sender.broadcast_direct(msg.clone()).await {
88                    // An error would be due to either of these:
89                    //
90                    // 1. the channel is closed.
91                    // 2. No active receivers.
92                    //
93                    // In either case, just log it unless this is the channel for the generic
94                    // unfiltered stream, where the channel is not created on-demand.
95                    if rule.is_some() {
96                        trace!(
97                            "Error broadcasting message to stream for `{:?}`: {:?}",
98                            rule, e
99                        );
100                    }
101                }
102            }
103            trace!("Broadcasted to all streams: {:?}", msg);
104
105            if msg.is_err() {
106                senders.clear();
107                trace!("Socket reading task stopped");
108
109                return;
110            }
111        }
112    }
113
114    fn dispatch_pending_reply(&self, msg: &Message) {
115        debug_assert!(matches!(
116            msg.message_type(),
117            Type::MethodReturn | Type::Error
118        ));
119
120        let reply_serial = match msg.header().reply_serial() {
121            Some(serial) => serial,
122            None => return,
123        };
124
125        let result = match msg.message_type() {
126            Type::MethodReturn => Ok(msg.clone()),
127            Type::Error => Err(msg.clone().into()),
128            Type::MethodCall | Type::Signal => return,
129        };
130        self.pending_method_calls
131            .complete_call(reply_serial, msg.recv_position(), result);
132    }
133
134    fn fail_pending_method_calls(&self, error: crate::Error) {
135        self.pending_method_calls.fail_all(error);
136    }
137
138    #[instrument(skip(self), level = "trace")]
139    async fn read_socket(&mut self) -> crate::Result<Message> {
140        self.activity_event.notify(usize::MAX);
141        let seq = self.prev_seq + 1;
142        let msg = self
143            .socket
144            .receive_message(
145                seq,
146                &mut self.already_received_bytes,
147                #[cfg(unix)]
148                &mut self.already_received_fds,
149            )
150            .await?;
151        self.prev_seq = seq;
152
153        Ok(msg)
154    }
155}