SSE-based directory size handling

This allows us to cancel stale clients and in general be much nicer and more modern.

3 files changed, 230 insertions, 100 deletions

diff --git a/src/handlers.rs b/src/handlers.rs
new file mode 100644
index 0000000..8b4d6e5
--- /dev/null
+++ b/src/handlers.rs
@@ -0,0 +1,183 @@
+use std::time::Duration;
+
+use actix_web::{HttpRequest, HttpResponse, Responder, http::header::ContentType, web};
+use actix_web_lab::sse;
+use bytesize::ByteSize;
+use dav_server::{
+    DavConfig, DavHandler,
+    actix::{DavRequest, DavResponse},
+};
+use log::{error, info, warn};
+use percent_encoding::percent_decode_str;
+use serde::{Deserialize, Serialize};
+use tokio::sync::Mutex;
+use tokio::task::JoinSet;
+
+use crate::{config::MiniserveConfig, errors::RuntimeError};
+use crate::{file_op::recursive_dir_size, file_utils};
+
+pub async fn dav_handler(req: DavRequest, davhandler: web::Data<DavHandler>) -> DavResponse {
+    if let Some(prefix) = req.prefix() {
+        let config = DavConfig::new().strip_prefix(prefix);
+        davhandler.handle_with(config, req.request).await.into()
+    } else {
+        davhandler.handle(req.request).await.into()
+    }
+}
+
+pub async fn error_404(req: HttpRequest) -> Result<HttpResponse, RuntimeError> {
+    Err(RuntimeError::RouteNotFoundError(req.path().to_string()))
+}
+
+pub async fn healthcheck() -> impl Responder {
+    HttpResponse::Ok().body("OK")
+}
+
+#[derive(Deserialize, Debug)]
+pub enum ApiCommand {
+    /// Request the size of a particular directory
+    CalculateDirSizes(Vec<String>),
+}
+
+pub type DirSizeJoinSet = JoinSet<Result<DirSize, RuntimeError>>;
+
+// Holds the result of a calculated dir size
+#[derive(Debug, Clone)]
+pub struct DirSize {
+    /// The web path of the dir (not the filesystem path)
+    pub web_path: String,
+
+    /// The calculcated recursive size of the dir
+    pub size: u64,
+}
+
+// Reply for a calculated dir size
+#[derive(Debug, Clone, Serialize)]
+pub struct DirSizeReply {
+    /// The web path of the dir (not the filesystem path)
+    pub web_path: String,
+
+    /// The formatted size of the dir
+    pub size: String,
+}
+
+// Reply to check whether the client is still connected
+//
+// If the client has disconnected, we can cancel all the tasks and save some compute.
+#[derive(Debug, Clone, Serialize)]
+pub struct HeartbeatReply;
+
+/// SSE API route that yields an event stream that clients can subscribe to
+pub async fn api_sse(
+    config: web::Data<MiniserveConfig>,
+    task_joinset: web::Data<Mutex<DirSizeJoinSet>>,
+) -> impl Responder {
+    let (sender, receiver) = tokio::sync::mpsc::channel(2);
+
+    actix_web::rt::spawn(async move {
+        loop {
+            let msg = match task_joinset.lock().await.try_join_next() {
+                Some(Ok(Ok(finished_task))) => {
+                    let dir_size = if config.show_exact_bytes {
+                        format!("{} B", finished_task.size)
+                    } else {
+                        ByteSize::b(finished_task.size).to_string()
+                    };
+
+                    let dir_size_reply = DirSizeReply {
+                        web_path: finished_task.web_path,
+                        size: dir_size,
+                    };
+
+                    sse::Data::new_json(dir_size_reply)
+                        .expect("Couldn't serialize as JSON")
+                        .event("dir-size")
+                }
+                Some(Ok(Err(e))) => {
+                    error!("Some error during dir size calculation: {e}");
+                    break;
+                }
+                Some(Err(e)) => {
+                    error!("Some error during dir size calculation joining: {e}");
+                    break;
+                }
+                None => sse::Data::new_json(HeartbeatReply)
+                    .expect("Couldn't serialize as JSON")
+                    .event("heartbeat"),
+            };
+
+            if sender.send(msg.into()).await.is_err() {
+                warn!("Client disconnected; could not send SSE message");
+                break;
+            }
+
+            tokio::time::sleep(Duration::from_secs(1)).await;
+        }
+    });
+
+    sse::Sse::from_infallible_receiver(receiver).with_keep_alive(Duration::from_secs(3))
+}
+
+async fn handle_dir_size_tasks(
+    dirs: Vec<String>,
+    config: &MiniserveConfig,
+    task_joinset: web::Data<Mutex<DirSizeJoinSet>>,
+) -> Result<(), RuntimeError> {
+    for dir in dirs {
+        // The dir argument might be percent-encoded so let's decode it just in case.
+        let decoded_path = percent_decode_str(&dir)
+            .decode_utf8()
+            .map_err(|e| RuntimeError::ParseError(dir.clone(), e.to_string()))?;
+
+        // Convert the relative dir to an absolute path on the system.
+        let sanitized_path =
+            file_utils::sanitize_path(&*decoded_path, true).expect("Expected a path to directory");
+
+        let full_path = config
+            .path
+            .canonicalize()
+            .expect("Couldn't canonicalize path")
+            .join(sanitized_path);
+        info!("Requested directory size for {full_path:?}");
+
+        let mut joinset = task_joinset.lock().await;
+        joinset.spawn(async move {
+            recursive_dir_size(&full_path).await.map(|dir_size| {
+                info!("Finished dir size calculation for {full_path:?}");
+                DirSize {
+                    web_path: dir,
+                    size: dir_size,
+                }
+            })
+        });
+    }
+    Ok(())
+}
+
+/// This "API" is pretty shitty but frankly miniserve doesn't really need a very fancy API. Or at
+/// least I hope so.
+pub async fn api_command(
+    command: web::Json<ApiCommand>,
+    config: web::Data<MiniserveConfig>,
+    task_joinset: web::Data<Mutex<DirSizeJoinSet>>,
+) -> Result<impl Responder, RuntimeError> {
+    match command.into_inner() {
+        ApiCommand::CalculateDirSizes(dirs) => {
+            handle_dir_size_tasks(dirs, &config, task_joinset).await?;
+            Ok("Directories are being calculated")
+        }
+    }
+}
+
+pub async fn favicon() -> impl Responder {
+    let logo = include_str!("../data/logo.svg");
+    HttpResponse::Ok()
+        .insert_header(ContentType(mime::IMAGE_SVG))
+        .body(logo)
+}
+
+pub async fn css(stylesheet: web::Data<String>) -> impl Responder {
+    HttpResponse::Ok()
+        .insert_header(ContentType(mime::TEXT_CSS))
+        .body(stylesheet.to_string())
+}
diff --git a/src/main.rs b/src/main.rs
index cd54bd8..3fe1410 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -6,25 +6,20 @@ use std::time::Duration;
 use actix_files::NamedFile;
 use actix_web::middleware::from_fn;
 use actix_web::{
-    App, HttpRequest, HttpResponse, Responder,
+    App,
     dev::{ServiceRequest, ServiceResponse, fn_service},
     guard,
-    http::{Method, header::ContentType},
+    http::Method,
     middleware, web,
 };
 use actix_web_httpauth::middleware::HttpAuthentication;
 use anyhow::Result;
