mod.rs

use anyhow::bail;
use arroyo_datastream::logical::LogicalProgram;
use arroyo_rpc::config::config;
use arroyo_rpc::connect_grpc;
use arroyo_rpc::grpc::rpc::node_grpc_client::NodeGrpcClient;
use arroyo_rpc::grpc::rpc::{
    HeartbeatNodeReq, RegisterNodeReq, StartWorkerReq, StopWorkerReq, StopWorkerStatus,
    WorkerFinishedReq,
};
use arroyo_types::{
    GENERATION_ENV, JOB_ID_ENV, JobId, MachineId, PIPELINE_ID_ENV, PipelineId, WorkerId,
};
use futures::future::join_all;
use lazy_static::lazy_static;
use prometheus::{Gauge, register_gauge};
use std::collections::HashMap;
use std::env::current_exe;
use std::ffi::OsString;
use std::path::PathBuf;
use std::str::FromStr;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};
use tokio::process::Command;
use tokio::sync::{Mutex, oneshot};
use tonic::transport::Channel;
use tonic::{Request, Status};
use tracing::{info, warn};

pub mod embedded;
pub mod kubernetes;

lazy_static! {
    static ref FREE_SLOTS: Gauge =
        register_gauge!("arroyo_controller_free_slots", "number of free task slots").unwrap();
    static ref REGISTERED_SLOTS: Gauge = register_gauge!(
        "arroyo_controller_registered_slots",
        "total number of registered task slots"
    )
    .unwrap();
    static ref REGISTERED_NODES: Gauge = register_gauge!(
        "arroyo_controller_registered_nodes",
        "total number of registered nodes"
    )
    .unwrap();
}

#[async_trait::async_trait]
pub trait Scheduler: Send + Sync {
    async fn start_workers(
        &self,
        start_pipeline_req: StartPipelineReq,
    ) -> Result<(), SchedulerError>;

    async fn register_node(&self, req: RegisterNodeReq);
    async fn heartbeat_node(&self, req: HeartbeatNodeReq) -> Result<(), Status>;
    async fn worker_finished(&self, req: WorkerFinishedReq);
    async fn stop_workers(
        &self,
        job_id: &str,
        generation: Option<u64>,
        force: bool,
    ) -> anyhow::Result<()>;
    async fn workers_for_job(
        &self,
        job_id: &str,
        generation: Option<u64>,
    ) -> anyhow::Result<Vec<WorkerId>>;

    async fn shutdown(&self) {}
}

pub struct ProcessWorker {
    pipeline_id: PipelineId,
    job_id: JobId,
    generation: u64,
    shutdown_tx: oneshot::Sender<()>,
    finished_rx: oneshot::Receiver<()>,
}

/// This Scheduler starts new processes to run the worker nodes
pub struct ProcessScheduler {
    workers: Arc<Mutex<HashMap<WorkerId, ProcessWorker>>>,
    worker_counter: AtomicU64,
}

impl ProcessScheduler {
    pub fn new() -> Self {
        Self {
            workers: Arc::new(Mutex::new(HashMap::new())),
            worker_counter: AtomicU64::new(100),
        }
    }
}

const PROCESS_WORKER_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(10);

pub struct StartPipelineReq {
    pub name: String,
    pub program: LogicalProgram,
    pub wasm_path: String,
    pub pipeline_id: PipelineId,
    pub organization_id: String,
    pub job_id: JobId,
    pub hash: String,
    pub generation: u64,
    pub slots: usize,
    pub env_vars: HashMap<String, String>,
    pub pipeline_tags: HashMap<String, String>,
    /// Per-job scheduler configuration overlay as raw JSON. An empty
    /// object means "use the controller's global scheduler config
    /// unchanged". The scheduler interprets the shape; the controller
    /// treats it as opaque and passes it through verbatim.
    pub scheduler_config: serde_json::Value,
}

