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 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<Build>) -> fairing::Result {
        let new_rocket = rocket.manage(ChromiumCoordinator::new().await);
        Ok(new_rocket)
    }
}
pub struct ChromiumCoordinator {
    instances: Arc<Mutex<Vec<BrowserHolder>>>,
}

impl ChromiumCoordinator {
   const NUMBER_OF_INSTANCES: usize = 3; // TODO: make this configurable

   pub async fn new() -> Self {
        let instances: Arc<Mutex<Vec<BrowserHolder>>> = 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<BrowserHolder> {
        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<BrowserHolder, ()> {
        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<BrowserHolder, ()> {
        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<Self, ()> {
        let coordinator = request.rocket().state::<ChromiumCoordinator>().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::<BrowserHolder, ()>::Success(browser),
                Err(_) => {
                    error!("Can't create new instance of browser");
                    Outcome::<BrowserHolder, ()>::Error((Status::InternalServerError, ()))
                }
            },
            Err(_) => {
                error!("Can't create new instance of browser (timeout)");
                Outcome::<BrowserHolder, ()>::Error((Status::InternalServerError, ()))
            }
        }
    }
}