fix: screenshots page lag + better virtualization pass

packages/app-lib/src/api/instance/screenshots/reconciliation.rs

@@ -8,12 +8,15 @@ use crate::state::instances::adapters::sqlite::{
     screenshot_rows::{self, ScreenshotRow},
 };
 use crate::util::fetch::sha1_file_async;
-use crate::util::io::{self, IOError};
+use crate::util::io::IOError;
 use chrono::{DateTime, Utc};
+use futures::stream::{self, StreamExt, TryStreamExt};
 use std::collections::HashMap;
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 use uuid::Uuid;
 
+const SCREENSHOT_HASH_CONCURRENCY: usize = 8;
+
 pub(super) struct ScannedScreenshot {
     file_name: String,
     created_at: DateTime<Utc>,
@@ -109,7 +112,14 @@ pub(super) async fn scan_source_screenshots(
     }
 
     let screenshots_dir = source_screenshots_dir(state, source).await?;
-    let mut entries = match io::read_dir(&screenshots_dir).await {
+    tokio::task::spawn_blocking(move || scan_screenshots_dir(&screenshots_dir))
+        .await?
+}
+
+fn scan_screenshots_dir(
+    screenshots_dir: &Path,
+) -> crate::Result<Vec<ScannedScreenshot>> {
+    let entries = match std::fs::read_dir(screenshots_dir) {
         Ok(entries) => entries,
         Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
             return Ok(Vec::new());
@@ -118,14 +128,11 @@ pub(super) async fn scan_source_screenshots(
     };
     let mut screenshots = Vec::new();
 
-    while let Some(entry) = entries
-        .next_entry()
-        .await
-        .map_err(|error| IOError::with_path(error, &screenshots_dir))?
-    {
+    for entry in entries {
+        let entry = entry
+            .map_err(|error| IOError::with_path(error, screenshots_dir))?;
         let file_type = entry
             .file_type()
-            .await
             .map_err(|error| IOError::with_path(error, entry.path()))?;
         if !file_type.is_file() {
             continue;
@@ -142,7 +149,6 @@ pub(super) async fn scan_source_screenshots(
         };
         let metadata = entry
             .metadata()
-            .await
             .map_err(|error| IOError::with_path(error, &path))?;
         let created_at = metadata
             .created()
@@ -179,53 +185,34 @@ pub(super) async fn reconcile_source_screenshots(
 ) -> crate::Result<Vec<InstanceScreenshot>> {
     let existing =
         screenshot_rows::list_screenshots(&source.id, &state.pool).await?;
-    let existing_by_name = existing
-        .iter()
-        .map(|row| (row.file_name.as_str(), row))
+    let mut unmatched_by_name = existing
+        .into_iter()
+        .map(|row| (row.file_name.clone(), row))
         .collect::<HashMap<_, _>>();
-    let metadata_matches = existing.len() == scanned.len()
-        && scanned.iter().all(|scanned| {
-            existing_by_name
-                .get(scanned.file_name.as_str())
-                .is_some_and(|row| {
-                    row.file_size == scanned.file_size
-                        && row.modified_at == scanned.modified_at
-                        && row.created_at
-                            == scanned.created_at.timestamp_millis()
-                })
-        });
-    let mut unmatched_by_hash =
-        HashMap::<(String, i64), Vec<ScreenshotRow>>::new();
     let mut resolved = Vec::with_capacity(scanned.len());
+    let mut needs_hash = Vec::new();
 
-    if metadata_matches {
-        let mut existing_by_name = existing
-            .into_iter()
-            .map(|row| (row.file_name.clone(), row))
-            .collect::<HashMap<_, _>>();
-        for scanned in scanned {
-            let row = existing_by_name.remove(&scanned.file_name).ok_or_else(
-                || {
-                    crate::ErrorKind::InputError(
-                        "Screenshot index changed during reconciliation"
-                            .to_string(),
-                    )
-                },
-            )?;
-            resolved.push(ResolvedScreenshot {
-                scanned,
-                row,
-                is_new: false,
-                changed: false,
-            });
+    for scanned in scanned {
+        match unmatched_by_name.remove(&scanned.file_name) {
+            Some(row)
+                if row.file_size == scanned.file_size
+                    && row.modified_at == scanned.modified_at
+                    && row.created_at
+                        == scanned.created_at.timestamp_millis() =>
+            {
+                resolved.push(ResolvedScreenshot {
+                    scanned,
+                    row,
+                    is_new: false,
+                    changed: false,
+                });
+            }
+            matched => needs_hash.push((scanned, matched)),
         }
-    } else {
-        let mut unmatched_by_name = existing
-            .into_iter()
-            .map(|row| (row.file_name.clone(), row))
-            .collect::<HashMap<_, _>>();
-        let mut hashed = Vec::with_capacity(scanned.len());
-        for mut scanned in scanned {
+    }
+
+    let hashed = stream::iter(needs_hash.into_iter().map(
+        |(mut scanned, matched)| async move {
             let (file_size, content_hash) =
                 sha1_file_async(&scanned.path).await?;
             scanned.file_size = i64::try_from(file_size).map_err(|_| {
@@ -233,54 +220,54 @@ pub(super) async fn reconcile_source_screenshots(
                     "Screenshot is too large to index".to_string(),
                 )
             })?;
-            hashed.push((scanned, content_hash));
-        }
+            Ok::<_, crate::Error>((scanned, content_hash, matched))
+        },
+    ))
+    .buffer_unordered(SCREENSHOT_HASH_CONCURRENCY)
+    .try_collect::<Vec<_>>()
+    .await?;
 
-        let mut renamed_or_new = Vec::new();
-        for (scanned, content_hash) in hashed {
-            if let Some(row) = unmatched_by_name.remove(&scanned.file_name) {
-                resolved.push(resolve_scanned_screenshot(
-                    source,
-                    scanned,
-                    content_hash,
-                    Some(row),
-                ));
-            } else {
-                renamed_or_new.push((scanned, content_hash));
-            }
-        }
+    let mut unmatched_by_hash =
+        HashMap::<(String, i64), Vec<ScreenshotRow>>::new();
+    for row in unmatched_by_name.into_values() {
+        unmatched_by_hash
+            .entry((row.content_hash.clone(), row.file_size))
+            .or_default()
+            .push(row);
+    }
 
-        for row in unmatched_by_name.into_values() {
-            unmatched_by_hash
-                .entry((row.content_hash.clone(), row.file_size))
-                .or_default()
-                .push(row);
-        }
-
-        for (scanned, content_hash) in renamed_or_new {
-            let hash_key = (content_hash.clone(), scanned.file_size);
-            let matched =
-                unmatched_by_hash.get_mut(&hash_key).and_then(|rows| {
-                    if rows.is_empty() {
-                        return None;
-                    }
-                    let created_at = scanned.created_at.timestamp_millis();
-                    let index = rows
-                        .iter()
-                        .position(|row| {
-                            row.modified_at == scanned.modified_at
-                                && row.created_at == created_at
-                        })
-                        .unwrap_or(rows.len() - 1);
-                    Some(rows.swap_remove(index))
-                });
+    for (scanned, content_hash, matched_by_name) in hashed {
+        if let Some(row) = matched_by_name {
             resolved.push(resolve_scanned_screenshot(
                 source,
                 scanned,
                 content_hash,
-                matched,
+                Some(row),
             ));
+            continue;
         }
+
+        let hash_key = (content_hash.clone(), scanned.file_size);
+        let matched = unmatched_by_hash.get_mut(&hash_key).and_then(|rows| {
+            if rows.is_empty() {
+                return None;
+            }
+            let created_at = scanned.created_at.timestamp_millis();
+            let index = rows
+                .iter()
+                .position(|row| {
+                    row.modified_at == scanned.modified_at
+                        && row.created_at == created_at
+                })
+                .unwrap_or(rows.len() - 1);
+            Some(rows.swap_remove(index))
+        });
+        resolved.push(resolve_scanned_screenshot(
+            source,
+            scanned,
+            content_hash,
+            matched,
+        ));
     }
 
     let mut tx = state.pool.begin().await?;