#[async_trait::async_trait]
impl Scheduler for ProcessScheduler {
    async fn start_workers(
        &self,
        start_pipeline_req: StartPipelineReq,
    ) -> Result<(), SchedulerError> {
        let workers = (start_pipeline_req.slots as f32
            / config().process_scheduler.slots_per_process as f32)
            .ceil() as usize;

        let mut slots_scheduled = 0;

        let base_path = PathBuf::from_str(&format!(
            "/tmp/arroyo-process/{}",
            start_pipeline_req.job_id
        ))
        .unwrap();

        for _ in 0..workers {
            let path = base_path.clone();

            let slots_here = (start_pipeline_req.slots - slots_scheduled)
                .min(config().process_scheduler.slots_per_process as usize);

            let worker_id = self.worker_counter.fetch_add(1, Ordering::SeqCst);

            let (tx, rx) = oneshot::channel();
            let (finished_tx, finished_rx) = oneshot::channel();

            {
                let mut workers = self.workers.lock().await;
                workers.insert(
                    WorkerId(worker_id),
                    ProcessWorker {
                        pipeline_id: start_pipeline_req.pipeline_id.clone(),
                        job_id: start_pipeline_req.job_id.clone(),
                        generation: start_pipeline_req.generation,
                        shutdown_tx: tx,
                        finished_rx,
                    },
                );
            }

            slots_scheduled += slots_here;
            let pipeline_id = start_pipeline_req.pipeline_id.clone();
            let job_id = start_pipeline_req.job_id.clone();
            let workers = Arc::downgrade(&self.workers);
            let env_map = start_pipeline_req.env_vars.clone();

            tokio::spawn(async move {
                let mut command =
                    Command::new(current_exe().expect("Could not get path of worker binary"));

                for (env, value) in env_map {
                    command.env(env, value);
                }

                let config = config();

                let mut args = vec![];
                if let Some(path) = &config.config_path {
                    args.push(OsString::from_str("-c").unwrap());
                    args.push(path.clone().into_os_string());
                }

                if let Some(path) = &config.config_dir {
                    args.push(OsString::from_str("--config-dir").unwrap());
                    args.push(path.clone().into_os_string());
                }

                args.push("worker".into());

                let mut child = match command
                    .args(args)
                    .env("ARROYO__ADMIN__HTTP_PORT", "0")
                    .env("ARROYO__WORKER__TASK_SLOTS", format!("{slots_here}"))
                    .env("ARROYO__WORKER__ID", format!("{worker_id}")) // start at 100 to make same length
                    .env("ARROYO__CONTROLLER_ENDPOINT", config.controller_endpoint())
                    .env("UNDER_PROCESS_SCHEDULER", "true")
                    .env(PIPELINE_ID_ENV, &*pipeline_id)
                    .env(JOB_ID_ENV, &*job_id)
                    .env(GENERATION_ENV, format!("{}", start_pipeline_req.generation))
                    .kill_on_drop(config.process_scheduler.shutdown_with_controller)
                    .spawn()
                {
                    Ok(child) => child,
                    Err(e) => {
                        warn!(
                            message = "failed to start process scheduler worker",
                            worker_id,
                            job_id = %job_id,
                            pipeline_id = %pipeline_id,
                            error = format!("{:?}", e),
                        );

                        if let Some(workers) = workers.upgrade() {
                            let mut state = workers.lock().await;
                            state.remove(&WorkerId(worker_id));
                        }
                        let _ = finished_tx.send(());
                        return;
                    }
                };

                tokio::select! {
                    status = child.wait() => {
                        info!(
                            job_id = %job_id,
                            pipeline_id = %pipeline_id,
                            "Child ({:?}) exited with status {:?}",
                            path,
                            status
                        );
                    }
                    _ = rx => {
                        if config.process_scheduler.shutdown_with_controller {
                            info!(
                                message = "Killing child",
                                worker_id,
                                job_id = %job_id,
                                pipeline_id = *pipeline_id
                            );
                            if let Err(e) = child.kill().await {
                                warn!(
                                    message = "failed to kill process scheduler worker",
                                    worker_id,
                                    job_id = %job_id,
                                    pipeline_id = %pipeline_id,
                                    error = format!("{:?}", e),
                                );
                            }
                        }
                    }
                }

                if let Some(workers) = workers.upgrade() {
                    let mut state = workers.lock().await;
                    state.remove(&WorkerId(worker_id));
                }
                let _ = finished_tx.send(());
            });
        }

        Ok(())
    }

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

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

