Skip to main content

headgate_shared/
lib.rs

1//! Dependency-light data types and utilities shared across headgate crates.
2//!
3//! This crate contains no store driver, network client, runtime, or exporter dependency,
4//! keeping it safe for core, adapters, and optional integrations to use as a leaf.
5
6use std::collections::BTreeMap;
7use std::time::Duration;
8
9pub mod log;
10
11/// Portable lifecycle result written by every worker runtime and store.
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum Outcome {
14    Success,
15    Retry,
16    Skip,
17    Revoke,
18    Snooze,
19    LeaseLost,
20    Undecodable,
21    RateLimited,
22}
23
24impl Outcome {
25    pub const fn as_str(self) -> &'static str {
26        match self {
27            Self::Success => "success",
28            Self::Retry => "retry",
29            Self::Skip => "skip",
30            Self::Revoke => "revoke",
31            Self::Snooze => "snooze",
32            Self::LeaseLost => "lease_lost",
33            Self::Undecodable => "undecodable",
34            Self::RateLimited => "rate_limited",
35        }
36    }
37
38    pub fn parse(value: &str) -> Option<Self> {
39        match value {
40            "success" => Some(Self::Success),
41            "retry" => Some(Self::Retry),
42            "skip" => Some(Self::Skip),
43            "revoke" => Some(Self::Revoke),
44            "snooze" => Some(Self::Snooze),
45            "lease_lost" => Some(Self::LeaseLost),
46            "undecodable" => Some(Self::Undecodable),
47            "rate_limited" => Some(Self::RateLimited),
48            _ => None,
49        }
50    }
51}
52
53/// What happens to periodic runs missed during downtime.
54#[derive(Clone, Copy, Debug, PartialEq, Eq)]
55pub enum MissedPolicy {
56    Skip,
57    RunOnce,
58    Backfill,
59}
60
61impl MissedPolicy {
62    pub const fn as_str(self) -> &'static str {
63        match self {
64            Self::Skip => "skip",
65            Self::RunOnce => "run_once",
66            Self::Backfill => "backfill",
67        }
68    }
69
70    pub fn parse(value: &str) -> Option<Self> {
71        match value {
72            "skip" => Some(Self::Skip),
73            "run_once" => Some(Self::RunOnce),
74            "backfill" => Some(Self::Backfill),
75            _ => None,
76        }
77    }
78}
79
80pub const DEFAULT_QUEUE: &str = "default";
81pub const DEFAULT_SCHEMA_VERSION: u32 = 1;
82pub const DEFAULT_MAX_ATTEMPTS: u32 = 25;
83pub const DEFAULT_WEIGHT: u32 = 1;
84pub const MAX_OPAQUE_SCHEMA_VERSION: u32 = i32::MAX as u32;
85
86pub fn normalize_queues(mut queues: Vec<String>) -> Vec<String> {
87    queues.sort();
88    queues.dedup();
89    queues
90}
91
92/// Validate the millisecond wire boundary without truncating a positive
93/// sub-millisecond duration or overflowing the signed store representation.
94pub fn duration_millis(duration: Duration) -> Option<i64> {
95    i64::try_from(duration.as_millis())
96        .ok()
97        .filter(|millis| *millis > 0)
98}
99
100pub fn effective_queue(queue: &str) -> &str {
101    if queue.is_empty() {
102        DEFAULT_QUEUE
103    } else {
104        queue
105    }
106}
107
108pub const fn effective_schema_version(version: u32) -> u32 {
109    if version == 0 {
110        DEFAULT_SCHEMA_VERSION
111    } else {
112        version
113    }
114}
115
116pub const fn effective_max_attempts(max_attempts: u32) -> u32 {
117    if max_attempts == 0 {
118        DEFAULT_MAX_ATTEMPTS
119    } else {
120        max_attempts
121    }
122}
123
124pub const fn effective_weight(weight: u32) -> u32 {
125    if weight == 0 { DEFAULT_WEIGHT } else { weight }
126}
127
128#[derive(Clone, Copy, Debug, Eq, PartialEq)]
129pub enum AckValidation {
130    Valid,
131    LeaseLost,
132    SnoozeDelayRequired,
133}
134
135pub const fn validate_ack(outcome: Outcome, delay_ms: Option<i64>) -> AckValidation {
136    if matches!(outcome, Outcome::LeaseLost) {
137        AckValidation::LeaseLost
138    } else if matches!(outcome, Outcome::Snooze) && !matches!(delay_ms, Some(delay) if delay > 0) {
139        AckValidation::SnoozeDelayRequired
140    } else {
141        AckValidation::Valid
142    }
143}
144
145#[derive(Clone, Copy, Debug, Eq, PartialEq)]
146pub enum OpaqueSchemaValidation {
147    Valid,
148    Zero,
149    TooLarge,
150}
151
152pub const fn validate_opaque_schema(version: u32) -> OpaqueSchemaValidation {
153    if version == 0 {
154        OpaqueSchemaValidation::Zero
155    } else if version > MAX_OPAQUE_SCHEMA_VERSION {
156        OpaqueSchemaValidation::TooLarge
157    } else {
158        OpaqueSchemaValidation::Valid
159    }
160}
161
162pub fn bulk_action_states(action: &str) -> Option<&'static [&'static str]> {
163    match action {
164        "retry" => Some(&["archived"]),
165        "cancel" => Some(&["scheduled", "available", "running"]),
166        "delete" => Some(&[
167            "scheduled",
168            "available",
169            "retryable",
170            "completed",
171            "archived",
172            "cancelled",
173            "quarantined",
174            "undecodable",
175        ]),
176        _ => None,
177    }
178}
179
180pub fn valid_worker_command(command: &str) -> bool {
181    matches!(
182        command,
183        "" | "quiet" | "resume" | "restart" | "terminate" | "resign"
184    )
185}
186
187pub fn format_generated_id(now_ms: u64, process_id: u32, sequence: u64) -> String {
188    format!(
189        "hg{now_ms:012x}{:05x}{:04x}",
190        process_id & 0xfffff,
191        sequence & 0xffff
192    )
193}
194
195#[derive(Clone, Debug, Default)]
196pub struct AdmissionFacts {
197    pub state: String,
198    pub now_ms: i64,
199    pub scheduled_at_ms: i64,
200    pub queue_paused: bool,
201    pub quarantined: bool,
202    pub fingerprint: String,
203    pub rate_class: String,
204    pub weight: i64,
205    pub tokens_available: Option<i64>,
206    pub tokens_ahead: i64,
207    pub limit_per_window: i64,
208    pub window_ms: i64,
209    pub max_concurrent: Option<i64>,
210    pub inflight: i64,
211    pub saturation: String,
212    pub position: i64,
213    pub deficit: i64,
214}
215
216#[derive(Clone, Debug, Eq, PartialEq)]
217pub struct AdmissionEvaluation {
218    pub admissible: bool,
219    pub blocked_by: Option<&'static str>,
220    pub detail: Vec<(String, String)>,
221    pub estimated_admission_ms: Option<i64>,
222}
223
224pub fn evaluate_admission(f: &AdmissionFacts) -> AdmissionEvaluation {
225    let mut result = AdmissionEvaluation {
226        admissible: false,
227        blocked_by: None,
228        detail: vec![("state".into(), f.state.clone())],
229        estimated_admission_ms: None,
230    };
231    let block = |mut value: AdmissionEvaluation, by, eta| {
232        value.blocked_by = Some(by);
233        value.estimated_admission_ms = eta;
234        value
235    };
236    match f.state.as_str() {
237        "running" => {
238            result.admissible = true;
239            result.estimated_admission_ms = Some(0);
240            return result;
241        }
242        "scheduled" | "retryable" => {
243            result
244                .detail
245                .push(("scheduled_at_ms".into(), f.scheduled_at_ms.to_string()));
246            return block(
247                result,
248                "schedule",
249                Some((f.scheduled_at_ms - f.now_ms).max(0)),
250            );
251        }
252        "quarantined" => return block(result, "quarantine", None),
253        "available" => {}
254        _ => return result,
255    }
256    if f.queue_paused {
257        return block(result, "queue_paused", None);
258    }
259    if f.scheduled_at_ms > f.now_ms {
260        result
261            .detail
262            .push(("scheduled_at_ms".into(), f.scheduled_at_ms.to_string()));
263        return block(result, "schedule", Some(f.scheduled_at_ms - f.now_ms));
264    }
265    if f.quarantined {
266        result
267            .detail
268            .push(("fingerprint".into(), f.fingerprint.clone()));
269        return block(result, "quarantine", None);
270    }
271    if !f.rate_class.is_empty() {
272        let weight = f.weight.max(1);
273        let required = f.tokens_ahead + weight;
274        result.detail.extend([
275            ("rate_class".into(), f.rate_class.clone()),
276            ("weight".into(), weight.to_string()),
277            ("tokens_ahead_in_class".into(), f.tokens_ahead.to_string()),
278        ]);
279        if let Some(available) = f.tokens_available {
280            result
281                .detail
282                .push(("tokens_available".into(), available.to_string()));
283            if available < required {
284                let eta = (f.limit_per_window > 0)
285                    .then(|| (required - available).max(1) * f.window_ms / f.limit_per_window);
286                return block(result, "rate_class", eta);
287            }
288        } else {
289            result.detail.push((
290                "tokens_available".into(),
291                "unlimited (no such rate class)".into(),
292            ));
293        }
294    }
295    if let Some(max_concurrent) = f.max_concurrent {
296        let strategy = if f.saturation.is_empty() {
297            "queue"
298        } else {
299            &f.saturation
300        };
301        result.detail.extend([
302            ("max_concurrent".into(), max_concurrent.to_string()),
303            ("inflight".into(), f.inflight.to_string()),
304            ("on_saturated".into(), strategy.into()),
305        ]);
306        if f.inflight >= max_concurrent && strategy != "cancel_running" {
307            return block(result, "concurrency_limit", None);
308        }
309    }
310    result.detail.extend([
311        ("position_in_partition".into(), f.position.to_string()),
312        ("partition_deficit".into(), f.deficit.to_string()),
313    ]);
314    result.admissible = true;
315    result.estimated_admission_ms = Some(0);
316    result
317}
318
319/// Durable progress within a resumable job.
320#[derive(Clone, Debug, Default, PartialEq)]
321pub struct Checkpoint {
322    pub last_completed_step: Option<String>,
323    /// Completed steps in execution order. Replay compares them positionally.
324    pub completed_steps: Vec<String>,
325    /// The step recorded before its side effects began.
326    pub in_progress_step: Option<String>,
327    pub cursor_step: Option<String>,
328    /// Stored outside checkpoint JSON to avoid base64 encoding native binary data.
329    pub cursor: Option<Vec<u8>>,
330    pub schema_version: u32,
331    pub step_set_hash: String,
332    pub crashes_by_step: Vec<(String, u32)>,
333}
334
335#[derive(Debug, Clone, Copy, PartialEq, Eq)]
336pub enum Resume {
337    Continue,
338    Remapped,
339    Undecodable,
340}
341
342impl Checkpoint {
343    /// Decide whether a checkpoint can continue against the current task definition.
344    pub fn resumability(&self, current_version: u32, current_step_set_hash: &str) -> Resume {
345        if self.step_set_hash.is_empty() || self.step_set_hash == current_step_set_hash {
346            Resume::Continue
347        } else if self.schema_version != current_version {
348            Resume::Remapped
349        } else {
350            Resume::Undecodable
351        }
352    }
353}
354
355pub mod inspection {
356    /// Largest row sample used by an aggregate inspection query.
357    pub const SAMPLE_LIMIT: i64 = 50_000;
358    /// Largest sample used to estimate a job's queue position.
359    pub const POSITION_LIMIT: i64 = 1_000;
360    /// Largest quiet-partition set inspected in one request.
361    pub const QUIET_PARTITION_LIMIT: i64 = 1_000;
362    /// Largest list page exposed by an inspection adapter.
363    pub const MAX_PAGE: u32 = 200;
364    /// Largest per-queue sample used for memory estimates.
365    pub const MEMORY_SAMPLE_LIMIT: u32 = 1_000;
366
367    pub const fn age_ms(now_ms: i64, at_ms: i64) -> i64 {
368        let age = now_ms - at_ms;
369        if age > 0 { age } else { 0 }
370    }
371
372    pub fn time_to_drain_ms(backlog: i64, arrival_rate: f64, drain_rate: f64) -> Option<i64> {
373        (drain_rate > arrival_rate && drain_rate > 0.0)
374            .then(|| (backlog as f64 / (drain_rate - arrival_rate) * 1000.0) as i64)
375    }
376}
377
378pub mod codec {
379    use super::{BTreeMap, Checkpoint};
380
381    pub fn encode_string_list(values: &[String]) -> String {
382        serde_json::to_string(values).unwrap_or_else(|_| "[]".into())
383    }
384
385    pub fn decode_string_list(encoded: &str) -> Vec<String> {
386        serde_json::from_str(encoded).unwrap_or_default()
387    }
388
389    pub fn encode_checkpoint_value(checkpoint: &Checkpoint) -> serde_json::Value {
390        let mut object = serde_json::Map::new();
391        if !checkpoint.completed_steps.is_empty() {
392            object.insert(
393                "completed".into(),
394                checkpoint.completed_steps.clone().into(),
395            );
396        }
397        if let Some(step) = &checkpoint.in_progress_step {
398            object.insert("in_progress".into(), step.clone().into());
399        }
400        if let Some(step) = &checkpoint.cursor_step {
401            object.insert("cursor_step".into(), step.clone().into());
402        }
403        if checkpoint.schema_version != 0 {
404            object.insert("version".into(), checkpoint.schema_version.into());
405        }
406        if !checkpoint.step_set_hash.is_empty() {
407            object.insert("hash".into(), checkpoint.step_set_hash.clone().into());
408        }
409        if !checkpoint.crashes_by_step.is_empty() {
410            let crashes = checkpoint
411                .crashes_by_step
412                .iter()
413                .map(|(step, count)| (step.clone(), (*count).into()))
414                .collect();
415            object.insert("crashes".into(), serde_json::Value::Object(crashes));
416        }
417        serde_json::Value::Object(object)
418    }
419
420    pub fn encode_checkpoint_json(checkpoint: &Checkpoint) -> String {
421        encode_checkpoint_value(checkpoint).to_string()
422    }
423
424    pub fn decode_checkpoint_value(
425        value: Option<serde_json::Value>,
426        cursor: Option<Vec<u8>>,
427    ) -> Checkpoint {
428        let mut checkpoint = Checkpoint {
429            cursor,
430            ..Default::default()
431        };
432        let Some(serde_json::Value::Object(object)) = value else {
433            return checkpoint;
434        };
435        if let Some(serde_json::Value::Array(completed)) = object.get("completed") {
436            checkpoint.completed_steps = completed
437                .iter()
438                .filter_map(|step| step.as_str().map(String::from))
439                .collect();
440            checkpoint.last_completed_step = checkpoint.completed_steps.last().cloned();
441        }
442        checkpoint.in_progress_step = object
443            .get("in_progress")
444            .and_then(|step| step.as_str())
445            .map(String::from);
446        checkpoint.cursor_step = object
447            .get("cursor_step")
448            .and_then(|step| step.as_str())
449            .map(String::from);
450        checkpoint.schema_version = object
451            .get("version")
452            .and_then(serde_json::Value::as_u64)
453            .unwrap_or(0) as u32;
454        checkpoint.step_set_hash = object
455            .get("hash")
456            .and_then(serde_json::Value::as_str)
457            .unwrap_or("")
458            .to_owned();
459        if let Some(serde_json::Value::Object(crashes)) = object.get("crashes") {
460            checkpoint.crashes_by_step = crashes
461                .iter()
462                .map(|(step, count)| (step.clone(), count.as_u64().unwrap_or(0) as u32))
463                .collect();
464        }
465        checkpoint
466    }
467
468    pub fn decode_checkpoint_bytes(json: Option<&[u8]>, cursor: Option<Vec<u8>>) -> Checkpoint {
469        let value = json.and_then(|bytes| serde_json::from_slice(bytes).ok());
470        decode_checkpoint_value(value, cursor)
471    }
472
473    pub fn decode_checkpoint_str(json: Option<&str>, cursor: Option<Vec<u8>>) -> Checkpoint {
474        decode_checkpoint_bytes(json.map(str::as_bytes), cursor)
475    }
476
477    pub fn encode_headers_value(headers: &BTreeMap<String, String>) -> serde_json::Value {
478        serde_json::Value::Object(
479            headers
480                .iter()
481                .map(|(key, value)| (key.clone(), serde_json::Value::String(value.clone())))
482                .collect(),
483        )
484    }
485
486    pub fn encode_headers_json(headers: &BTreeMap<String, String>, omit_empty: bool) -> String {
487        if omit_empty && headers.is_empty() {
488            String::new()
489        } else {
490            encode_headers_value(headers).to_string()
491        }
492    }
493
494    pub fn decode_headers_value(value: Option<serde_json::Value>) -> BTreeMap<String, String> {
495        let Some(serde_json::Value::Object(object)) = value else {
496            return BTreeMap::new();
497        };
498        object
499            .into_iter()
500            .filter_map(|(key, value)| match value {
501                serde_json::Value::String(text) => Some((key, text)),
502                _ => None,
503            })
504            .collect()
505    }
506
507    pub fn decode_headers_bytes(json: Option<&[u8]>) -> BTreeMap<String, String> {
508        let value = json.and_then(|bytes| serde_json::from_slice(bytes).ok());
509        decode_headers_value(value)
510    }
511
512    pub fn decode_headers_str(json: Option<&str>) -> BTreeMap<String, String> {
513        decode_headers_bytes(json.map(str::as_bytes))
514    }
515}
516
517#[cfg(test)]
518mod tests {
519    use super::{AdmissionFacts, Checkpoint, Outcome, Resume, codec};
520
521    #[test]
522    fn checkpoint_codec_has_a_stable_wire_shape() {
523        let checkpoint = Checkpoint {
524            completed_steps: vec!["fetch".into(), "transform".into()],
525            in_progress_step: Some("publish".into()),
526            cursor_step: Some("transform".into()),
527            cursor: Some(b"opaque".to_vec()),
528            schema_version: 2,
529            step_set_hash: "steps-v2".into(),
530            crashes_by_step: vec![("publish".into(), 1)],
531            ..Default::default()
532        };
533        let encoded = codec::encode_checkpoint_json(&checkpoint);
534        assert_eq!(
535            encoded,
536            r#"{"completed":["fetch","transform"],"crashes":{"publish":1},"cursor_step":"transform","hash":"steps-v2","in_progress":"publish","version":2}"#
537        );
538        let mut expected = checkpoint.clone();
539        expected.last_completed_step = Some("transform".into());
540        assert_eq!(
541            codec::decode_checkpoint_str(Some(&encoded), checkpoint.cursor.clone()),
542            expected
543        );
544    }
545
546    #[test]
547    fn malformed_checkpoint_preserves_cursor() {
548        let checkpoint = codec::decode_checkpoint_str(Some("{"), Some(b"cursor".to_vec()));
549        assert_eq!(checkpoint.cursor.as_deref(), Some(b"cursor".as_slice()));
550        assert!(checkpoint.completed_steps.is_empty());
551    }
552
553    #[test]
554    fn resumability_remains_conservative() {
555        let checkpoint = Checkpoint {
556            schema_version: 1,
557            step_set_hash: "old".into(),
558            ..Default::default()
559        };
560        assert_eq!(checkpoint.resumability(1, "old"), Resume::Continue);
561        assert_eq!(checkpoint.resumability(2, "new"), Resume::Remapped);
562        assert_eq!(checkpoint.resumability(1, "new"), Resume::Undecodable);
563    }
564
565    #[test]
566    fn policy_and_admission_rules_are_shared() {
567        for raw in [
568            "success",
569            "retry",
570            "skip",
571            "revoke",
572            "snooze",
573            "lease_lost",
574            "undecodable",
575            "rate_limited",
576        ] {
577            let outcome = Outcome::parse(raw).expect("known outcome");
578            assert_eq!(outcome.as_str(), raw);
579        }
580        assert_eq!(
581            super::bulk_action_states("cancel"),
582            Some(["scheduled", "available", "running"].as_slice())
583        );
584        let evaluation = super::evaluate_admission(&AdmissionFacts {
585            state: "available".into(),
586            rate_class: "api".into(),
587            weight: 3,
588            tokens_available: Some(2),
589            limit_per_window: 1,
590            window_ms: 1_000,
591            ..Default::default()
592        });
593        assert_eq!(evaluation.blocked_by, Some("rate_class"));
594        assert_eq!(evaluation.estimated_admission_ms, Some(1_000));
595        assert_eq!(
596            super::inspection::time_to_drain_ms(10, 2.0, 4.0),
597            Some(5_000)
598        );
599    }
600}