lib.rs

#![allow(clippy::new_without_default)]
// TODO: factor out complex types
#![allow(clippy::type_complexity)]
// https://github.com/rust-lang/rust-clippy/issues/12908
#![allow(clippy::needless_lifetimes)]

use anyhow::Result;
use arroyo_rpc::config::config;
use arroyo_rpc::grpc::rpc;
use arroyo_rpc::grpc::rpc::controller_grpc_server::{ControllerGrpc, ControllerGrpcServer};
use arroyo_rpc::grpc::rpc::job_controller_grpc_server::{
    JobControllerGrpc, JobControllerGrpcServer,
};
use arroyo_rpc::grpc::rpc::{
    GrpcOutputSubscription, HeartbeatNodeReq, HeartbeatNodeResp, HeartbeatReq, HeartbeatResp,
    JobMetricsReq, JobMetricsResp, NonfatalErrorReq, OutputData, RegisterNodeReq, RegisterNodeResp,
    RegisterWorkerReq, RegisterWorkerResp, SinkDataReq, SinkDataResp, TaskCheckpointCompletedReq,
    TaskCheckpointCompletedResp, TaskCheckpointEventReq, TaskCheckpointEventResp, TaskFailedReq,
    TaskFailedResp, TaskFinishedReq, TaskFinishedResp, TaskStartedReq, TaskStartedResp,
    WorkerErrorRes, WorkerFinishedReq, WorkerFinishedResp, WorkerInitializationCompleteReq,
    WorkerInitializationCompleteResp,
};
use arroyo_rpc::public_ids::{IdTypes, generate_id};
use arroyo_rpc::worker_types::{RunningMessage, TaskFailedEvent};
use arroyo_rpc::{StateContext, config, errors};
use arroyo_server_common::shutdown::ShutdownGuard;
use arroyo_server_common::wrap_start;
use arroyo_types::{MachineId, PipelineId, WorkerId, from_micros};
use arroyo_worker::job_controller::job_metrics::JobMetrics;
use cornucopia_async::DatabaseSource;
use lazy_static::lazy_static;
use prometheus::{IntGaugeVec, register_int_gauge_vec};
use states::{Created, State, StateMachine};
use std::collections::{HashMap, HashSet};
use std::env;
use std::net::SocketAddr;
use std::str::FromStr;
use std::sync::Arc;
use std::time::{Duration, Instant};
use time::OffsetDateTime;
use tokio::net::TcpListener;
use tokio::sync::RwLock;
use tokio::sync::mpsc::Sender;
use tokio::sync::mpsc::error::TrySendError;
use tokio_stream::wrappers::{ReceiverStream, TcpListenerStream};
use tonic::codec::CompressionEncoding;
use tonic::{Request, Response, Status};
use tracing::{debug, info, warn};

//pub mod compiler;
pub mod job_controller;
pub mod schedulers;
mod states;

const TTL_PIPELINE_CLEANUP_TIME: Duration = Duration::from_secs(60 * 60);

lazy_static! {
    static ref JOBS_BY_STATE: IntGaugeVec = register_int_gauge_vec!(
        "arroyo_controller_jobs",
        "Current number of jobs by controller state",
        &["state"]
    )
    .unwrap();
}

fn metric_job_state<'a>(state: Option<&'a str>, failure_domain: Option<&str>) -> &'a str {
    let state = state.unwrap_or("Created");
    if state == "Failed" && failure_domain == Some("user") {
        "UserFailed"
    } else {
        state
    }
}

fn job_state_counts<'a>(
    jobs: impl Iterator<Item = (Option<&'a str>, Option<&'a str>)>,
) -> HashMap<&'a str, i64> {
    let mut counts = HashMap::new();
    for (state, failure_domain) in jobs {
        *counts
            .entry(metric_job_state(state, failure_domain))
            .or_default() += 1;
    }
    counts
}

fn update_job_state_metrics(counts: &HashMap<&str, i64>) {
    JOBS_BY_STATE.reset();
    for (state, count) in counts {
        JOBS_BY_STATE.with_label_values(&[state]).set(*count);
    }
}

include!(concat!(env!("OUT_DIR"), "/controller-sql.rs"));

