S3-FIFO caching

stalwartlabs · Dec 27, 2024 · f2b00cc · f2b00cc
1 parent 3530b66
commit f2b00cc
Show file tree

Hide file tree

Showing 33 changed files with 631 additions and 383 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/common/src/auth/access_token.rs b/crates/common/src/auth/access_token.rs
@@ -6,7 +6,7 @@
 
 use directory::{
     backend::internal::{lookup::DirectoryStore, PrincipalField},
-    Permission, Principal, QueryBy,
+    Permission, Principal, QueryBy, PERMISSIONS_BITSET_SIZE,
 };
 use jmap_proto::{
     request::RequestMethod,
@@ -15,14 +15,15 @@ use jmap_proto::{
 use std::{
     hash::{DefaultHasher, Hash, Hasher},
     sync::Arc,
-    time::Instant,
 };
 use store::query::acl::AclQuery;
 use trc::AddContext;
-use utils::map::{
-    bitmap::{Bitmap, BitmapItem},
-    ttl_dashmap::TtlMap,
-    vec_map::VecMap,
+use utils::{
+    cache::TtlEntry,
+    map::{
+        bitmap::{Bitmap, BitmapItem},
+        vec_map::VecMap,
+    },
 };
 
 use crate::Server;
@@ -109,6 +110,7 @@ impl Server {
                 .unwrap_or_default(),
             quota: principal.quota(),
             permissions,
+            obj_size: 0,
         })
     }
 
@@ -183,27 +185,52 @@ impl Server {
             }
         }
 
-        Ok(access_token)
+        Ok(access_token.update_size())
     }
 
-    pub fn cache_access_token(&self, access_token: Arc<AccessToken>) {
-        self.inner.data.access_tokens.insert_with_ttl(
-            access_token.primary_id(),
-            access_token,
-            Instant::now() + self.core.jmap.session_cache_ttl,
-        );
+    pub async fn get_or_build_access_token(
+        &self,
+        principal: Principal,
+    ) -> trc::Result<Arc<AccessToken>> {
+        match self
+            .inner
+            .cache
+            .access_tokens
+            .get_value_or_guard_async(&principal.id())
+            .await
+        {
+            Ok(token) => Ok(token),
+            Err(guard) => {
+                let token = Arc::new(
+                    self.update_access_token(self.build_access_token(principal).await?)
+                        .await?,
+                );
+                let _ = guard.insert(TtlEntry::new(
+                    token.clone(),
+                    self.core.jmap.session_cache_ttl,
+                ));
+                Ok(token)
+            }
+        }
     }
 
     pub async fn get_cached_access_token(&self, primary_id: u32) -> trc::Result<Arc<AccessToken>> {
-        if let Some(access_token) = self.inner.data.access_tokens.get_with_ttl(&primary_id) {
-            Ok(access_token)
-        } else {
-            // Refresh ACL token
-            self.get_access_token(primary_id).await.map(|access_token| {
-                let access_token = Arc::new(access_token);
-                self.cache_access_token(access_token.clone());
-                access_token
-            })
+        match self
+            .inner
+            .cache
+            .access_tokens
+            .get_value_or_guard_async(&primary_id)
+            .await
+        {
+            Ok(token) => Ok(token),
+            Err(guard) => {
+                let token = Arc::new(self.get_access_token(primary_id).await?);
+                let _ = guard.insert(TtlEntry::new(
+                    token.clone(),
+                    self.core.jmap.session_cache_ttl,
+                ));
+                Ok(token)
+            }
         }
     }
 }
@@ -467,4 +494,17 @@ impl AccessToken {
             tenant: self.tenant,
         }
     }
+
+    pub fn update_size(mut self) -> Self {
+        self.obj_size = ((std::mem::size_of::<u32>() * 2)
+            + (std::mem::size_of::<u64>() * 3)
+            + (self.member_of.len() * std::mem::size_of::<u32>())
+            + (self.access_to.len() * (std::mem::size_of::<u32>() + std::mem::size_of::<u64>()))
+            + self.name.len()
+            + self.description.as_ref().map_or(0, |v| v.len())
+            + self.emails.iter().map(|v| v.len()).sum::<usize>()
+            + (PERMISSIONS_BITSET_SIZE * std::mem::size_of::<usize>()))
+            as u64;
+        self
+    }
 }
diff --git a/crates/common/src/auth/mod.rs b/crates/common/src/auth/mod.rs
@@ -4,15 +4,18 @@
  * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
  */
 
-use std::{net::IpAddr, sync::Arc, time::Instant};
+use std::{net::IpAddr, sync::Arc};
 
 use directory::{
     core::secret::verify_secret_hash, Directory, Permission, Permissions, Principal, QueryBy,
 };
 use jmap_proto::types::collection::Collection;
 use mail_send::Credentials;
 use oauth::GrantType;
