AttachmentStore now has its own runloop, can download attachments

kordophoned/src/daemon/attachment_store.rs

@@ -7,7 +7,20 @@ use anyhow::Result;
 use futures_util::StreamExt;
 use kordophone::APIInterface;
 use thiserror::Error;
+
+use kordophone_db::database::Database;
+use kordophone_db::database::DatabaseAccess;
+
+use crate::daemon::events::Event;
+use crate::daemon::events::Reply;
+use crate::daemon::Daemon;
+
+use std::sync::Arc;
+use tokio::sync::Mutex;
+use tokio::sync::mpsc::{Receiver, Sender};
+
 use tokio::pin;
+use tokio::time::Duration;
 
 mod target {
     pub static ATTACHMENTS: &str = "attachments";
@@ -16,8 +29,31 @@ mod target {
 #[derive(Debug, Clone)]
 pub struct Attachment {
     pub guid: String,
-    pub path: PathBuf,
-    pub downloaded: bool,
+    pub base_path: PathBuf,
+}
+
+impl Attachment {
+    pub fn get_path(&self, preview: bool) -> PathBuf {
+        self.base_path.with_extension(if preview { "preview" } else { "full" })
+    }
+
+    pub fn is_downloaded(&self, preview: bool) -> bool {
+        std::fs::exists(&self.get_path(preview))
+            .expect(format!("Wasn't able to check for the existence of an attachment file path at {}", &self.get_path(preview).display()).as_str())
+    }
+}
+
+#[derive(Debug)]
+pub enum AttachmentStoreEvent {
+    // Get the attachment info for a given attachment guid.
+    // Args: attachment guid, reply channel.
+    GetAttachmentInfo(String, Reply<Attachment>),
+
+    // Queue a download for a given attachment guid.
+    // Args: 
+    // - attachment guid
+    // - preview: whether to download the preview (true) or full attachment (false)
+    QueueDownloadAttachment(String, bool),
 }
 
 #[derive(Debug, Error)]
@@ -31,10 +67,15 @@ enum AttachmentStoreError {
 
 pub struct AttachmentStore {
     store_path: PathBuf,
+    database: Arc<Mutex<Database>>,
+    daemon_event_sink: Sender<Event>,
+
+    event_source: Receiver<AttachmentStoreEvent>,
+    event_sink: Option<Sender<AttachmentStoreEvent>>,
 }
 
 impl AttachmentStore {
-    pub fn new(data_dir: &PathBuf) -> AttachmentStore {
+    pub fn new(data_dir: &PathBuf, database: Arc<Mutex<Database>>, daemon_event_sink: Sender<Event>) -> AttachmentStore {
         let store_path = data_dir.join("attachments");
         log::info!(target: target::ATTACHMENTS, "Attachment store path: {}", store_path.display());
 
@@ -42,39 +83,31 @@ impl AttachmentStore {
         std::fs::create_dir_all(&store_path)
             .expect("Wasn't able to create the attachment store path");
 
+        let (event_sink, event_source) = tokio::sync::mpsc::channel(100);
+
         AttachmentStore {
             store_path: store_path,
+            database: database,
+            daemon_event_sink: daemon_event_sink,
+            event_source: event_source,
+            event_sink: Some(event_sink),
         }
     }
 
-    pub fn get_attachment(&self, guid: &String) -> Attachment {
-        let path = self.store_path.join(guid);
-        let path_exists = std::fs::exists(&path).expect(
-            format!(
-                "Wasn't able to check for the existence of an attachment file path at {}",
-                &path.display()
-            )
-            .as_str(),
-        );
+    pub fn get_event_sink(&mut self) -> Sender<AttachmentStoreEvent> {
+        self.event_sink.take().unwrap()
+    }
 
+    fn get_attachment(&self, guid: &String, preview: bool) -> Attachment {
+        let base_path = self.store_path.join(guid);
         Attachment {
             guid: guid.to_owned(),
-            path: path,
-            downloaded: path_exists,
+            base_path: base_path,
         }
     }
-
-    pub async fn download_attachment<C, F, Fut>(
-        &mut self,
-        attachment: &Attachment,
-        mut client_factory: F,
-    ) -> Result<()>
-    where
-        C: APIInterface,
-        F: FnMut() -> Fut,
-        Fut: std::future::Future<Output = Result<C>>,
-    {
-        if attachment.downloaded {
+    
+    async fn download_attachment(&mut self, attachment: &Attachment, preview: bool) -> Result<()> {
+        if attachment.is_downloaded(preview) {
             log::info!(target: target::ATTACHMENTS, "Attachment already downloaded: {}", attachment.guid);
             return Err(AttachmentStoreError::AttachmentAlreadyDownloaded.into());
         }
@@ -82,15 +115,15 @@ impl AttachmentStore {
         log::info!(target: target::ATTACHMENTS, "Starting download for attachment: {}", attachment.guid);
 
         // Create temporary file first, we'll atomically swap later.
-        assert!(!std::fs::exists(&attachment.path).unwrap());
-        let file = std::fs::File::create(&attachment.path)?;
+        assert!(!std::fs::exists(&attachment.get_path(preview)).unwrap());
+        let file = std::fs::File::create(&attachment.get_path(preview))?;
         let mut writer = BufWriter::new(&file);
 
-        log::trace!(target: target::ATTACHMENTS, "Created attachment file at {}", &attachment.path.display());
+        log::trace!(target: target::ATTACHMENTS, "Created attachment file at {}", &attachment.get_path(preview).display());
 
-        let mut client = client_factory().await?;
+        let mut client = Daemon::get_client_impl(&mut self.database).await?;
         let stream = client
-            .fetch_attachment_data(&attachment.guid)
+            .fetch_attachment_data(&attachment.guid, preview)
             .await
             .map_err(|e| AttachmentStoreError::APIClientError(format!("{:?}", e)))?;
 
@@ -106,9 +139,31 @@ impl AttachmentStore {
         Ok(())
     }
     
-    /// Check if an attachment should be downloaded
-    pub fn should_download(&self, attachment_id: &str) -> bool {
-        let attachment = self.get_attachment(&attachment_id.to_string());
-        !attachment.downloaded
+    pub async fn run(&mut self) {
+        loop {
+            tokio::select! {
+                Some(event) = self.event_source.recv() => {
+                    log::debug!(target: target::ATTACHMENTS, "Received attachment store event: {:?}", event);
+
+                    match event {
+                        AttachmentStoreEvent::QueueDownloadAttachment(guid, preview) => {
+                            let attachment = self.get_attachment(&guid, preview);
+                            if !attachment.is_downloaded(preview) {
+                                self.download_attachment(&attachment, preview).await.unwrap_or_else(|e| {
+                                    log::error!(target: target::ATTACHMENTS, "Error downloading attachment: {}", e);
+                                });
+                            } else {
+                                log::info!(target: target::ATTACHMENTS, "Attachment already downloaded: {}", guid);
+                            }
+                        }
+
+                        AttachmentStoreEvent::GetAttachmentInfo(guid, reply) => {
+                            let attachment = self.get_attachment(&guid, false);
+                            reply.send(attachment).unwrap();
+                        }
+                    }
+                }
+            }
+        }
     }
 }