1use std::collections::BTreeMap;
7use std::time::Duration;
8
9#[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#[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
90pub 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#[derive(Clone, Debug, Default, PartialEq)]
319pub struct Checkpoint {
320 pub last_completed_step: Option<String>,
321 pub completed_steps: Vec<String>,
323 pub in_progress_step: Option<String>,
325 pub cursor_step: Option<String>,
326 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 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 pub const SAMPLE_LIMIT: i64 = 50_000;
356 pub const POSITION_LIMIT: i64 = 1_000;
358 pub const QUIET_PARTITION_LIMIT: i64 = 1_000;
360 pub const MAX_PAGE: u32 = 200;
362 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}