Privatize access to internal session map

This way, RwLockReadGuard isn't being juggled anywhere, and the bounds of the session MutexGuard is limited to the FnOnce explicitly passed to with_session (less room for accidents, nice!)

This approach was inspired by the source code of RegionMap of Bamboo (https://gitlab.com/macmv/bamboo ref 8e32b54a5e7a69986339fea465d4afcc20425dd7), which uses a similar RwLock<HashMap<_, Mutex<_>>>. I like this approach because it allows me to provide access to the Session without having to clone it and allows the MutexGuard and RwLockReadGuard to be held as the user needs without allowing the user to access the RwLockReadGuard and poke around the rest of the entries in the HashMap.

4 files changed, 74 insertions, 63 deletions

diff --git a/src/app.rs b/src/app.rs
new file mode 100644
index 0000000..43e2ac6
--- /dev/null
+++ b/src/app.rs
@@ -0,0 +1,37 @@
+//! Internal app logic.
+
+use crate::session::Session;
+use std::collections::HashMap;
+use std::sync::{Mutex, MutexGuard, RwLock};
+
+/// Provider of the state that Tabletop Ambulator needs to keep track of.
+///
+/// Provides the app state, including runtime data such as the active game
+/// sessions. It provides access to relevant data through synchronization
+/// primitives, but does not handle shared ownership; for that, you may
+/// want to use a solution like [`std::sync::Arc`].
+#[derive(Debug, Default)]
+pub struct AppState {
+    sessions: RwLock<HashMap<String, Mutex<Session>>>,
+}
+
+impl AppState {
+    /// Creates a new `AppState`.
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Creates and stores a new [`Session`], which can later be accessed with [`with_session`].
+    pub fn create_session(&self, id: String, name: String) {
+        let session = Mutex::new(Session::new(name));
+        self.sessions.write().unwrap().insert(id, session);
+    }
+
+    /// Finds a [`Session`] by the provided `id`. If the session exists, it is provided within a
+    /// [`MutexGuard`] to the supplied function. Returns the result of the function, or `None` if
+    /// there was no session with the supplied id.
+    pub fn with_session<T>(&self, id: &str, f: impl FnOnce(MutexGuard<Session>) -> T) -> Option<T> {
+        let sessions = self.sessions.read().unwrap();
+        sessions.get(id).map(|session| f(session.lock().unwrap()))
+    }
+}
diff --git a/src/main.rs b/src/main.rs
index b423f6b..38bbbf2 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -6,12 +6,14 @@
 
 #![warn(missing_docs, missing_debug_implementations)]
 
+mod app;
 mod play;
 mod session;
 mod template;
 
+use crate::app::AppState;
 use crate::play::handle_play;
-use crate::session::{HandObject, PlayerColor, Session};
+use crate::session::{HandObject, PlayerColor};
 use crate::template::{IndexTemplate, SessionTemplate};
 use askama::Template;
 use axum::extract::{Path, Query, State, WebSocketUpgrade};
@@ -23,23 +25,12 @@ use rust_embed::Embed;
 use std::array;
 use std::collections::HashMap;
 use std::net::SocketAddr;
-use std::sync::{Arc, Mutex, RwLock};
+use std::sync::Arc;
 
 #[derive(Embed)]
 #[folder = "assets/"]
 struct EmbedAsset;
 
-#[derive(Default)]
-struct AppState {
-    sessions: RwLock<HashMap<String, Mutex<Session>>>,
-}
-
-impl AppState {
-    fn new() -> Self {
-        Self::default()
-    }
-}
-
 #[tokio::main]
 async fn main() {
     let app = Router::new()
@@ -57,8 +48,7 @@ async fn main() {
 }
 
 async fn serve_index() -> axum::response::Result<Html<String>> {
-    let template = IndexTemplate;
-    Template::render(&template)
+    Template::render(&IndexTemplate)
         .map(Html)
         .inspect_err(|e| eprintln!("Template render error: {}", e))
         .map_err(|_| ErrorResponse::from(StatusCode::INTERNAL_SERVER_ERROR))
@@ -89,19 +79,16 @@ async fn visit_session(
     Path(id): Path<String>,
     State(state): State<Arc<AppState>>,
 ) -> axum::response::Result<Html<String>> {
-    let sessions = state.sessions.read().unwrap();
-    let session = sessions
-        .get(&id)
-        .ok_or((StatusCode::NOT_FOUND, "Session does not exist"))?
-        .lock()
-        .unwrap();
-
-    let template = SessionTemplate {
-        id: id.as_str(),
-        session: &session,
-    };
-    Template::render(&template)
-        .map(Html)
+    let rendered = state
+        .with_session(id.as_str(), |session| {
+            Template::render(&SessionTemplate {
+                id: id.as_str(),
+                session: &session,
+            })
+            .map(Html)
+        })
+        .ok_or(StatusCode::NOT_FOUND)?;
+    rendered
         .inspect_err(|e| eprintln!("Template render error: {}", e))
         .map_err(|_| ErrorResponse::from(StatusCode::INTERNAL_SERVER_ERROR))
 }
@@ -111,12 +98,12 @@ async fn create_session(
     Query(query): Query<HashMap<String, String>>,
     State(state): State<Arc<AppState>>,
 ) -> StatusCode {
-    let name = query.get("name").cloned().unwrap_or("Unknown".to_string());
+    let name = query
+        .get("name")
+        .cloned()
+        .unwrap_or("Unknown User".to_string());
 
-    let mut sessions = state.sessions.write().unwrap();
-
-    let session = Session::new(name);
-    sessions.insert(id, Mutex::new(session));
+    state.create_session(id, name);
 
     StatusCode::CREATED
 }
@@ -124,32 +111,24 @@ async fn create_session(
 async fn update_hands(
     Path(id): Path<String>,
     State(state): State<Arc<AppState>>,
-    Json(payload): Json<HashMap<String, Vec<HandObject>>>,
+    Json(mut payload): Json<HashMap<String, Vec<HandObject>>>,
 ) -> StatusCode {
-    let sessions = state.sessions.read().unwrap();
-
     let hand = array::from_fn(|i| {
         let color = PlayerColor::try_from(i);
         let color = color.as_ref().map(AsRef::as_ref);
         if let Ok(color) = color {
-            // It would not be necessary to clone here if the vector could be moved, which could be
-            // done if payload was mutable
-            payload.get(color).cloned().unwrap_or_else(Vec::new)
+            payload.remove(color).unwrap_or_else(Vec::new)
         } else {
             Vec::new()
         }
     });
 
-    match sessions.get(&id) {
-        Some(session) => {
-            let mut session = session.lock().unwrap();
-
+    state
+        .with_session(id.as_str(), |mut session| {
             session.update_hands(hand);
-
             StatusCode::NO_CONTENT
-        }
-        None => StatusCode::NOT_FOUND,
-    }
+        })
+        .unwrap_or(StatusCode::NOT_FOUND)
 }
 
 async fn upgrade_play(ws: WebSocketUpgrade, State(state): State<Arc<AppState>>) -> Response {
diff --git a/src/play.rs b/src/play.rs
index 998ac29..0ec9758 100644
--- a/src/play.rs
+++ b/src/play.rs
@@ -144,16 +144,7 @@ async fn handle_play_message(
 ) -> Result<(), Error> {
     match message {
         IncomingPlayMessage::Initialize { id } => {
-            let (colors, update_rx) = {
-                let sessions = app_state.sessions.read().unwrap();
-                let session = sessions.get(&id).map(|session| session.lock().unwrap());
-                // The Option is unwrapped with let else here instead of ok_or and propagating
-                // because the error moves id which would be used later, and the compiler does not
-                // reason about propagation returning early.
-                let Some(session) = session else {
-                    return Err(Error::InvalidSession(id));
-                };
-
+            let data_opt = app_state.with_session(id.as_str(), |session| {
                 let colors: Vec<String> = session
                     .seats
                     .iter()
@@ -165,6 +156,11 @@ async fn handle_play_message(
                 let update_rx = session.update_tx.subscribe();
 
                 (colors, update_rx)
+            });
+            // let else used instead of propagating Option::ok_or_else because compiler wouldn't
+            // know about early return when moving id
+            let Some((colors, update_rx)) = data_opt else {
+                return Err(Error::InvalidSession(id));
             };
 
             let update_cancel_rx = {
@@ -194,16 +190,15 @@ async fn handle_play_message(
                 let mut state = state.lock().unwrap();
                 state.color =
                     PlayerColor::try_from(color.as_str()).map_err(|_| Error::InvalidColor)?;
+
                 let name = state
                     .session
                     .as_ref()
                     .ok_or_else(|| Error::InvalidSession(String::default()))?;
-                let sessions = app_state.sessions.read().unwrap();
-                let session = sessions
-                    .get(name)
-                    .ok_or(Error::InvalidSession(name.clone()))?;
 
-                session.lock().unwrap().seats[&state.color].clone()
+                app_state
+                    .with_session(name.as_str(), |session| session.seats[&state.color].clone())
+                    .ok_or_else(|| Error::InvalidSession(name.clone()))?
             };
 
             sender_tx
diff --git a/src/session.rs b/src/session.rs
index f4c8fe3..c15bfda 100644
--- a/src/session.rs
+++ b/src/session.rs
@@ -50,7 +50,7 @@ pub struct CustomDeck {
 ///
 /// See [Player Colors](https://api.tabletopsimulator.com/player/colors/) from the Tabletop
 /// Simulator API knowledge base.
-#[derive(Clone, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
+#[derive(Debug, Clone, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub enum PlayerColor {
     White,
     Brown,