trz_gateway_server/connection/
mod.rs

1//! Cache of client connections to the Terrazzo Gateway.
2
3use std::sync::Arc;
4use std::time::Duration;
5use std::time::Instant;
6
7use connection_id::ConnectionId;
8use dashmap::DashMap;
9use nameth::NamedEnumValues as _;
10use nameth::nameth;
11use scopeguard::defer;
12use tokio::time::error::Elapsed;
13use tokio::time::timeout;
14use tonic::transport::Channel;
15use tracing::Instrument as _;
16use tracing::info;
17use tracing::info_span;
18use tracing::warn;
19use trz_gateway_common::consts::HEALTH_CHECK_PERIOD;
20use trz_gateway_common::consts::HEALTH_CHECK_TIMEOUT;
21use trz_gateway_common::id::ClientName;
22use trz_gateway_common::protos::terrazzo::remote::health::Ping;
23use trz_gateway_common::protos::terrazzo::remote::health::Pong;
24use trz_gateway_common::protos::terrazzo::remote::health::health_service_client::HealthServiceClient;
25
26use self::balance::IncomingClients;
27use crate::auth_code::AuthCode;
28
29mod balance;
30pub mod connection_id;
31mod pending_requests;
32
33/// The cache of all the channels connected to the Terrazzo Gateway,
34/// grouped by [ClientName].
35#[derive(Default)]
36pub struct Connections {
37    cache: DashMap<ClientName, IncomingClients<Channel>>,
38}
39
40impl Connections {
41    /// Adds a connection to the cache and schedules the periodic keep-alive.
42    ///
43    /// The connection is removed from the cache, and closed, if the keep-alive fails.
44    pub fn add(self: &Arc<Self>, client_name: ClientName, channel: Channel) {
45        let connection_id = ConnectionId::next();
46        let _span = info_span!("Connection", %connection_id).entered();
47        match self.cache.entry(client_name.clone()) {
48            dashmap::Entry::Occupied(mut entry) => {
49                self.add_channel(entry.get_mut(), client_name, connection_id, channel);
50            }
51            dashmap::Entry::Vacant(entry) => {
52                let mut connections = IncomingClients::new();
53                self.add_channel(&mut connections, client_name, connection_id, channel);
54                entry.insert(connections);
55            }
56        }
57    }
58
59    fn add_channel(
60        self: &Arc<Self>,
61        connections: &mut IncomingClients<Channel>,
62        client_name: ClientName,
63        connection_id: ConnectionId,
64        channel: Channel,
65    ) {
66        connections.add_channel(connection_id, channel.clone());
67        tokio::spawn(
68            self.clone()
69                .channel_health_check(client_name, connection_id, channel)
70                .in_current_span(),
71        );
72    }
73
74    /// Runs the keep-alive.
75    ///   1. First, run a quick ping/pong check
76    ///   2. Then, send a ping and expect a response after [PERIOD]
77    ///   3. go back to step 1
78    async fn channel_health_check(
79        self: Arc<Connections>,
80        client_name: ClientName,
81        connection_id: ConnectionId,
82        channel: Channel,
83    ) -> Result<(), ChannelHealthError> {
84        defer!(self.remove(client_name, connection_id));
85        let mut health_client = HealthServiceClient::new(channel);
86        let health_check_loop = async move {
87            loop {
88                let start = Instant::now();
89                let pong = health_client.ping_pong(Ping {
90                    connection_id: connection_id.to_string(),
91                    ..Ping::default()
92                });
93                let Pong { .. } = timeout(HEALTH_CHECK_TIMEOUT, pong).await??.into_inner();
94                let latency = Instant::now() - start;
95                info!(?latency, "Ping");
96
97                let start = Instant::now();
98                let pong = health_client.ping_pong(Ping {
99                    connection_id: connection_id.to_string(),
100                    delay: Some(HEALTH_CHECK_PERIOD.try_into()?),
101                    auth_code: AuthCode::current().to_string(),
102                });
103                let Pong { .. } = timeout(HEALTH_CHECK_PERIOD + HEALTH_CHECK_TIMEOUT, pong)
104                    .await??
105                    .into_inner();
106                let elapsed = start.elapsed();
107                if elapsed < HEALTH_CHECK_PERIOD {
108                    return Err(ChannelHealthError::TooSoon(elapsed));
109                }
110            }
111        };
112        health_check_loop
113            .await
114            .inspect(|()| info!("Health check loop DONE"))
115            .inspect_err(|error| warn!("Health check loop FAILED: {error}"))
116    }
117
118    fn remove(self: &Arc<Self>, client_name: ClientName, connection_id: ConnectionId) {
119        let dashmap::Entry::Occupied(mut connections) = self.cache.entry(client_name) else {
120            return;
121        };
122        connections.get_mut().remove_channel(connection_id);
123        if connections.get_mut().is_empty() {
124            connections.remove();
125        }
126    }
127}
128
129#[nameth]
130#[derive(thiserror::Error, Debug)]
131pub enum ChannelHealthError {
132    #[error("[{n}] {0}", n = self.name())]
133    GrpcError(#[from] tonic::Status),
134
135    #[error("[{n}] {0}", n = self.name())]
136    Timeout(#[from] Elapsed),
137
138    #[error("[{n}] {0}", n = self.name())]
139    DurationError(#[from] prost_types::DurationError),
140
141    #[error("[{n}] Client slept for {0:?}, should have been {HEALTH_CHECK_PERIOD:?}", n = self.name())]
142    TooSoon(Duration),
143}
144
145impl Connections {
146    /// Returns the list of connected clients.
147    pub fn clients(&self) -> Vec<ClientName> {
148        self.cache.iter().map(|entry| entry.key().clone()).collect()
149    }
150
151    /// Returns a connection for the given client.
152    ///
153    /// Multiple connections for the same [ClientName] are load-balanced.
154    pub fn get_client(
155        &self,
156        client_name: &ClientName,
157    ) -> Option<pending_requests::PendingRequests<Channel>> {
158        self.cache
159            .get_mut(client_name)
160            .and_then(|mut c| c.get_channel())
161    }
162}
trz_gateway_server/connection/mod.rs

trz_gateway_server/connection/
mod.rs
