create BotManager

src/bot_manager.rs

@@ -1,6 +1,7 @@
-use std::{collections::HashMap, sync::RwLock, thread::JoinHandle};
+use std::{collections::HashMap, future::Future, sync::RwLock, thread::JoinHandle, time::Duration};
 
 use lazy_static::lazy_static;
+use log::{error, info};
 use teloxide::{
     dispatching::dialogue::serializer::Json,
     dptree,
@@ -18,7 +19,7 @@ use crate::{
 pub struct BotRunner {
     controller: BotController,
     info: BotInfo,
-    thread: JoinHandle<BotResult<()>>,
+    thread: Option<JoinHandle<BotResult<()>>>,
 }
 
 unsafe impl Sync for BotRunner {}
@@ -33,9 +34,141 @@ lazy_static! {
     static ref BOT_POOL: RwLock<HashMap<String, BotRunner>> = RwLock::new(HashMap::new());
 }
 
-static DEFAULT_SCRIPT: &str =
+pub static DEFAULT_SCRIPT: &str =
     include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/default_script.js"));
 
+pub struct BotManager<BIG, HM, BIS, HI, FBIS, FHI>
+where
+    BIG: FnMut() -> FBIS,
+    FBIS: Future<Output = BIS>,
+    BIS: Iterator<Item = BotInstance>,
+    HM: FnMut(BotInstance) -> FHI,
+    FHI: Future<Output = HI>,
+    HI: Iterator<Item = BotHandler>,
+{
+    bot_pool: HashMap<String, BotRunner>,
+    bi_getter: BIG,
+    h_mapper: HM,
+}
+
+impl<BIG, HM, BIS, HI, FBIS, FHI> BotManager<BIG, HM, BIS, HI, FBIS, FHI>
+where
+    BIG: FnMut() -> FBIS,
+    FBIS: Future<Output = BIS>,
+    BIS: Iterator<Item = BotInstance>,
+    HM: FnMut(BotInstance) -> FHI,
+    FHI: Future<Output = HI>,
+    HI: Iterator<Item = BotHandler>,
+{
+    /// bi_getter - fnmut that returns iterator over BotInstance
+    /// h_map     - fnmut that returns iterator over handlers by BotInstance
+    pub fn with(bi_getter: BIG, h_mapper: HM) -> Self {
+        Self {
+            bot_pool: Default::default(),
+            bi_getter,
+            h_mapper,
+        }
+    }
+
+    pub async fn dispatch(mut self, db: &mut DB) -> ! {
+        loop {
+            'biter: for bi in (self.bi_getter)().await {
+                // removing handler to force restart
+                // TODO: wait till all updates are processed in bot
+                // Temporarly disabling code, because it's free of js runtime
+                // spreads panic
+                if bi.restart_flag {
+                    info!(
+                        "Trying to restart bot `{}`, new script: {}",
+                        bi.name, bi.script
+                    );
+                    let runner = self.bot_pool.remove(&bi.name);
+                };
+                // start, if not started
+                let mut bot_runner = match self.bot_pool.remove(&bi.name) {
+                    Some(br) => br,
+                    None => {
+                        let handlers = (self.h_mapper)(bi.clone()).await;
+                        info!("NEW INSTANCE: Starting new instance! bot name: {}", bi.name);
+                        self.start_bot(bi, db, handlers.collect()).await.unwrap();
+                        continue 'biter;
+                    }
+                };
+
+                // checking if thread is not finished, otherwise clearing handler
+                bot_runner.thread = match bot_runner.thread {
+                    Some(thread) => {
+                        if thread.is_finished() {
+                            let err = thread.join();
+                            error!("Thread bot `{}` finished with error: {:?}", bi.name, err);
+                            None
+                        } else {
+                            Some(thread)
+                        }
+                    }
+                    None => None,
+                };
+
+                // checking if thread is running, otherwise start thread
+                bot_runner.thread = match bot_runner.thread {
+                    Some(thread) => Some(thread),
+                    None => {
+                        let handlers = (self.h_mapper)(bi.clone()).await;
+                        let handler =
+                            script_handler_gen(bot_runner.controller.clone(), handlers.collect())
+                                .await;
+                        Some(
+                            spawn_bot_thread(bot_runner.controller.clone(), db, handler)
+                                .await
+                                .unwrap(),
+                        )
+                    }
+                };
+                self.bot_pool.insert(bi.name.clone(), bot_runner);
+            }
+            tokio::time::sleep(Duration::from_secs(1)).await;
+        }
+    }
+
+    pub async fn start_bot(
+        &mut self,
+        bi: BotInstance,
+        db: &mut DB,
+        plug_handlers: Vec<BotHandler>,
+    ) -> BotResult<BotInfo> {
+        let mut db = db.clone().with_name(bi.name.clone());
+        let controller = BotController::with_db(db.clone(), &bi.token, &bi.script).await?;
+
+        let handler = script_handler_gen(controller.clone(), plug_handlers).await;
+
+        let thread = spawn_bot_thread(controller.clone(), &mut db, handler).await?;
+
+        let info = BotInfo {
+            name: bi.name.clone(),
+        };
+        let runner = BotRunner {
+            controller,
+            info: info.clone(),
+            thread: Some(thread),
+        };
+
+        self.bot_pool.insert(bi.name.clone(), runner);
+
+        Ok(info)
+    }
+}
+
+async fn script_handler_gen(c: BotController, plug_handlers: Vec<BotHandler>) -> BotHandler {
+    let handler = script_handler(c.rc.clone());
+    // each handler will be added to dptree::entry()
+    let handler = plug_handlers
+        .into_iter()
+        // as well as the script handler at the end
+        .chain(std::iter::once(handler))
+        .fold(dptree::entry(), |h, plug| h.branch(plug));
+    handler
+}
+
 pub async fn create_bot(db: &mut DB, token: &str) -> BotResult<BotInstance> {
     let bot = Bot::new(token);
     let name = bot.get_me().await?.username().to_string();
@@ -71,7 +204,7 @@ pub async fn start_bot(
     let runner = BotRunner {
         controller,
         info: info.clone(),
-        thread,
+        thread: Some(thread),
     };
 
     BOT_POOL