cronoq_server/
server.rs

1//! Server assembly: listeners, TLS, store, scheduler and leadership.
2//!
3//! Listeners are bound and served from the start. Until this instance holds leadership (always
4//! immediately with the memory store), the API answers `UNAVAILABLE` / `503` and `/readyz`
5//! reports `503`; it switches to serving without rebinding once the scheduler is built.
6
7use std::future::Future;
8use std::net::SocketAddr;
9use std::path::Path;
10use std::sync::atomic::{AtomicBool, Ordering};
11use std::sync::{Arc, OnceLock};
12use std::time::Duration;
13
14use anyhow::Context;
15use bytes::Bytes;
16use cronoq::{ExternalScaler, Pool, Scheduler};
17use cronoq_proto::v1 as pb;
18use rustc_hash::FxHashMap;
19use rustls_pki_types::pem::PemObject;
20use rustls_pki_types::{CertificateDer, PrivateKeyDer};
21use tokio::net::TcpListener;
22use tokio::task::JoinHandle;
23use tokio_util::sync::CancellationToken;
24
25use crate::auth::Auth;
26use crate::config::{Config, ScalerSection, StoreConfig, TlsConfig};
27use crate::convert::unix_ns;
28use crate::grpc::{AdminService, TasksService, WorkersService};
29use crate::info::TaskInfoCache;
30use crate::remote::{OutputCache, RemoteExecutor, WorkerRegistry};
31use crate::scaler::WebhookScaler;
32use crate::state::{AppState, PoolInfo, Shared};
33use crate::store::Opened;
34
35/// How long listeners may take to drain after the server stops serving.
36const DRAIN: Duration = Duration::from_secs(5);
37
38/// How long workers get to read their `Shutdown` message before their streams are cut.
39const WORKER_GRACE: Duration = Duration::from_secs(1);
40
41/// Listening addresses, reported once this instance serves as the leader.
42#[derive(Clone, Debug)]
43pub struct Bound {
44    /// gRPC listener.
45    pub grpc: SocketAddr,
46    /// HTTP listener.
47    pub http: SocketAddr,
48    /// Unauthenticated metrics listener, if configured.
49    pub metrics: Option<SocketAddr>,
50}
51
52/// Open the configured store; `retention` is how far back finished tasks are recovered.
53async fn open_store(cfg: &StoreConfig, retention: Duration) -> anyhow::Result<Opened> {
54    match cfg {
55        StoreConfig::Memory => Ok(Opened::memory()),
56        StoreConfig::Postgres {
57            url: inline,
58            url_env,
59            schema,
60            ..
61        } => {
62            let url = match (inline, url_env) {
63                (Some(u), _) => u.clone(),
64                (None, env) => {
65                    let var = env.as_deref().unwrap_or("DATABASE_URL");
66                    std::env::var(var).map_err(|_| {
67                        anyhow::anyhow!("environment variable {var} (store URL) is not set")
68                    })?
69                }
70            };
71            crate::pg::warn_insecure_url(&url, inline.is_some());
72            crate::pg::open(&url, schema.as_deref(), retention).await
73        }
74        StoreConfig::Sqlite { path, .. } => {
75            let cfg = cronoq_store_sqlite::SqliteStoreConfig {
76                retention,
77                ..Default::default()
78            };
79            let store = cronoq_store_sqlite::SqliteStore::open_with(path, cfg)
80                .await
81                .map_err(|e| anyhow::anyhow!("sqlite {}: {e}", path.display()))?;
82            Ok(Opened::durable(store))
83        }
84    }
85}
86
87fn read(what: &str, path: &Path) -> anyhow::Result<Vec<u8>> {
88    std::fs::read(path).with_context(|| format!("reading TLS {what} {}", path.display()))
89}
90
91/// TLS material, loaded once at startup.
92struct Tls {
93    grpc: tonic::transport::ServerTlsConfig,
94    http: tokio_rustls::TlsAcceptor,
95}
96
97fn load_tls(cfg: &TlsConfig) -> anyhow::Result<Tls> {
98    let cert = read("certificate", &cfg.cert)?;
99    let key = read("key", &cfg.key)?;
100    let ca = cfg
101        .client_ca
102        .as_ref()
103        .map(|p| read("client CA", p))
104        .transpose()?;
105
106    let chain: Vec<CertificateDer<'static>> = CertificateDer::pem_slice_iter(&cert)
107        .collect::<Result<_, _>>()
108        .with_context(|| format!("parsing TLS certificate {}", cfg.cert.display()))?;
109    anyhow::ensure!(
110        !chain.is_empty(),
111        "no certificate found in {}",
112        cfg.cert.display()
113    );
114    let private = PrivateKeyDer::from_pem_slice(&key)
115        .with_context(|| format!("parsing TLS key {}", cfg.key.display()))?;
116    let provider = Arc::new(rustls::crypto::ring::default_provider());
117    let builder = rustls::ServerConfig::builder_with_provider(provider.clone())
118        .with_safe_default_protocol_versions()?;
119    let builder = match &ca {
120        Some(ca) => {
121            let mut roots = rustls::RootCertStore::empty();
122            for c in CertificateDer::pem_slice_iter(ca) {
123                let c = c.with_context(|| "parsing TLS client CA")?;
124                roots.add(c).context("adding TLS client CA")?;
125            }
126            anyhow::ensure!(
127                !roots.is_empty(),
128                "no certificate found in the client CA file"
129            );
130            let verifier = rustls::server::WebPkiClientVerifier::builder_with_provider(
131                Arc::new(roots),
132                provider,
133            )
134            .build()
135            .context("building the client certificate verifier")?;
136            builder.with_client_cert_verifier(verifier)
137        }
138        None => builder.with_no_client_auth(),
139    };
140    let mut http = builder
141        .with_single_cert(chain, private)
142        .context("TLS certificate and key do not match")?;
143    http.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
144
145    let mut grpc = tonic::transport::ServerTlsConfig::new()
146        .identity(tonic::transport::Identity::from_pem(cert, key));
147    if let Some(ca) = ca {
148        grpc = grpc.client_ca_root(tonic::transport::Certificate::from_pem(ca));
149    }
150    Ok(Tls {
151        grpc,
152        http: tokio_rustls::TlsAcceptor::from(Arc::new(http)),
153    })
154}
155
156async fn bind(what: &str, addr: SocketAddr) -> anyhow::Result<TcpListener> {
157    TcpListener::bind(addr)
158        .await
159        .with_context(|| format!("binding the {what} listener on {addr}"))
160}
161
162/// Take over a listener the caller bound.
163fn adopt(l: std::net::TcpListener) -> anyhow::Result<TcpListener> {
164    l.set_nonblocking(true)?;
165    Ok(TcpListener::from_std(l)?)
166}
167
168/// gRPC and HTTP listening sockets bound by the caller, for [`run_with_listeners`].
169#[derive(Debug)]
170pub struct Listeners {
171    /// Serves gRPC (in place of `server.grpc_addr`).
172    pub grpc: std::net::TcpListener,
173    /// Serves HTTP (in place of `server.http_addr`).
174    pub http: std::net::TcpListener,
175}
176
177/// Run until `shutdown` resolves, or until the scheduler stops (e.g. leadership is lost, which
178/// returns an error). `on_bound` is called once this instance serves as the leader.
179pub async fn run(
180    cfg: Config,
181    shutdown: impl Future<Output = ()>,
182    on_bound: impl FnOnce(Bound),
183) -> anyhow::Result<()> {
184    serve(cfg, None, shutdown, on_bound).await
185}
186
187/// [`run`] on listeners the caller has already bound, so their addresses are known before
188/// this instance becomes the leader (a standby is reachable, but `on_bound` is not called
189/// until it leads). The metrics listener, if configured, is still bound from
190/// `server.metrics_addr`.
191pub async fn run_with_listeners(
192    mut cfg: Config,
193    listeners: Listeners,
194    shutdown: impl Future<Output = ()>,
195    on_bound: impl FnOnce(Bound),
196) -> anyhow::Result<()> {
197    cfg.server.grpc_addr = listeners.grpc.local_addr()?;
198    cfg.server.http_addr = listeners.http.local_addr()?;
199    serve(cfg, Some(listeners), shutdown, on_bound).await
200}
201
202async fn serve(
203    cfg: Config,
204    listeners: Option<Listeners>,
205    shutdown: impl Future<Output = ()>,
206    on_bound: impl FnOnce(Bound),
207) -> anyhow::Result<()> {
208    cfg.validate()?;
209    let shards = cfg
210        .shard_plan()
211        .map(|p| Arc::new(crate::shard::ShardMap::new(&p)));
212    // From here on: this shard's pools and store only.
213    let cfg = cfg.for_own_shard();
214    if let Some(s) = &shards {
215        tracing::info!(shard = %s.own().name, number = s.own().number, "serving one shard");
216    }
217    let auth = Auth::from_config(&cfg.auth)?;
218    cfg.check_auth_exposure(auth.enabled())?;
219    if !auth.enabled() {
220        tracing::warn!(
221            "authentication is DISABLED: configure [[auth.tokens]] for any shared deployment"
222        );
223    }
224    let tls = cfg.server.tls.as_ref().map(load_tls).transpose()?;
225    let peers = match &shards {
226        Some(_) => Some(Arc::new(
227            crate::peer::Peers::new(
228                cfg.server.peer_tls.as_ref(),
229                auth.enabled(),
230                cfg.server.insecure_peer_http,
231                &cfg.limits,
232            )
233            .context("configuring calls to other shards")?,
234        )),
235        None => None,
236    };
237    if let (Some(peers), Some(map)) = (&peers, &shards) {
238        for sh in map.others() {
239            if !sh.endpoint.starts_with("https://") && auth.enabled() {
240                if peers.may_carry_tokens(sh) && cfg.server.insecure_peer_http {
241                    tracing::warn!(shard = %sh.name, endpoint = %sh.endpoint,
242                        "forwarding callers' tokens to this shard over plain http (insecure_peer_http)");
243                } else if !peers.may_carry_tokens(sh) {
244                    tracing::warn!(shard = %sh.name, endpoint = %sh.endpoint,
245                        "calls for this shard are not forwarded, because its endpoint is plain http \
246                         on another host; clients are told to connect to it directly");
247                }
248            }
249        }
250    }
251    let limits = cfg.limits.clone();
252
253    let (grpc_listener, http_listener) = match listeners {
254        Some(l) => (adopt(l.grpc)?, adopt(l.http)?),
255        None => (
256            bind("gRPC", cfg.server.grpc_addr).await?,
257            bind("HTTP", cfg.server.http_addr).await?,
258        ),
259    };
260    let metrics_listener = match cfg.server.metrics_addr {
261        Some(a) => Some(bind("metrics", a).await?),
262        None => None,
263    };
264    let bound = Bound {
265        grpc: grpc_listener.local_addr()?,
266        http: http_listener.local_addr()?,
267        metrics: metrics_listener
268            .as_ref()
269            .map(|l| l.local_addr())
270            .transpose()?,
271    };
272
273    let shared = Arc::new(Shared {
274        loopback_hosts_only: !auth.enabled() && !cfg.server.insecure_no_auth,
275        auth,
276        limits: limits.clone(),
277        ready: Arc::new(AtomicBool::new(false)),
278        app: OnceLock::new(),
279        stop: CancellationToken::new(),
280    });
281
282    // Listeners serve from the start; the API reports UNAVAILABLE until leadership.
283    let mut servers: Vec<JoinHandle<()>> = Vec::new();
284    {
285        let max_msg = limits.max_message_bytes;
286        let mut grpc = tonic::transport::Server::builder()
287            .http2_keepalive_interval(Some(limits.keepalive_interval()))
288            .http2_keepalive_timeout(Some(limits.keepalive_timeout()))
289            .tcp_keepalive(Some(limits.keepalive_interval()))
290            .max_concurrent_streams(limits.max_concurrent_streams)
291            .timeout(limits.max_wait() + limits.request_timeout());
292        if let Some(t) = &tls {
293            grpc = grpc
294                .tls_config(t.grpc.clone())
295                .context("configuring gRPC TLS")?;
296        }
297        let web = cfg
298            .server
299            .grpc_web
300            .as_ref()
301            .map(|w| crate::grpc_web::layers(w, shared.loopback_hosts_only));
302        if web.is_some() {
303            // Browsers without TLS (or without HTTP/2) send gRPC-Web over HTTP/1.1.
304            grpc = grpc.accept_http1(true);
305        }
306        let semaphore = Arc::new(tokio::sync::Semaphore::new(limits.max_concurrent_requests));
307        let router = grpc
308            .layer(tower::ServiceBuilder::new().option_layer(web).layer(
309                tower::limit::GlobalConcurrencyLimitLayer::with_semaphore(semaphore),
310            ))
311            .add_service(
312                pb::tasks_server::TasksServer::new(TasksService {
313                    shared: shared.clone(),
314                })
315                .max_decoding_message_size(max_msg)
316                .max_encoding_message_size(max_msg),
317            )
318            .add_service(
319                pb::workers_server::WorkersServer::new(WorkersService {
320                    shared: shared.clone(),
321                })
322                .max_decoding_message_size(max_msg)
323                .max_encoding_message_size(max_msg),
324            )
325            .add_service(
326                pb::admin_server::AdminServer::new(AdminService {
327                    shared: shared.clone(),
328                })
329                .max_decoding_message_size(max_msg)
330                .max_encoding_message_size(max_msg),
331            );
332        let standby = shared.clone();
333        let incoming = tokio_stream::StreamExt::map(
334            tokio_stream::wrappers::TcpListenerStream::new(grpc_listener),
335            move |c| c.map(|tcp| crate::standby::CloseOnStandby::new(tcp, standby.clone())),
336        );
337        let stop = shared.stop.clone();
338        servers.push(tokio::spawn(async move {
339            if let Err(e) = router
340                .serve_with_incoming_shutdown(incoming, stop.cancelled_owned())
341                .await
342            {
343                tracing::error!(error = %e, "gRPC server failed");
344            }
345        }));
346    }
347    servers.push(tokio::spawn(crate::http::serve(
348        http_listener,
349        crate::http::router(shared.clone()),
350        tls.as_ref().map(|t| t.http.clone()),
351        shared.clone(),
352        shared.stop.clone(),
353        DRAIN,
354    )));
355    if let Some(l) = metrics_listener {
356        servers.push(tokio::spawn(crate::http::serve(
357            l,
358            crate::http::metrics_router(shared.clone()),
359            None,
360            shared.clone(),
361            shared.stop.clone(),
362            DRAIN,
363        )));
364    }
365
366    let result = lead(&cfg, shards, peers, &shared, shutdown, bound, on_bound).await;
367    shared.ready.store(false, Ordering::SeqCst);
368    shared.stop.cancel();
369    for mut s in servers {
370        if tokio::time::timeout(DRAIN + Duration::from_secs(1), &mut s)
371            .await
372            .is_err()
373        {
374            tracing::warn!("a listener did not stop in time; aborting it");
375            s.abort();
376        }
377    }
378    result
379}
380
381/// Acquire leadership, serve, and stop. Returns when `shutdown` resolves or leadership is lost.
382async fn lead(
383    cfg: &Config,
384    shards: Option<Arc<crate::shard::ShardMap>>,
385    peers: Option<Arc<crate::peer::Peers>>,
386    shared: &Arc<Shared>,
387    shutdown: impl Future<Output = ()>,
388    bound: Bound,
389    on_bound: impl FnOnce(Bound),
390) -> anyhow::Result<()> {
391    let limits = &cfg.limits;
392    tokio::pin!(shutdown);
393    let Opened { store, queries } = tokio::select! {
394        s = open_store(&cfg.store, cfg.limits.key_retention()) => s?,
395        _ = &mut shutdown => return Ok(()),
396    };
397    let registry = WorkerRegistry::new();
398    let outputs = OutputCache::new(cfg.outputs.cache_entries, cfg.outputs.cache_bytes);
399    let info = TaskInfoCache::new(cfg.outputs.info_entries);
400    let mut builder = Scheduler::<Bytes, Bytes>::builder()
401        .store_arc(store.clone())
402        // Outputs are cached by the server (bounded by count and bytes) and the store; the
403        // runtime's count-only cache would hold up to 1024 maximum-size outputs.
404        .output_cache(0)
405        .emit_transitions(true)
406        .configure(|c| {
407            c.global_bounds.max_pending = Some(limits.max_pending_total);
408            c.retention = limits.key_retention();
409            c.retention_max = limits.key_retention_max;
410            c.id_shard = shards.as_ref().map(|s| s.own().number);
411        });
412    let mut pools = FxHashMap::default();
413    for p in &cfg.pools {
414        let pc = p.to_pool_config(limits);
415        let name: Arc<str> = pc.name.clone();
416        let lease_ttl = pc.lease_ttl.unwrap_or(Duration::from_secs(30));
417        pools.insert(
418            name.clone(),
419            PoolInfo {
420                slot_cap: p.slot_cap(limits),
421                lease_ttl,
422                max_attempts: pc.retry.max_attempts,
423            },
424        );
425        let mut pool = Pool::new(name.clone())
426            .configure(|c| *c = pc.clone())
427            .executor(RemoteExecutor::new(
428                name.clone(),
429                lease_ttl,
430                registry.clone(),
431                outputs.clone(),
432            ));
433        pool = match &p.scaler {
434            ScalerSection::External => pool.scaler(ExternalScaler),
435            ScalerSection::Webhook { url, token_env } => {
436                let token = match token_env {
437                    Some(env) => Some(std::env::var(env).map_err(|_| {
438                        anyhow::anyhow!(
439                            "pool '{name}': environment variable {env} (webhook token) is not set"
440                        )
441                    })?),
442                    None => None,
443                };
444                pool.scaler(WebhookScaler::new(url.clone(), token)?)
445            }
446        };
447        builder = builder.pool(pool);
448    }
449    for r in &cfg.resources {
450        builder = builder.resource(r.to_config());
451    }
452    if queries.is_some() {
453        tracing::info!("waiting for leadership");
454    }
455    let sched = tokio::select! {
456        s = builder.build() => s.context("starting the scheduler")?,
457        _ = &mut shutdown => {
458            tracing::info!("shutdown requested while standing by");
459            return Ok(());
460        }
461    };
462    registry.attach(sched.clone());
463    let pruner = match (&queries, &cfg.store) {
464        (
465            Some(q),
466            StoreConfig::Postgres {
467                keep_finished_s, ..
468            }
469            | StoreConfig::Sqlite {
470                keep_finished_s, ..
471            },
472        ) => Some(crate::store::spawn_pruner(
473            q.clone(),
474            Duration::from_secs_f64(keep_finished_s.unwrap_or(7.0 * 86400.0)),
475            Duration::from_secs(3600),
476        )),
477        _ => None,
478    };
479    let feeder = spawn_info_feeder(&sched, info.clone(), shared.stop.clone());
480    let app = Arc::new(AppState {
481        sched: sched.clone(),
482        registry: registry.clone(),
483        outputs,
484        info,
485        pools,
486        limits: limits.clone(),
487        queries: queries.clone(),
488        shards,
489        peers,
490    });
491    let _ = shared.app.set(app);
492    shared.ready.store(true, Ordering::SeqCst);
493    tracing::info!(grpc = %bound.grpc, http = %bound.http, "leader: serving");
494    on_bound(bound);
495
496    let lost_leadership = tokio::select! {
497        _ = &mut shutdown => false,
498        _ = sched.stopped() => true,
499        _ = crate::store::leadership_lost(queries.clone()) => true,
500    };
501    shared.ready.store(false, Ordering::SeqCst);
502    if let Some(p) = pruner {
503        p.abort();
504    }
505    let result = if lost_leadership {
506        tracing::error!("leadership lost or scheduler stopped: no longer serving");
507        registry.close_all("this server is no longer the leader", true, WORKER_GRACE);
508        let _ = sched.shutdown(Duration::ZERO).await;
509        Err(anyhow::anyhow!("lost leadership or the scheduler stopped"))
510    } else {
511        let grace = Duration::from_secs_f64(cfg.server.shutdown_grace_s);
512        // Workers keep reporting outcomes while running attempts finish.
513        let r = sched.shutdown(grace).await;
514        registry.close_all("server shutting down", false, WORKER_GRACE);
515        // Hand over at once instead of letting a standby wait for the lease to run out.
516        store.release_leadership().await;
517        r.context("scheduler shutdown")
518    };
519    // Ends Watch streams and pending waits before the listeners drain.
520    shared.stop.cancel();
521    feeder.abort();
522    result
523}
524
525/// Keep the task-info cache up to date from the scheduler's status events.
526fn spawn_info_feeder(
527    sched: &Scheduler<Bytes, Bytes>,
528    info: Arc<TaskInfoCache>,
529    stop: CancellationToken,
530) -> JoinHandle<()> {
531    let mut events = sched.events();
532    let names: Arc<[Arc<str>]> = sched.pools().into();
533    let clock = *sched.clock();
534    tokio::spawn(async move {
535        loop {
536            let e = tokio::select! {
537                e = events.recv() => e,
538                _ = stop.cancelled() => return,
539            };
540            match e {
541                Ok(e) => {
542                    let Some(pool) = names.get(e.pool.0 as usize) else {
543                        continue;
544                    };
545                    info.on_event(&e, pool, unix_ns(clock.to_system(e.at)));
546                }
547                Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
548                    tracing::debug!(dropped = n, "task info cache missed events");
549                }
550                Err(tokio::sync::broadcast::error::RecvError::Closed) => return,
551            }
552        }
553    })
554}
cronoq_server/server.rs

cronoq_server/
server.rs
