use std::sync::Arc;
use chronon_core::context::{ContextFactory, NoOpContextFactory};
use chronon_core::error::{ChrononError, Result};
use chronon_core::store::SchedulerStore;
use chronon_executor::{Executor, ScriptRegistry};
use chronon_scheduler::{Scheduler, SchedulerConfig};
use chronon_telemetry::{NoOpSink, TelemetrySink};
use tokio::sync::{mpsc, Notify};
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub enum DeploymentShape {
#[default]
Embedded,
CoordinatorOnly,
Worker(String),
RemoteClient(String),
}
pub struct ChrononBuilder {
store: Option<Arc<dyn SchedulerStore>>,
context_factory: Option<Arc<dyn ContextFactory>>,
telemetry: Option<Arc<dyn TelemetrySink>>,
registry: Option<Arc<ScriptRegistry>>,
deployment: DeploymentShape,
auto_registry: bool,
tick_interval_ms: u64,
instance_id: Option<String>,
}
impl ChrononBuilder {
pub fn new() -> Self {
Self {
store: None,
context_factory: None,
telemetry: None,
registry: None,
deployment: DeploymentShape::Embedded,
auto_registry: false,
tick_interval_ms: chronon_scheduler::tick_interval_ms_from_env(),
instance_id: None,
}
}
pub fn scheduler_store(mut self, store: Arc<dyn SchedulerStore>) -> Self {
self.store = Some(store);
self
}
pub fn scheduler_store_from_global(mut self) -> Result<Self> {
self.store = Some(chronon_core::default_store_from_global()?);
Ok(self)
}
pub fn context_factory(mut self, factory: Arc<dyn ContextFactory>) -> Self {
self.context_factory = Some(factory);
self
}
pub fn telemetry_sink(mut self, sink: Arc<dyn TelemetrySink>) -> Self {
self.telemetry = Some(sink);
self
}
pub fn script_registry(mut self, registry: Arc<ScriptRegistry>) -> Self {
self.registry = Some(registry);
self
}
pub fn instance_id(mut self, id: impl Into<String>) -> Self {
self.instance_id = Some(id.into());
self
}
pub fn embedded(mut self) -> Self {
self.deployment = DeploymentShape::Embedded;
self
}
pub fn coordinator_only(mut self) -> Self {
self.deployment = DeploymentShape::CoordinatorOnly;
self
}
pub fn worker(mut self, pool_id: impl Into<String>) -> Self {
self.deployment = DeploymentShape::Worker(pool_id.into());
self
}
pub fn remote_coordinator(mut self, base_url: impl Into<String>) -> Self {
self.deployment = DeploymentShape::RemoteClient(base_url.into());
self
}
pub fn auto_registry(mut self) -> Self {
self.auto_registry = true;
self
}
pub fn tick_interval_ms(mut self, ms: u64) -> Self {
self.tick_interval_ms = ms;
self
}
pub fn build(self) -> Result<super::Chronon> {
let store = self
.store
.ok_or_else(|| ChrononError::Internal("scheduler_store is required".into()))?;
let context_factory = self
.context_factory
.unwrap_or_else(|| Arc::new(NoOpContextFactory));
let telemetry = self
.telemetry
.unwrap_or_else(|| Arc::new(NoOpSink) as Arc<dyn TelemetrySink>);
let registry = match self.registry {
Some(registry) => registry,
None if self.auto_registry => Arc::new(ScriptRegistry::from_inventory()),
None => Arc::new(ScriptRegistry::new()),
};
let embedded_partitions = matches!(self.deployment, DeploymentShape::Embedded);
let instance_id = self
.instance_id
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
let scheduler = Arc::new(Scheduler::new(
SchedulerConfig {
tick_interval_ms: self.tick_interval_ms,
instance_id,
num_partitions: chronon_scheduler::num_partitions_from_env(),
embedded: embedded_partitions,
},
store.clone(),
telemetry.clone(),
));
let (event_tx, event_rx) = mpsc::unbounded_channel();
let executor = Arc::new(Executor::new(
registry,
context_factory,
telemetry,
event_tx,
));
Ok(super::Chronon::new(
store,
scheduler,
executor,
self.deployment,
Arc::new(Notify::new()),
event_rx,
))
}
}
impl Default for ChrononBuilder {
fn default() -> Self {
Self::new()
}
}
pub fn builder() -> ChrononBuilder {
ChrononBuilder::new()
}
#[cfg(test)]
mod tests {
use super::*;
use chronon_backend_mem::InMemorySchedulerStore;
use chronon_telemetry::ConsoleSink;
#[tokio::test]
async fn builder_embedded_compiles() {
let store = Arc::new(InMemorySchedulerStore::new());
let chronon = ChrononBuilder::new()
.scheduler_store(store)
.telemetry_sink(Arc::new(ConsoleSink))
.embedded()
.build()
.expect("build");
assert_eq!(chronon.deployment, DeploymentShape::Embedded);
}
#[tokio::test]
async fn builder_auto_registry_from_inventory() {
let store = Arc::new(InMemorySchedulerStore::new());
let chronon = ChrononBuilder::new()
.scheduler_store(store)
.embedded()
.auto_registry()
.build()
.expect("build");
let _ = chronon.executor().script_count();
}
#[tokio::test]
async fn builder_scheduler_store_from_global() {
let _installed = chronon_backend_mem::install_default_mem_store();
let chronon = ChrononBuilder::new()
.scheduler_store_from_global()
.expect("global store")
.embedded()
.build()
.expect("build");
assert_eq!(chronon.deployment, DeploymentShape::Embedded);
}
}