Implementation:ArroyoSystems Arroyo Scheduler Trait

Overview

Scheduler is the core trait defining the worker lifecycle management interface for the Arroyo controller. It abstracts over different deployment backends including process-based, node-based, and Kubernetes scheduling, providing a uniform API for starting, stopping, and monitoring workers.

Description

The module defines three scheduler implementations:

• Scheduler trait: An async trait with methods for starting workers, registering nodes, handling heartbeats, processing worker completion, stopping workers, and querying workers for a job.
• ProcessScheduler: Spawns worker processes on the local machine using tokio::process::Command. Each worker runs as a child process with configured environment variables and task slots. Supports graceful shutdown via oneshot channels.
• NodeScheduler: Manages workers distributed across remote nodes via gRPC. Tracks node status (free slots, scheduled slots, heartbeat time) and worker state. Implements slot allocation using a largest-free-slots-first strategy with node expiration after 30 seconds without heartbeat.

Supporting types include:

• StartPipelineReq: Contains pipeline name, program, WASM path, job ID, hash, run ID, slot count, and environment variables.
• SchedulerError: Error enum with NotEnoughSlots and Other variants.
• NodeStatus: Tracks per-node free slots, scheduled slots, address, and last heartbeat time with Prometheus gauge integration.
• NodeSchedulerState: Holds the map of nodes and workers with node expiration logic.

Usage

The controller selects a scheduler implementation at startup based on configuration. All scheduler operations are accessed through the Scheduler trait interface.

Code Reference

Source Location

crates/arroyo-controller/src/schedulers/mod.rs

Signature

#[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, run_id: Option<u64>, force: bool) -> anyhow::Result<()>;
    async fn workers_for_job(&self, job_id: &str, run_id: Option<u64>) -> anyhow::Result<Vec<WorkerId>>;
}

pub struct StartPipelineReq {
    pub name: String,
    pub program: LogicalProgram,
    pub wasm_path: String,
    pub job_id: Arc<String>,
    pub hash: String,
    pub run_id: u64,
    pub slots: usize,
    pub env_vars: HashMap<String, String>,
}

pub enum SchedulerError {
    NotEnoughSlots { slots_needed: usize },
    Other(String),
}

pub struct ProcessScheduler { /* ... */ }
pub struct NodeScheduler { /* ... */ }

Import

use crate::schedulers::{Scheduler, SchedulerError, StartPipelineReq};
use crate::schedulers::{ProcessScheduler, NodeScheduler};

I/O Contract

Inputs

Name	Type	Description
StartPipelineReq	struct	Pipeline parameters including slots, job ID, run ID, and environment
RegisterNodeReq	gRPC message	Node registration with machine ID, task slots, and address
HeartbeatNodeReq	gRPC message	Node heartbeat for liveness tracking
WorkerFinishedReq	gRPC message	Worker completion notification with worker info and slot count

Outputs

Name	Type	Description
workers	Vec<WorkerId>	Active worker identifiers for a job
SchedulerError	enum	Error indicating insufficient slots or other scheduling failure

Usage Examples

// Using ProcessScheduler for local development
let scheduler = ProcessScheduler::new();
let req = StartPipelineReq {
    name: "test_pipeline".to_string(),
    program: logical_program,
    wasm_path: "file:///wasm".to_string(),
    job_id: Arc::new("job_1".to_string()),
    hash: "abc123".to_string(),
    run_id: 1,
    slots: 4,
    env_vars: HashMap::new(),
};
scheduler.start_workers(req).await?;

// Query active workers
let workers = scheduler.workers_for_job("job_1", Some(1)).await?;

// Stop workers gracefully
scheduler.stop_workers("job_1", Some(1), false).await?;

Related Pages

• ArroyoSystems_Arroyo_Kubernetes_Scheduler - Kubernetes-based Scheduler implementation
• ArroyoSystems_Arroyo_Controller_Server - The controller that uses the Scheduler trait
• ArroyoSystems_Arroyo_Datastream_Types - LogicalProgram type used in StartPipelineReq