    async fn stop_workers(
        &self,
        job_id: &str,
        run_id: Option<u64>,
        _force: bool,
    ) -> anyhow::Result<()> {
        for worker_id in self.workers_for_job(job_id, run_id).await? {
            let worker = {
                let mut state = self.workers.lock().await;
                let Some(worker) = state.remove(&worker_id) else {
                    return Ok(());
                };
                worker
            };

            let _ = worker.shutdown_tx.send(());
        }

        Ok(())
    }

    async fn shutdown(&self) {
        let workers: Vec<_> = self.workers.lock().await.drain().collect();

        if workers.is_empty() || !config().process_scheduler.shutdown_with_controller {
            return;
        }

        let worker_count = workers.len();
        info!(
            message = "shutting down process scheduler workers",
            workers = worker_count,
        );

        let waiters = workers.into_iter().map(|(worker_id, worker)| async move {
            let job_id = worker.job_id.clone();
            let pipeline_id = worker.pipeline_id.clone();
            let _ = worker.shutdown_tx.send(());
            (worker_id, job_id, pipeline_id, worker.finished_rx.await)
        });

        match tokio::time::timeout(PROCESS_WORKER_SHUTDOWN_TIMEOUT, join_all(waiters)).await {
            Ok(results) => {
                for (worker_id, job_id, pipeline_id, result) in results {
                    if result.is_err() {
                        warn!(
                            message = "process scheduler worker exited without completion signal",
                            worker_id = worker_id.0,
                            job_id = %job_id,
                            pipeline_id = *pipeline_id,
                        );
                    }
                }
            }
            Err(_) => {
                warn!(
                    message = "timed out waiting for process scheduler workers to stop",
                    workers = worker_count,
                    timeout_secs = PROCESS_WORKER_SHUTDOWN_TIMEOUT.as_secs(),
                );
            }
        }
    }
}

/// A "manual" scheduler that relies on the user to manage worker processes by executing commands
/// printed out to the terminal. This is mostly useful for testing scheduling behavior.
pub struct ManualScheduler {}

impl ManualScheduler {
    pub fn new() -> Self {
        Self {}
    }
}

#[async_trait::async_trait]
impl Scheduler for ManualScheduler {
    async fn start_workers(
        &self,
        start_pipeline_req: StartPipelineReq,
    ) -> Result<(), SchedulerError> {
        let config = config();
        let slots_per_process = config.manual_scheduler.slots_per_process as usize;

        let workers = (start_pipeline_req.slots as f32 / slots_per_process as f32).ceil() as usize;

        let exe = current_exe().map_err(|e| {
            SchedulerError::Other(format!("Could not get path of worker binary: {e:?}"))
        })?;

        let mut slots_scheduled = 0;

        for _ in 0..workers {
            let slots_here = (start_pipeline_req.slots - slots_scheduled).min(slots_per_process);
            slots_scheduled += slots_here;

            let mut envs: Vec<(String, String)> = vec![
                ("ARROYO__ADMIN__HTTP_PORT".to_string(), "0".to_string()),
                (
                    "ARROYO__WORKER__TASK_SLOTS".to_string(),
                    slots_here.to_string(),
                ),
                (
                    "ARROYO__CONTROLLER_ENDPOINT".to_string(),
                    config.controller_endpoint(),
                ),
                (
                    PIPELINE_ID_ENV.to_string(),
                    start_pipeline_req.pipeline_id.to_string(),
                ),
                (
                    JOB_ID_ENV.to_string(),
                    start_pipeline_req.job_id.to_string(),
                ),
                (
                    GENERATION_ENV.to_string(),
                    start_pipeline_req.generation.to_string(),
                ),
            ];

            for (env, value) in &start_pipeline_req.env_vars {
                envs.push((env.clone(), value.clone()));
            }

            // Build the command-line arguments, mirroring the process scheduler.
            let mut args: Vec<String> = vec![];
            if let Some(path) = &config.config_path {
                args.push("-c".to_string());
                args.push(path.to_string_lossy().to_string());
            }
            if let Some(path) = &config.config_dir {
                args.push("--config-dir".to_string());
                args.push(path.to_string_lossy().to_string());
            }
            args.push("worker".to_string());

            // Assemble the copy-pasteable command line.
            let mut cmdline = String::new();
            for (env, value) in &envs {
                cmdline.push_str(&format!("{}={} ", env, value));
            }
            cmdline.push_str(&exe.to_string_lossy());
            for arg in &args {
                cmdline.push(' ');
                cmdline.push_str(arg);
            }

            println!();
            println!(
                "════════════════════════════════════════════════════════════════════════════════"
            );
            println!(
                "[manual scheduler] Run worker {} of {} (worker id {} slots) for job {} in a \
                 separate terminal:",
                slots_scheduled / slots_per_process.max(1),
                workers,
                slots_here,
                start_pipeline_req.job_id,
            );
            println!();
            println!("{cmdline}");
            println!(
                "════════════════════════════════════════════════════════════════════════════════"
            );
            println!();
        }

        Ok(())
    }

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

