cronoq_server/
worker.rs

1//! Rust client for remote workers.
2//!
3//! [`run_worker`] connects to a cronoq server, registers for a pool and runs a handler for each
4//! lease, with at most `slots` handlers at a time. It heartbeats running leases, reports
5//! outcomes, reconnects with backoff when the connection drops or the server steps down, and
6//! drains running handlers when `shutdown` resolves.
7//!
8//! ```no_run
9//! # async fn demo() -> anyhow::Result<()> {
10//! use cronoq_server::worker::{WorkerOptions, run_worker};
11//! let opts = WorkerOptions::new("http://127.0.0.1:7070", "default").slots(8);
12//! run_worker(opts, |lease| async move { Ok(lease.payload) }, std::future::pending()).await
13//! # }
14//! ```
15//!
16//! A handler that panics fails its attempt as retryable; an output larger than the server
17//! accepts fails it permanently.
18//!
19//! When the pool belongs to another shard, the server refuses the worker and names that
20//! shard's endpoint; the worker then connects there instead (see [`MAX_REDIRECTS`]). It never
21//! follows a redirect from `https://` to `http://`, and never sends its token over plain http
22//! to another host unless [`WorkerOptions::insecure_redirect_http`] allows it.
23
24use std::collections::HashMap;
25use std::fmt;
26use std::future::Future;
27use std::panic::AssertUnwindSafe;
28use std::sync::atomic::{AtomicUsize, Ordering};
29use std::sync::{Arc, Mutex};
30use std::time::{Duration, SystemTime, UNIX_EPOCH};
31
32use bytes::Bytes;
33use cronoq::TaskError;
34use cronoq_proto::v1 as pb;
35use futures_util::FutureExt;
36use pb::worker_message::Msg;
37use pb::workers_client::WorkersClient;
38use rustc_hash::FxHashMap;
39use tokio::sync::{Notify, mpsc};
40use tokio_stream::wrappers::ReceiverStream;
41use tokio_util::sync::CancellationToken;
42use tonic::Code;
43use tonic::metadata::MetadataValue;
44use tonic::transport::{Certificate, ClientTlsConfig, Endpoint, Identity};
45
46/// TLS settings for the worker connection.
47#[derive(Clone, Default)]
48pub struct WorkerTls {
49    /// PEM CA certificate(s) to trust (default: the system roots).
50    pub ca_pem: Option<Vec<u8>>,
51    /// PEM client certificate and key, for mTLS.
52    pub identity: Option<(Vec<u8>, Vec<u8>)>,
53    /// Server name to verify (default: the endpoint host).
54    pub domain: Option<String>,
55}
56
57impl fmt::Debug for WorkerTls {
58    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
59        f.debug_struct("WorkerTls")
60            .field("ca_pem", &self.ca_pem.as_ref().map(|_| "<pem>"))
61            .field("identity", &self.identity.as_ref().map(|_| "<redacted>"))
62            .field("domain", &self.domain)
63            .finish()
64    }
65}
66
67/// Worker settings.
68#[derive(Clone)]
69pub struct WorkerOptions {
70    /// Server gRPC endpoint (`http://` or `https://`).
71    pub endpoint: String,
72    /// Pool to serve.
73    pub pool: String,
74    /// Leases run concurrently (the server may grant fewer).
75    pub slots: u32,
76    /// Bearer token (role `worker`).
77    pub token: Option<String>,
78    /// Name shown in server logs.
79    pub name: String,
80    /// TLS settings; they need an `https://` endpoint (an `https://` endpoint without them uses
81    /// the system roots).
82    pub tls: Option<WorkerTls>,
83    /// Reconnect after the connection drops or the server steps down (default `true`).
84    pub reconnect: bool,
85    /// First reconnect delay (doubles up to `max_backoff`).
86    pub initial_backoff: Duration,
87    /// Longest reconnect delay.
88    pub max_backoff: Duration,
89    /// How long running handlers may take to finish after `shutdown` resolves.
90    pub drain_timeout: Duration,
91    /// Largest lease message accepted from the server.
92    pub max_message_bytes: usize,
93    /// Labels of this worker: tasks that require labels are leased only to workers that have
94    /// all of them.
95    pub labels: Vec<(String, String)>,
96    /// Follow redirects to plain `http://` endpoints on other hosts with the token (default
97    /// `false`: only loopback); for trusted networks without TLS.
98    pub insecure_redirect_http: bool,
99}
100
101impl fmt::Debug for WorkerOptions {
102    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
103        f.debug_struct("WorkerOptions")
104            .field("endpoint", &self.endpoint)
105            .field("pool", &self.pool)
106            .field("slots", &self.slots)
107            .field("token", &self.token.as_ref().map(|_| "<redacted>"))
108            .field("name", &self.name)
109            .field("tls", &self.tls)
110            .field("reconnect", &self.reconnect)
111            .field("drain_timeout", &self.drain_timeout)
112            .field("labels", &self.labels)
113            .field("insecure_redirect_http", &self.insecure_redirect_http)
114            .finish_non_exhaustive()
115    }
116}
117
118impl WorkerOptions {
119    /// Options for serving `pool` on the server at `endpoint`, with one slot.
120    pub fn new(endpoint: impl Into<String>, pool: impl Into<String>) -> Self {
121        WorkerOptions {
122            endpoint: endpoint.into(),
123            pool: pool.into(),
124            slots: 1,
125            token: None,
126            name: String::new(),
127            tls: None,
128            reconnect: true,
129            initial_backoff: Duration::from_millis(200),
130            max_backoff: Duration::from_secs(10),
131            drain_timeout: Duration::from_secs(30),
132            max_message_bytes: 64 << 20,
133            labels: Vec::new(),
134            insecure_redirect_http: false,
135        }
136    }
137    /// Allow redirects to plain `http://` endpoints on other hosts with the token.
138    pub fn insecure_redirect_http(mut self, on: bool) -> Self {
139        self.insecure_redirect_http = on;
140        self
141    }
142    /// Add a label (`name = value`) to this worker.
143    pub fn label(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
144        self.labels.push((name.into(), value.into()));
145        self
146    }
147    /// Concurrent leases.
148    pub fn slots(mut self, n: u32) -> Self {
149        self.slots = n;
150        self
151    }
152    /// Bearer token.
153    pub fn token(mut self, t: impl Into<String>) -> Self {
154        self.token = Some(t.into());
155        self
156    }
157    /// Name shown in server logs.
158    pub fn name(mut self, n: impl Into<String>) -> Self {
159        self.name = n.into();
160        self
161    }
162    /// TLS settings.
163    pub fn tls(mut self, tls: WorkerTls) -> Self {
164        self.tls = Some(tls);
165        self
166    }
167    /// Whether to reconnect after the connection drops.
168    pub fn reconnect(mut self, on: bool) -> Self {
169        self.reconnect = on;
170        self
171    }
172    /// Time running handlers get to finish on shutdown.
173    pub fn drain_timeout(mut self, d: Duration) -> Self {
174        self.drain_timeout = d;
175        self
176    }
177}
178
179/// One leased attempt.
180#[derive(Clone, Debug)]
181pub struct LeaseInfo {
182    /// Lease id (unique per attempt).
183    pub lease: u64,
184    /// Task id.
185    pub task_id: u64,
186    /// Attempt number (1-based).
187    pub attempt: u32,
188    /// Pool the task runs in.
189    pub pool: String,
190    /// Task payload.
191    pub payload: Bytes,
192    /// Payload content type (may be empty).
193    pub content_type: String,
194    /// Task kind.
195    pub kind: String,
196    /// Task key (may be empty).
197    pub key: String,
198    /// Task tags.
199    pub tags: HashMap<String, String>,
200    /// Attempt timeout, if any.
201    pub timeout: Option<Duration>,
202    /// Task deadline, if any.
203    pub deadline: Option<SystemTime>,
204    /// Cancelled when the server aborts this attempt or the worker stops.
205    pub cancel: CancellationToken,
206}
207
208fn duration(s: Option<f64>) -> Option<Duration> {
209    s.and_then(|s| Duration::try_from_secs_f64(s).ok())
210}
211
212/// Redirects to other shards followed in a row before the worker gives up.
213pub const MAX_REDIRECTS: u32 = 3;
214
215/// Metadata naming the endpoint of the shard that serves a refused pool.
216const ENDPOINT_HEADER: &str = "cronoq-shard-endpoint";
217
218/// The endpoint a refusal sends the worker to, if it is a shard redirect: `Ok(None)` when it is
219/// not one, an error when the endpoint may not be followed from `current`.
220fn redirect(
221    opts: &WorkerOptions,
222    current: &str,
223    s: &tonic::Status,
224) -> anyhow::Result<Option<String>> {
225    if s.code() != Code::FailedPrecondition {
226        return Ok(None);
227    }
228    let Some(target) = s
229        .metadata()
230        .get(ENDPOINT_HEADER)
231        .and_then(|v| v.to_str().ok())
232    else {
233        return Ok(None);
234    };
235    let https = target.starts_with("https://");
236    anyhow::ensure!(
237        https || target.starts_with("http://"),
238        "the server redirected the worker to '{target}', which is not an http(s) endpoint"
239    );
240    anyhow::ensure!(
241        https || !current.starts_with("https://"),
242        "the server redirected the worker from https to plain http ('{target}'); refusing"
243    );
244    let uri = target.parse::<tonic::transport::Uri>()?;
245    anyhow::ensure!(
246        uri.authority().is_some_and(|a| !a.as_str().contains('@')),
247        "the server redirected the worker to '{target}', which is not a plain host and port"
248    );
249    let plain_remote = target.starts_with("http://") && !is_loopback_endpoint(target);
250    if opts.token.is_some() && plain_remote {
251        anyhow::ensure!(
252            opts.insecure_redirect_http,
253            "the server redirected the worker to '{target}', which would send its token over \
254             plain http to another host; refusing (insecure_redirect_http allows it)"
255        );
256        tracing::warn!(endpoint = %target, "following a redirect that sends the worker token over plain http");
257    }
258    Ok(Some(target.to_string()))
259}
260
261/// Errors that retrying cannot fix.
262fn fatal(code: Code) -> bool {
263    matches!(
264        code,
265        Code::Unauthenticated
266            | Code::PermissionDenied
267            | Code::NotFound
268            | Code::InvalidArgument
269            | Code::FailedPrecondition
270            | Code::Unimplemented
271    )
272}
273
274/// Running handlers, across connections.
275#[derive(Default)]
276struct Running {
277    leases: Mutex<FxHashMap<u64, Lease>>,
278    count: AtomicUsize,
279    idle: Notify,
280}
281
282struct Lease {
283    cancel: CancellationToken,
284    /// Set when the server aborted the attempt: no outcome is reported.
285    aborted: Arc<std::sync::atomic::AtomicBool>,
286}
287
288/// Decrements the running count when a handler task ends, however it ends.
289struct Slot {
290    running: Arc<Running>,
291    lease: u64,
292}
293
294impl Running {
295    /// Free a lease's slot (once): when its handler ends, or as soon as the server aborts it,
296    /// so the server may reuse the slot right away.
297    fn release(&self, lease: u64) -> Option<Lease> {
298        let l = self.leases.lock().unwrap().remove(&lease)?;
299        if self.count.fetch_sub(1, Ordering::SeqCst) == 1 {
300            self.idle.notify_waiters();
301        }
302        Some(l)
303    }
304}
305
306impl Drop for Slot {
307    fn drop(&mut self) {
308        self.running.release(self.lease);
309    }
310}
311
312/// Aborts a spawned task when dropped.
313struct AbortOnDrop(tokio::task::JoinHandle<()>);
314
315impl Drop for AbortOnDrop {
316    fn drop(&mut self) {
317        self.0.abort();
318    }
319}
320
321enum SessionEnd {
322    /// The caller asked to stop.
323    Shutdown,
324    /// The server closed the stream or asked us to go away.
325    Ended(String),
326    /// A gRPC error.
327    Failed(tonic::Status),
328}
329
330/// Connect, register and process leases until `shutdown` resolves (then drain running
331/// handlers), reconnecting with backoff when the connection is lost. Returns an error for
332/// failures that retrying cannot fix (bad credentials, unknown pool, ...), or for any failure
333/// when reconnecting is disabled.
334pub async fn run_worker<F, Fut>(
335    opts: WorkerOptions,
336    handler: F,
337    shutdown: impl Future<Output = ()>,
338) -> anyhow::Result<()>
339where
340    F: Fn(LeaseInfo) -> Fut + Send + Sync + 'static,
341    Fut: Future<Output = Result<Bytes, TaskError>> + Send + 'static,
342{
343    endpoint(&opts, &opts.endpoint)?;
344    let auth: Option<MetadataValue<_>> = match &opts.token {
345        Some(t) => {
346            let mut v: MetadataValue<_> = format!("Bearer {t}").parse()?;
347            v.set_sensitive(true);
348            Some(v)
349        }
350        None => None,
351    };
352    let handler = Arc::new(handler);
353    let running = Arc::new(Running::default());
354    let stop = CancellationToken::new();
355    // Handlers are cancelled if this future is dropped, too.
356    let _stop_on_drop = stop.clone().drop_guard();
357    tokio::pin!(shutdown);
358    let mut backoff = opts.initial_backoff;
359    // The endpoint connected to: the configured one, or the shard it redirected the worker to.
360    let mut current = opts.endpoint.clone();
361    let mut redirects = 0;
362    let result = loop {
363        let mut connect_failed = false;
364        let session = async {
365            let channel = endpoint(&opts, &current)
366                .map_err(|e| tonic::Status::invalid_argument(e.to_string()))?
367                .connect()
368                .await
369                .map_err(|e| {
370                    connect_failed = true;
371                    tonic::Status::unavailable(format!("connect: {e}"))
372                })?;
373            let auth = auth.clone();
374            let client =
375                WorkersClient::with_interceptor(channel, move |mut req: tonic::Request<()>| {
376                    if let Some(a) = &auth {
377                        req.metadata_mut().insert("authorization", a.clone());
378                    }
379                    Ok(req)
380                })
381                .max_decoding_message_size(opts.max_message_bytes);
382            Ok::<_, tonic::Status>(client)
383        };
384        let end = tokio::select! {
385            c = session => match c {
386                Ok(client) => serve(&opts, client, handler.clone(), running.clone(), &stop, &mut shutdown, &mut backoff).await,
387                Err(s) => SessionEnd::Failed(s),
388            },
389            _ = &mut shutdown => SessionEnd::Shutdown,
390        };
391        if let SessionEnd::Failed(s) = &end {
392            match redirect(&opts, &current, s) {
393                Ok(Some(target)) => {
394                    redirects += 1;
395                    if redirects > MAX_REDIRECTS {
396                        break Err(anyhow::anyhow!(
397                            "worker redirected {redirects} times in a row, more than the {MAX_REDIRECTS} it follows: {s}"
398                        ));
399                    }
400                    tracing::info!(from = %current, to = %target, "pool is served by another shard; reconnecting there");
401                    current = target;
402                    continue;
403                }
404                Ok(None) => {}
405                Err(e) => break Err(e),
406            }
407        }
408        redirects = 0;
409        if connect_failed && current != opts.endpoint {
410            // The shard is unreachable: ask the configured endpoint again where the pool is.
411            tracing::warn!(endpoint = %current, "shard unreachable; returning to the configured endpoint");
412            current = opts.endpoint.clone();
413        }
414        match end {
415            SessionEnd::Shutdown => break Ok(()),
416            SessionEnd::Failed(s) if fatal(s.code()) => {
417                break Err(anyhow::anyhow!("worker rejected: {s}"));
418            }
419            SessionEnd::Failed(s) if !opts.reconnect => {
420                break Err(anyhow::anyhow!("worker stream failed: {s}"));
421            }
422            SessionEnd::Ended(why) if !opts.reconnect => {
423                break Err(anyhow::anyhow!("worker stream ended: {why}"));
424            }
425            SessionEnd::Failed(s) => {
426                tracing::warn!(error = %s, retry_in = ?backoff, "worker connection failed")
427            }
428            SessionEnd::Ended(why) => {
429                tracing::info!(reason = %why, retry_in = ?backoff, "worker stream ended; reconnecting")
430            }
431        }
432        tokio::select! {
433            _ = tokio::time::sleep(backoff) => {}
434            _ = &mut shutdown => break Ok(()),
435        }
436        backoff = (backoff * 2).min(opts.max_backoff);
437    };
438    // Whatever is still running after the drain (or after a fatal error) is cancelled.
439    stop.cancel();
440    for (_, l) in running.leases.lock().unwrap().drain() {
441        l.cancel.cancel();
442    }
443    result
444}
445
446/// The endpoint at `url`: the configured one, or a shard the server redirected the worker to
447/// (which keeps the TLS settings except a fixed server name).
448fn endpoint(opts: &WorkerOptions, url: &str) -> anyhow::Result<Endpoint> {
449    let mut ep = Endpoint::from_shared(url.to_string())?
450        .connect_timeout(Duration::from_secs(10))
451        .tcp_keepalive(Some(Duration::from_secs(20)))
452        .http2_keep_alive_interval(Duration::from_secs(20))
453        .keep_alive_timeout(Duration::from_secs(20))
454        .keep_alive_while_idle(true);
455    let https = url.starts_with("https://");
456    // tonic only applies TLS to https:// endpoints; anything else would silently be plaintext.
457    anyhow::ensure!(
458        https || opts.tls.is_none(),
459        "worker TLS settings need an https:// endpoint, not '{url}'"
460    );
461    if !https && opts.token.is_some() && !is_loopback_endpoint(url) {
462        tracing::warn!(endpoint = %url, "the worker token is sent over plain http; use https outside a trusted network");
463    }
464    if https {
465        let tls = opts.tls.clone().unwrap_or_default();
466        let mut cfg = ClientTlsConfig::new();
467        cfg = match &tls.ca_pem {
468            Some(ca) => cfg.ca_certificate(Certificate::from_pem(ca)),
469            None => cfg.with_native_roots(),
470        };
471        if let Some((cert, key)) = &tls.identity {
472            cfg = cfg.identity(Identity::from_pem(cert, key));
473        }
474        if let Some(d) = &tls.domain
475            && url == opts.endpoint
476        {
477            cfg = cfg.domain_name(d.clone());
478        }
479        ep = ep.tls_config(cfg)?;
480    }
481    Ok(ep)
482}
483
484fn is_loopback_endpoint(endpoint: &str) -> bool {
485    // Parsed as tonic connects, so userinfo (`http://127.0.0.1:1@host`) cannot hide the host.
486    let Ok(uri) = endpoint.parse::<tonic::transport::Uri>() else {
487        return false;
488    };
489    let host = uri.host().unwrap_or("");
490    let host = host.trim_start_matches('[').trim_end_matches(']');
491    host.eq_ignore_ascii_case("localhost")
492        || host
493            .parse::<std::net::IpAddr>()
494            .is_ok_and(|ip| ip.is_loopback())
495}
496
497#[allow(clippy::too_many_arguments)]
498async fn serve<F, Fut, I>(
499    opts: &WorkerOptions,
500    mut client: WorkersClient<
501        tonic::service::interceptor::InterceptedService<tonic::transport::Channel, I>,
502    >,
503    handler: Arc<F>,
504    running: Arc<Running>,
505    stop: &CancellationToken,
506    shutdown: &mut std::pin::Pin<&mut impl Future<Output = ()>>,
507    backoff: &mut Duration,
508) -> SessionEnd
509where
510    F: Fn(LeaseInfo) -> Fut + Send + Sync + 'static,
511    Fut: Future<Output = Result<Bytes, TaskError>> + Send + 'static,
512    I: tonic::service::Interceptor,
513{
514    let cap = (opts.slots as usize).saturating_mul(4).clamp(64, 65_536);
515    let (tx, rx) = mpsc::channel::<pb::WorkerMessage>(cap);
516    let register = pb::WorkerMessage {
517        msg: Some(Msg::Register(pb::Register {
518            pool: opts.pool.clone(),
519            slots: opts.slots,
520            worker_name: opts.name.clone(),
521            protocol_version: cronoq_proto::PROTOCOL_VERSION,
522            labels: opts.labels.iter().cloned().collect(),
523        })),
524    };
525    if tx.send(register).await.is_err() {
526        return SessionEnd::Ended("local channel closed".into());
527    }
528    let mut inbound = match client.work(ReceiverStream::new(rx)).await {
529        Ok(r) => r.into_inner(),
530        Err(s) => return SessionEnd::Failed(s),
531    };
532    let registered = match inbound.message().await {
533        Ok(Some(pb::ServerMessage {
534            msg: Some(pb::server_message::Msg::Registered(r)),
535        })) => r,
536        Ok(Some(pb::ServerMessage {
537            msg: Some(pb::server_message::Msg::Shutdown(s)),
538        })) => {
539            return SessionEnd::Ended(format!("server shutting down: {}", s.reason));
540        }
541        Ok(other) => {
542            return SessionEnd::Failed(tonic::Status::internal(format!(
543                "unexpected first server message: {other:?}"
544            )));
545        }
546        Err(s) => return SessionEnd::Failed(s),
547    };
548    *backoff = opts.initial_backoff;
549    let ttl = duration(registered.lease_ttl_s);
550    let slots = registered.slots as usize;
551    // Leave room for the message envelope.
552    let max_output = (registered.max_message_bytes as usize).saturating_sub(1024);
553    tracing::info!(worker = registered.worker_id, pool = %opts.pool, slots, "worker registered");
554    let mut draining: Option<tokio::time::Instant> = None;
555    loop {
556        let drain_done = async {
557            match draining {
558                Some(deadline) => loop {
559                    let idle = running.idle.notified();
560                    if running.count.load(Ordering::SeqCst) == 0 {
561                        return;
562                    }
563                    tokio::select! {
564                        _ = idle => {}
565                        _ = tokio::time::sleep_until(deadline) => return,
566                    }
567                },
568                None => std::future::pending().await,
569            }
570        };
571        let msg = tokio::select! {
572            m = inbound.message() => m,
573            _ = shutdown.as_mut(), if draining.is_none() => {
574                tracing::info!(running = running.count.load(Ordering::SeqCst), "worker draining");
575                draining = Some(tokio::time::Instant::now() + opts.drain_timeout);
576                continue;
577            }
578            _ = drain_done => {
579                // Handlers still running after the drain timeout report a retryable failure.
580                for l in running.leases.lock().unwrap().values() {
581                    l.cancel.cancel();
582                }
583                // Half-close: once every outcome is sent, the server ends the stream.
584                drop(tx);
585                let flush = async { while let Ok(Some(_)) = inbound.message().await {} };
586                let _ = tokio::time::timeout(Duration::from_secs(5), flush).await;
587                return SessionEnd::Shutdown;
588            }
589        };
590        let msg = match msg {
591            Ok(Some(m)) => m.msg,
592            Ok(None) => {
593                return if draining.is_some() {
594                    SessionEnd::Shutdown
595                } else {
596                    SessionEnd::Ended("server closed the stream".into())
597                };
598            }
599            Err(s) => {
600                return if draining.is_some() {
601                    SessionEnd::Shutdown
602                } else {
603                    SessionEnd::Failed(s)
604                };
605            }
606        };
607        match msg {
608            Some(pb::server_message::Msg::Lease(l)) => {
609                let busy = running.count.load(Ordering::SeqCst);
610                if draining.is_some() || busy >= slots {
611                    let reason = if draining.is_some() {
612                        "worker is draining"
613                    } else {
614                        "no free slot"
615                    };
616                    let _ = tx.try_send(pb::WorkerMessage {
617                        msg: Some(Msg::Nack(pb::Nack {
618                            lease: l.lease,
619                            reason: reason.into(),
620                        })),
621                    });
622                    continue;
623                }
624                start_lease(
625                    l,
626                    handler.clone(),
627                    running.clone(),
628                    tx.clone(),
629                    ttl,
630                    max_output,
631                    stop,
632                );
633            }
634            Some(pb::server_message::Msg::Abort(a)) => {
635                if let Some(l) = running.release(a.lease) {
636                    l.aborted.store(true, Ordering::SeqCst);
637                    l.cancel.cancel();
638                }
639            }
640            Some(pb::server_message::Msg::Shutdown(s)) => {
641                if draining.is_some() {
642                    return SessionEnd::Shutdown;
643                }
644                let why = if s.not_leader {
645                    format!("server is no longer the leader: {}", s.reason)
646                } else {
647                    s.reason
648                };
649                return SessionEnd::Ended(why);
650            }
651            Some(pb::server_message::Msg::Registered(_)) | None => {}
652        }
653    }
654}
655
656fn start_lease<F, Fut>(
657    l: pb::Lease,
658    handler: Arc<F>,
659    running: Arc<Running>,
660    tx: mpsc::Sender<pb::WorkerMessage>,
661    ttl: Option<Duration>,
662    max_output: usize,
663    stop: &CancellationToken,
664) where
665    F: Fn(LeaseInfo) -> Fut + Send + Sync + 'static,
666    Fut: Future<Output = Result<Bytes, TaskError>> + Send + 'static,
667{
668    let lease = l.lease;
669    let cancel = stop.child_token();
670    let aborted = Arc::new(std::sync::atomic::AtomicBool::new(false));
671    running.count.fetch_add(1, Ordering::SeqCst);
672    running.leases.lock().unwrap().insert(
673        lease,
674        Lease {
675            cancel: cancel.clone(),
676            aborted: aborted.clone(),
677        },
678    );
679    let slot = Slot { running, lease };
680    let info = LeaseInfo {
681        lease,
682        task_id: l.task_id,
683        attempt: l.attempt,
684        pool: l.pool,
685        payload: Bytes::from(l.payload),
686        content_type: l.content_type,
687        kind: l.kind,
688        key: l.key,
689        tags: l.tags,
690        timeout: duration(l.timeout_s),
691        deadline: l
692            .deadline_unix_ns
693            .filter(|ns| *ns >= 0)
694            .map(|ns| UNIX_EPOCH + Duration::from_nanos(ns as u64)),
695        cancel: cancel.clone(),
696    };
697    let ttl = duration(l.lease_ttl_s).or(ttl);
698    tokio::spawn(async move {
699        let _slot = slot;
700        let _beat = ttl.map(|t| {
701            let tx = tx.clone();
702            AbortOnDrop(tokio::spawn(async move {
703                let mut iv = tokio::time::interval((t / 3).max(Duration::from_millis(10)));
704                iv.tick().await;
705                loop {
706                    iv.tick().await;
707                    if tx
708                        .send(pb::WorkerMessage {
709                            msg: Some(Msg::Heartbeat(pb::Heartbeat { lease })),
710                        })
711                        .await
712                        .is_err()
713                    {
714                        break;
715                    }
716                }
717            }))
718        });
719        let run = AssertUnwindSafe(handler(info)).catch_unwind();
720        let result = tokio::select! {
721            r = run => Some(r),
722            _ = cancel.cancelled() => None,
723        };
724        let msg = match result {
725            None if aborted.load(Ordering::SeqCst) => return,
726            None => Msg::Fail(pb::Fail {
727                lease,
728                message: "worker stopped before the attempt finished".into(),
729                permanent: false,
730                retry_after_s: None,
731            }),
732            Some(Err(panic)) => {
733                let what = panic
734                    .downcast_ref::<&str>()
735                    .map(|s| s.to_string())
736                    .or_else(|| panic.downcast_ref::<String>().cloned())
737                    .unwrap_or_else(|| "unknown panic".into());
738                tracing::error!(lease, panic = %what, "handler panicked");
739                Msg::Fail(pb::Fail {
740                    lease,
741                    message: format!("handler panicked: {what}"),
742                    permanent: false,
743                    retry_after_s: None,
744                })
745            }
746            Some(Ok(Ok(output))) if max_output > 0 && output.len() > max_output => {
747                Msg::Fail(pb::Fail {
748                    lease,
749                    message: format!(
750                        "output of {} bytes exceeds the server limit of {max_output} bytes",
751                        output.len()
752                    ),
753                    permanent: true,
754                    retry_after_s: None,
755                })
756            }
757            Some(Ok(Ok(output))) => Msg::Complete(pb::Complete {
758                lease,
759                output: output.to_vec(),
760                content_type: String::new(),
761            }),
762            Some(Ok(Err(TaskError::Permanent { message }))) => Msg::Fail(pb::Fail {
763                lease,
764                message: message.to_string(),
765                permanent: true,
766                retry_after_s: None,
767            }),
768            Some(Ok(Err(TaskError::Retryable {
769                message,
770                retry_after,
771            }))) => Msg::Fail(pb::Fail {
772                lease,
773                message: message.to_string(),
774                permanent: false,
775                retry_after_s: retry_after.map(|d| d.as_secs_f64()),
776            }),
777            Some(Ok(Err(TaskError::Cancelled))) => Msg::Cancelled(pb::Cancelled { lease }),
778        };
779        let _ = tx.send(pb::WorkerMessage { msg: Some(msg) }).await;
780    });
781}
782
783#[cfg(test)]
784mod tests {
785    use super::*;
786
787    fn refusal(code: Code, endpoint: Option<&'static str>) -> tonic::Status {
788        let mut s = tonic::Status::new(code, "pool 'llm' is served by shard 'gpu'");
789        if let Some(e) = endpoint {
790            s.metadata_mut()
791                .insert(ENDPOINT_HEADER, MetadataValue::from_static(e));
792        }
793        s
794    }
795
796    #[test]
797    fn redirects_never_weaken_the_connection() {
798        let opts = WorkerOptions::new("https://cpu.internal:7070", "llm").token("t");
799        let follow = |current: &str, s| redirect(&opts, current, &s);
800        let fp = Code::FailedPrecondition;
801        assert_eq!(
802            follow(
803                current(&opts),
804                refusal(fp, Some("https://gpu.internal:7070"))
805            )
806            .unwrap(),
807            Some("https://gpu.internal:7070".into())
808        );
809        // Not a redirect: another code, or no endpoint.
810        assert_eq!(
811            follow(current(&opts), refusal(Code::NotFound, Some("https://x:1"))).unwrap(),
812            None
813        );
814        assert_eq!(follow(current(&opts), refusal(fp, None)).unwrap(), None);
815        // No downgrade from https, not even to loopback.
816        assert!(follow(current(&opts), refusal(fp, Some("http://127.0.0.1:7070"))).is_err());
817        assert!(follow(current(&opts), refusal(fp, Some("unix:/tmp/s"))).is_err());
818        assert!(
819            follow(
820                current(&opts),
821                refusal(fp, Some("https://u:p@gpu.internal:7070"))
822            )
823            .is_err()
824        );
825
826        // From plain http on loopback, a token never goes over plain http to another host.
827        let local = WorkerOptions::new("http://127.0.0.1:7070", "llm").token("t");
828        let r = |t| redirect(&local, "http://127.0.0.1:7070", &refusal(fp, Some(t)));
829        assert!(r("http://127.0.0.1:7071").unwrap().is_some());
830        assert!(r("https://gpu.internal:7070").unwrap().is_some());
831        assert!(r("http://gpu.internal:7070").is_err());
832        // Userinfo does not make another host look like loopback.
833        assert!(r("http://127.0.0.1:1@gpu.internal:7070").is_err());
834        assert!(!is_loopback_endpoint(
835            "http://127.0.0.1:1@gpu.internal:7070"
836        ));
837        assert!(is_loopback_endpoint("http://[::1]:7070"));
838        // Without a token there is nothing to leak.
839        let anon = WorkerOptions::new("http://127.0.0.1:7070", "llm");
840        let to_remote = refusal(fp, Some("http://gpu.internal:7070"));
841        assert!(
842            redirect(&anon, "http://127.0.0.1:7070", &to_remote)
843                .unwrap()
844                .is_some()
845        );
846        // A plain-http endpoint on another host does not by itself allow more of them.
847        let remote = WorkerOptions::new("http://cpu.internal:7070", "llm").token("t");
848        assert!(redirect(&remote, "http://cpu.internal:7070", &to_remote).is_err());
849        // Only the opt-in does.
850        let trusted = remote.insecure_redirect_http(true);
851        assert!(
852            redirect(&trusted, "http://cpu.internal:7070", &to_remote)
853                .unwrap()
854                .is_some()
855        );
856    }
857
858    fn current(opts: &WorkerOptions) -> &str {
859        &opts.endpoint
860    }
861}
cronoq_server/worker.rs

cronoq_server/
worker.rs
