add missing files

.gitignore

@@ -1,6 +1,6 @@
 /target
 .env
-cache
+/cache
 flamegraph*.svg
 perf.data*
 dhat.out.*

src/cache/disk_cache.rs

@@ -0,0 +1,227 @@
+//! Low memory caching stuff
+
+use std::path::PathBuf;
+use std::str::FromStr;
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use bytes::Bytes;
+use futures::StreamExt;
+use log::{warn, LevelFilter};
+use sqlx::sqlite::SqliteConnectOptions;
+use sqlx::{ConnectOptions, SqlitePool};
+use tokio::fs::remove_file;
+use tokio::sync::mpsc::{channel, Receiver, Sender};
+use tokio_stream::wrappers::ReceiverStream;
+
+use super::{BoxedImageStream, Cache, CacheError, CacheKey, CacheStream, ImageMetadata};
+
+pub struct LowMemCache {
+    disk_path: PathBuf,
+    disk_cur_size: AtomicU64,
+    db_update_channel_sender: Sender<DbMessage>,
+}
+
+enum DbMessage {
+    Get(Arc<PathBuf>),
+    Put(Arc<PathBuf>, u32),
+}
+
+impl LowMemCache {
+    /// Constructs a new low memory cache at the provided path and capaci ty.
+    /// This internally spawns a task that will wait for filesystem
+    /// notifications when a file has been written.
+    #[allow(clippy::new_ret_no_self)]
+    pub async fn new(disk_max_size: u64, disk_path: PathBuf) -> Arc<Box<dyn Cache>> {
+        let (db_tx, db_rx) = channel(128);
+        let db_pool = {
+            let db_url = format!("sqlite:{}/metadata.sqlite", disk_path.to_str().unwrap());
+            let mut options = SqliteConnectOptions::from_str(&db_url)
+                .unwrap()
+                .create_if_missing(true);
+            options.log_statements(LevelFilter::Trace);
+            let db = SqlitePool::connect_with(options).await.unwrap();
+
+            // Run db init
+            sqlx::query_file!("./db_queries/init.sql")
+                .execute(&mut db.acquire().await.unwrap())
+                .await
+                .unwrap();
+
+            db
+        };
+
+        let new_self: Arc<Box<dyn Cache>> = Arc::new(Box::new(Self {
+            disk_path,
+            disk_cur_size: AtomicU64::new(0),
+            db_update_channel_sender: db_tx,
+        }));
+
+        tokio::spawn(db_listener(
+            Arc::clone(&new_self),
+            db_rx,
+            db_pool,
+            disk_max_size / 20 * 19,
+        ));
+
+        new_self
+    }
+}
+
+/// Spawn a new task that will listen for updates to the db, pruning if the size
+/// becomes too large.
+async fn db_listener(
+    cache: Arc<Box<dyn Cache>>,
+    db_rx: Receiver<DbMessage>,
+    db_pool: SqlitePool,
+    max_on_disk_size: u64,
+) {
+    let mut recv_stream = ReceiverStream::new(db_rx).ready_chunks(128);
+    while let Some(messages) = recv_stream.next().await {
+        let now = chrono::Utc::now();
+        let mut transaction = db_pool.begin().await.unwrap();
+        for message in messages {
+            match message {
+                DbMessage::Get(entry) => {
+                    let key = entry.as_os_str().to_str();
+                    let query =
+                        sqlx::query!("update Images set accessed = ? where id = ?", now, key)
+                            .execute(&mut transaction)
+                            .await;
+                    if let Err(e) = query {
+                        warn!("Failed to update timestamp in db for {:?}: {}", key, e);
+                    }
+                }
+                DbMessage::Put(entry, size) => {
+                    let key = entry.as_os_str().to_str();
+                    let query = sqlx::query!(
+                        "insert into Images (id, size, accessed) values (?, ?, ?) on conflict do nothing",
+                        key,
+                        size,
+                        now,
+                    )
+                    .execute(&mut transaction)
+                    .await;
+                    if let Err(e) = query {
+                        warn!("Failed to add {:?} to db: {}", key, e);
+                    }
+
+                    cache.increase_usage(size);
+                }
+            }
+        }
+        transaction.commit().await.unwrap();
+
+        if cache.on_disk_size() >= max_on_disk_size {
+            let mut conn = db_pool.acquire().await.unwrap();
+            let items =
+                sqlx::query!("select id, size from Images order by accessed asc limit 1000")
+                    .fetch_all(&mut conn)
+                    .await
+                    .unwrap();
+
+            let mut size_freed = 0;
+            for item in items {
+                // Can't be helped, SQLite doesn't support unsigned integers
+                #[allow(clippy::cast_sign_loss)]
+                {
+                    size_freed += item.size as u64;
+                }
+                tokio::spawn(remove_file(item.id));
+            }
+
+            cache.decrease_usage(size_freed);
+        }
+    }
+}
+
+#[async_trait]
+impl Cache for LowMemCache {
+    async fn get(
+        &self,
+        key: &CacheKey,
+    ) -> Option<Result<(CacheStream, ImageMetadata), CacheError>> {
+        let channel = self.db_update_channel_sender.clone();
+
+        let path = Arc::new(self.disk_path.clone().join(PathBuf::from(key)));
+        let path_0 = Arc::clone(&path);
+
+        tokio::spawn(async move { channel.send(DbMessage::Get(path_0)).await });
+
+        super::fs::read_file(&path)
+            .await
+            .map(|res| res.map_err(Into::into))
+    }
+
+    async fn put(
+        &self,
+        key: CacheKey,
+        image: BoxedImageStream,
+        metadata: ImageMetadata,
+    ) -> Result<CacheStream, CacheError> {
+        let channel = self.db_update_channel_sender.clone();
+
+        let path = Arc::new(self.disk_path.clone().join(PathBuf::from(&key)));
+        let path_0 = Arc::clone(&path);
+
+        let db_callback = |size: u32| async move {
+            let _ = channel.send(DbMessage::Put(path_0, size)).await;
+        };
+
+        super::fs::write_file(&path, key, image, metadata, db_callback, None)
+            .await
+            .map_err(Into::into)
+    }
+
+    #[inline]
+    fn increase_usage(&self, amt: u32) {
+        self.disk_cur_size
+            .fetch_add(u64::from(amt), Ordering::Release);
+    }
+
+    #[inline]
+    fn decrease_usage(&self, amt: u64) {
+        self.disk_cur_size.fetch_sub(amt, Ordering::Release);
+    }
+
+    #[inline]
+    fn on_disk_size(&self) -> u64 {
+        (self.disk_cur_size.load(Ordering::Acquire) + 4095) / 4096 * 4096
+    }
+
+    fn mem_size(&self) -> u64 {
+        0
+    }
+
+    async fn put_with_on_completed_callback(
+        &self,
+        key: CacheKey,
+        image: BoxedImageStream,
+        metadata: ImageMetadata,
+        on_complete: Sender<(CacheKey, Bytes, ImageMetadata, usize)>,
+    ) -> Result<CacheStream, CacheError> {
+        let channel = self.db_update_channel_sender.clone();
+
+        let path = Arc::new(self.disk_path.clone().join(PathBuf::from(&key)));
+        let path_0 = Arc::clone(&path);
+
+        let db_callback = |size: u32| async move {
+            let _ = channel.send(DbMessage::Put(path_0, size)).await;
+        };
+
+        super::fs::write_file(&path, key, image, metadata, db_callback, Some(on_complete))
+            .await
+            .map_err(Into::into)
+    }
+
+    #[inline]
+    async fn put_internal(&self, _: CacheKey, _: Bytes, _: ImageMetadata, _: usize) {
+        unimplemented!()
+    }
+
+    #[inline]
+    async fn pop_memory(&self) -> Option<(CacheKey, Bytes, ImageMetadata, usize)> {
+        None
+    }
+}