    async fn workers_for_job(
        &self,
        _job_id: &str,
        _run_id: Option<u64>,
    ) -> anyhow::Result<Vec<WorkerId>> {
        Ok(vec![])
    }

    async fn stop_workers(
        &self,
        job_id: &str,
        _run_id: Option<u64>,
        _force: bool,
    ) -> anyhow::Result<()> {
        println!("[manual scheduler] Stop workers for job {}", job_id);

        Ok(())
    }
}

#[derive(Debug, Clone)]
struct NodeStatus {
    id: MachineId,
    free_slots: usize,
    scheduled_slots: HashMap<WorkerId, usize>,
    addr: String,
    last_heartbeat: Instant,
}

impl NodeStatus {
    fn new(id: MachineId, slots: usize, addr: String) -> NodeStatus {
        FREE_SLOTS.add(slots as f64);
        REGISTERED_SLOTS.add(slots as f64);

        NodeStatus {
            id,
            free_slots: slots,
            scheduled_slots: HashMap::new(),
            addr,
            last_heartbeat: Instant::now(),
        }
    }

    fn take_slots(&mut self, worker: WorkerId, slots: usize) {
        if let Some(v) = self.free_slots.checked_sub(slots) {
            FREE_SLOTS.sub(slots as f64);
            self.free_slots = v;
            self.scheduled_slots.insert(worker, slots);
        } else {
            panic!(
                "Attempted to schedule more slots than are available on node {} ({} < {})",
                self.addr, self.free_slots, slots
            );
        }
    }

    fn release_slots(&mut self, worker_id: WorkerId, slots: usize) {
        if let Some(freed) = self.scheduled_slots.remove(&worker_id) {
            assert_eq!(
                freed, slots,
                "Controller and node disagree about how many slots are scheduled for worker {worker_id:?} ({freed} != {slots})"
            );

            self.free_slots += slots;

            FREE_SLOTS.add(slots as f64);
        } else {
            warn!(
                "Received release request for unknown worker {:?}",
                worker_id
            );
        }
    }
}

#[derive(Clone)]
struct NodeWorker {
    pipeline_id: PipelineId,
    job_id: JobId,
    node_id: MachineId,
    generation: u64,
    running: bool,
}

#[derive(Default)]
pub struct NodeSchedulerState {
    nodes: HashMap<MachineId, NodeStatus>,
    workers: HashMap<WorkerId, NodeWorker>,
}

impl NodeSchedulerState {
    fn expire_nodes(&mut self, expiration_time: Instant) {
        let expired_nodes: Vec<_> = self
            .nodes
            .iter()
            .filter_map(|(node_id, status)| {
                if status.last_heartbeat >= expiration_time {
                    None
                } else {
                    Some(node_id.clone())
                }
            })
            .collect();
        for node_id in expired_nodes {
            warn!("expiring node {:?} from scheduler state", node_id);
            self.nodes.remove(&node_id);
        }
    }
}

pub struct NodeScheduler {
    state: Arc<Mutex<NodeSchedulerState>>,
}

#[derive(Debug)]
pub enum SchedulerError {
    NotEnoughSlots {
        slots_needed: usize,
    },
    Other(String),
    /// Non-retryable scheduler failure. The state machine treats
    /// this as fatal and transitions the job to `Failed` without retrying.
    Fatal(String),
}

pub fn is_empty_overlay(v: &serde_json::Value) -> bool {
    match v {
        serde_json::Value::Null => true,
        serde_json::Value::Object(m) => m.is_empty(),
        _ => false,
    }
}

impl NodeScheduler {
    pub fn new() -> Self {
        Self {
            state: Arc::new(Mutex::new(NodeSchedulerState::default())),
        }
    }