use crate::schedulers::{ManualScheduler, NodeScheduler, ProcessScheduler, Scheduler};
use types::public::LogLevel;
use types::public::{RestartMode, StopMode};

pub const CHECKPOINTS_TO_KEEP: u32 = 5;

#[derive(PartialEq, Clone, Debug)]
pub struct JobConfig {
    id: Arc<String>,
    organization_id: String,
    pipeline_name: String,
    pipeline_id: i64,
    stop_mode: StopMode,
    checkpoint_interval: Duration,
    ttl: Option<Duration>,
    parallelism_overrides: HashMap<u32, usize>,
    restart_nonce: i32,
    restart_mode: RestartMode,
    /// Minimum checkpoint epoch in controller mode; generation to start without state in leader
    /// mode.
    ignore_state_before_epoch: Option<i32>,
    /// Per-job environment variables forwarded to workers at scheduling time.
    env_vars: serde_json::Value,
    /// Per-job scheduler configuration overlay as raw JSON (same
    /// shape as the controller-wide scheduler config). The scheduler
    /// interprets this; the controller treats it as opaque. An empty
    /// object is the no-override case.
    scheduler_config: serde_json::Value,
}

/// Per-pipeline data that doesn't change for the lifetime of a job.
#[derive(Clone, Debug)]
pub struct PipelineInfo {
    pub pipeline_id: PipelineId,
    pub state_url: Option<String>,
    pub tags: HashMap<String, String>,
}

#[derive(Clone, Debug)]
pub struct JobStatus {
    id: Arc<String>,
    generation: u64,
    state: String,
    start_time: Option<OffsetDateTime>,
    finish_time: Option<OffsetDateTime>,
    tasks: Option<i32>,
    failure_message: Option<String>,
    failure_domain: Option<String>,
    restarts: i32,
    pipeline_path: Option<String>,
    wasm_path: Option<String>,
    restart_nonce: i32,
    state_context: StateContext,
}

impl JobStatus {
    pub async fn update_db(&self, database: &DatabaseSource) -> Result<(), String> {
        let c = database.client().await.map_err(|e| format!("{e:?}"))?;
        let res = queries::controller_queries::execute_update_job_status(
            &c,
            &self.state,
            &self.start_time,
            &self.finish_time,
            &self.tasks,
            &self.failure_message,
            &self.failure_domain,
            &self.restarts,
            &self.pipeline_path,
            &self.wasm_path,
            &(self.generation as i64),
            &self.restart_nonce,
            &serde_json::to_value(&self.state_context).expect("failed to serialize"),
            &*self.id,
        )
        .await
        .map_err(|e| format!("{e:?}"))?;

        if res == 0 {
            Err("Job status does not exist".to_string())
        } else {
            Ok(())
        }
    }
}

fn job_in_final_state(config: &JobConfig, status: &JobStatus) -> bool {
    match status.state.as_str() {
        "Stopped" | "Finished" => config.stop_mode != StopMode::none,
        "Failed" => config.restart_nonce == status.restart_nonce,
        _ => false,
    }
}

#[derive(Debug)]
pub enum JobMessage {
    ConfigUpdate(JobConfig),
    WorkerConnect {
        worker_id: WorkerId,
        machine_id: MachineId,
        generation: u64,
        rpc_address: String,
        data_address: String,
        slots: usize,
    },
    WorkerInitializationComplete {
        worker_id: WorkerId,
        success: bool,
        error_message: Option<String>,
    },
    TaskStarted {
        worker_id: WorkerId,
        task_id: u32,
        subtask_idx: u32,
    },
    RunningMessage(RunningMessage),
}

#[derive(Clone)]
pub struct ControllerServer {
    job_state: Arc<tokio::sync::Mutex<HashMap<String, StateMachine>>>,
    data_txs: Arc<tokio::sync::Mutex<HashMap<String, Vec<Sender<Result<OutputData, Status>>>>>>,
    scheduler: Arc<dyn Scheduler>,
    metrics: Arc<RwLock<HashMap<Arc<String>, JobMetrics>>>,
    db: DatabaseSource,
}

