Skip to main content

zbus/connection/
socket_reader.rs

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