Instrument caches

src/api/server_server.rs

@@ -66,7 +66,9 @@ use tracing::{debug, error, warn};
 use crate::{
     api::client_server::{self, claim_keys_helper, get_keys_helper},
     service::pdu::{gen_event_id_canonical_json, PduBuilder},
-    services, utils, Ar, Error, PduEvent, Ra, Result,
+    services,
+    utils::{self, FoundIn},
+    Ar, Error, PduEvent, Ra, Result,
 };
 
 /// Wraps either an literal IP address plus port, or a hostname plus complement
@@ -125,7 +127,7 @@ impl FedDest {
     }
 }
 
-#[tracing::instrument(skip(request))]
+#[tracing::instrument(skip(request), fields(destination_cache_result))]
 pub(crate) async fn send_request<T: OutgoingRequest>(
     destination: &ServerName,
     request: T,
@@ -156,6 +158,7 @@ where
         .cloned();
 
     let (actual_destination, host) = if let Some(result) = cached_result {
+        FoundIn::Cache.record("destination_cache_result");
         result
     } else {
         write_destination_to_cache = true;
@@ -294,6 +297,7 @@ where
                 let response =
                     T::IncomingResponse::try_from_http_response(http_response);
                 if response.is_ok() && write_destination_to_cache {
+                    FoundIn::Remote.record("destination_cache_result");
                     services()
                         .globals
                         .actual_destination_cache

src/database/key_value/rooms/auth_chain.rs

@@ -1,8 +1,14 @@
 use std::{collections::HashSet, mem::size_of, sync::Arc};
 
-use crate::{database::KeyValueDatabase, service, utils, Result};
+use crate::{
+    database::KeyValueDatabase,
+    service,
+    utils::{self, FoundIn},
+    Result,
+};
 
 impl service::rooms::auth_chain::Data for KeyValueDatabase {
+    #[tracing::instrument(skip(self, key), fields(cache_result))]
     fn get_cached_eventid_authchain(
         &self,
         key: &[u64],
@@ -10,6 +16,7 @@ impl service::rooms::auth_chain::Data for KeyValueDatabase {
         // Check RAM cache
         if let Some(result) = self.auth_chain_cache.lock().unwrap().get_mut(key)
         {
+            FoundIn::Cache.record("cache_result");
             return Ok(Some(Arc::clone(result)));
         }
 
@@ -30,6 +37,7 @@ impl service::rooms::auth_chain::Data for KeyValueDatabase {
                 });
 
             if let Some(chain) = chain {
+                FoundIn::Database.record("cache_result");
                 let chain = Arc::new(chain);
 
                 // Cache in RAM
@@ -42,6 +50,7 @@ impl service::rooms::auth_chain::Data for KeyValueDatabase {
             }
         }
 
+        FoundIn::Nothing.record("cache_result");
         Ok(None)
     }
 

src/database/key_value/rooms/short.rs

@@ -3,24 +3,33 @@ use std::sync::Arc;
 use ruma::{events::StateEventType, EventId, RoomId};
 
 use crate::{
-    database::KeyValueDatabase, service, services, utils, Error, Result,
+    database::KeyValueDatabase,
+    service, services,
+    utils::{self, FoundIn},
+    Error, Result,
 };
 
 impl service::rooms::short::Data for KeyValueDatabase {
+    #[tracing::instrument(skip(self), fields(cache_result))]
     fn get_or_create_shorteventid(&self, event_id: &EventId) -> Result<u64> {
         if let Some(short) =
             self.eventidshort_cache.lock().unwrap().get_mut(event_id)
         {
+            FoundIn::Cache.record("cache_result");
             return Ok(*short);
         }
 
         let short = if let Some(shorteventid) =
             self.eventid_shorteventid.get(event_id.as_bytes())?
         {
+            FoundIn::Database.record("cache_result");
+
             utils::u64_from_bytes(&shorteventid).map_err(|_| {
                 Error::bad_database("Invalid shorteventid in db.")
             })?
         } else {
+            FoundIn::Nothing.record("cache_result");
+
             let shorteventid = services().globals.next_count()?;
             self.eventid_shorteventid
                 .insert(event_id.as_bytes(), &shorteventid.to_be_bytes())?;
@@ -37,6 +46,7 @@ impl service::rooms::short::Data for KeyValueDatabase {
         Ok(short)
     }
 
+    #[tracing::instrument(skip(self), fields(cache_result))]
     fn get_shortstatekey(
         &self,
         event_type: &StateEventType,
@@ -48,6 +58,7 @@ impl service::rooms::short::Data for KeyValueDatabase {
             .unwrap()
             .get_mut(&(event_type.clone(), state_key.to_owned()))
         {
+            FoundIn::Cache.record("cache_result");
             return Ok(Some(*short));
         }
 
@@ -66,15 +77,20 @@ impl service::rooms::short::Data for KeyValueDatabase {
             .transpose()?;
 
         if let Some(s) = short {
+            FoundIn::Database.record("cache_result");
+
             self.statekeyshort_cache
                 .lock()
                 .unwrap()
                 .insert((event_type.clone(), state_key.to_owned()), s);
+        } else {
+            FoundIn::Nothing.record("cache_result");
         }
 
         Ok(short)
     }
 
+    #[tracing::instrument(skip(self), fields(cache_result))]
     fn get_or_create_shortstatekey(
         &self,
         event_type: &StateEventType,
@@ -86,6 +102,7 @@ impl service::rooms::short::Data for KeyValueDatabase {
             .unwrap()
             .get_mut(&(event_type.clone(), state_key.to_owned()))
         {
+            FoundIn::Cache.record("cache_result");
             return Ok(*short);
         }
 
@@ -96,10 +113,14 @@ impl service::rooms::short::Data for KeyValueDatabase {
         let short = if let Some(shortstatekey) =
             self.statekey_shortstatekey.get(&db_key)?
         {
+            FoundIn::Database.record("cache_result");
+
             utils::u64_from_bytes(&shortstatekey).map_err(|_| {
                 Error::bad_database("Invalid shortstatekey in db.")
             })?
         } else {
+            FoundIn::Nothing.record("cache_result");
+
             let shortstatekey = services().globals.next_count()?;
             self.statekey_shortstatekey
                 .insert(&db_key, &shortstatekey.to_be_bytes())?;
@@ -116,6 +137,7 @@ impl service::rooms::short::Data for KeyValueDatabase {
         Ok(short)
     }
 
+    #[tracing::instrument(skip(self), fields(cache_result))]
     fn get_eventid_from_short(
         &self,
         shorteventid: u64,
@@ -123,6 +145,7 @@ impl service::rooms::short::Data for KeyValueDatabase {
         if let Some(id) =
             self.shorteventid_cache.lock().unwrap().get_mut(&shorteventid)
         {
+            FoundIn::Cache.record("cache_result");
             return Ok(Arc::clone(id));
         }
 
@@ -144,6 +167,8 @@ impl service::rooms::short::Data for KeyValueDatabase {
             Error::bad_database("EventId in shorteventid_eventid is invalid.")
         })?;
 
+        FoundIn::Database.record("cache_result");
+
         self.shorteventid_cache
             .lock()
             .unwrap()
@@ -152,6 +177,7 @@ impl service::rooms::short::Data for KeyValueDatabase {
         Ok(event_id)
     }
 
+    #[tracing::instrument(skip(self), fields(cache_result))]
     fn get_statekey_from_short(
         &self,
         shortstatekey: u64,
@@ -159,6 +185,7 @@ impl service::rooms::short::Data for KeyValueDatabase {
         if let Some(id) =
             self.shortstatekey_cache.lock().unwrap().get_mut(&shortstatekey)
         {
+            FoundIn::Cache.record("cache_result");
             return Ok(id.clone());
         }
 
@@ -193,6 +220,8 @@ impl service::rooms::short::Data for KeyValueDatabase {
 
         let result = (event_type, state_key);
 
+        FoundIn::Database.record("cache_result");
+
         self.shortstatekey_cache
             .lock()
             .unwrap()
@@ -202,6 +231,7 @@ impl service::rooms::short::Data for KeyValueDatabase {
     }
 
     /// Returns `(shortstatehash, already_existed)`
+    #[tracing::instrument(skip(self))]
     fn get_or_create_shortstatehash(
         &self,
         state_hash: &[u8],

src/database/key_value/rooms/state_cache.rs

@@ -9,7 +9,9 @@ use ruma::{
 use crate::{
     database::KeyValueDatabase,
     service::{self, appservice::RegistrationInfo},
-    services, utils, Error, Result,
+    services,
+    utils::{self, FoundIn},
+    Error, Result,
 };
 
 impl service::rooms::state_cache::Data for KeyValueDatabase {
@@ -170,7 +172,7 @@ impl service::rooms::state_cache::Data for KeyValueDatabase {
         Ok(())
     }
 
-    #[tracing::instrument(skip(self, room_id))]
+    #[tracing::instrument(skip(self), fields(cache_result))]
     fn get_our_real_users(
         &self,
         room_id: &RoomId,
@@ -178,16 +180,21 @@ impl service::rooms::state_cache::Data for KeyValueDatabase {
         let maybe =
             self.our_real_users_cache.read().unwrap().get(room_id).cloned();
         if let Some(users) = maybe {
+            FoundIn::Cache.record("cache_result");
             Ok(users)
         } else {
             self.update_joined_count(room_id)?;
+            FoundIn::Database.record("cache_result");
             Ok(Arc::clone(
                 self.our_real_users_cache.read().unwrap().get(room_id).unwrap(),
             ))
         }
     }
 
-    #[tracing::instrument(skip(self, room_id, appservice))]
+    #[tracing::instrument(
+        skip(self, appservice),
+        fields(cache_result, appservice_id = appservice.registration.id),
+    )]
     fn appservice_in_room(
         &self,
         room_id: &RoomId,
@@ -202,6 +209,7 @@ impl service::rooms::state_cache::Data for KeyValueDatabase {
             .copied();
 
         if let Some(b) = maybe {
+            FoundIn::Cache.record("cache_result");
             Ok(b)
         } else {
             let bridge_user_id = UserId::parse_with_server_name(
@@ -218,6 +226,7 @@ impl service::rooms::state_cache::Data for KeyValueDatabase {
                 })
             });
 
+            FoundIn::Database.record("cache_result");
             self.appservice_in_room_cache
                 .write()
                 .unwrap()

src/database/key_value/rooms/timeline.rs

@@ -8,11 +8,14 @@ use service::rooms::timeline::PduCount;
 use tracing::error;
 
 use crate::{
-    database::KeyValueDatabase, service, services, utils, Error, PduEvent,
+    database::KeyValueDatabase,
-    Result,
+    service, services,
+    utils::{self, FoundIn},
+    Error, PduEvent, Result,
 };
 
 impl service::rooms::timeline::Data for KeyValueDatabase {
+    #[tracing::instrument(skip(self), fields(cache_result))]
     fn last_timeline_count(
         &self,
         sender_user: &UserId,
@@ -34,12 +37,17 @@ impl service::rooms::timeline::Data for KeyValueDatabase {
                         r.ok()
                     })
                 {
+                    FoundIn::Database.record("cache_result");
                     Ok(*v.insert(last_count.0))
                 } else {
+                    FoundIn::Nothing.record("cache_result");
                     Ok(PduCount::Normal(0))
                 }
             }
-            hash_map::Entry::Occupied(o) => Ok(*o.get()),
+            hash_map::Entry::Occupied(o) => {
+                FoundIn::Cache.record("cache_result");
+                Ok(*o.get())
+            }
         }
     }
 
@@ -119,8 +127,10 @@ impl service::rooms::timeline::Data for KeyValueDatabase {
     /// Returns the pdu.
     ///
     /// Checks the `eventid_outlierpdu` Tree if not found in the timeline.
+    #[tracing::instrument(skip(self), fields(cache_result))]
     fn get_pdu(&self, event_id: &EventId) -> Result<Option<Arc<PduEvent>>> {
         if let Some(p) = self.pdu_cache.lock().unwrap().get_mut(event_id) {
+            FoundIn::Cache.record("cache_result");
             return Ok(Some(Arc::clone(p)));
         }
 
@@ -141,12 +151,14 @@ impl service::rooms::timeline::Data for KeyValueDatabase {
             )?
             .map(Arc::new)
         {
+            FoundIn::Database.record("cache_result");
             self.pdu_cache
                 .lock()
                 .unwrap()
                 .insert(event_id.to_owned(), Arc::clone(&pdu));
             Ok(Some(pdu))
         } else {
+            FoundIn::Nothing.record("cache_result");
             Ok(None)
         }
     }

src/service/rooms/state_accessor.rs

@@ -27,7 +27,9 @@ use serde_json::value::to_raw_value;
 use tokio::sync::MutexGuard;
 use tracing::{error, warn};
 
-use crate::{service::pdu::PduBuilder, services, Error, PduEvent, Result};
+use crate::{
+    service::pdu::PduBuilder, services, utils::FoundIn, Error, PduEvent, Result,
+};
 
 pub(crate) struct Service {
     pub(crate) db: &'static dyn Data,
@@ -121,7 +123,7 @@ impl Service {
 
     /// Whether a server is allowed to see an event through federation, based on
     /// the room's history_visibility at that event's state.
-    #[tracing::instrument(skip(self, origin, room_id, event_id))]
+    #[tracing::instrument(skip(self), fields(cache_result))]
     pub(crate) fn server_can_see_event(
         &self,
         origin: &ServerName,
@@ -138,6 +140,7 @@ impl Service {
             .unwrap()
             .get_mut(&(origin.to_owned(), shortstatehash))
         {
+            FoundIn::Cache.record("cache_result");
             return Ok(*visibility);
         }
 
@@ -188,6 +191,7 @@ impl Service {
             }
         };
 
+        FoundIn::Database.record("cache_result");
         self.server_visibility_cache
             .lock()
             .unwrap()
@@ -198,7 +202,7 @@ impl Service {
 
     /// Whether a user is allowed to see an event, based on
     /// the room's history_visibility at that event's state.
-    #[tracing::instrument(skip(self, user_id, room_id, event_id))]
+    #[tracing::instrument(skip(self), fields(cache_result))]
     pub(crate) fn user_can_see_event(
         &self,
         user_id: &UserId,
@@ -215,6 +219,7 @@ impl Service {
             .unwrap()
             .get_mut(&(user_id.to_owned(), shortstatehash))
         {
+            FoundIn::Cache.record("cache_result");
             return Ok(*visibility);
         }
 
@@ -257,6 +262,7 @@ impl Service {
             }
         };
 
+        FoundIn::Database.record("cache_result");
         self.user_visibility_cache
             .lock()
             .unwrap()

src/service/rooms/state_compressor.rs

@@ -10,7 +10,11 @@ use lru_cache::LruCache;
 use ruma::{EventId, RoomId};
 
 use self::data::StateDiff;
-use crate::{services, utils, Result};
+use crate::{
+    services,
+    utils::{self, FoundIn},
+    Result,
+};
 
 #[derive(Clone)]
 pub(crate) struct CompressedStateLayer {
@@ -33,7 +37,7 @@ impl Service {
     /// Returns a stack with info on shortstatehash, full state, added diff and
     /// removed diff for the selected shortstatehash and each parent layer.
     #[allow(clippy::type_complexity)]
-    #[tracing::instrument(skip(self))]
+    #[tracing::instrument(skip(self), fields(cache_result))]
     pub(crate) fn load_shortstatehash_info(
         &self,
         shortstatehash: u64,
@@ -41,6 +45,7 @@ impl Service {
         if let Some(r) =
             self.stateinfo_cache.lock().unwrap().get_mut(&shortstatehash)
         {
+            FoundIn::Cache.record("cache_result");
             return Ok(r.clone());
         }
 
@@ -50,7 +55,7 @@ impl Service {
             removed,
         } = self.db.get_statediff(shortstatehash)?;
 
-        if let Some(parent) = parent {
+        let response = if let Some(parent) = parent {
             let mut response = self.load_shortstatehash_info(parent)?;
             let mut state = (*response.last().unwrap().full_state).clone();
             state.extend(added.iter().copied());
@@ -65,26 +70,23 @@ impl Service {
                 added,
                 removed: Arc::new(removed),
             });
-
+            response
-            self.stateinfo_cache
-                .lock()
-                .unwrap()
-                .insert(shortstatehash, response.clone());
-
-            Ok(response)
         } else {
-            let response = vec![CompressedStateLayer {
+            vec![CompressedStateLayer {
                 shortstatehash,
                 full_state: added.clone(),
                 added,
                 removed,
-            }];
+            }]
-            self.stateinfo_cache
+        };
-                .lock()
+
-                .unwrap()
+        FoundIn::Database.record("cache_result");
-                .insert(shortstatehash, response.clone());
+        self.stateinfo_cache
-            Ok(response)
+            .lock()
-        }
+            .unwrap()
+            .insert(shortstatehash, response.clone());
+
+        Ok(response)
     }
 
     // Allowed because this function uses `services()`

src/utils.rs

@@ -231,6 +231,35 @@ pub(crate) fn truncate_str_for_debug(
     }
 }
 
+/// Type to record cache performance in a tracing span field.
+pub(crate) enum FoundIn {
+    /// Found in cache
+    Cache,
+    /// Cache miss, but it was in the database. The cache has been updated.
+    Database,
+    /// Cache and database miss, but another server had it. The cache has been
+    /// updated.
+    Remote,
+    /// The entry could not be found anywhere.
+    Nothing,
+}
+
+impl FoundIn {
+    fn value(&self) -> &'static str {
+        match self {
+            FoundIn::Cache => "hit",
+            FoundIn::Database => "miss-database",
+            FoundIn::Remote => "miss-remote",
+            FoundIn::Nothing => "not-found",
+        }
+    }
+
+    // TODO: use tracing::Value instead if it ever becomes accessible
+    pub(crate) fn record(&self, field: &str) {
+        tracing::Span::current().record(field, self.value());
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::utils::truncate_str_for_debug;