refactor: improve error handling and metrics in relay components

clippy fixes

Author: sameh-farouk

src/bins/rmb-relay.rs

@@ -3,6 +3,7 @@ use std::time::Duration;
 
 use anyhow::{Context, Result};
 use clap::{builder::ArgAction, Parser};
+use prometheus::default_registry;
 use rmb::cache::RedisCache;
 use rmb::events;
 use rmb::redis;
@@ -170,6 +171,7 @@ async fn app(args: Args) -> Result<JoinHandle<()>> {
         )
     };
     let ranker = relay::ranker::RelayRanker::new(Duration::from_secs(args.ranker_period));
+    rmb::cache::register_cache_metrics(default_registry());
     let r = relay::Relay::new(
         args.domains.iter().cloned().collect(),
         twins,

src/cache/memory.rs

@@ -3,7 +3,9 @@ use std::sync::Arc;
 use std::time::Duration;
 use tokio::sync::RwLock;
 
-use super::Cache;
+#[cfg(feature = "tracker")]
+use super::RMB_CACHE_HITS_TOTAL;
+use super::{Cache, RMB_CACHE_ENTRIES, RMB_CACHE_FLUSHES_TOTAL, RMB_CACHE_MISSES_TOTAL};
 use anyhow::Result;
 use ttl_cache::TtlCache;
 
@@ -38,21 +40,35 @@ where
 {
     async fn set<K: ToString + Send + Sync>(&self, key: K, obj: T) -> Result<()> {
         let mut mem = self.mem.write().await;
-        mem.insert(key.to_string(), obj, self.ttl);
+        let prev = mem.insert(key.to_string(), obj, self.ttl);
+        if prev.is_none() {
+            RMB_CACHE_ENTRIES.inc();
+        }
 
         Ok(())
     }
 
     async fn get<K: ToString + Send + Sync>(&self, key: K) -> Result<Option<T>> {
         let mem = self.mem.read().await;
         match mem.get(&key.to_string()) {
-            None => Ok(None),
-            Some(v) => Ok(Some(v.clone())),
+            None => {
+                RMB_CACHE_MISSES_TOTAL.inc();
+                Ok(None)
+            }
+            Some(v) => {
+                #[cfg(feature = "tracker")]
+                {
+                    RMB_CACHE_HITS_TOTAL.inc();
+                }
+                Ok(Some(v.clone()))
+            }
         }
     }
     async fn flush(&self) -> Result<()> {
         let mut mem = self.mem.write().await;
         mem.clear();
+        RMB_CACHE_ENTRIES.set(0);
+        RMB_CACHE_FLUSHES_TOTAL.inc();
 
         Ok(())
     }

src/cache/mod.rs

@@ -9,6 +9,45 @@ use std::{
     marker::{Send, Sync},
 };
 
+use lazy_static::lazy_static;
+use prometheus::{IntCounter, IntGauge, Registry};
+
+#[cfg(feature = "tracker")]
+lazy_static! {
+    static ref RMB_CACHE_HITS_TOTAL: IntCounter = IntCounter::new(
+        "rmb_cache_hits_total",
+        "Number of cache hits when fetching twins",
+    )
+    .unwrap();
+}
+
+lazy_static! {
+    static ref RMB_CACHE_MISSES_TOTAL: IntCounter = IntCounter::new(
+        "rmb_cache_misses_total",
+        "Number of cache misses when fetching twins (RPC fallback likely)",
+    )
+    .unwrap();
+    static ref RMB_CACHE_FLUSHES_TOTAL: IntCounter = IntCounter::new(
+        "rmb_cache_flushes_total",
+        "Number of times the twin cache has been flushed",
+    )
+    .unwrap();
+    static ref RMB_CACHE_ENTRIES: IntGauge = IntGauge::new(
+        "rmb_cache_entries",
+        "Current number of entries in the twin cache",
+    )
+    .unwrap();
+}
+
+pub fn register_cache_metrics(registry: &Registry) {
+    // Best-effort register into provided registry; ignore duplicate registration errors.
+    #[cfg(feature = "tracker")]
+    let _ = registry.register(Box::new(RMB_CACHE_HITS_TOTAL.clone()));
+    let _ = registry.register(Box::new(RMB_CACHE_MISSES_TOTAL.clone()));
+    let _ = registry.register(Box::new(RMB_CACHE_FLUSHES_TOTAL.clone()));
+    let _ = registry.register(Box::new(RMB_CACHE_ENTRIES.clone()));
+}
+
 pub trait Cache<T>: Send + Sync + 'static {
     fn set<S: ToString + Send + Sync>(
         &self,

src/cache/redis.rs

@@ -1,4 +1,6 @@
-use super::Cache;
+#[cfg(feature = "tracker")]
+use super::RMB_CACHE_HITS_TOTAL;
+use super::{Cache, RMB_CACHE_ENTRIES, RMB_CACHE_FLUSHES_TOTAL, RMB_CACHE_MISSES_TOTAL};
 
 use anyhow::{Context, Result};
 use bb8_redis::{
@@ -47,12 +49,15 @@ where
     async fn set<S: ToString + Send + Sync>(&self, key: S, obj: T) -> Result<()> {
         let mut conn = self.get_connection().await?;
         let obj = serde_json::to_vec(&obj).context("unable to serialize twin object for redis")?;
-        let _: () = cmd("HSET")
+        let added: i64 = cmd("HSET")
             .arg(&self.prefix)
             .arg(key.to_string())
             .arg(obj)
-            .query_async::<_, ()>(&mut *conn)
+            .query_async::<_, i64>(&mut *conn)
             .await?;
+        if added == 1 {
+            RMB_CACHE_ENTRIES.inc();
+        }
 
         Ok(())
     }
@@ -70,14 +75,23 @@ where
                 let ret: T = serde_json::from_slice(&val)
                     .context("unable to deserialize redis value to twin object")?;
 
+                #[cfg(feature = "tracker")]
+                {
+                    RMB_CACHE_HITS_TOTAL.inc();
+                }
                 Ok(Some(ret))
             }
-            None => Ok(None),
+            None => {
+                RMB_CACHE_MISSES_TOTAL.inc();
+                Ok(None)
+            }
         }
     }
     async fn flush(&self) -> Result<()> {
         let mut conn = self.get_connection().await?;
         let _: () = cmd("DEL").arg(&self.prefix).query_async(&mut *conn).await?;
+        RMB_CACHE_ENTRIES.set(0);
+        RMB_CACHE_FLUSHES_TOTAL.inc();
 
         Ok(())
     }

src/identity/ed25519.rs

@@ -103,7 +103,7 @@ mod tests {
         Ed25519Signer::try_from(WORDS).expect("key must be loaded");
 
         let err = Ed25519Signer::try_from("invalid words");
-        assert_eq!(err.is_err(), true);
+        assert!(err.is_err());
     }
 
     #[test]
@@ -115,6 +115,6 @@ mod tests {
         let _: [u8; 32] = result.as_slice().try_into().expect("key of size 32");
 
         let err = Ed25519Signer::try_from("0xinvalidseed");
-        assert_eq!(err.is_err(), true);
+        assert!(err.is_err());
     }
 }