    async fn client(node: &NodeStatus) -> anyhow::Result<NodeGrpcClient<Channel>> {
        let channel = connect_grpc(
            "controller",
            format!("http://{}", node.addr),
            &config().controller.tls,
            &config().node.tls,
        )
        .await?;
        Ok(NodeGrpcClient::new(channel))
    }

    async fn stop_worker(
        &self,
        job_id: &str,
        worker_id: WorkerId,
        force: bool,
    ) -> anyhow::Result<Option<WorkerId>> {
        let state = self.state.lock().await;

        let Some(worker) = state.workers.get(&worker_id) else {
            // assume it's already finished
            return Ok(Some(worker_id));
        };

        let Some(node) = state.nodes.get(&worker.node_id) else {
            warn!(
                message = "node not found for stop worker",
                node_id = *worker.node_id.0,
                job_id = %worker.job_id,
                pipeline_id = *worker.pipeline_id
            );
            return Ok(Some(worker_id));
        };

        let worker = worker.clone();
        let node = node.clone();
        drop(state);

        info!(
            message = "stopping worker",
            job_id = %worker.job_id,
            pipeline_id = *worker.pipeline_id,
            node_id = *worker.node_id.0,
            node_addr = node.addr,
            worker_id = worker_id.0
        );

        let Ok(mut client) = Self::client(&node).await else {
            warn!(
                job_id = %worker.job_id,
                pipeline_id = *worker.pipeline_id,
                "Failed to connect to worker to stop; this likely means it is dead"
            );
            return Ok(Some(worker_id));
        };

        let Ok(resp) = client
            .stop_worker(Request::new(StopWorkerReq {
                job_id: job_id.to_string(),
                worker_id: worker_id.0,
                force,
            }))
            .await
        else {
            warn!(
                job_id = %worker.job_id,
                pipeline_id = *worker.pipeline_id,
                "Failed to connect to worker to stop; this likely means it is dead"
            );
            return Ok(Some(worker_id));
        };

        match (resp.get_ref().status(), force) {
            (StopWorkerStatus::NotFound, false) => {
                bail!("couldn't find worker, will only continue if force")
            }
            (StopWorkerStatus::StopFailed, _) => bail!("tried to kill and couldn't"),
            _ => Ok(None),
        }
    }
}

#[async_trait::async_trait]
impl Scheduler for NodeScheduler {
    async fn register_node(&self, req: RegisterNodeReq) {
        let mut state = self.state.lock().await;
        if let std::collections::hash_map::Entry::Vacant(e) =
            state.nodes.entry(MachineId(req.machine_id.clone().into()))
        {
            e.insert(NodeStatus::new(
                MachineId(req.machine_id.into()),
                req.task_slots as usize,
                req.addr,
            ));
        }
    }

    async fn heartbeat_node(&self, req: HeartbeatNodeReq) -> Result<(), Status> {
        let mut state = self.state.lock().await;
        if let Some(node) = state
            .nodes
            .get_mut(&MachineId(req.machine_id.clone().into()))
        {
            node.last_heartbeat = Instant::now();
            Ok(())
        } else {
            warn!(
                "Received heartbeat for unregistered node {}, failing request",
                req.machine_id
            );
            Err(Status::not_found(format!(
                "node {} not in scheduler's collection of nodes",
                req.machine_id
            )))
        }
    }

    async fn worker_finished(&self, req: WorkerFinishedReq) {
        let mut state = self.state.lock().await;
        let Some(worker_context) = req.worker_context else {
            warn!("Got worker finished with no worker context");
            return;
        };

        let worker_id = WorkerId(worker_context.worker_id);
        let job_id = worker_context.job_id;
        let pipeline_id = worker_context.pipeline_id;
        let machine_id = MachineId(Arc::new(worker_context.machine_id));

        if let Some(node) = state.nodes.get_mut(&machine_id) {
            node.release_slots(worker_id, req.slots as usize);
        } else {
            warn!(
                %job_id,
                pipeline_id, "Got worker finished message for unknown node {}", machine_id
            );
        }

        if state.workers.remove(&worker_id).is_none() {
            warn!(
                %job_id,
                pipeline_id, "Got worker finished message for unknown worker {}", worker_id.0
            );
        }
    }

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

