embedded.rs

use crate::schedulers::{Scheduler, SchedulerError, StartPipelineReq};
use arroyo_rpc::grpc::rpc::{HeartbeatNodeReq, RegisterNodeReq, WorkerFinishedReq};
use arroyo_server_common::shutdown::{Shutdown, SignalBehavior};
use arroyo_types::{JobId, MachineId, WorkerId};
use arroyo_worker::WorkerServer;
use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use tokio::sync::Mutex;
use tokio::task::JoinHandle;
use tonic::Status;
use tracing::{error, info};

pub struct EmbeddedWorker {
    job_id: JobId,
    generation: u64,
    shutdown: Shutdown,
    handle: JoinHandle<()>,
}

/// The EmbedddedScheduler runs workers within the controller process
pub struct EmbeddedScheduler {
    tasks: Arc<Mutex<HashMap<WorkerId, EmbeddedWorker>>>,
    worker_counter: AtomicU64,
}

impl EmbeddedScheduler {
    pub fn new() -> Self {
        Self {
            tasks: Arc::new(Default::default()),
            worker_counter: AtomicU64::new(100),
        }
    }

    async fn clear_finished(&self) {
        self.tasks
            .lock()
            .await
            .retain(|_, w| !w.handle.is_finished());
    }
}

#[async_trait]
impl Scheduler for EmbeddedScheduler {
    async fn start_workers(&self, req: StartPipelineReq) -> Result<(), SchedulerError> {
        self.clear_finished().await;

        let shutdown = Shutdown::new("embedded-worker", SignalBehavior::None);
        let guard = shutdown.guard("embedded-worker");

        let job_id = req.job_id.clone();
        let pipeline_id = req.pipeline_id.clone();
        let log_job_id = job_id.clone();
        let generation = req.generation;
        let worker_id = WorkerId(self.worker_counter.fetch_add(1, Ordering::SeqCst));
        let handle = tokio::task::spawn(async move {
            let server = WorkerServer::new(
                MachineId(Arc::new("embedded".to_string())),
                worker_id,
                req.pipeline_id.clone(),
                req.job_id.clone(),
                req.generation,
                guard,
            );
            let worker_job_id = log_job_id.clone();
            let worker_pipeline_id = pipeline_id.clone();

            match tokio::task::spawn(async move {
                if let Err(e) = server.start_async().await {
                    error!(
                        job_id = %worker_job_id,
                        pipeline_id = *worker_pipeline_id,
                        "Failed to start worker {:?}: {:?}",
                        worker_id,
                        e
                    );
                }
            })
            .await
            {
                Ok(_) => {
                    info!(
                        job_id = %log_job_id,
                        pipeline_id = *pipeline_id,
                        "Worker {:?} finished",
                        worker_id
                    );
                }
                Err(err) => {
                    error!(
                        job_id = %log_job_id,
                        pipeline_id = *pipeline_id,
                        "Worker {:?} panicked: {:?}",
                        worker_id,
                        err
                    );
                }
            }
        });

        self.tasks.lock().await.insert(
            worker_id,
            EmbeddedWorker {
                job_id,
                generation,
                shutdown,
                handle,
            },
        );

        Ok(())
    }

    async fn register_node(&self, _: RegisterNodeReq) {}

    async fn heartbeat_node(&self, _: HeartbeatNodeReq) -> Result<(), Status> {
        Ok(())
    }

    async fn worker_finished(&self, _: WorkerFinishedReq) {}

    async fn stop_workers(
        &self,
        job_id: &str,
        generation: Option<u64>,
        _: bool,
    ) -> anyhow::Result<()> {
        for w in self.workers_for_job(job_id, generation).await? {
            let state = self.tasks.lock().await;
            if let Some(worker) = state.get(&w) {
                worker.shutdown.token().cancel();
            }
        }

        Ok(())
    }

    async fn workers_for_job(
        &self,
        job_id: &str,
        generation: Option<u64>,
    ) -> anyhow::Result<Vec<WorkerId>> {
        let state = self.tasks.lock().await;
        Ok(state
            .iter()
            .filter(|(_, t)| {
                *t.job_id == job_id && (generation.is_none() || generation.unwrap() == t.generation)
            })
            .map(|(k, _)| *k)
            .collect())
    }
}