src/identity/sr25519.rs

@@ -100,14 +100,14 @@ mod tests {
         Sr25519Signer::try_from(WORDS).expect("key must be loaded");
 
         let err = Sr25519Signer::try_from("invalid words");
-        assert_eq!(err.is_err(), true);
+        assert!(err.is_err());
     }
 
     #[test]
     fn test_load_from_seed() {
         Sr25519Signer::try_from(SEED).expect("key must be loaded");
 
         let err = Sr25519Signer::try_from("0xinvalidseed");
-        assert_eq!(err.is_err(), true);
+        assert!(err.is_err());
     }
 }

src/peer/e2e/mod.rs

@@ -141,8 +141,8 @@ mod test {
                 .parse()
                 .unwrap();
 
-        let shared1 = sk1.shared(&sk2.public()).unwrap();
-        let shared2 = sk2.shared(&sk1.public()).unwrap();
+        let shared1 = sk1.shared(sk2.public()).unwrap();
+        let shared2 = sk2.shared(sk1.public()).unwrap();
 
         assert_eq!(shared1, shared2);
 

src/relay/api.rs

@@ -209,7 +209,7 @@ async fn federation<D: TwinDB, R: RateLimiter>(
 
     let dst: SessionID = (&envelope.destination).into();
 
-    // fast-fail: if requested and destination is not local, return error immediately
+    // fast-fail path: if requested and destination is not local, return error immediately
     if has_fast_fail(&envelope) && !data.switch.is_local(&dst).await {
         // destination session doesn’t exist on this relay
         return Response::builder()
@@ -376,18 +376,15 @@ impl<M: Metrics, D: TwinDB> Session<M, D> {
         }
 
         let is_local = self.switch.is_local(&dst).await;
-        // fast-fail: if requested and destination is not local, return error immediately
-        if has_fast_fail(envelope) && !is_local {
-            anyhow::bail!("destination offline");
-        }
 
         // check if the dst twin id is already connected locally
         // if so, we don't have to check federation and directly
         // switch the message
         if is_local {
-            log::debug!("found local session for '{}' , forwarding message", dst);
+            log::trace!("found local session for '{}' , forwarding message", dst);
             if let Err(err) = self.switch.send(&dst, &msg).await {
                 log::error!("failed to route message to peer '{}' : {}", dst, err);
+                anyhow::bail!("failed to route message to peer '{}': {}", dst, err);
             }
             return Ok(());
         }
@@ -414,13 +411,18 @@ impl<M: Metrics, D: TwinDB> Session<M, D> {
             // push message to the (relay.federation) queue
             return Ok(self.federator.send(&msg).await?);
         }
-
+        // fast-fail path: we confirmed that this is not an foreign message, and if it is fast-fail taged and destination session is not connected
+        // then the peer must be offline, return error immediately
+        if has_fast_fail(envelope) && !is_local {
+            anyhow::bail!("destination offline");
+        }
         // we don't return an error because when we return an error
         // we will send this error back to the sender user. Hence
         // calling the switch.send again
         // this is an internal error anyway, and should not happen
         if let Err(err) = self.switch.send(&dst, &msg).await {
             log::error!("failed to route message to peer '{}': {}", dst, err);
+            anyhow::bail!("failed to route message to peer '{}': {}", dst, err);
         }
 
         Ok(())
@@ -474,7 +476,7 @@ impl<M: Metrics, D: TwinDB> Session<M, D> {
                     // TODO: throttling to avoid sending too many messages in short time!
                     match message {
                         Message::Text(_) => {
-                            log::trace!("received unsupported (text) message. disconnecting!");
+                            log::debug!("received unsupported (text) message. disconnecting!");
                             break;
                         }
                         Message::Binary(msg) => {
@@ -487,16 +489,21 @@ impl<M: Metrics, D: TwinDB> Session<M, D> {
                             super::switch::MESSAGE_RX_TWIN.with_label_values(&[&format!("{}", envelope.source.twin)]).inc();
 
                             if !self.metrics.measure(msg.len()).await {
-                                log::trace!("twin with stream id {} exceeded its request limits, dropping message", self.id);
-                                self.send_error(envelope, "exceeded rate limits, dropping message").await;
-                                continue
+                                log::debug!("twin with stream id {} exceeded its request limits, dropping message", self.id);
+                                if envelope.has_request() {
+                                    self.send_error(envelope, "exceeded rate limits, dropping message").await;
+                                }
+                                continue;
                             }
 
                             // if we failed to route back the message to the user
                             // for any reason we send an error message back
                             // server error has no source address set.
                             if let Err(err) = self.route(&mut envelope, msg).await {
-                                self.send_error(envelope, err).await;
+                                // check if the envolope is request or response
+                                if envelope.has_request() {
+                                    self.send_error(envelope, err).await;
+                                }
                             }
                         }
                         Message::Ping(_) => {

src/relay/switch/mod.rs

@@ -27,10 +27,7 @@ use tokio::time::Duration;
 use tokio_util::sync::{CancellationToken, DropGuard};
 use worker::WorkerJob;
 
-use prometheus::{IntCounter, IntGaugeVec, Opts, Registry};
-
-#[cfg(feature = "tracker")]
-use prometheus::IntCounterVec;
+use prometheus::{IntCounter, IntCounterVec, IntGaugeVec, Opts, Registry};
 
 lazy_static::lazy_static! {
     static ref CON_PER_WORKER: IntGaugeVec = IntGaugeVec::new(
@@ -45,6 +42,16 @@ lazy_static::lazy_static! {
     static ref MESSAGE_RX_BYTES: IntCounter = IntCounter::new("relay_message_rx_bytes", "size of messages received by relay in bytes").unwrap();
 
     static ref MESSAGE_TX_BYTES: IntCounter = IntCounter::new("relay_message_tx_bytes", "size of messages forwarded by relay in bytes").unwrap();
+
+    // Total number of session evictions, labeled by reason
+    // reason ∈ { "closed", "back_pressure_timeout", "cancelled" }
+    static ref RELAY_SESSION_EVICTIONS_TOTAL: IntCounterVec = IntCounterVec::new(
+        Opts::new(
+            "relay_session_evictions_total",
+            "Total number of session evictions from the relay, labeled by reason",
+        ),
+        &["reason"],
+    ).unwrap();
 }
 
 #[cfg(feature = "tracker")]
@@ -201,6 +208,8 @@ where
         opts.registry.register(Box::new(MESSAGE_TX.clone()))?;
         opts.registry.register(Box::new(MESSAGE_RX_BYTES.clone()))?;
         opts.registry.register(Box::new(MESSAGE_TX_BYTES.clone()))?;
+        opts.registry
+            .register(Box::new(RELAY_SESSION_EVICTIONS_TOTAL.clone()))?;
         #[cfg(feature = "tracker")]
         opts.registry.register(Box::new(MESSAGE_RX_TWIN.clone()))?;
 
@@ -277,6 +286,9 @@ where
     }
 
     /// checks if a session is connected locally
+    /// it doesn't verify if the twin is belongs to this relay
+    /// so should be treated as cheap way to route to locals online twins bypassing twin lookup
+    /// if the session is not found, further routing verification should be done
     pub async fn is_local(&self, id: &SessionID) -> bool {
         self.sessions.lock().await.contains_key(id)
     }
@@ -382,7 +394,7 @@ mod test {
             .expect("can connect");
 
         let switch = Switch::<MessageSender>::new(super::SwitchOptions {
-            pool: pool,
+            pool,
             workers: 2,
             max_users: 100,
             registry: Registry::new(),