#[allow(clippy::result_large_err)]
fn job_id_from_context(worker_context: &Option<rpc::WorkerContext>) -> Result<String, Status> {
    Ok(worker_context
        .as_ref()
        .ok_or_else(|| Status::invalid_argument("missing worker_context"))?
        .job_id
        .clone())
}

#[tonic::async_trait]
impl ControllerGrpc for ControllerServer {
    async fn register_worker(
        &self,
        request: Request<RegisterWorkerReq>,
    ) -> Result<Response<RegisterWorkerResp>, Status> {
        let remote_addr = request.remote_addr();
        let req = request.into_inner();
        let worker = req
            .worker_context
            .ok_or_else(|| Status::invalid_argument("missing worker_context"))?;
        info!(
            job_id = worker.job_id,
            pipeline_id = worker.pipeline_id,
            "Worker registered: {:?} -- {:?}",
            worker,
            remote_addr
        );

        self.send_to_job_queue(
            &worker.job_id,
            JobMessage::WorkerConnect {
                worker_id: WorkerId(worker.worker_id),
                machine_id: MachineId(worker.machine_id.into()),
                generation: worker.generation,
                rpc_address: req.rpc_address,
                data_address: req.data_address,
                slots: req.slots as usize,
            },
        )
        .await?;

        Ok(Response::new(RegisterWorkerResp {}))
    }

    async fn task_started(
        &self,
        request: Request<TaskStartedReq>,
    ) -> Result<Response<TaskStartedResp>, Status> {
        let req = request.into_inner();
        let ctx = req
            .worker_context
            .ok_or_else(|| Status::invalid_argument("missing worker_context"))?;
        info!(
            job_id = ctx.job_id,
            pipeline_id = ctx.pipeline_id,
            worker_id = ctx.worker_id,
            task_id = req.task_id,
            subtask_idx = req.subtask_idx,
            "task started"
        );

        self.send_to_job_queue(
            &ctx.job_id,
            JobMessage::TaskStarted {
                worker_id: WorkerId(ctx.worker_id),
                task_id: req.task_id,
                subtask_idx: req.subtask_idx,
            },
        )
        .await?;

        Ok(Response::new(TaskStartedResp {}))
    }

    async fn register_node(
        &self,
        request: Request<RegisterNodeReq>,
    ) -> Result<Response<RegisterNodeResp>, Status> {
        let req = request.into_inner();
        info!(
            "Received node registration from {} at {} with {} slots",
            req.machine_id, req.addr, req.task_slots
        );

        self.scheduler.register_node(req).await;

        Ok(Response::new(RegisterNodeResp {}))
    }

    async fn heartbeat_node(
        &self,
        request: Request<HeartbeatNodeReq>,
    ) -> Result<Response<HeartbeatNodeResp>, Status> {
        self.scheduler.heartbeat_node(request.into_inner()).await?;
        Ok(Response::new(HeartbeatNodeResp {}))
    }

    async fn worker_finished(
        &self,
        request: Request<WorkerFinishedReq>,
    ) -> Result<Response<WorkerFinishedResp>, Status> {
        self.scheduler.worker_finished(request.into_inner()).await;
        Ok(Response::new(WorkerFinishedResp {}))
    }

    async fn send_sink_data(
        &self,
        request: Request<SinkDataReq>,
    ) -> Result<Response<SinkDataResp>, Status> {
        let req = request.into_inner();
        let mut data_txs = self.data_txs.lock().await;
        if let Some(v) = data_txs.get_mut(&req.job_id) {
            let output = OutputData {
                operator_id: req.operator_id,
                subtask_idx: req.subtask_idx,
                timestamps: req.timestamps,
                batch: req.batch,
                start_id: req.start_id,
                done: req.done,
            };

            let mut remove = HashSet::new();
            for (i, tx) in v.iter().enumerate() {
                match tx.try_send(Ok(output.clone())) {
                    Ok(_) => {}
                    Err(TrySendError::Closed(_)) => {
                        remove.insert(i);
                    }
                    Err(TrySendError::Full(_)) => {
                        debug!("queue full");
                    }
                }
            }

            let mut i = 0;
            v.retain(|_tx| {
                i += 1;
                !remove.contains(&(i - 1))
            });
        }
        Ok(Response::new(SinkDataResp::default()))
    }

