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