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