Merge branch '5-migrate-from-tokio-loop-to-futures-stream' into 'master'

refactor: use futures instead of tokio for media upload

Closes #5

See merge request veretcle/scootaloo!28

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "scootaloo"
-version = "0.8.0"
+version = "0.9.0"
 authors = ["VC <veretcle+framagit@mateu.be>"]
 edition = "2021"
 
@@ -10,11 +10,11 @@ edition = "2021"
 chrono = "^0.4"
 serde = { version = "1.0", features = ["derive"] }
 toml = "^0.5"
-clap = "^2.34"
-futures = "^0.3"
+clap = "^4"
 egg-mode = "^0.16"
 rusqlite = "^0.27"
-tokio = { version = "1", features = ["full"]}
+tokio = { version = "^1", features = ["full"]}
+futures = "^0.3"
 elefren = "^0.22"
 html-escape = "^0.2"
 reqwest = "^0.11"

README.md

@@ -18,6 +18,7 @@ First up, create a configuration file (default path is `/usr/local/etc/scootaloo
 [scootaloo]
 db_path = "/var/lib/scootaloo/scootaloo.sqlite" ## file containing the SQLite Tweet corresponding Toot DB, must be writeable
 cache_path = "/tmp/scootaloo" ## a dir where the temporary files will be download, must be writeable
+rate_limiting = 4 ## optional, default 4, number of accounts handled simultaneously
 
 [twitter]
 ## Consumer/Access key for Twitter (can be generated at https://developer.twitter.com/en/apps)

src/config.rs

@@ -32,6 +32,7 @@ pub struct MastodonConfig {
 pub struct ScootalooConfig {
     pub db_path: String,
     pub cache_path: String,
+    pub rate_limit: Option<usize>,
 }
 
 /// Parses the TOML file into a Config Struct

src/error.rs

@@ -1,8 +1,12 @@
 use std::{
+    boxed::Box,
+    convert::From,
     error::Error,
     fmt::{Display, Formatter, Result},
 };
 
+use elefren::Error as elefrenError;
+
 #[derive(Debug)]
 pub struct ScootalooError {
     details: String,
@@ -23,3 +27,15 @@ impl Display for ScootalooError {
         write!(f, "{}", self.details)
     }
 }
+
+impl From<Box<dyn Error>> for ScootalooError {
+    fn from(error: Box<dyn Error>) -> Self {
+        ScootalooError::new(&format!("Error in a subset crate: {}", error))
+    }
+}
+
+impl From<elefrenError> for ScootalooError {
+    fn from(error: elefrenError) -> Self {
+        ScootalooError::new(&format!("Error in elefren crate: {}", error))
+    }
+}

src/lib.rs

@@ -25,12 +25,13 @@ use rusqlite::Connection;
 use std::sync::Arc;
 use tokio::{spawn, sync::Mutex};
 
+use futures::StreamExt;
+
+const DEFAULT_RATE_LIMIT: usize = 4;
+
 /// This is where the magic happens
 #[tokio::main]
 pub async fn run(config: Config) {
-    // create the task vector for handling multiple accounts
-    let mut mtask = vec![];
-
     // open the SQLite connection
     let conn = Arc::new(Mutex::new(
         Connection::open(&config.scootaloo.db_path).unwrap_or_else(|e| {
@@ -41,121 +42,113 @@ pub async fn run(config: Config) {
         }),
     ));
 
-    for mastodon_config in config.mastodon.into_values() {
-        // create temporary value for each task
-        let scootaloo_cache_path = config.scootaloo.cache_path.clone();
-        let token = get_oauth2_token(&config.twitter);
-        let task_conn = conn.clone();
+    let mut stream = futures::stream::iter(config.mastodon.into_values())
+        .map(|mastodon_config| {
+            // create temporary value for each task
+            let scootaloo_cache_path = config.scootaloo.cache_path.clone();
+            let token = get_oauth2_token(&config.twitter);
+            let task_conn = conn.clone();
 
-        let task = spawn(async move {
-            info!("Starting treating {}", &mastodon_config.twitter_screen_name);
+            spawn(async move {
+                info!("Starting treating {}", &mastodon_config.twitter_screen_name);
 
-            // retrieve the last tweet ID for the username
-            let lconn = task_conn.lock().await;
-            let last_tweet_id = read_state(&lconn, &mastodon_config.twitter_screen_name, None)
-                .unwrap_or_else(|e| panic!("Cannot retrieve last_tweet_id: {}", e))
-                .map(|s| s.tweet_id);
-            drop(lconn);
+                // retrieve the last tweet ID for the username
+                let lconn = task_conn.lock().await;
+                let last_tweet_id = read_state(&lconn, &mastodon_config.twitter_screen_name, None)?
+                    .map(|r| r.tweet_id);
+                drop(lconn);
 
-            // get Mastodon instance
-            let mastodon = get_mastodon_token(&mastodon_config);
+                // get user timeline feed (Vec<tweet>)
+                let mut feed =
+                    get_user_timeline(&mastodon_config.twitter_screen_name, &token, last_tweet_id)
+                        .await?;
 
-            // get user timeline feed (Vec<tweet>)
-            let mut feed =
-                get_user_timeline(&mastodon_config.twitter_screen_name, &token, last_tweet_id)
-                    .await
-                    .unwrap_or_else(|e| {
-                        panic!(
-                            "Something went wrong when trying to retrieve {}’s timeline: {}",
-                            &mastodon_config.twitter_screen_name, e
-                        )
-                    });
-
-            // empty feed -> exiting
-            if feed.is_empty() {
-                info!("Nothing to retrieve since last time, exiting…");
-                return;
-            }
-
-            // order needs to be chronological
-            feed.reverse();
-
-            for tweet in &feed {
-                info!("Treating Tweet {} inside feed", tweet.id);
-                // initiate the toot_reply_id var
-                let mut toot_reply_id: Option<String> = None;
-                // determine if the tweet is part of a thread (response to self) or a standard response
-                if let Some(r) = &tweet.in_reply_to_screen_name {
-                    if r.to_lowercase() != mastodon_config.twitter_screen_name.to_lowercase() {
-                        // we are responding not threading
-                        info!("Tweet is a direct response, skipping");
-                        continue;
-                    }
-                    info!("Tweet is a thread");
-                    // get the corresponding toot id
-                    let lconn = task_conn.lock().await;
-                    toot_reply_id = read_state(
-                        &lconn,
-                        &mastodon_config.twitter_screen_name,
-                        tweet.in_reply_to_status_id,
-                    )
-                    .unwrap_or(None)
-                    .map(|s| s.toot_id);
-                    drop(lconn);
-                };
-
-                // build basic status by just yielding text and dereferencing contained urls
-                let mut status_text = build_basic_status(tweet);
-
-                // building associative media list
-                let (media_url, status_medias) =
-                    generate_media_ids(tweet, &scootaloo_cache_path, &mastodon).await;
-
-                status_text = status_text.replace(&media_url, "");
-
-                info!("Building corresponding Mastodon status");
-
-                let mut status_builder = StatusBuilder::new();
-
-                status_builder.status(&status_text).media_ids(status_medias);
-
-                if let Some(i) = toot_reply_id {
-                    status_builder.in_reply_to(&i);
+                // empty feed -> exiting
+                if feed.is_empty() {
+                    info!("Nothing to retrieve since last time, exiting…");
+                    return Ok(());
                 }
 
-                // can be activated for test purposes
-                // status_builder.visibility(elefren::status_builder::Visibility::Private);
+                // get Mastodon instance
+                let mastodon = get_mastodon_token(&mastodon_config);
 
-                let status = status_builder
-                    .build()
-                    .unwrap_or_else(|_| panic!("Cannot build status with text {}", &status_text));
+                // order needs to be chronological
+                feed.reverse();
 
-                // publish status
-                // again unwrap is safe here as we are in the main thread
-                let published_status = mastodon.new_status(status).unwrap();
-                // this will panic if it cannot publish the status, which is a good thing, it allows the
-                // last_tweet gathered not to be written
+                for tweet in &feed {
+                    info!("Treating Tweet {} inside feed", tweet.id);
+                    // initiate the toot_reply_id var
+                    let mut toot_reply_id: Option<String> = None;
+                    // determine if the tweet is part of a thread (response to self) or a standard response
+                    if let Some(r) = &tweet.in_reply_to_screen_name {
+                        if r.to_lowercase() != mastodon_config.twitter_screen_name.to_lowercase() {
+                            // we are responding not threading
+                            info!("Tweet is a direct response, skipping");
+                            continue;
+                        }
+                        info!("Tweet is a thread");
+                        // get the corresponding toot id
+                        let lconn = task_conn.lock().await;
+                        toot_reply_id = read_state(
+                            &lconn,
+                            &mastodon_config.twitter_screen_name,
+                            tweet.in_reply_to_status_id,
+                        )
+                        .unwrap_or(None)
+                        .map(|s| s.toot_id);
+                        drop(lconn);
+                    };
 
-                let ttt_towrite = TweetToToot {
-                    twitter_screen_name: mastodon_config.twitter_screen_name.clone(),
-                    tweet_id: tweet.id,
-                    toot_id: published_status.id,
-                };
+                    // build basic status by just yielding text and dereferencing contained urls
+                    let mut status_text = build_basic_status(tweet);
 
-                // write the current state (tweet ID and toot ID) to avoid copying it another time
-                let lconn = task_conn.lock().await;
-                write_state(&lconn, ttt_towrite)
-                    .unwrap_or_else(|e| panic!("Can’t write the last tweet retrieved: {}", e));
-                drop(lconn);
-            }
-        });
+                    // building associative media list
+                    let (media_url, status_medias) =
+                        generate_media_ids(tweet, &scootaloo_cache_path, &mastodon).await;
 
-        // push each task into the vec task
-        mtask.push(task);
-    }
+                    status_text = status_text.replace(&media_url, "");
+
+                    info!("Building corresponding Mastodon status");
+
+                    let mut status_builder = StatusBuilder::new();
+
+                    status_builder.status(&status_text).media_ids(status_medias);
+
+                    if let Some(i) = toot_reply_id {
+                        status_builder.in_reply_to(&i);
+                    }
+
+                    // can be activated for test purposes
+                    // status_builder.visibility(elefren::status_builder::Visibility::Private);
+
+                    let status = status_builder.build()?;
+
+                    let published_status = mastodon.new_status(status)?;
+                    // this will return if it cannot publish the status preventing the last_tweet from
+                    // being written into db
+
+                    let ttt_towrite = TweetToToot {
+                        twitter_screen_name: mastodon_config.twitter_screen_name.clone(),
+                        tweet_id: tweet.id,
+                        toot_id: published_status.id,
+                    };
+
+                    // write the current state (tweet ID and toot ID) to avoid copying it another time
+                    let lconn = task_conn.lock().await;
+                    write_state(&lconn, ttt_towrite)?;
+                    drop(lconn);
+                }
+                Ok::<(), ScootalooError>(())
+            })
+        })
+        .buffer_unordered(config.scootaloo.rate_limit.unwrap_or(DEFAULT_RATE_LIMIT));
 
     // launch and wait for every handle
-    for handle in mtask {
-        handle.await.unwrap();
+    while let Some(result) = stream.next().await {
+        match result {
+            Ok(Err(e)) => eprintln!("Error within thread: {}", e),
+            Err(e) => eprintln!("Error with thread: {}", e),
+            _ => (),
+        }
     }
 }

src/main.rs

@@ -1,5 +1,5 @@
-use clap::{App, Arg, SubCommand};
-use log::{error, LevelFilter};
+use clap::{Arg, Command};
+use log::LevelFilter;
 use scootaloo::*;
 use simple_logger::SimpleLogger;
 use std::str::FromStr;
@@ -7,115 +7,121 @@ use std::str::FromStr;
 const DEFAULT_CONFIG_PATH: &str = "/usr/local/etc/scootaloo.toml";
 
 fn main() {
-    let matches = App::new(env!("CARGO_PKG_NAME"))
+    let matches = Command::new(env!("CARGO_PKG_NAME"))
         .version(env!("CARGO_PKG_VERSION"))
         .about("A Twitter to Mastodon bot")
         .arg(
-            Arg::with_name("config")
-                .short("c")
+            Arg::new("config")
+                .short('c')
                 .long("config")
                 .value_name("CONFIG_FILE")
                 .help(&format!(
                     "TOML config file for scootaloo (default {})",
                     DEFAULT_CONFIG_PATH
                 ))
-                .takes_value(true)
+                .num_args(1)
+                .default_value(DEFAULT_CONFIG_PATH)
                 .display_order(1),
         )
         .arg(
-            Arg::with_name("log_level")
-                .short("l")
+            Arg::new("log_level")
+                .short('l')
                 .long("loglevel")
                 .value_name("LOGLEVEL")
                 .help("Log level. Valid values are: Off, Warn, Error, Info, Debug")
-                .takes_value(true)
+                .num_args(1)
+                .value_parser(["Off", "Warn", "Error", "Info", "Debug"])
                 .display_order(2),
         )
         .subcommand(
-            SubCommand::with_name("register")
+            Command::new("register")
                 .version(env!("CARGO_PKG_VERSION"))
                 .about("Command to register to a Mastodon Instance")
                 .arg(
-                    Arg::with_name("host")
-                        .short("H")
+                    Arg::new("host")
+                        .short('H')
                         .long("host")
                         .value_name("HOST")
                         .help("Base URL of the Mastodon instance to register to (no default)")
-                        .takes_value(true)
+                        .num_args(1)
                         .required(true)
                         .display_order(1)
                     )
                 .arg(
-                    Arg::with_name("name")
-                        .short("n")
+                    Arg::new("name")
+                        .short('n')
                         .long("name")
                         .help("Twitter Screen Name (like https://twitter.com/screen_name, no default)")
-                        .takes_value(true)
+                        .num_args(1)
                         .required(true)
                         .display_order(2)
                 ),
         )
         .subcommand(
-            SubCommand::with_name("init")
+            Command::new("init")
                 .version(env!("CARGO_PKG_VERSION"))
                 .about("Command to init Scootaloo DB")
                 .arg(
-                    Arg::with_name("config")
-                        .short("c")
+                    Arg::new("config")
+                        .short('c')
                         .long("config")
                         .value_name("CONFIG_FILE")
                         .help(&format!(
                             "TOML config file for scootaloo (default {})",
                             DEFAULT_CONFIG_PATH
                         ))
-                        .takes_value(true)
+                        .default_value(DEFAULT_CONFIG_PATH)
+                        .num_args(1)
                         .display_order(1),
                 ),
         )
         .subcommand(
-            SubCommand::with_name("migrate")
+            Command::new("migrate")
             .version(env!("CARGO_PKG_VERSION"))
             .about("Command to migrate Scootaloo DB")
             .arg(
-                Arg::with_name("config")
-                    .short("c")
+                Arg::new("config")
+                    .short('c')
                     .long("config")
                     .value_name("CONFIG_FILE")
                     .help(&format!("TOML config file for scootaloo (default {})", DEFAULT_CONFIG_PATH))
-                    .takes_value(true)
+                    .default_value(DEFAULT_CONFIG_PATH)
+                    .num_args(1)
                     .display_order(1),
             )
             .arg(
-                    Arg::with_name("name")
-                        .short("n")
+                    Arg::new("name")
+                        .short('n')
                         .long("name")
                         .help("Twitter Screen Name (like https://twitter.com/screen_name, no default)")
-                        .takes_value(true)
+                        .num_args(1)
                         .display_order(2)
             )
         )
         .get_matches();
 
     match matches.subcommand() {
-        ("register", Some(sub_m)) => {
+        Some(("register", sub_m)) => {
             register(
-                sub_m.value_of("host").unwrap(),
-                sub_m.value_of("name").unwrap(),
+                sub_m.get_one::<String>("host").unwrap(),
+                sub_m.get_one::<String>("name").unwrap(),
             );
             return;
         }
-        ("init", Some(sub_m)) => {
-            let config = parse_toml(sub_m.value_of("config").unwrap_or(DEFAULT_CONFIG_PATH));
+        Some(("init", sub_m)) => {
+            let config = parse_toml(sub_m.get_one::<String>("config").unwrap());
             init_db(&config.scootaloo.db_path).unwrap();
             return;
         }
-        ("migrate", Some(sub_m)) => {
-            let config = parse_toml(sub_m.value_of("config").unwrap_or(DEFAULT_CONFIG_PATH));
+        Some(("migrate", sub_m)) => {
+            let config = parse_toml(sub_m.get_one::<String>("config").unwrap());
             let config_twitter_screen_name =
                 &config.mastodon.values().next().unwrap().twitter_screen_name;
             migrate_db(
                 &config.scootaloo.db_path,
-                sub_m.value_of("name").unwrap_or(config_twitter_screen_name),
+                sub_m
+                    .get_one::<String>("name")
+                    .unwrap_or(config_twitter_screen_name),
             )
             .unwrap();
             return;
@@ -123,20 +129,14 @@ fn main() {
         _ => (),
     }
 
-    if matches.is_present("log_level") {
-        match LevelFilter::from_str(matches.value_of("log_level").unwrap()) {
-            Ok(level) => SimpleLogger::new().with_level(level).init().unwrap(),
-            Err(e) => {
-                SimpleLogger::new()
-                    .with_level(LevelFilter::Error)
-                    .init()
-                    .unwrap();
-                error!("Unknown log level filter: {}", e);
-            }
-        };
+    if let Some(level) = matches.get_one::<String>("log_level") {
+        SimpleLogger::new()
+            .with_level(LevelFilter::from_str(level).unwrap())
+            .init()
+            .unwrap();
     }
 
-    let config = parse_toml(matches.value_of("config").unwrap_or(DEFAULT_CONFIG_PATH));
+    let config = parse_toml(matches.get_one::<String>("config").unwrap());
 
     run(config);
 }

src/state.rs

@@ -87,7 +87,7 @@ pub fn migrate_db(d: &str, s: &str) -> Result<(), Box<dyn Error>> {
     match res {
         Err(e) => match e.to_string().as_str() {
             "duplicate column name: twitter_screen_name" => Ok(()),
-            _ => Err(Box::new(e)),
+            _ => Err(e.into()),
         },
         _ => Ok(()),
     }

src/util.rs

@@ -15,6 +15,8 @@ use tokio::{
     io::copy,
 };
 
+use futures::{stream, stream::StreamExt};
+
 /// Generate associative table between media ids and tweet extended entities
 pub async fn generate_media_ids(
     tweet: &Tweet,
@@ -25,63 +27,41 @@ pub async fn generate_media_ids(
     let mut media_ids: Vec<String> = vec![];
 
     if let Some(m) = &tweet.extended_entities {
-        // create tasks list
-        let mut tasks = vec![];
-
-        // size of media_ids vector, should be equal to the media vector
-        media_ids.resize(m.media.len(), String::new());
-
         info!("{} medias in tweet", m.media.len());
 
-        for (i, media) in m.media.iter().enumerate() {
-            // attribute media url
-            media_url = media.url.clone();
+        let medias = m.media.clone();
 
-            // clone everything we need
-            let cache_path = String::from(cache_path);
-            let media = media.clone();
-            let mastodon = mastodon.clone();
+        let mut stream = stream::iter(medias)
+            .map(|media| {
+                // attribute media url
+                media_url = media.url.clone();
 
-            let task = tokio::task::spawn(async move {
-                info!("Start treating {}", media.media_url_https);
-                // get the tweet embedded media
-                let local_tweet_media_path = match get_tweet_media(&media, &cache_path).await {
-                    Ok(l) => l,
-                    Err(e) => {
-                        return Err(ScootalooError::new(&format!(
-                            "Cannot get tweet media for {}: {}",
-                            &media.url, e
-                        )))
-                    }
-                };
+                // clone everything we need
+                let cache_path = String::from(cache_path);
+                let mastodon = mastodon.clone();
 
-                // upload media to Mastodon
-                let mastodon_media = mastodon.media(Cow::from(local_tweet_media_path.to_owned()));
-                // at this point, we can safely erase the original file
-                // it doesn’t matter if we can’t remove, cache_media fn is idempotent
-                remove_file(&local_tweet_media_path).await.ok();
+                tokio::task::spawn(async move {
+                    info!("Start treating {}", media.media_url_https);
+                    // get the tweet embedded media
+                    let local_tweet_media_path = get_tweet_media(&media, &cache_path).await?;
 
-                let mastodon_media = match mastodon_media {
-                    Ok(m) => m,
-                    Err(e) => {
-                        return Err(ScootalooError::new(&format!(
-                            "Attachment {} cannot be uploaded to Mastodon Instance: {}",
-                            &local_tweet_media_path, e
-                        )))
-                    }
-                };
+                    // upload media to Mastodon
+                    let mastodon_media =
+                        mastodon.media(Cow::from(local_tweet_media_path.to_owned()))?;
+                    // at this point, we can safely erase the original file
+                    // it doesn’t matter if we can’t remove, cache_media fn is idempotent
+                    remove_file(&local_tweet_media_path).await.ok();
 
-                Ok((i, mastodon_media.id))
-            });
+                    Ok::<String, ScootalooError>(mastodon_media.id)
+                })
+            })
+            .buffered(4); // there are max four medias per tweet and they need to be treated in
+                          // order
 
-            tasks.push(task);
-        }
-
-        for task in tasks {
-            match task.await {
-                // insert the media at the right place
-                Ok(Ok((i, v))) => media_ids[i] = v,
-                Ok(Err(e)) => warn!("{}", e),
+        while let Some(result) = stream.next().await {
+            match result {
+                Ok(Ok(v)) => media_ids.push(v),
+                Ok(Err(e)) => warn!("Cannot treat media: {}", e),
                 Err(e) => error!("Something went wrong when joining the main thread: {}", e),
             }
         }