ovlach_pdf/src/chromium/rocket.rs

use std::{sync::{Arc}, time::Duration};
use async_trait::async_trait;
use async_mutex::Mutex;
use headless_chrome::{Browser, LaunchOptions};
use log::{warn, error};
use rocket::{fairing::{Fairing, self}, Rocket, Build, catcher::Result, Request, Response, request::{FromRequest, Outcome}};

#[derive(Default)]
pub struct Chromium {
}

impl Chromium {
    pub fn ignite() -> Self {
        Self::default()
    }

    pub fn default() -> Self {
        Self {
        }
    }
}

#[async_trait]
impl Fairing for Chromium {
    fn info(&self) -> fairing::Info {
        fairing::Info {
            name: "Chromium",
            kind: fairing::Kind::Ignite,
        }
    }

    async fn on_ignite(&self, rocket: Rocket<Build>) -> fairing::Result {
        let new_rocket = rocket.manage(ChromiumCoordinator::new().await);
        let coordinator = new_rocket.state::<ChromiumCoordinator>().unwrap();
        error!("{:p}", &coordinator);
        Ok(new_rocket)
    }
}
pub struct ChromiumCoordinator {
    instances: Arc<Mutex<Vec<BrowserHolder>>>,
}

impl ChromiumCoordinator {
   const NUMBER_OF_INSTANCES: usize = 1;
   
   pub async fn new() -> Self {
        error!("ChromiumCoordinator::new()");
        let instances: Arc<Mutex<Vec<BrowserHolder>>> = Arc::new(Mutex::new(Vec::with_capacity(Self::NUMBER_OF_INSTANCES)));
        while instances.lock().await.len() < Self::NUMBER_OF_INSTANCES {
            instances.lock().await.push(BrowserHolder { browser: Browser::new(LaunchOptions {
                idle_browser_timeout: Duration::from_secs(1), // Wait inifinity for commands
                ..LaunchOptions::default()
            }).unwrap() });
        }

        Self { instances }
   }

   fn spawn_browser(&self) {
    let instances = self.instances.clone(); 
    tokio::spawn(async move {
        error!("Spawn new instance of browser");
        // Create new instance
        let browser = Browser::new(LaunchOptions::default()).unwrap();
        instances.lock().await.push( BrowserHolder {browser });
    });
   }

   fn try_get_instance(&self) -> Option<BrowserHolder> {
        let mut instances = self.instances.lock().await;
        if instances.len() == 0 {
            return None;
        }
        Some(instances.remove(0))
   }

   pub async fn try_get_browser(&self) -> std::result::Result<BrowserHolder, ()> {
        let instances = self.instances.clone(); 
        
        let mut loop_count = 0;
        
        while instances.lock().await.len() == 0 {
            warn!("Waiting for Chromium instances to start...");
            std::thread::sleep(std::time::Duration::from_millis(100));
            if loop_count > 100 * 10 * 60 {
                panic!("Can't start Chromium instances");
            }
            loop_count += 1;
        }

        let browser =  self.instances.lock().await.remove(0);
        self.spawn_browser();
        // test connection state
        match browser.browser.get_version() {
            Ok(_) => Ok(browser),
            Err(_) => {
                Err(())
            }
        }
    }

    pub async fn get_browser(&self) -> std::result::Result<BrowserHolder, ()> {
        loop {
            match self.try_get_browser().await {
                Ok(browser) => return Ok(browser),
                Err(_) => {} // all instances may be dead ... we must wait for new instance
            }
        }
    }
}

pub struct BrowserHolder {
    pub browser: Browser,
}

#[async_trait]
impl<'r> FromRequest<'r> for BrowserHolder {
    type Error = ();

    async fn from_request(request: &'r Request<'_>) -> Outcome<Self, ()> {
        let coordinator = request.rocket().state::<ChromiumCoordinator>().unwrap();
        error!("{:p}", &coordinator);
        let browser = coordinator.get_browser().await.unwrap();
        Outcome::Success(browser)
    }
}