    type SubscribeToOutputStream = ReceiverStream<Result<OutputData, Status>>;

    async fn subscribe_to_output(
        &self,
        request: Request<GrpcOutputSubscription>,
    ) -> Result<Response<Self::SubscribeToOutputStream>, Status> {
        let job_id = request.into_inner().job_id;
        if self
            .job_state
            .lock()
            .await
            .get(&job_id)
            .ok_or_else(|| Status::not_found(format!("Job {job_id} does not exist")))?
            .state
            .read()
            .unwrap()
            .as_str()
            != "Running"
        {
            return Err(Status::failed_precondition(
                "Job must be running to read output",
            ));
        }

        let (tx, rx) = tokio::sync::mpsc::channel(32);

        let mut data_txs = self.data_txs.lock().await;
        data_txs.entry(job_id).or_default().push(tx);

        Ok(Response::new(ReceiverStream::new(rx)))
    }

    async fn worker_initialization_complete(
        &self,
        request: Request<WorkerInitializationCompleteReq>,
    ) -> Result<Response<WorkerInitializationCompleteResp>, Status> {
        let req = request.into_inner();
        let ctx = req
            .worker_context
            .ok_or_else(|| Status::invalid_argument("missing worker_context"))?;
        info!(
            job_id = ctx.job_id,
            pipeline_id = ctx.pipeline_id,
            "Worker {} initialization completed: success={}, error={:?}",
            ctx.worker_id,
            req.success,
            req.error_message
        );

        self.send_to_job_queue(
            &ctx.job_id,
            JobMessage::WorkerInitializationComplete {
                worker_id: WorkerId(ctx.worker_id),
                success: req.success,
                error_message: req.error_message,
            },
        )
        .await?;

        Ok(Response::new(WorkerInitializationCompleteResp {}))
    }
}

#[tonic::async_trait]
impl JobControllerGrpc for ControllerServer {
    async fn task_checkpoint_event(
        &self,
        request: Request<TaskCheckpointEventReq>,
    ) -> Result<Response<TaskCheckpointEventResp>, Status> {
        let req = request.into_inner();

        let ctx = req
            .worker_context
            .as_ref()
            .ok_or_else(|| Status::invalid_argument("missing worker_context"))?;
        debug!(
            job_id = ctx.job_id,
            pipeline_id = ctx.pipeline_id,
            "received task checkpoint event {:?}",
            req
        );
        let job_id = ctx.job_id.clone();

        self.send_to_job_queue(
            &job_id,
            JobMessage::RunningMessage(RunningMessage::TaskCheckpointEvent(req)),
        )
        .await?;

        Ok(Response::new(TaskCheckpointEventResp {}))
    }

    async fn task_checkpoint_completed(
        &self,
        request: Request<TaskCheckpointCompletedReq>,
    ) -> Result<Response<TaskCheckpointCompletedResp>, Status> {
        let req = request.into_inner();

        let ctx = req
            .worker_context
            .as_ref()
            .ok_or_else(|| Status::invalid_argument("missing worker_context"))?;
        debug!(
            job_id = ctx.job_id,
            pipeline_id = ctx.pipeline_id,
            "received task checkpoint completed {:?}",
            req
        );
        let job_id = ctx.job_id.clone();

        self.send_to_job_queue(
            &job_id,
            JobMessage::RunningMessage(RunningMessage::TaskCheckpointFinished(req)),
        )
        .await?;

        Ok(Response::new(TaskCheckpointCompletedResp {}))
    }

    async fn task_finished(
        &self,
        request: Request<TaskFinishedReq>,
    ) -> Result<Response<TaskFinishedResp>, Status> {
        let req = request.into_inner();

        let ctx = req
            .worker_context
            .ok_or_else(|| Status::invalid_argument("missing worker_context"))?;

        self.send_to_job_queue(
            &ctx.job_id,
            JobMessage::RunningMessage(RunningMessage::TaskFinished {
                worker_id: WorkerId(ctx.worker_id),
                time: from_micros(req.time),
                task_id: req.task_id,
                subtask_idx: req.subtask_idx,
            }),
        )
        .await?;

        Ok(Response::new(TaskFinishedResp {}))
    }

