Skip to main content

runledger_runtime/
reaper.rs

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(&registry, &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(&registry, &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(&registry, &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(&registry, &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}