-use bytesize::ByteSize;
 use clap::{CommandFactory, Parser, crate_version};
 use colored::*;
-use dav_server::{
-    DavConfig, DavHandler, DavMethodSet,
-    actix::{DavRequest, DavResponse},
-};
+use dav_server::{DavHandler, DavMethodSet};
 use fast_qr::QRBuilder;
-use log::{error, info, warn};
-use percent_encoding::percent_decode_str;
-use serde::Deserialize;
+use log::{error, warn};
+use tokio::sync::Mutex;
 
 mod archive;
 mod args;
@@ -34,14 +29,17 @@ mod consts;
 mod errors;
 mod file_op;
 mod file_utils;
+mod handlers;
 mod listing;
 mod pipe;
 mod renderer;
 mod webdav_fs;
 
 use crate::config::MiniserveConfig;
-use crate::errors::{RuntimeError, StartupError};
-use crate::file_op::recursive_dir_size;
+use crate::errors::StartupError;
+use crate::handlers::{
+    DirSizeJoinSet, api_command, api_sse, css, dav_handler, error_404, favicon, healthcheck,
+};
 use crate::webdav_fs::RestrictedFs;
 
 static STYLESHEET: &str = grass::include!("data/style.scss");
@@ -216,10 +214,13 @@ async fn run(miniserve_config: MiniserveConfig) -> Result<(), StartupError> {
         .join("\n"),
     );
 
+    let dir_size_join_set = web::Data::new(Mutex::new(DirSizeJoinSet::new()));
+
     let srv = actix_web::HttpServer::new(move || {
         App::new()
             .wrap(configure_header(&inside_config.clone()))
             .app_data(web::Data::new(inside_config.clone()))
+            .app_data(dir_size_join_set.clone())
             .app_data(stylesheet.clone())
             .wrap(from_fn(errors::error_page_middleware))
             .wrap(middleware::Logger::default())
@@ -228,7 +229,8 @@ async fn run(miniserve_config: MiniserveConfig) -> Result<(), StartupError> {
                 middleware::Compress::default(),
             ))
             .route(&inside_config.healthcheck_route, web::get().to(healthcheck))
-            .route(&inside_config.api_route, web::post().to(api))
+            .route(&inside_config.api_route, web::post().to(api_command))
+            .route(&inside_config.api_route, web::get().to(api_sse))
             .route(&inside_config.favicon_route, web::get().to(favicon))
             .route(&inside_config.css_route, web::get().to(css))
             .service(
@@ -425,74 +427,3 @@ fn configure_app(app: &mut web::ServiceConfig, conf: &MiniserveConfig) {
         );
     }
 }
