use std::{
collections::HashMap,
error::Error,
fs,
path::Path,
sync::Arc,
time::{Duration, Instant},
};
use datum::{Keep, Sink, Source};
use datum_agent::{
Agent, AgentConfig, DesiredJobState, JobMat, JobSpec, JobState,
dcp::{DcpJobFactories, DcpQuicServerConfig, DcpServer, DcpServerConfig, DcpTcpServerConfig},
};
use datum_net::quic::{
crypto::rustls::QuicServerConfig,
quinn,
rustls::{
RootCertStore, ServerConfig as RustlsServerConfig,
pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer},
server::WebPkiClientVerifier,
},
};
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
let config_values = load_config_values()?;
let agent = Agent::start_with_config(AgentConfig::default())?;
let factories = registered_factories();
let server_config = dcp_server_config(&config_values)?;
let server = DcpServer::from_agent(&agent, factories, server_config);
let handle = server.start().await?;
if let Some(addr) = handle.tcp_addr() {
eprintln!("datum-agent DCP TCP listening on {addr}");
}
if let Some(addr) = handle.quic_addr() {
eprintln!("datum-agent DCP QUIC listening on {addr}");
}
tokio::signal::ctrl_c().await?;
drain_jobs(agent.registry().clone(), Duration::from_secs(30)).await;
handle.shutdown().await;
agent.registry().shutdown()?;
Ok(())
}
fn load_config_values() -> Result<HashMap<String, String>, Box<dyn Error>> {
let mut values = HashMap::new();
if let Some(path) = std::env::args()
.nth(1)
.or_else(|| std::env::var("DATUM_AGENT_CONFIG").ok())
{
let content = fs::read_to_string(path)?;
for line in content.lines() {
let line = line.trim();
if line.is_empty() || line.starts_with('#') {
continue;
}
if let Some((key, value)) = line.split_once('=') {
values.insert(key.trim().to_owned(), value.trim().to_owned());
}
}
}
for (key, value) in std::env::vars() {
if key.starts_with("DATUM_AGENT_") {
values.insert(key, value);
}
}
Ok(values)
}
fn dcp_server_config(values: &HashMap<String, String>) -> Result<DcpServerConfig, Box<dyn Error>> {
let mut config = DcpServerConfig::default();
if let Some(node_id) = values.get("DATUM_AGENT_DCP_NODE_ID") {
config.node_id = node_id.clone();
}
if let Some(token) = values.get("DATUM_AGENT_DCP_AUTH_TOKEN") {
config.auth_token = Some(token.clone());
}
if let Some(interval) = values.get("DATUM_AGENT_DCP_METRICS_INTERVAL_MS") {
config.metrics_interval = Duration::from_millis(interval.parse()?);
}
let tcp_enabled = values
.get("DATUM_AGENT_DCP_TCP")
.map(|value| value != "0" && value != "false")
.unwrap_or(true);
config.tcp = if tcp_enabled {
let addr = values
.get("DATUM_AGENT_DCP_TCP_ADDR")
.map(String::as_str)
.unwrap_or("127.0.0.1:9555")
.parse()?;
Some(DcpTcpServerConfig { addr })
} else {
None
};
config.quic = match values.get("DATUM_AGENT_DCP_QUIC_ADDR") {
Some(addr) => Some(DcpQuicServerConfig {
addr: addr.parse()?,
server_config: load_quic_server_config(values)?,
}),
None => None,
};
Ok(config)
}
fn load_quic_server_config(
values: &HashMap<String, String>,
) -> Result<quinn::ServerConfig, Box<dyn Error>> {
let cert = read_required_der(values, "DATUM_AGENT_DCP_QUIC_CERT_DER")?;
let key = read_required_der(values, "DATUM_AGENT_DCP_QUIC_KEY_DER")?;
let client_ca = read_required_der(values, "DATUM_AGENT_DCP_QUIC_CLIENT_CA_DER")?;
let mut client_roots = RootCertStore::empty();
client_roots.add(CertificateDer::from(client_ca))?;
let verifier = WebPkiClientVerifier::builder(Arc::new(client_roots)).build()?;
let server_crypto = RustlsServerConfig::builder()
.with_client_cert_verifier(verifier)
.with_single_cert(
vec![CertificateDer::from(cert)],
PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(key)),
)?;
Ok(quinn::ServerConfig::with_crypto(Arc::new(
QuicServerConfig::try_from(server_crypto)?,
)))
}
fn read_required_der(
values: &HashMap<String, String>,
key: &str,
) -> Result<Vec<u8>, Box<dyn Error>> {
let path = values
.get(key)
.ok_or_else(|| format!("missing required config key: {key}"))?;
Ok(fs::read(Path::new(path))?)
}
fn registered_factories() -> DcpJobFactories {
let factories = DcpJobFactories::new();
factories
.register("empty", |instance_name, _params| {
Ok(JobSpec::new(instance_name, |context| {
let control = context.control();
Ok(Source::<u64>::empty()
.instrumented(
format!("{}:{}", context.name(), context.generation()),
context.instrumentation_registry(),
)
.via_mat(context.drain_flow(), Keep::right)
.to_mat(Sink::ignore(), move |_switch, completion| {
JobMat::new(completion, control.clone())
}))
}))
})
.expect("built-in empty factory registers");
factories
}
async fn drain_jobs(registry: datum_agent::JobRegistryHandle, timeout: Duration) {
let _ = tokio::task::spawn_blocking({
let registry = registry.clone();
move || {
if let Ok(statuses) = registry.list() {
for status in statuses {
if matches!(
status.desired_state,
DesiredJobState::Running | DesiredJobState::Draining
) {
let _ = registry.drain(status.name);
}
}
}
}
})
.await;
let deadline = Instant::now() + timeout;
loop {
let registry = registry.clone();
let done = tokio::task::spawn_blocking(move || {
registry
.list()
.map(|statuses| {
statuses.iter().all(|status| {
matches!(
status.state,
JobState::Completed
| JobState::Drained
| JobState::Stopped
| JobState::Failed
)
})
})
.unwrap_or(true)
})
.await
.unwrap_or(true);
if done || Instant::now() >= deadline {
break;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
}