    async fn task_failed(
        &self,
        request: Request<TaskFailedReq>,
    ) -> Result<Response<TaskFailedResp>, Status> {
        let req = request.into_inner();
        let ctx = req
            .worker_context
            .ok_or_else(|| Status::invalid_argument("TaskFailedReq missing worker_context"))?;
        let err = req
            .error
            .ok_or_else(|| Status::invalid_argument("TaskFailedReq missing error"))?;

        self.send_to_job_queue(
            &ctx.job_id,
            JobMessage::RunningMessage(RunningMessage::TaskFailed(TaskFailedEvent {
                worker_id: WorkerId(ctx.worker_id),
                task_id: err.task_id,
                subtask_idx: err.subtask_idx,
                error_domain: err.error_domain().into(),
                retry_hint: err.retry_hint().into(),
                operator_id: err.operator_id,
                reason: err.error,
                details: err.details,
            })),
        )
        .await?;

        Ok(Response::new(TaskFailedResp {}))
    }

    async fn heartbeat(
        &self,
        request: Request<HeartbeatReq>,
    ) -> Result<Response<HeartbeatResp>, Status> {
        let req = request.into_inner();

        let job_id = job_id_from_context(&req.worker_context)?;

        self.send_to_job_queue(
            &job_id,
            JobMessage::RunningMessage(RunningMessage::WorkerHeartbeat {
                worker_id: WorkerId(req.worker_context.as_ref().unwrap().worker_id),
                time: Instant::now(),
            }),
        )
        .await?;

        return Ok(Response::new(HeartbeatResp {}));
    }

    async fn nonfatal_error(
        &self,
        request: Request<NonfatalErrorReq>,
    ) -> Result<Response<WorkerErrorRes>, Status> {
        let req = request.into_inner();
        let ctx = req
            .worker_context
            .ok_or_else(|| Status::invalid_argument("NonfatalErrorReq missing worker_context"))?;
        let err = req
            .error
            .ok_or_else(|| Status::invalid_argument("NonfatalErrorReq missing error"))?;

        info!(
            job_id = ctx.job_id,
            pipeline_id = ctx.pipeline_id,
            operator_id = err.operator_id,
            message = "operator error",
            error_message = err.error,
            error_details = err.details
        );

        let client = self.db.client().await.unwrap();
        match queries::controller_queries::execute_create_job_log_message(
            &client,
            &generate_id(IdTypes::JobLogMessage),
            &ctx.job_id,
            &err.operator_id,
            &(err.subtask_idx as i64),
            &LogLevel::error,
            &err.error,
            &err.details,
            &errors::ErrorDomain::from(err.error_domain()).as_str(),
            &errors::RetryHint::from(err.retry_hint()).as_str(),
        )
        .await
        {
            Ok(_) => Ok(Response::new(WorkerErrorRes {})),
            Err(db_err) => Err(Status::from_error(Box::new(db_err))),
        }
    }

    async fn job_metrics(
        &self,
        request: Request<JobMetricsReq>,
    ) -> Result<Response<JobMetricsResp>, Status> {
        let job_id = request.into_inner().job_id;
        let metrics = self
            .metrics
            .read()
            .await
            .get(&job_id)
            .ok_or_else(|| Status::not_found("No metrics for job"))?
            .clone();

        // TODO: send this over in a more efficient format like protobuf
        Ok(Response::new(JobMetricsResp {
            metrics: serde_json::to_string(&metrics.get_groups()).unwrap(),
        }))
    }
}

impl ControllerServer {
    pub async fn new(database: DatabaseSource) -> Self {
        let scheduler: Arc<dyn Scheduler> = match &config().controller.scheduler {
            config::Scheduler::Node => {
                info!("Using node scheduler");
                Arc::new(NodeScheduler::new())
            }
            config::Scheduler::Kubernetes => {
                info!("Using kubernetes scheduler");
                Arc::new(schedulers::kubernetes::KubernetesScheduler::new().await)
            }
            config::Scheduler::Embedded => {
                info!("Using embedded scheduler");
                Arc::new(schedulers::embedded::EmbeddedScheduler::new())
            }
            config::Scheduler::Process => {
                info!("Using process scheduler");
                Arc::new(ProcessScheduler::new())
            }
            config::Scheduler::Manual => {
                info!("Using manual scheduler");
                Arc::new(ManualScheduler::new())
            }
        };

        Self {
            scheduler,
            data_txs: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
            job_state: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
            db: database,
            metrics: Default::default(),
        }
    }

