fix session saving and loading data race

This fixes a bug where session data could be lost if a session is loaded
from the store while it's also being saved back to the store. Because
sessions are loaded and stored in memory at the beginning of a request,
it was previously possible for a session's inner value to be correctly
altered but for that modification to be lost between the moment the
session is saved to the store and when it's loaded. In other words, the
data loaded would not reflect the in-progress update and instead be
stale.

To address this we could hold a lock for the duration of the request.
However, this imposes a performance penalty we would like to avoid.
Alternatively we can manage a registry of sessions, a la tower-cookie's
management of `CookieJar`: with this, we load a session first from the
service's cache of sessions and fallback to the store only when the
session is not present.

Currently sessions are managed like this via DashMap. However,
alternatively caching facilities, such as Moka, might prove to be a
better fit. Further, the current implementation does not bound the total
number of cached sessions--this should be addressed in future
updates. At present, this means memory will continue to grow, albeit
typically at a very small rate; this is a memory leak!

Because sessions are safely written to and read from this cache, we
prevent stale reads and thus a data race which could lead to data loss
in highly concurrent environments, especially where multiple requests
are made using the same session cookie simultaneously.

Note that `replace_if_equal` is still required to prevent data races in
application code.

Cargo.toml

@@ -27,6 +27,7 @@ tokio = ["dep:tokio"]
 
 [dependencies]
 async-trait = "0.1.73"
+dashmap = { version = "5.5.3", default-features = false }
 http = "0.2.9"
 futures = { version = "0.3.28", default-features = false, features = [
     "async-await",

src/service.rs

@@ -2,9 +2,11 @@
 use std::{
     future::Future,
     pin::Pin,
+    sync::Arc,
     task::{Context, Poll},
 };
 
+use dashmap::{mapref::entry::Entry, DashMap};
 use http::{Request, Response};
 use time::Duration;
 use tower_cookies::{cookie::SameSite, Cookie, CookieManager, Cookies};
@@ -20,15 +22,21 @@ use crate::{
 #[derive(Debug, Clone)]
 pub struct SessionManager<S, Store: SessionStore> {
     inner: S,
+    // TODO: This is currently unbounded, meaning sessions grow in memory indefinitely.
+    // Alternatively we could set some configurable capacity. Moka is possibly a better candidate
+    // for this use case.
+    sessions: Arc<DashMap<SessionId, Session>>,
     session_store: Store,
     cookie_config: CookieConfig,
 }
 
 impl<S, Store: SessionStore> SessionManager<S, Store> {
     /// Create a new [`SessionManager`].
     pub fn new(inner: S, session_store: Store, cookie_config: CookieConfig) -> Self {
+        let sessions = Arc::new(DashMap::new());
         Self {
             inner,
+            sessions,
             session_store,
             cookie_config,
         }
@@ -54,6 +62,7 @@ where
     }
 
     fn call(&mut self, mut req: Request<ReqBody>) -> Self::Future {
+        let sessions = self.sessions.clone();
         let session_store = self.session_store.clone();
         let cookie_config = self.cookie_config.clone();
 
@@ -70,18 +79,28 @@ where
             let mut session = if let Some(session_cookie) =
                 cookies.get(&cookie_config.name).map(Cookie::into_owned)
             {
-                // We do have a session cookie, so let's see if our store has the associated
-                // session.
+                // We do have a session cookie, so we load it either directly from our
+                // in-memory cache or the backing session store.
                 //
                 // N.B.: Our store will *not* have the session if the session is empty.
                 let session_id = session_cookie.value().try_into()?;
-                let session = session_store.load(&session_id).await?;
+                let session = match sessions.entry(session_id) {
+                    Entry::Vacant(entry) => {
+                        if let Some(session) = session_store.load(&session_id).await? {
+                            session_loaded_from_store = true;
+                            entry.insert(session.clone());
+                            Some(session)
+                        } else {
+                            None
+                        }
+                    }
+
+                    Entry::Occupied(entry) => Some(entry.get()).cloned(),
+                };
 
                 // If the store does not know about this session, we should remove the cookie.
                 if session.is_none() {
                     cookies.remove(session_cookie);
-                } else {
-                    session_loaded_from_store = true;
                 }
 
                 session
@@ -113,6 +132,7 @@ where
                             session_store.delete(&session.id()).await?;
                         }
 
+                        sessions.remove(&session.id());
                         cookies.remove(cookie_config.build_cookie(&session));
 
                         // Since the session has been deleted, there's no need for further
@@ -124,8 +144,9 @@ where
                         if session_loaded_from_store {
                             session_store.delete(&deleted_id).await?;
                         }
-
+                        sessions.remove(&session.id());
                         session.id = SessionId::default();
+                        sessions.insert(session.id(), session.clone());
                     }
                 }
             };
@@ -139,6 +160,7 @@ where
             // the `Set-Cookie` header whenever modified or if some "always save" marker is
             // set.
             if session.modified() {
+                sessions.insert(session.id(), session.clone());
                 let session_record = (&session).into();
                 session_store.save(&session_record).await?;
                 cookies.add(cookie_config.build_cookie(&session))
@@ -280,6 +302,7 @@ impl<S, Store: SessionStore> Layer<S> for SessionManagerLayer<Store> {
     fn layer(&self, inner: S) -> Self::Service {
         let session_manager = SessionManager {
             inner,
+            sessions: Arc::new(DashMap::new()),
             session_store: self.session_store.clone(),
             cookie_config: self.cookie_config.clone(),
         };