src/cache/mem_cache.rs

@@ -0,0 +1,145 @@
+use std::path::PathBuf;
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::Arc;
+
+use crate::cache::LowMemCache;
+
+use super::{BoxedImageStream, Cache, CacheKey, CacheStream, ImageMetadata, MemStream};
+use async_trait::async_trait;
+use bytes::Bytes;
+use futures::FutureExt;
+use lru::LruCache;
+use tokio::sync::mpsc::{channel, Sender};
+use tokio::sync::Mutex;
+
+/// Memory accelerated disk cache. Uses an LRU in memory to speed up reads.
+pub struct MemoryLruCache {
+    inner: Arc<Box<dyn Cache>>,
+    cur_mem_size: AtomicU64,
+    mem_cache: Mutex<LruCache<CacheKey, (Bytes, ImageMetadata, usize)>>,
+    master_sender: Sender<(CacheKey, Bytes, ImageMetadata, usize)>,
+}
+
+impl MemoryLruCache {
+    #[allow(clippy::new_ret_no_self)]
+    pub async fn new(
+        disk_max_size: u64,
+        disk_path: PathBuf,
+        max_mem_size: u64,
+    ) -> Arc<Box<dyn Cache>> {
+        let (tx, mut rx) = channel(100);
+        let new_self = Arc::new(Box::new(Self {
+            inner: LowMemCache::new(disk_max_size, disk_path).await,
+            cur_mem_size: AtomicU64::new(0),
+            mem_cache: Mutex::new(LruCache::unbounded()),
+            master_sender: tx,
+        }) as Box<dyn Cache>);
+
+        let new_self_0 = Arc::clone(&new_self);
+        tokio::spawn(async move {
+            let new_self = new_self_0;
+            let max_mem_size = max_mem_size / 20 * 19;
+            while let Some((key, bytes, metadata, size)) = rx.recv().await {
+                new_self.increase_usage(size as u32);
+                new_self.put_internal(key, bytes, metadata, size).await;
+                while new_self.mem_size() >= max_mem_size {
+                    if let Some((_, _, _, size)) = new_self.pop_memory().await {
+                        new_self.decrease_usage(size as u64);
+                    } else {
+                        break;
+                    }
+                }
+            }
+        });
+
+        new_self
+    }
+}
+
+#[async_trait]
+impl Cache for MemoryLruCache {
+    #[inline]
+    async fn get(
+        &self,
+        key: &CacheKey,
+    ) -> Option<Result<(CacheStream, ImageMetadata), super::CacheError>> {
+        match self.mem_cache.lock().now_or_never() {
+            Some(mut mem_cache) => match mem_cache.get(key).map(|(bytes, metadata, _)| {
+                Ok((CacheStream::Memory(MemStream(bytes.clone())), *metadata))
+            }) {
+                Some(v) => Some(v),
+                None => self.inner.get(key).await,
+            },
+            None => self.inner.get(key).await,
+        }
+    }
+
+    #[inline]
+    async fn put(
+        &self,
+        key: CacheKey,
+        image: BoxedImageStream,
+        metadata: ImageMetadata,
+    ) -> Result<CacheStream, super::CacheError> {
+        self.inner
+            .put_with_on_completed_callback(key, image, metadata, self.master_sender.clone())
+            .await
+    }
+
+    #[inline]
+    fn increase_usage(&self, amt: u32) {
+        self.cur_mem_size
+            .fetch_add(u64::from(amt), Ordering::Release);
+    }
+
+    #[inline]
+    fn decrease_usage(&self, amt: u64) {
+        self.cur_mem_size.fetch_sub(amt, Ordering::Release);
+    }
+
+    #[inline]
+    fn on_disk_size(&self) -> u64 {
+        self.inner.on_disk_size()
+    }
+
+    #[inline]
+    fn mem_size(&self) -> u64 {
+        self.cur_mem_size.load(Ordering::Acquire)
+    }
+
+    #[inline]
+    async fn put_with_on_completed_callback(
+        &self,
+        key: CacheKey,
+        image: BoxedImageStream,
+        metadata: ImageMetadata,
+        on_complete: Sender<(CacheKey, Bytes, ImageMetadata, usize)>,
+    ) -> Result<CacheStream, super::CacheError> {
+        self.inner
+            .put_with_on_completed_callback(key, image, metadata, on_complete)
+            .await
+    }
+
+    #[inline]
+    async fn put_internal(
+        &self,
+        key: CacheKey,
+        image: Bytes,
+        metadata: ImageMetadata,
+        size: usize,
+    ) {
+        self.mem_cache
+            .lock()
+            .await
+            .put(key, (image, metadata, size));
+    }
+
+    #[inline]
+    async fn pop_memory(&self) -> Option<(CacheKey, Bytes, ImageMetadata, usize)> {
+        self.mem_cache
+            .lock()
+            .await
+            .pop_lru()
+            .map(|(key, (bytes, metadata, size))| (key, bytes, metadata, size))
+    }
+}