1use runledger_core::jobs::{JobContext, JobDeadLetterInfo, JobDeadLetterReason, JobFailure};
2use std::collections::HashMap;
3use std::sync::Arc;
4use tokio::sync::watch;
5use tokio::task::{Id, JoinSet};
6use tokio::time::{Duration, sleep, timeout};
7use tracing::{info, warn};
8
9use crate::ReaperError;
10use crate::RuntimeLoopExit;
11use crate::config::JobsConfig;
12use crate::registry::JobRegistry;
13
14const REAPER_WORKER_ID: &str = "reaper";
15const LEASE_EXPIRED_CODE: &str = "job.lease_expired";
16const LEASE_EXPIRED_MESSAGE: &str = "Job lease expired before completion.";
17const REAPER_TERMINAL_HOOK_MAX_CONCURRENCY: usize = 8;
18#[cfg(test)]
19const TERMINAL_HOOK_TIMEOUT: Duration = Duration::from_millis(100);
20#[cfg(not(test))]
21const TERMINAL_HOOK_TIMEOUT: Duration = Duration::from_secs(10);
22#[cfg(test)]
23const TERMINAL_HOOK_ABORT_DRAIN_TIMEOUT: Duration = Duration::from_millis(50);
24#[cfg(not(test))]
25const TERMINAL_HOOK_ABORT_DRAIN_TIMEOUT: Duration = Duration::from_millis(250);
26
27pub async fn run_reaper_loop(
28 pool: runledger_postgres::DbPool,
29 registry: JobRegistry,
30 config: JobsConfig,
31 mut shutdown: watch::Receiver<bool>,
32) -> RuntimeLoopExit {
33 if let Err(error) = config.validate_reaper_loop() {
34 warn!(%error, "invalid jobs config; stopping reaper loop");
35 return RuntimeLoopExit::InvalidConfig(error);
36 }
37
38 let registry = Arc::new(registry);
39
40 loop {
41 if shutdown_requested_or_closed(&shutdown) {
42 return reaper_shutdown_complete();
43 }
44
45 match runledger_postgres::jobs::reap_expired_leases_with_diagnostics(
46 &pool,
47 config.claim_batch_size,
48 config.reaper_retry_delay_ms,
49 )
50 .await
51 {
52 Ok(result) => {
53 if result.summary.processed > 0 {
54 info!(
55 reaped = result.summary.processed,
56 "reaper reclaimed expired leases"
57 );
58 }
59 log_deferred_row_errors(&result);
60
61 let fanout_result = notify_handlers_of_terminal_lease_expirations(
62 registry.as_ref(),
63 &result.summary.terminal_dead_lettered,
64 &mut shutdown,
65 )
66 .await;
67 if matches!(
68 fanout_result,
69 TerminalHookFanoutResult::InterruptedByShutdown
70 ) {
71 return reaper_shutdown_complete();
72 }
73 }
74 Err(error) => {
75 let error = ReaperError::ReapExpiredLeases {
76 batch_size: config.claim_batch_size,
77 retry_delay_ms: config.reaper_retry_delay_ms,
78 source: error,
79 };
80 warn!(%error, "reaper iteration failed");
81 }
82 }
83
84 if wait_for_shutdown_or_poll(&mut shutdown, config.reaper_interval).await {
85 return reaper_shutdown_complete();
86 }
87 }
88}
89
90fn log_deferred_row_errors(result: &runledger_postgres::jobs::ReapExpiredLeasesDetailedResult) {
91 for error in &result.deferred_row_errors {
92 warn!(
93 job_id = %error.job_id,
94 run_number = error.run_number,
95 attempt = error.attempt,
96 error_code = %error.error_code,
97 error_message = %error.error_message,
98 error_sqlstate = error.sqlstate.as_deref().unwrap_or(""),
99 deferred_row_error_count = result.deferred_row_error_count,
100 logged_deferred_row_errors = result.deferred_row_errors.len(),
101 "reaper deferred expired leased job after row-level processing error"
102 );
103 }
104
105 if result.deferred_row_error_count > result.deferred_row_errors.len() {
106 warn!(
107 deferred_row_error_count = result.deferred_row_error_count,
108 logged_deferred_row_errors = result.deferred_row_errors.len(),
109 "reaper deferred additional expired leased jobs after row-level processing errors"
110 );
111 }
112}
113
114fn reaper_shutdown_complete() -> RuntimeLoopExit {
115 info!("reaper shutdown complete");
116 RuntimeLoopExit::Shutdown
117}
118
119fn shutdown_requested_or_closed(shutdown: &watch::Receiver<bool>) -> bool {
120 *shutdown.borrow() || shutdown.has_changed().is_err()
121}
122
123async fn wait_for_shutdown_or_poll(
124 shutdown: &mut watch::Receiver<bool>,
125 poll_interval: Duration,
126) -> bool {
127 tokio::select! {
128 changed = shutdown.changed() => changed.is_err() || *shutdown.borrow(),
129 _ = sleep(poll_interval) => false,
130 }
131}
132
133async fn notify_handlers_of_terminal_lease_expirations(
134 registry: &JobRegistry,
135 jobs: &[runledger_postgres::jobs::ReapedTerminalLeaseRecord],
136 shutdown: &mut watch::Receiver<bool>,
137) -> TerminalHookFanoutResult {
138 let mut in_flight = JoinSet::new();
139 let mut metadata: HashMap<Id, HookMetadata> = HashMap::new();
140
141 for job in jobs {
142 if *shutdown.borrow() {
143 abort_terminal_hook_fanout(&mut in_flight, &mut metadata).await;
144 return TerminalHookFanoutResult::InterruptedByShutdown;
145 }
146
147 let Some(handler) = registry.get(job.job_type.as_borrowed()) else {
148 continue;
149 };
150
151 let context = JobContext {
152 job_id: job.job_id,
153 run_number: job.run_number,
154 attempt: job.attempt,
155 organization_id: job.organization_id,
156 worker_id: REAPER_WORKER_ID.to_string(),
157 };
158 let payload = job.payload.clone();
159 let dead_letter = JobDeadLetterInfo::new(
160 lease_expired_failure(),
161 JobDeadLetterReason::LeaseExpired,
162 None,
163 );
164 let hook_meta = HookMetadata {
165 job_id: job.job_id.to_string(),
166 job_type: job.job_type.to_string(),
167 run_number: job.run_number,
168 attempt: job.attempt,
169 };
170
171 let abort_handle = in_flight.spawn(async move {
172 if tokio::time::timeout(
173 TERMINAL_HOOK_TIMEOUT,
174 handler.on_dead_letter(context, payload, dead_letter),
175 )
176 .await
177 .is_ok()
178 {
179 HookOutcome::Completed
180 } else {
181 HookOutcome::TimedOut
182 }
183 });
184 metadata.insert(abort_handle.id(), hook_meta);
185
186 if in_flight.len() >= REAPER_TERMINAL_HOOK_MAX_CONCURRENCY
187 && matches!(
188 await_next_hook_result_or_shutdown(&mut in_flight, &mut metadata, shutdown).await,
189 NextHookResult::InterruptedByShutdown
190 )
191 {
192 abort_terminal_hook_fanout(&mut in_flight, &mut metadata).await;
193 return TerminalHookFanoutResult::InterruptedByShutdown;
194 }
195 }
196
197 while !in_flight.is_empty() {
198 if matches!(
199 await_next_hook_result_or_shutdown(&mut in_flight, &mut metadata, shutdown).await,
200 NextHookResult::InterruptedByShutdown
201 ) {
202 abort_terminal_hook_fanout(&mut in_flight, &mut metadata).await;
203 return TerminalHookFanoutResult::InterruptedByShutdown;
204 }
205 }
206
207 if !metadata.is_empty() {
208 warn!(
209 stale_hook_metadata_entries = metadata.len(),
210 "reaper hook metadata diverged from in-flight task set; clearing stale metadata"
211 );
212 metadata.clear();
213 }
214
215 TerminalHookFanoutResult::Completed
216}
217
218fn lease_expired_failure() -> JobFailure {
219 JobFailure::lease_expired(LEASE_EXPIRED_CODE, LEASE_EXPIRED_MESSAGE)
220}
221
222#[derive(Debug)]
223enum HookOutcome {
224 Completed,
225 TimedOut,
226}
227
228#[derive(Debug, PartialEq, Eq)]
229enum TerminalHookFanoutResult {
230 Completed,
231 InterruptedByShutdown,
232}
233
234#[derive(Debug, PartialEq, Eq)]
235enum NextHookResult {
236 HookSettled,
237 InterruptedByShutdown,
238}
239
240#[derive(Debug)]
241struct HookMetadata {
242 job_id: String,
243 job_type: String,
244 run_number: i32,
245 attempt: i32,
246}
247
248type HookJoinResult = std::result::Result<(Id, HookOutcome), tokio::task::JoinError>;
249
250async fn await_next_hook_result_or_shutdown(
251 in_flight: &mut JoinSet<HookOutcome>,
252 metadata: &mut HashMap<Id, HookMetadata>,
253 shutdown: &mut watch::Receiver<bool>,
254) -> NextHookResult {
255 if *shutdown.borrow() {
256 return NextHookResult::InterruptedByShutdown;
257 }
258
259 loop {
260 tokio::select! {
261 changed = shutdown.changed() => {
262 if changed.is_err() || *shutdown.borrow() {
263 return NextHookResult::InterruptedByShutdown;
264 }
265 }
266 result = in_flight.join_next_with_id() => {
267 if let Some(result) = result {
268 handle_next_hook_result(result, metadata);
269 }
270 return NextHookResult::HookSettled;
271 }
272 }
273 }
274}
275
276async fn abort_terminal_hook_fanout(
277 in_flight: &mut JoinSet<HookOutcome>,
278 metadata: &mut HashMap<Id, HookMetadata>,
279) {
280 if in_flight.is_empty() {
281 if !metadata.is_empty() {
282 warn!(
283 stale_hook_metadata_entries = metadata.len(),
284 "reaper hook metadata diverged while shutdown fanout abort had no in-flight tasks"
285 );
286 metadata.clear();
287 }
288 return;
289 }
290
291 warn!(
292 in_flight_terminal_hooks = in_flight.len(),
293 drain_timeout_ms = TERMINAL_HOOK_ABORT_DRAIN_TIMEOUT.as_millis(),
294 "shutdown requested while terminal failure hooks are running; waiting briefly for in-flight hooks to finish"
295 );
296
297 let drain_result = timeout(
298 TERMINAL_HOOK_ABORT_DRAIN_TIMEOUT,
299 drain_terminal_hook_results(in_flight, metadata),
300 )
301 .await;
302 match drain_result {
303 Ok(()) => {}
304 Err(_) => {
305 warn!(
306 remaining_in_flight_hooks = in_flight.len(),
307 drain_timeout_ms = TERMINAL_HOOK_ABORT_DRAIN_TIMEOUT.as_millis(),
308 "terminal hooks did not finish before shutdown deadline; aborting remaining hooks"
309 );
310 in_flight.abort_all();
311
312 let abort_drain_result = timeout(
313 TERMINAL_HOOK_ABORT_DRAIN_TIMEOUT,
314 drain_aborted_terminal_hook_results(in_flight, metadata),
315 )
316 .await;
317
318 match abort_drain_result {
319 Ok(cancelled_hook_count) => {
320 if cancelled_hook_count > 0 {
321 info!(
322 cancelled_hook_count,
323 "terminal failure hooks cancelled due to shutdown request"
324 );
325 }
326 }
327 Err(_) => {
328 warn!(
329 remaining_in_flight_hooks = in_flight.len(),
330 undrained_hook_metadata_entries = metadata.len(),
331 drain_timeout_ms = TERMINAL_HOOK_ABORT_DRAIN_TIMEOUT.as_millis(),
332 "terminal hook abort drain timed out during shutdown; dropping unresolved hook tasks"
333 );
334 }
335 }
336 }
337 }
338
339 if !metadata.is_empty() {
340 warn!(
341 stale_hook_metadata_entries = metadata.len(),
342 "reaper hook metadata remains after shutdown fanout abort; clearing stale metadata"
343 );
344 metadata.clear();
345 }
346}
347
348async fn drain_terminal_hook_results(
349 in_flight: &mut JoinSet<HookOutcome>,
350 metadata: &mut HashMap<Id, HookMetadata>,
351) {
352 while let Some(result) = in_flight.join_next_with_id().await {
353 handle_next_hook_result(result, metadata);
354 }
355}
356
357async fn drain_aborted_terminal_hook_results(
358 in_flight: &mut JoinSet<HookOutcome>,
359 metadata: &mut HashMap<Id, HookMetadata>,
360) -> usize {
361 let mut cancelled_hook_count: usize = 0;
362
363 while let Some(result) = in_flight.join_next_with_id().await {
364 match result {
365 Err(error) if error.is_cancelled() => {
366 let id = error.id();
367 if metadata.remove(&id).is_none() {
368 warn!(
369 "terminal failure hook cancellation observed during shutdown; metadata missing in reaper loop"
370 );
371 }
372 cancelled_hook_count += 1;
373 }
374 other => handle_next_hook_result(other, metadata),
375 }
376 }
377
378 cancelled_hook_count
379}
380
381fn handle_next_hook_result(result: HookJoinResult, metadata: &mut HashMap<Id, HookMetadata>) {
382 match result {
383 Ok((id, HookOutcome::Completed)) => {
384 if metadata.remove(&id).is_none() {
385 warn!("terminal failure hook completed; metadata missing in reaper loop");
386 }
387 }
388 Ok((id, HookOutcome::TimedOut)) => {
389 if let Some(meta) = metadata.remove(&id) {
390 warn!(
391 job_id = meta.job_id,
392 job_type = meta.job_type,
393 run_number = meta.run_number,
394 attempt = meta.attempt,
395 timeout_ms = TERMINAL_HOOK_TIMEOUT.as_millis(),
396 "terminal failure hook timed out; continuing reaper loop"
397 );
398 } else {
399 warn!(
400 timeout_ms = TERMINAL_HOOK_TIMEOUT.as_millis(),
401 "terminal failure hook timed out; metadata missing in reaper loop"
402 );
403 }
404 }
405 Err(error) => {
406 let id = error.id();
407 if let Some(meta) = metadata.remove(&id) {
408 if error.is_panic() {
409 warn!(
410 job_id = meta.job_id,
411 job_type = meta.job_type,
412 run_number = meta.run_number,
413 attempt = meta.attempt,
414 error = %error,
415 "terminal failure hook panicked; continuing reaper loop"
416 );
417 } else if error.is_cancelled() {
418 warn!(
419 job_id = meta.job_id,
420 job_type = meta.job_type,
421 run_number = meta.run_number,
422 attempt = meta.attempt,
423 error = %error,
424 "terminal failure hook was cancelled; continuing reaper loop"
425 );
426 } else {
427 warn!(
428 job_id = meta.job_id,
429 job_type = meta.job_type,
430 run_number = meta.run_number,
431 attempt = meta.attempt,
432 error = %error,
433 "terminal failure hook join failed; continuing reaper loop"
434 );
435 }
436 } else if error.is_panic() {
437 warn!(
438 error = %error,
439 "terminal failure hook panicked; metadata missing in reaper loop"
440 );
441 } else if error.is_cancelled() {
442 warn!(
443 error = %error,
444 "terminal failure hook was cancelled; metadata missing in reaper loop"
445 );
446 } else {
447 warn!(
448 error = %error,
449 "terminal failure hook join failed; metadata missing in reaper loop"
450 );
451 }
452 }
453 }
454}
455
456#[cfg(test)]
457mod tests {
458 use std::future::pending;
459 use std::sync::atomic::{AtomicUsize, Ordering};
460 use std::sync::{Arc, Mutex};
461 use std::time::Duration;
462
463 use runledger_core::jobs::{
464 JobCompletion, JobContext, JobDeadLetterInfo, JobDeadLetterReason, JobFailure, JobType,
465 };
466 use runledger_postgres::jobs::ReapedTerminalLeaseRecord;
467 use serde_json::{Value, json};
468 use sqlx::types::Uuid;
469 use tokio::sync::{Notify, watch};
470 use tokio::time::{sleep, timeout};
471
472 use crate::registry::{JobHandler, JobRegistry};
473
474 use super::{TerminalHookFanoutResult, notify_handlers_of_terminal_lease_expirations};
475
476 struct HangingTerminalHookHandler;
477 struct ControlledTerminalHookHandler {
478 started: Arc<Notify>,
479 release: Arc<Notify>,
480 completions: Arc<AtomicUsize>,
481 }
482
483 struct RecordingDeadLetterHandler {
484 dead_letters: Arc<Mutex<Vec<JobDeadLetterInfo>>>,
485 }
486
487 #[async_trait::async_trait]
488 impl JobHandler for HangingTerminalHookHandler {
489 fn job_type(&self) -> JobType<'static> {
490 JobType::new("jobs.test.reaper.hook.hang")
491 }
492
493 async fn execute(
494 &self,
495 _context: JobContext,
496 _payload: Value,
497 ) -> Result<JobCompletion, JobFailure> {
498 Ok(JobCompletion::success())
499 }
500
501 async fn on_dead_letter(
502 &self,
503 _context: JobContext,
504 _payload: Value,
505 _dead_letter: JobDeadLetterInfo,
506 ) {
507 pending::<()>().await;
508 }
509 }
510
511 #[async_trait::async_trait]
512 impl JobHandler for ControlledTerminalHookHandler {
513 fn job_type(&self) -> JobType<'static> {
514 JobType::new("jobs.test.reaper.hook.controlled")
515 }
516
517 async fn execute(
518 &self,
519 _context: JobContext,
520 _payload: Value,
521 ) -> Result<JobCompletion, JobFailure> {
522 Ok(JobCompletion::success())
523 }
524
525 async fn on_dead_letter(
526 &self,
527 _context: JobContext,
528 _payload: Value,
529 _dead_letter: JobDeadLetterInfo,
530 ) {
531 self.started.notify_one();
532 self.release.notified().await;
533 self.completions.fetch_add(1, Ordering::SeqCst);
534 }
535 }
536
537 #[async_trait::async_trait]
538 impl JobHandler for RecordingDeadLetterHandler {
539 fn job_type(&self) -> JobType<'static> {
540 JobType::new("jobs.test.reaper.dead_letter.record")
541 }
542
543 async fn execute(
544 &self,
545 _context: JobContext,
546 _payload: Value,
547 ) -> Result<JobCompletion, JobFailure> {
548 Ok(JobCompletion::success())
549 }
550
551 async fn on_dead_letter(
552 &self,
553 _context: JobContext,
554 _payload: Value,
555 dead_letter: JobDeadLetterInfo,
556 ) {
557 self.dead_letters
558 .lock()
559 .expect("dead-letter list lock should not be poisoned")
560 .push(dead_letter);
561 }
562 }
563
564 #[tokio::test]
565 async fn notify_handlers_survives_terminal_hook_timeout() {
566 let mut registry = JobRegistry::new();
567 registry.register(HangingTerminalHookHandler);
568 let (_shutdown_tx, mut shutdown_rx) = watch::channel(false);
569
570 let jobs = vec![ReapedTerminalLeaseRecord {
571 job_id: Uuid::now_v7(),
572 job_type: runledger_core::jobs::JobTypeName::from_static("jobs.test.reaper.hook.hang"),
573 organization_id: None,
574 run_number: 1,
575 attempt: 1,
576 payload: json!({ "kind": "hook-timeout" }),
577 }];
578
579 let result = timeout(
580 Duration::from_secs(2),
581 notify_handlers_of_terminal_lease_expirations(®istry, &jobs, &mut shutdown_rx),
582 )
583 .await
584 .expect("notification pass should return even when terminal hook hangs");
585 assert_eq!(result, TerminalHookFanoutResult::Completed);
586 }
587
588 #[tokio::test]
589 async fn notify_handlers_shutdown_interrupts_terminal_hook_fanout() {
590 let mut registry = JobRegistry::new();
591 registry.register(HangingTerminalHookHandler);
592 let (shutdown_tx, mut shutdown_rx) = watch::channel(false);
593
594 let jobs: Vec<_> = (0..256)
595 .map(|idx| ReapedTerminalLeaseRecord {
596 job_id: Uuid::now_v7(),
597 job_type: runledger_core::jobs::JobTypeName::from_static(
598 "jobs.test.reaper.hook.hang",
599 ),
600 organization_id: None,
601 run_number: 1,
602 attempt: 1,
603 payload: json!({ "kind": "hook-shutdown", "index": idx }),
604 })
605 .collect();
606
607 tokio::spawn(async move {
608 sleep(Duration::from_millis(20)).await;
609 let _ = shutdown_tx.send(true);
610 });
611
612 let result = timeout(
613 Duration::from_secs(2),
614 notify_handlers_of_terminal_lease_expirations(®istry, &jobs, &mut shutdown_rx),
615 )
616 .await
617 .expect("notification pass should return quickly after shutdown signal");
618 assert_eq!(result, TerminalHookFanoutResult::InterruptedByShutdown);
619 }
620
621 #[tokio::test]
622 async fn notify_handlers_ignores_non_shutdown_watch_updates() {
623 let started = Arc::new(Notify::new());
624 let release = Arc::new(Notify::new());
625 let completions = Arc::new(AtomicUsize::new(0));
626 let mut registry = JobRegistry::new();
627 registry.register(ControlledTerminalHookHandler {
628 started: started.clone(),
629 release: release.clone(),
630 completions: completions.clone(),
631 });
632 let (shutdown_tx, mut shutdown_rx) = watch::channel(false);
633 let update_tx = shutdown_tx.clone();
634
635 let jobs = vec![ReapedTerminalLeaseRecord {
636 job_id: Uuid::now_v7(),
637 job_type: runledger_core::jobs::JobTypeName::from_static(
638 "jobs.test.reaper.hook.controlled",
639 ),
640 organization_id: None,
641 run_number: 1,
642 attempt: 1,
643 payload: json!({ "kind": "non-shutdown-watch-update" }),
644 }];
645
646 tokio::spawn(async move {
647 started.notified().await;
648 let _ = update_tx.send(false);
649 sleep(Duration::from_millis(20)).await;
650 release.notify_waiters();
651 });
652
653 let result =
654 notify_handlers_of_terminal_lease_expirations(®istry, &jobs, &mut shutdown_rx).await;
655 drop(shutdown_tx);
656
657 assert_eq!(result, TerminalHookFanoutResult::Completed);
658 assert_eq!(completions.load(Ordering::SeqCst), 1);
659 }
660
661 #[tokio::test]
662 async fn notify_handlers_reports_lease_expiration_dead_letter_reason() {
663 let dead_letters = Arc::new(Mutex::new(Vec::new()));
664 let mut registry = JobRegistry::new();
665 registry.register(RecordingDeadLetterHandler {
666 dead_letters: dead_letters.clone(),
667 });
668 let (_shutdown_tx, mut shutdown_rx) = watch::channel(false);
669
670 let jobs = vec![ReapedTerminalLeaseRecord {
671 job_id: Uuid::now_v7(),
672 job_type: runledger_core::jobs::JobTypeName::from_static(
673 "jobs.test.reaper.dead_letter.record",
674 ),
675 organization_id: None,
676 run_number: 1,
677 attempt: 1,
678 payload: json!({ "kind": "lease-expired" }),
679 }];
680
681 let result =
682 notify_handlers_of_terminal_lease_expirations(®istry, &jobs, &mut shutdown_rx).await;
683 assert_eq!(result, TerminalHookFanoutResult::Completed);
684
685 let dead_letters = dead_letters
686 .lock()
687 .expect("dead-letter list lock should not be poisoned");
688 assert_eq!(dead_letters.len(), 1);
689 let dead_letter = &dead_letters[0];
690 assert_eq!(dead_letter.reason, JobDeadLetterReason::LeaseExpired);
691 assert_eq!(
692 dead_letter.failure.kind,
693 runledger_core::jobs::JobFailureKind::LeaseExpired
694 );
695 assert_eq!(dead_letter.max_attempts, None);
696 }
697}