    #[allow(unreachable_code, unused)]
    async fn start_workers(
        &self,
        start_pipeline_req: StartPipelineReq,
    ) -> Result<(), SchedulerError> {
        // TODO: make this locking more fine-grained
        let mut state = self.state.lock().await;

        state.expire_nodes(Instant::now() - Duration::from_secs(30));

        let free_slots = state.nodes.values().map(|n| n.free_slots).sum::<usize>();
        let slots = start_pipeline_req.slots;
        if slots > free_slots {
            return Err(SchedulerError::NotEnoughSlots {
                slots_needed: slots - free_slots,
            });
        }

        let mut to_schedule = slots;
        let mut slots_assigned = vec![];
        while to_schedule > 0 {
            // find the node with the most free slots and fill it
            let node = {
                if let Some(status) = state
                    .nodes
                    .values()
                    .filter(|n| {
                        n.free_slots > 0 && n.last_heartbeat.elapsed() < Duration::from_secs(30)
                    })
                    .max_by_key(|n| n.free_slots)
                    .cloned()
                {
                    status
                } else {
                    unreachable!();
                }
            };

            let slots_for_this_one = node.free_slots.min(to_schedule);
            info!(
                job_id = %start_pipeline_req.job_id,
                pipeline_id = *start_pipeline_req.pipeline_id,
                "Scheduling {} slots on node {}",
                slots_for_this_one,
                node.addr
            );

            let mut client = Self::client(&node)
                .await
                // TODO: handle this issue more gracefully by moving trying other nodes
                .map_err(|e| {
                    // release back slots already scheduled.
                    slots_assigned
                        .iter()
                        .for_each(|(node_id, worker_id, slots)| {
                            state
                                .nodes
                                .get_mut(node_id)
                                .unwrap()
                                .release_slots(*worker_id, *slots);
                        });
                    SchedulerError::Other(format!(
                        "Failed to connect to node {}: {:?}",
                        node.addr, e
                    ))
                })?;

            let req = StartWorkerReq {
                name: start_pipeline_req.name.clone(),
                pipeline_id: (*start_pipeline_req.pipeline_id).clone(),
                job_id: (*start_pipeline_req.job_id).clone(),
                slots: slots_for_this_one as u64,
                machine_id: node.id.to_string(),
                generation: start_pipeline_req.generation,
                env_vars: start_pipeline_req.env_vars.clone(),
            };

            let res = client
                .start_worker(Request::new(req))
                .await
                .map_err(|e| {
                    // release back slots already scheduled.
                    slots_assigned
                        .iter()
                        .for_each(|(node_id, worker_id, slots)| {
                            state
                                .nodes
                                .get_mut(node_id)
                                .unwrap()
                                .release_slots(*worker_id, *slots);
                        });
                    SchedulerError::Other(format!(
                        "Failed to start worker on node {}: {:?}",
                        node.addr, e
                    ))
                })?
                .into_inner();

            state
                .nodes
                .get_mut(&node.id)
                .unwrap()
                .take_slots(WorkerId(res.worker_id), slots_for_this_one);

            state.workers.insert(
                WorkerId(res.worker_id),
                NodeWorker {
                    pipeline_id: start_pipeline_req.pipeline_id.clone(),
                    job_id: start_pipeline_req.job_id.clone(),
                    generation: start_pipeline_req.generation,
                    node_id: node.id.clone(),
                    running: true,
                },
            );

            slots_assigned.push((node.id, WorkerId(res.worker_id), slots_for_this_one));

            to_schedule -= slots_for_this_one;
        }
        Ok(())
    }

    async fn stop_workers(
        &self,
        job_id: &str,
        run_id: Option<u64>,
        force: bool,
    ) -> anyhow::Result<()> {
        // iterate through all of the workers from workers_for_job and stop them in parallel
        let workers = self.workers_for_job(job_id, run_id).await?;
        let mut futures = vec![];
        for worker_id in workers {
            futures.push(self.stop_worker(job_id, worker_id, force));
        }

        for f in futures {
            match f.await? {
                Some(worker_id) => {
                    let mut state = self.state.lock().await;
                    if let Some(worker) = state.workers.get_mut(&worker_id) {
                        worker.running = false;
                    }
                }
                None => {
                    bail!("Failed to stop worker");
                }
            }
        }

        Ok(())
    }
}