-
-async fn dav_handler(req: DavRequest, davhandler: web::Data<DavHandler>) -> DavResponse {
-    if let Some(prefix) = req.prefix() {
-        let config = DavConfig::new().strip_prefix(prefix);
-        davhandler.handle_with(config, req.request).await.into()
-    } else {
-        davhandler.handle(req.request).await.into()
-    }
-}
-
-async fn error_404(req: HttpRequest) -> Result<HttpResponse, RuntimeError> {
-    Err(RuntimeError::RouteNotFoundError(req.path().to_string()))
-}
-
-async fn healthcheck() -> impl Responder {
-    HttpResponse::Ok().body("OK")
-}
-
-#[derive(Deserialize, Debug)]
-enum ApiCommand {
-    /// Request the size of a particular directory
-    DirSize(String),
-}
-
-/// This "API" is pretty shitty but frankly miniserve doesn't really need a very fancy API. Or at
-/// least I hope so.
-async fn api(
-    command: web::Json<ApiCommand>,
-    config: web::Data<MiniserveConfig>,
-) -> Result<impl Responder, RuntimeError> {
-    match command.into_inner() {
-        ApiCommand::DirSize(path) => {
-            // The dir argument might be percent-encoded so let's decode it just in case.
-            let decoded_path = percent_decode_str(&path)
-                .decode_utf8()
-                .map_err(|e| RuntimeError::ParseError(path.clone(), e.to_string()))?;
-
-            // Convert the relative dir to an absolute path on the system.
-            let sanitized_path = file_utils::sanitize_path(&*decoded_path, true)
-                .expect("Expected a path to directory");
-
-            let full_path = config
-                .path
-                .canonicalize()
-                .expect("Couldn't canonicalize path")
-                .join(sanitized_path);
-            info!("Requested directory listing for {full_path:?}");
-
-            let dir_size = recursive_dir_size(&full_path).await?;
-            if config.show_exact_bytes {
-                Ok(format!("{dir_size} B"))
-            } else {
-                let dir_size = ByteSize::b(dir_size);
-                Ok(dir_size.to_string())
-            }
-        }
-    }
-}
-
-async fn favicon() -> impl Responder {
-    let logo = include_str!("../data/logo.svg");
-    HttpResponse::Ok()
-        .insert_header(ContentType(mime::IMAGE_SVG))
-        .body(logo)
-}
-
-async fn css(stylesheet: web::Data<String>) -> impl Responder {
-    HttpResponse::Ok()
-        .insert_header(ContentType(mime::TEXT_CSS))
-        .body(stylesheet.to_string())
-}
diff --git a/src/renderer.rs b/src/renderer.rs
index 7846388..bcbbdfe 100644
--- a/src/renderer.rs
+++ b/src/renderer.rs
@@ -674,10 +674,8 @@ fn page_header(
             script {
                 (format!("const API_ROUTE = '{api_route}';"))
                 (PreEscaped(r#"
-                    let dirSizeCache = {};
-
                     // Query the directory size from the miniserve API
-                    function fetchDirSize(dir) {
+                    function requestDirSizeCalculation(dirs) {
                         return fetch(API_ROUTE, {
                             headers: {
                                 'Accept': 'application/json',
@@ -685,32 +683,50 @@ fn page_header(
                             },
                             method: 'POST',
                             body: JSON.stringify({
-                                DirSize: dir
+                                CalculateDirSizes: dirs
                             })
                         }).then(resp => resp.ok ? resp.text() : "~")
                     }
 
-                    function updateSizeCells() {
+                    function updateSizeCells(event_data) {
+                        console.log("Received dir-size object", event_data);
                         const directoryCells = document.querySelectorAll('tr.entry-type-directory .size-cell');
-
                         directoryCells.forEach(cell => {
                             // Get the dir from the sibling anchor tag.
                             const href = cell.parentNode.querySelector('a').href;
                             const target = new URL(href).pathname;
 
-                            // First check our local cache
-                            if (target in dirSizeCache) {
-                                cell.dataset.size = dirSizeCache[target];
-                            } else {
-                                fetchDirSize(target).then(dir_size => {
-                                    cell.dataset.size = dir_size;
-                                    dirSizeCache[target] = dir_size;
-                                })
-                                .catch(error => console.error("Error fetching dir size:", error));
+                            if (target === event_data.web_path) {
+                                cell.dataset.size = event_data.size
                             }
+                        });
+                    }
+
+                    function requestDirSizes() {
+                        // Subscribe to the SSE stream
+                        let dir_sizes_events = new EventSource(API_ROUTE)
+
+                        dir_sizes_events.addEventListener('dir-size', (event) => {
+                            updateSizeCells(JSON.parse(event.data));
                         })
+
+                        const directoryCells = document.querySelectorAll('tr.entry-type-directory .size-cell');
+
+                        // The list of dirs we'll request to have calculated.
+                        let dirs = [];
+
+                        directoryCells.forEach(cell => {
+                            // Get the dir from the sibling anchor tag.
+                            const href = cell.parentNode.querySelector('a').href;
+                            const target = new URL(href).pathname;
+
+                            dirs.push(target);
+                        });
+
+                        requestDirSizeCalculation(dirs);
                     }
-                    setInterval(updateSizeCells, 1000);
+
+                    window.addEventListener('load', requestDirSizes);
                 "#))
             }