Saving

src/stores/message_store.rs

@@ -1,15 +1,17 @@
 use crate::message::Message;
 use crate::stores::{IMessageSearcher, IMessageStorage, IMessageStore, MessageStoreError};
+use async_trait::async_trait;
 use chrono::{DateTime, Utc};
+use futures::future::join_all;
 use log::error;
 use maildir::{MailEntry, Maildir};
 use pbr::{MultiBar, ProgressBar};
-use rayon::prelude::*;
 use std::collections::HashSet;
 use std::path::PathBuf;
-use std::sync::mpsc;
 use std::thread;
 use std::time::Duration;
+use tokio::sync::mpsc;
+use tokio::task;
 
 pub struct MessageStore {
     pub searcher: Box<dyn IMessageSearcher>,
@@ -17,6 +19,7 @@ pub struct MessageStore {
     progress: Option<ProgressBar<pbr::Pipe>>,
     display_progress: bool,
 }
+#[async_trait]
 impl IMessageStore for MessageStore {
     fn get_message(&self, id: String) -> Result<Option<Message>, MessageStoreError> {
         self.searcher.get_message(id)
@@ -32,8 +35,8 @@ impl IMessageStore for MessageStore {
         Ok(id)
     }
 
-    fn add_maildir(&mut self, path: PathBuf, all: bool) -> Result<usize, MessageStoreError> {
-        self.index_mails(path, all)
+    async fn add_maildir(&mut self, path: PathBuf, all: bool) -> Result<usize, MessageStoreError> {
+        self.index_mails(path, all).await
     }
     fn tag_message_id(
         &mut self,
@@ -135,7 +138,28 @@ impl MessageStore {
         }
     }
 
-    fn do_index_mails(&mut self, maildir: Maildir, full: bool) -> Result<usize, MessageStoreError> {
+    async fn parse_and_add_message(
+        mail: MailEntry,
+    ) -> Result<(Message, String), MessageStoreError> {
+        let message = Message::from_mailentry(mail);
+        match message {
+            Ok(msg) => {
+                let parsed_body = msg.get_body().as_text();
+                Ok((msg, parsed_body))
+            }
+            Err(err) => {
+                error!("A message could not be parsed: {}", err.message);
+                Err(MessageStoreError::CouldNotAddMessage(
+                    "Failed to parse email".to_string(),
+                ))
+            }
+        }
+    }
+    async fn do_index_mails(
+        &mut self,
+        maildir: Maildir,
+        full: bool,
+    ) -> Result<usize, MessageStoreError> {
         let mails: Vec<Result<MailEntry, _>> = Self::mail_iterator(&maildir, full).collect();
         let count = mails.len();
         self.start_indexing_process(count)?;
@@ -143,37 +167,26 @@ impl MessageStore {
         if self.display_progress {
             progress_handle = Some(self.init_progress(count));
         }
-        let (tx, rx) = mpsc::channel();
-
-        let t = thread::spawn(move || {
-            mails
-                .into_par_iter()
-                .for_each_with(tx, |tx, msg| match msg {
-                    Ok(unparsed_msg) => {
-                        let message = Message::from_mailentry(unparsed_msg);
-                        match message {
-                            Ok(msg) => {
-                                let parsed_body = msg.get_body().as_text();
-                                tx.send((msg, parsed_body))
-                                    .expect("Could not send to channel?")
-                            }
-                            Err(err) => {
-                                error!("A message could not be parsed: {}", err.message);
-                            }
-                        }
+        let (tx, mut rx) = mpsc::channel(100);
+        let handles = mails
+            .into_iter()
+            .map(|m| {
+                let mut tx = tx.clone();
+                task::spawn(async move {
+                    if let Ok(msg) = m {
+                        if let Ok((msg, body)) = MessageStore::parse_and_add_message(msg).await {
+                            tx.send((msg, body)).await;
+                        };
                     }
-                    Err(e) => {
-                        error!("Failed to get message {}", e);
-                    }
-                });
-        });
-
-        while let Ok((msg, parsed_body)) = rx.recv() {
+                })
+            })
+            .collect::<Vec<task::JoinHandle<_>>>();
+        while let Some((msg, parsed_body)) = rx.recv().await {
             self.add_message(msg, parsed_body)?;
-            self.inc_progress();
         }
+        join_all(handles).await;
         self.finish_indexing_process()?;
-        t.join().expect("Unable to join threads for some reason");
+        self.inc_progress();
         self.finish_progress();
         if let Some(handle) = progress_handle {
             handle
@@ -183,11 +196,15 @@ impl MessageStore {
         Ok(count)
     }
 
-    pub fn index_mails(&mut self, path: PathBuf, full: bool) -> Result<usize, MessageStoreError> {
+    pub async fn index_mails(
+        &mut self,
+        path: PathBuf,
+        full: bool,
+    ) -> Result<usize, MessageStoreError> {
         let maildir = self.maildir(path);
         match maildir {
             Ok(maildir) => {
-                self.do_index_mails(maildir, full)?;
+                self.do_index_mails(maildir, full).await?;
                 Ok(1)
             }
             Err(_) => Err(MessageStoreError::CouldNotOpenMaildir(

src/stores/mod.rs

@@ -1,4 +1,5 @@
 use crate::message::Message;
+use async_trait::async_trait;
 use chrono::{DateTime, Utc};
 use std::collections::HashSet;
 use std::path::PathBuf;
@@ -192,6 +193,7 @@ pub trait IMessageStorage {
     ) -> Result<Vec<Message>, MessageStoreError>;
 }
 
+#[async_trait]
 pub trait IMessageStore {
     fn get_message(&self, id: String) -> Result<Option<Message>, MessageStoreError>;
     fn add_message(
@@ -199,7 +201,7 @@ pub trait IMessageStore {
         msg: Message,
         parsed_body: String,
     ) -> Result<String, MessageStoreError>;
-    fn add_maildir(&mut self, path: PathBuf, all: bool) -> Result<usize, MessageStoreError>;
+    async fn add_maildir(&mut self, path: PathBuf, all: bool) -> Result<usize, MessageStoreError>;
     fn tag_message_id(
         &mut self,
         id: String,