use std::{sync::Arc, time::Duration}; use async_trait::async_trait; use async_mutex::Mutex; use headless_chrome::{Browser, LaunchOptions}; use tracing::{error, info, debug, trace, warn, trace_span, info_span}; use rocket::{fairing::{Fairing, self}, Rocket, Build, Request, request::{FromRequest, Outcome}, futures::pin_mut, http::Status}; use tokio::time::sleep; #[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) -> fairing::Result { let new_rocket = rocket.manage(ChromiumCoordinator::new().await); Ok(new_rocket) } } pub struct ChromiumCoordinator { instances: Arc>>, } impl ChromiumCoordinator { const NUMBER_OF_INSTANCES: usize = 1; // TODO: make this configurable pub async fn new() -> Self { let instances: Arc>> = Arc::new(Mutex::new(Vec::with_capacity(Self::NUMBER_OF_INSTANCES))); trace!("creating {} Chromium instances", Self::NUMBER_OF_INSTANCES); while instances.lock().await.len() < Self::NUMBER_OF_INSTANCES { instances.lock().await.push(Self::create_browser_instance()); } trace!("done creating {} Chromium instances", Self::NUMBER_OF_INSTANCES); Self { instances } } fn create_browser_instance() -> BrowserHolder { trace_span!("Creating Chromium instance").in_scope(|| { BrowserHolder { browser: Browser::new(LaunchOptions { idle_browser_timeout: Duration::MAX, // Wait inifinity for commands sandbox: false, ..LaunchOptions::default() }).unwrap() } }) } fn spawn_browser(&self) { let instances = self.instances.clone(); tokio::spawn(async move { debug!("spawn new instance of browser"); // Create new instance instances.lock().await.push(Self::create_browser_instance()); debug!("spawned new instance of browser"); }); } async fn try_get_instance(&self) -> Option { let mut instances = self.instances.lock().await; if instances.len() == 0 { return None; } self.spawn_browser(); Some(instances.remove(0)) } async fn loop_get_instance(&self) -> BrowserHolder { loop { match self.try_get_instance().await { Some(instance) => { return instance } None => { sleep(std::time::Duration::from_millis(100)).await; } } } } pub async fn try_get_browser(&self) -> std::result::Result { let browser = self.loop_get_instance().await; match browser.browser.get_version() { Ok(_) => Ok(browser), Err(_) => { warn!("found dead browser instance"); Err(()) }, } } pub async fn get_browser(&self) -> std::result::Result { let span = info_span!("get_browser"); let _enter = span.enter(); loop { info!("trying to get a browser instance"); match self.try_get_browser().await { Ok(browser) => return Ok(browser), Err(_) => { warn!("all instances of browser are dead... waiting for new instance"); // 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 { let coordinator = request.rocket().state::().unwrap(); let get_instance = coordinator.get_browser(); pin_mut!(get_instance); // TODO: wth? match tokio::time::timeout(std::time::Duration::from_secs(10), &mut get_instance).await { Ok(maybebrowser) => match maybebrowser { Ok(browser) => Outcome::::Success(browser), Err(_) => { error!("Can't create new instance of browser"); Outcome::::Error((Status::InternalServerError, ())) } }, Err(_) => { error!("Can't create new instance of browser (timeout)"); Outcome::::Error((Status::InternalServerError, ())) } } } }