    async fn send_to_job_queue(&self, job_id: &str, msg: JobMessage) -> Result<(), Status> {
        // Keep per-job backpressure from holding the global job map lock.
        let tx = {
            let jobs = self.job_state.lock().await;
            let Some(sm) = jobs.get(job_id) else {
                warn!(message = "Received message for unknown job id", job_id);
                return Err(Status::failed_precondition(format!(
                    "No job with id {job_id}"
                )));
            };

            sm.sender().ok_or_else(|| {
                Status::failed_precondition(format!(
                    "Cannot handle message for {job_id}: State machine is inactive"
                ))
            })?
        };

        tx.send(msg).await.map_err(|_| {
            Status::failed_precondition(format!(
                "Cannot handle message for {job_id}: State machine is inactive"
            ))
        })
    }

    fn start_updater(&self, guard: ShutdownGuard) {
        let db = self.db.clone();
        let jobs = Arc::clone(&self.job_state);
        let scheduler = Arc::clone(&self.scheduler);
        let metrics = Arc::clone(&self.metrics);

        let token = guard.token();

        let mut cleaned_at = Instant::now();

        let our_guard = guard.child("update-thread");
        our_guard.into_spawn_task(async move {
            while !token.is_cancelled() {
                let client = db.client().await?;
                let res = queries::controller_queries::fetch_all_jobs(&client).await?;
                let state_counts = job_state_counts(
                    res.iter()
                        .map(|p| (p.state.as_deref(), p.failure_domain.as_deref())),
                );
                update_job_state_metrics(&state_counts);

                for p in res {
                    let id = Arc::new(p.id);
                    let config = JobConfig {
                        id: id.clone(),
                        organization_id: p.org_id,
                        pipeline_id: p.pipeline_id,
                        pipeline_name: p.pipeline_name,
                        stop_mode: p.stop,
                        checkpoint_interval: Duration::from_micros(
                            p.checkpoint_interval_micros as u64,
                        ),
                        ttl: p.ttl_micros.map(|t| Duration::from_micros(t as u64)),
                        parallelism_overrides: p
                            .parallelism_overrides
                            .as_object()
                            .unwrap()
                            .into_iter()
                            .filter_map(|(k, v)| {
                                Some((u32::from_str(k).ok()?, v.as_u64()? as usize))
                            })
                            .collect(),
                        restart_nonce: p.config_restart_nonce,
                        restart_mode: p.restart_mode,
                        ignore_state_before_epoch: p.ignore_state_before_epoch,
                        env_vars: p.env_vars,
                        scheduler_config: p.scheduler_config,
                    };

                    let mut jobs = jobs.lock().await;

                    let state_context: StateContext =
                        serde_json::from_value(p.state_context.clone()).unwrap_or_else(|e| {
                            warn!(job_id = *id, original =? p.state_context, error =? e,
                                "failed to deserialize state context");
                            StateContext {
                                version: 1,
                                leader: None,
                            }
                        });

                    let status = JobStatus {
                        id: id.clone(),
                        generation: p.run_id.unwrap_or(0).max(0) as u64,
                        state: p.state.unwrap_or_else(|| Created {}.name().to_string()),
                        start_time: p.start_time,
                        finish_time: p.finish_time,
                        tasks: p.tasks,
                        failure_message: p.failure_message,
                        failure_domain: p.failure_domain,
                        restarts: p.restarts,
                        pipeline_path: p.pipeline_path,
                        wasm_path: p.wasm_path,
                        restart_nonce: p.status_restart_nonce,
                        state_context,
                    };

                    if let Some(sm) = jobs.get_mut(&*id) {
                        sm.update(config, status, &guard).await;
                    } else if !job_in_final_state(&config, &status) {
                        jobs.insert(
                            (*id).clone(),
                            StateMachine::new(
                                config,
                                status,
                                db.clone(),
                                scheduler.clone(),
                                guard.clone_temporary(),
                                metrics.clone(),
                            )
                            .await,
                        );
                    }
                }

                if cleaned_at.elapsed() > Duration::from_secs(5) {
                    let res = queries::controller_queries::execute_clean_preview_pipelines(
                        &client,
                        &(OffsetDateTime::now_utc() - TTL_PIPELINE_CLEANUP_TIME),
                    )
                    .await?;
                    if res > 0 {
                        info!("Cleaned {res} preview pipelines from database");
                    }
                    cleaned_at = Instant::now();
                }

                tokio::time::sleep(Duration::from_millis(500)).await;
            }
            Ok(())
        });
    }