-use utils::map::{bitmap::Bitmap, ttl_dashmap::TtlMap, vec_map::VecMap};
+use utils::{
+    cache::CacheItemWeight,
+    map::{bitmap::Bitmap, vec_map::VecMap},
+};
 
 use crate::Server;
 
@@ -32,6 +35,7 @@ pub struct AccessToken {
     pub quota: u64,
     pub permissions: Permissions,
     pub tenant: Option<TenantInfo>,
+    pub obj_size: u64,
 }
 
 #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
@@ -72,21 +76,7 @@ impl Server {
                 }
             }
             _ => match self.authenticate_credentials(req, directory).await {
-                Ok(principal) => {
-                    if let Some(access_token) =
-                        self.inner.data.access_tokens.get_with_ttl(&principal.id())
-                    {
-                        Ok(access_token)
-                    } else {
-                        self.build_access_token(principal)
-                            .await
-                            .map(|access_token| {
-                                let access_token = Arc::new(access_token);
-                                self.cache_access_token(access_token.clone());
-                                access_token
-                            })
-                    }
-                }
+                Ok(principal) => self.get_or_build_access_token(principal).await,
                 Err(err) => Err(err),
             },
         }
@@ -195,14 +185,6 @@ impl Server {
                 ))
         }
     }
-
-    pub fn cache_session(&self, session_id: String, access_token: &AccessToken) {
-        self.inner.data.http_auth_cache.insert_with_ttl(
-            session_id,
-            access_token.primary_id(),
-            Instant::now() + self.core.jmap.session_cache_ttl,
-        );
-    }
 }
 
 impl<'x> AuthRequest<'x> {
@@ -247,6 +229,12 @@ impl<'x> AuthRequest<'x> {
     }
 }
 
+impl CacheItemWeight for AccessToken {
+    fn weight(&self) -> u64 {
+        self.obj_size
+    }
+}
+
 pub(crate) trait CredentialsUsername {
     fn login(&self) -> Option<&str>;
 }

diff --git a/crates/common/src/auth/roles.rs b/crates/common/src/auth/roles.rs
@@ -9,9 +9,11 @@ use std::sync::{Arc, LazyLock};
 use ahash::AHashSet;
 use directory::{
     backend::internal::{lookup::DirectoryStore, PrincipalField},
-    Permission, Permissions, QueryBy, ROLE_ADMIN, ROLE_TENANT_ADMIN, ROLE_USER,
+    Permission, Permissions, QueryBy, PERMISSIONS_BITSET_SIZE, ROLE_ADMIN, ROLE_TENANT_ADMIN,
+    ROLE_USER,
 };
 use trc::AddContext;
+use utils::cache::CacheItemWeight;
 
 use crate::Server;
 
@@ -33,10 +35,19 @@ impl Server {
             ROLE_ADMIN => Ok(ADMIN_PERMISSIONS.clone()),
             ROLE_TENANT_ADMIN => Ok(TENANT_ADMIN_PERMISSIONS.clone()),
             role_id => {
-                if let Some(role_permissions) = self.inner.data.permissions.get(&role_id) {
-                    Ok(role_permissions.clone())
-                } else {
-                    self.build_role_permissions(role_id).await
+                match self
+                    .inner
+                    .cache
+                    .permissions
+                    .get_value_or_guard_async(&role_id)
+                    .await
+                {
+                    Ok(permissions) => Ok(permissions),
+                    Err(guard) => {
+                        let permissions = self.build_role_permissions(role_id).await?;
+                        let _ = guard.insert(permissions.clone());
+                        Ok(permissions)
+                    }
                 }
             }
         }
@@ -81,7 +92,7 @@ impl Server {
                     }
                     role_id => {
                         // Try with the cache
-                        if let Some(role_permissions) = self.inner.data.permissions.get(&role_id) {
+                        if let Some(role_permissions) = self.inner.cache.permissions.get(&role_id) {
                             return_permissions.union(role_permissions.as_ref());
                         } else {
                             let mut role_permissions = RolePermissions::default();
@@ -133,7 +144,7 @@ impl Server {
                             } else {
                                 // Cache role
                                 self.inner
-                                    .data
+                                    .cache
                                     .permissions
                                     .insert(role_id, Arc::new(role_permissions));
                             }
@@ -147,13 +158,7 @@ impl Server {
             }
         }
 
-        // Cache role
-        let return_permissions = Arc::new(return_permissions);
-        self.inner
-            .data
-            .permissions
-            .insert(role_id, return_permissions.clone());
-        Ok(return_permissions)
+        Ok(Arc::new(return_permissions))
     }
 }
 
@@ -213,3 +218,9 @@ fn admin_permissions() -> Arc<RolePermissions> {
         disabled: Permissions::new(),
     })
 }
+
+impl CacheItemWeight for RolePermissions {
+    fn weight(&self) -> u64 {
+        (PERMISSIONS_BITSET_SIZE * std::mem::size_of::<usize>() * 2) as u64
+    }
+}