Extract internal cache listener to function

src/cache/disk.rs

@@ -125,6 +125,8 @@ async fn db_listener(
     db_pool: SqlitePool,
     max_on_disk_size: u64,
 ) {
+    // This is in a receiver stream to process up to 128 simultaneous db updates
+    // in one transaction
     let mut recv_stream = ReceiverStream::new(db_rx).ready_chunks(128);
     while let Some(messages) = recv_stream.next().await {
         let mut transaction = match db_pool.begin().await {

src/cache/mem.rs

@@ -7,7 +7,7 @@ use bytes::Bytes;
 use futures::FutureExt;
 use lfu_cache::LfuCache;
 use lru::LruCache;
-use tokio::sync::mpsc::{channel, Sender};
+use tokio::sync::mpsc::{channel, Receiver, Sender};
 use tokio::sync::Mutex;
 
 type CacheValue = (Bytes, ImageMetadata, u64);
@@ -84,7 +84,7 @@ where
     ColdCache: 'static + Cache,
 {
     pub async fn new(inner: ColdCache, max_mem_size: crate::units::Bytes) -> Arc<Self> {
-        let (tx, mut rx) = channel(100);
+        let (tx, rx) = channel(100);
         let new_self = Arc::new(Self {
             inner,
             cur_mem_size: AtomicU64::new(0),
@@ -92,43 +92,50 @@ where
             master_sender: tx,
         });
 
-        let new_self_0 = Arc::clone(&new_self);
+        tokio::spawn(internal_cache_listener(
-        tokio::spawn(async move {
+            Arc::clone(&new_self),
-            let new_self = new_self_0;
+            max_mem_size,
+            rx,
+        ));
+
+        new_self
+    }
+}
+
+async fn internal_cache_listener<MemoryCacheImpl, ColdCache>(
+    cache: Arc<MemoryCache<MemoryCacheImpl, ColdCache>>,
+    max_mem_size: crate::units::Bytes,
+    mut rx: Receiver<(CacheKey, Bytes, ImageMetadata, u64)>,
+) where
+    MemoryCacheImpl: InternalMemoryCache,
+    ColdCache: Cache,
+{
     let max_mem_size = max_mem_size.get() / 20 * 19;
     while let Some((key, bytes, metadata, size)) = rx.recv().await {
         // Add to memory cache
         // We can add first because we constrain our memory usage to 95%
-                new_self
+        cache.cur_mem_size.fetch_add(size as u64, Ordering::Release);
-                    .cur_mem_size
+        cache
-                    .fetch_add(size as u64, Ordering::Release);
-                new_self
             .mem_cache
             .lock()
             .await
             .push(key, (bytes, metadata, size));
 
         // Pop if too large
-                while new_self.cur_mem_size.load(Ordering::Acquire) >= max_mem_size as u64 {
+        while cache.cur_mem_size.load(Ordering::Acquire) >= max_mem_size as u64 {
-                    let popped = new_self
+            let popped = cache
                 .mem_cache
                 .lock()
                 .await
                 .pop()
                 .map(|(key, (bytes, metadata, size))| (key, bytes, metadata, size));
             if let Some((_, _, _, size)) = popped {
-                        new_self
+                cache.cur_mem_size.fetch_sub(size as u64, Ordering::Release);
-                            .cur_mem_size
-                            .fetch_sub(size as u64, Ordering::Release);
             } else {
                 break;
             }
         }
     }
-        });
-
-        new_self
-    }
 }
 
 #[async_trait]