vv_agent/app_server/
server.rs

1use std::collections::HashSet;
2
3use tokio::sync::mpsc;
4
5use crate::app_server::outgoing::OutgoingEnvelope;
6use crate::app_server::processor::MessageProcessor;
7use crate::app_server::protocol::{AppServerError, RequestId};
8use crate::app_server::transport::{
9    AppServerTransport, ConnectionId, TransportConnectionMode, TransportEvent,
10};
11
12pub struct AppServer<T: AppServerTransport> {
13    transport: T,
14    processor: MessageProcessor,
15    outgoing: mpsc::Receiver<OutgoingEnvelope>,
16}
17
18impl<T: AppServerTransport> AppServer<T> {
19    pub fn new(
20        transport: T,
21        processor: MessageProcessor,
22        outgoing: mpsc::Receiver<OutgoingEnvelope>,
23    ) -> Self {
24        Self {
25            transport,
26            processor,
27            outgoing,
28        }
29    }
30
31    pub async fn run(&mut self) -> Result<(), AppServerError> {
32        let mut open_connections = HashSet::new();
33        let mut disconnected_connections = HashSet::new();
34        loop {
35            tokio::select! {
36                biased;
37                envelope = self.outgoing.recv() => {
38                    let Some(envelope) = envelope else {
39                        self.disconnect_all(&mut open_connections).await;
40                        return Ok(());
41                    };
42                    if !self
43                        .processor
44                        .outgoing()
45                        .is_connection_registered(envelope.connection_id)
46                        .await
47                    {
48                        continue;
49                    }
50                    if let Err(error) = self.transport
51                        .send(envelope.connection_id, envelope.message)
52                        .await
53                    {
54                        disconnected_connections.insert(envelope.connection_id);
55                        self.disconnect_connection(&mut open_connections, envelope.connection_id)
56                            .await;
57                        if self.transport.connection_mode() == TransportConnectionMode::Single {
58                            self.disconnect_all(&mut open_connections).await;
59                            return Err(error);
60                        }
61                    }
62                }
63                event = self.transport.next_event() => {
64                    match event {
65                        Some(Ok(TransportEvent::Opened { connection_id })) => {
66                            disconnected_connections.remove(&connection_id);
67                            open_connections.insert(connection_id);
68                            self.processor
69                                .outgoing()
70                                .register_connection(connection_id)
71                                .await;
72                        }
73                        Some(Ok(TransportEvent::Message {
74                            connection_id,
75                            message,
76                        })) => {
77                            if disconnected_connections.contains(&connection_id) {
78                                continue;
79                            }
80                            open_connections.insert(connection_id);
81                            self.processor.process_message(connection_id, message).await;
82                        }
83                        Some(Ok(TransportEvent::ProtocolError {
84                            connection_id,
85                            error,
86                        })) => {
87                            if disconnected_connections.contains(&connection_id) {
88                                continue;
89                            }
90                            open_connections.insert(connection_id);
91                            self.processor
92                                .outgoing()
93                                .register_connection(connection_id)
94                                .await;
95                            let _ = self
96                                .processor
97                                .outgoing()
98                                .send_error(connection_id, RequestId::Null, error)
99                                .await;
100                        }
101                        Some(Ok(TransportEvent::Closed { connection_id })) => {
102                            disconnected_connections.insert(connection_id);
103                            self.disconnect_connection(&mut open_connections, connection_id).await;
104                        }
105                        Some(Err(error)) => {
106                            self.disconnect_all(&mut open_connections).await;
107                            return Err(error);
108                        }
109                        None => {
110                            self.disconnect_all(&mut open_connections).await;
111                            return Ok(());
112                        }
113                    }
114                }
115            }
116        }
117    }
118
119    pub fn processor(&self) -> &MessageProcessor {
120        &self.processor
121    }
122
123    async fn disconnect_connection(
124        &mut self,
125        open_connections: &mut HashSet<ConnectionId>,
126        connection_id: ConnectionId,
127    ) {
128        open_connections.remove(&connection_id);
129        self.processor.disconnect_connection(connection_id).await;
130    }
131
132    async fn disconnect_all(&mut self, open_connections: &mut HashSet<ConnectionId>) {
133        let mut connection_ids = std::mem::take(open_connections);
134        connection_ids.extend(self.processor.connection_ids());
135        for connection_id in connection_ids {
136            self.processor.disconnect_connection(connection_id).await;
137        }
138    }
139}
vv_agent/app_server/server.rs

vv_agent/app_server/
server.rs
