use less-terrible way to keep track of count state (#26)

* use less-terrible way to keep track of count state

* bmp

* add resync check

* comments and unset dirty if resync

server/Cargo.toml

@@ -7,6 +7,7 @@ edition = "2021"
 strip = "symbols"
 
 [dependencies]
+anyhow = "1.0.79"
 axum = "0.5.17" # The webapp framework of choice.
 dotenv = "0.15.0" # For accessing env vars stored in .env
 serde = { version = "1.0.137", features = ["derive"] } # For JSON serialization.

server/clean_log.sh

@@ -2,4 +2,4 @@
 
 set -eu
 
-find /home/pyreko/htv-server/ -mtime +7 -name "volume.log*" -print -exec /bin/rm {} \;
+find ~/yc-server/ -mtime +7 -name "volume.log*" -print -exec /bin/rm {} \;

server/restart.sh

@@ -1,8 +1,8 @@
 #!/bin/bash
 
-set -eu
+set -eux
 
-pkill -2 hows-the-volume
+pkill -2 hows-the-volume || true
 while pgrep -u $UID -x hows-the-volume >/dev/null; do sleep 1; done
 cp ./target/release/hows-the-volume-server ~/htv-server/hows-the-volume-server
 cp -r ./assets ~/htv-server/

server/src/main.rs

@@ -1,5 +1,8 @@
-use std::{env, fs, sync::Arc};
+mod state;
 
+use std::{env, fs, sync::Arc, time::Duration};
+
+use anyhow::Result;
 use axum::{
     body::Body,
     extract::{rejection::PathRejection, Path},
@@ -11,74 +14,16 @@ use axum::{
 };
 use dotenv::dotenv;
 use serde::Serialize;
+
 use sqlx::{Pool, Sqlite, SqlitePool};
+use state::State;
+use tokio::time::timeout;
 use tower::util::ServiceExt;
 use tower_http::{cors::CorsLayer, services::ServeDir};
-use tracing::{error, info, warn};
+use tracing::{error, info};
 use tracing_subscriber::filter::EnvFilter;
 
-#[tokio::main]
-async fn main() {
-    dotenv().ok().unwrap();
-
-    let file_appender = tracing_appender::rolling::daily("./", "volume.log");
-    let (non_blocking, _guard) = tracing_appender::non_blocking(file_appender);
-
-    tracing_subscriber::fmt()
-        .with_env_filter(EnvFilter::from_default_env())
-        .with_writer(non_blocking)
-        .with_ansi(false)
-        .init();
-
-    info!("Starting up HTV...");
-    let pool = SqlitePool::connect(&env::var("DATABASE_URL").unwrap())
-        .await
-        .unwrap();
-
-    let pool_rc = Arc::new(pool);
-
-    let origins = [
-        "http://localhost:3000".parse::<HeaderValue>().unwrap(),
-        "https://howsthevolu.me".parse::<HeaderValue>().unwrap(),
-    ];
-    let cors = CorsLayer::new()
-        .allow_methods(vec![Method::GET, Method::POST])
-        .allow_origin(origins);
-
-    let app = Router::new()
-        .route("/sound/:id", get(sound))
-        .route("/count", get(count))
-        .route("/increment", post(increment))
-        .route("/num-files", get(num_audio_tracks))
-        .layer(cors)
-        .layer(Extension(pool_rc.clone()))
-        .fallback(not_found_handler.into_service());
-
-    let addr = "127.0.0.1:8080".parse().unwrap();
-
-    info!("Listening on {}", addr);
-
-    if std::path::Path::new("assets/").exists() {
-        let num_files = fs::read_dir("assets/").unwrap().count();
-        info!("Found {} files in assets!", num_files);
-    } else {
-        error!("Warning - no asset/ folder found! There should be one located near the binary!");
-    }
-
-    axum::Server::bind(&addr)
-        .serve(app.into_make_service())
-        .with_graceful_shutdown(async {
-            tokio::signal::ctrl_c().await.unwrap();
-            info!("Shutdown signal received.");
-        })
-        .await
-        .unwrap();
-
-    info!("Closing SQLite connection.");
-    pool_rc.close().await;
-
-    info!("Shutting down server.");
-}
+use crate::state::init_state;
 
 #[derive(Serialize)]
 struct EmptyJson {}
@@ -91,8 +36,7 @@ fn not_found() -> StatusCode {
     StatusCode::NOT_FOUND
 }
 
-type PoolExt = Arc<Pool<Sqlite>>;
-
+/// A JSON message containing a count.
 #[derive(Serialize)]
 struct Count {
     count: u64,
@@ -104,26 +48,9 @@ impl Count {
     }
 }
 
-/// Returns the current global count, stored in the DB.
-async fn count(Extension(pool): Extension<PoolExt>) -> Json<Count> {
-    match pool.acquire().await {
-        Ok(mut conn) => {
-            match sqlx::query!("SELECT count FROM counts WHERE name = 'volume'")
-                .fetch_one(&mut conn)
-                .await
-            {
-                Ok(query) => Json(Count::new(query.count as u64)),
-                Err(err) => {
-                    warn!("SQLite query for count failed - err: {}", err);
-                    Json(Count::new(0))
-                }
-            }
-        }
-        Err(err) => {
-            error!("Failed to get pool connection to SQLite DB - err: {}", err);
-            Json(Count::new(0))
-        }
-    }
+async fn count(Extension(state): Extension<State>) -> Json<Count> {
+    let val = state.lock().await;
+    Json(Count::new(val.count))
 }
 
 /// Returns the selected sound file if it exists.
@@ -132,7 +59,7 @@ async fn sound(id_result: Result<Path<u32>, PathRejection>) -> Response {
         const PREFIX: &str = "volume_";
         const SUFFIX: &str = ".mp3";
 
-        let uri = format!("/{}{:0>2}{}", PREFIX, id, SUFFIX);
+        let uri = format!("/{PREFIX}{id:0>2}{SUFFIX}");
         match Request::builder().uri(&uri).body(Body::empty()) {
             Ok(req) => match ServeDir::new("assets/").oneshot(req).await {
                 Ok(resp) => {
@@ -141,11 +68,11 @@ async fn sound(id_result: Result<Path<u32>, PathRejection>) -> Response {
                     }
                 }
                 Err(err) => {
-                    error!("Failed to get a response for file {} - err: {}", uri, err);
+                    error!("Failed to get a response for file {uri} - err: {err}");
                 }
             },
             Err(err) => {
-                error!("Failed to build a request for file {} - err: {}", uri, err);
+                error!("Failed to build a request for file {uri} - err: {err}");
             }
         }
     }
@@ -154,17 +81,19 @@ async fn sound(id_result: Result<Path<u32>, PathRejection>) -> Response {
 }
 
 /// Increments the count.
-async fn increment(Extension(pool): Extension<PoolExt>) {
-    let mut conn = pool.acquire().await.unwrap();
+async fn increment(Extension(state): Extension<State>) {
+    let mut val = state.lock().await;
+    val.count += 1;
+    val.dirty = true;
+}
 
-    match sqlx::query!("UPDATE counts SET count = count + 1 WHERE name = 'volume'")
-        .execute(&mut conn)
-        .await
+/// Open the SQLite pool.
+async fn open_pool() -> Result<Arc<Pool<Sqlite>>> {
     {
-        Ok(_) => {}
-        Err(err) => {
-            error!("Failed to increment in DB - err: {}", err);
-        }
+        let url = env::var("DATABASE_URL")?;
+        let pool = SqlitePool::connect(&url).await?;
+
+        Ok(Arc::new(pool))
     }
 }
 
@@ -185,3 +114,70 @@ async fn num_audio_tracks() -> Json<Count> {
 
     Json(Count::new(num_tracks))
 }
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    dotenv().ok().unwrap();
+
+    let file_appender = tracing_appender::rolling::daily("./", "volume.log");
+    let (non_blocking, _guard) = tracing_appender::non_blocking(file_appender);
+
+    tracing_subscriber::fmt()
+        .with_env_filter(EnvFilter::from_default_env())
+        .with_writer(non_blocking)
+        .with_ansi(false)
+        .init();
+
+    info!("Starting up HTV...");
+
+    let pool = open_pool().await?;
+    let (state, sync_task, shutdown) = init_state(pool.clone()).await?;
+
+    let origins = [
+        "http://localhost:3000".parse::<HeaderValue>()?,
+        "https://howsthevolu.me".parse::<HeaderValue>()?,
+    ];
+
+    let cors = CorsLayer::new()
+        .allow_methods(vec![Method::GET, Method::POST])
+        .allow_origin(origins);
+
+    let app = Router::new()
+        .route("/sound/:id", get(sound))
+        .route("/count", get(count))
+        .route("/increment", post(increment))
+        .route("/num-files", get(num_audio_tracks))
+        .layer(cors)
+        .layer(Extension(state))
+        .fallback(not_found_handler.into_service());
+
+    let addr = "127.0.0.1:8080".parse()?;
+
+    info!("Listening on {addr}");
+
+    if std::path::Path::new("assets/").exists() {
+        let num_files = fs::read_dir("assets/").unwrap().count();
+        info!("Found {num_files} files in assets!");
+    } else {
+        error!("Warning - no asset/ folder found! There should be one located near the binary!");
+    }
+
+    axum::Server::bind(&addr)
+        .serve(app.into_make_service())
+        .with_graceful_shutdown(async {
+            tokio::signal::ctrl_c().await.unwrap();
+            info!("Shutdown signal received.");
+        })
+        .await?;
+
+    info!("Stopping sync task.");
+    let _ = shutdown.send(());
+    let _ = timeout(Duration::from_secs(15), sync_task).await;
+
+    info!("Closing SQLite connection.");
+    pool.close().await;
+
+    info!("Cleanup complete, shutting down HTV server.");
+
+    Ok(())
+}