    pub async fn start(self, guard: ShutdownGuard) -> anyhow::Result<u16> {
        // let reflection = tonic_reflection::server::Builder::configure()
        //     .register_encoded_file_descriptor_set(arroyo_rpc::grpc::API_FILE_DESCRIPTOR_SET)
        //     .build_v1()
        //     .unwrap();

        let config = config();
        let addr = SocketAddr::new(config.controller.bind_address, config.controller.rpc_port);

        let listener = TcpListener::bind(addr).await?;
        let local_addr = listener.local_addr()?;

        let controller_service = ControllerGrpcServer::new(self.clone())
            .send_compressed(CompressionEncoding::Zstd)
            .accept_compressed(CompressionEncoding::Zstd);

        let job_controller_service = JobControllerGrpcServer::new(self.clone())
            .send_compressed(CompressionEncoding::Zstd)
            .accept_compressed(CompressionEncoding::Zstd);

        let scheduler_shutdown_guard = guard.child("scheduler-shutdown");
        let scheduler_shutdown_token = scheduler_shutdown_guard.token();
        let scheduler = Arc::clone(&self.scheduler);
        tokio::spawn(async move {
            scheduler_shutdown_token.cancelled().await;
            scheduler.shutdown().await;
            drop(scheduler_shutdown_guard);
        });

        self.start_updater(guard.child("updater"));

        if let Some(tls_config) = config.get_tls_config(&config.controller.tls) {
            info!("Starting arroyo-controller with TLS on {}", local_addr);

            let server = arroyo_server_common::grpc_server_with_tls(tls_config)
                .await?
                .accept_http1(true)
                .add_service(controller_service)
                .add_service(job_controller_service);

            guard.into_spawn_task(wrap_start(
                "controller",
                local_addr,
                server.serve_with_incoming(TcpListenerStream::new(listener)),
            ));
        } else {
            info!("Starting arroyo-controller on {}", local_addr);

            guard.into_spawn_task(wrap_start(
                "controller",
                local_addr,
                arroyo_server_common::grpc_server()
                    .accept_http1(true)
                    .add_service(controller_service)
                    .add_service(job_controller_service)
                    .serve_with_incoming(TcpListenerStream::new(listener)),
            ));
        }

        Ok(local_addr.port())
    }
}

#[cfg(test)]
mod tests {
    use prometheus::core::Collector;

    use super::*;

    #[test]
    fn metric_job_states_preserve_raw_states() {
        for state in ["Created", "Running", "Finished", "Unexpected"] {
            assert_eq!(metric_job_state(Some(state), None), state);
        }
        assert_eq!(metric_job_state(None, None), "Created");

        assert_eq!(metric_job_state(Some("Failed"), Some("user")), "UserFailed");
        assert_eq!(metric_job_state(Some("Failed"), Some("internal")), "Failed");
    }

    #[test]
    fn job_state_metrics_clear_absent_states() {
        let counts = job_state_counts(
            [
                (Some("Running"), None),
                (Some("Running"), None),
                (Some("Failed"), Some("user")),
            ]
            .into_iter(),
        );
        update_job_state_metrics(&counts);
        assert_eq!(JOBS_BY_STATE.with_label_values(&["Running"]).get(), 2);
        assert_eq!(JOBS_BY_STATE.with_label_values(&["UserFailed"]).get(), 1);

        update_job_state_metrics(&HashMap::new());
        assert!(JOBS_BY_STATE.collect()[0].get_metric().is_empty());
    }
}