haste_worker/async_jobs/
mod.rs

1//! Runs asynchronous jobs: claims queued rows from `async_jobs`, hands each to
2//! the [`JobHandler`] registered for its kind, keeps the job's lease renewed
3//! while the handler works, and records the result.
4//!
5//! Any number of these workers may run at once, on any number of hosts: a
6//! claim is one row-locking statement, so no two take the same job, and a
7//! worker that dies simply stops renewing, which hands its jobs back to the
8//! queue (up to each job's `max_attempts`). The same worker also runs the
9//! maintenance sweeps: failing jobs abandoned past their last attempt, and
10//! removing expired ones, letting their handlers remove what they produced.
11
12use crate::traits::Worker;
13use haste_fhir_model::r4::generated::{resources::Resource, terminology::IssueType};
14use haste_fhir_operation_error::OperationOutcomeError;
15use haste_repository::{
16    async_job::{AsyncJob, AsyncJobProvider, JobLease, JobResult, LeaseRenewal},
17    utilities::generate_id,
18};
19use std::{
20    collections::HashMap,
21    pin::Pin,
22    sync::{
23        Arc,
24        atomic::{AtomicBool, Ordering},
25    },
26    time::{Duration, Instant},
27};
28use tokio::{
29    sync::{Mutex, Semaphore},
30    task::JoinHandle,
31};
32
33/// What a finished handler hands back, and how the status endpoint serves it.
34#[derive(Debug)]
35pub enum JobOutput {
36    /// Served as `application/fhir+json`: the response the synchronous form
37    /// of the request would have had.
38    Resource(Box<Resource>),
39    /// Served as `application/json`: a Bulk Data manifest, for one.
40    Json(serde_json::Value),
41}
42
43impl JobOutput {
44    fn into_result(self) -> Result<JobResult, OperationOutcomeError> {
45        let (output, content_type) = match self {
46            JobOutput::Resource(resource) => (
47                serde_json::to_value(&resource).map_err(|e| {
48                    OperationOutcomeError::fatal(
49                        IssueType::exception(),
50                        format!("Failed to serialize the job's output: {e}"),
51                    )
52                })?,
53                "application/fhir+json",
54            ),
55            JobOutput::Json(value) => (value, "application/json"),
56        };
57
58        Ok(JobResult::Completed {
59            output,
60            content_type: content_type.to_string(),
61        })
62    }
63}
64
65/// What a running handler sees of its job, and how it talks back: progress
66/// for the client's `X-Progress`, and whether the client has cancelled.
67pub struct JobContext {
68    pub job: AsyncJob,
69    pending_progress: Mutex<Option<String>>,
70    cancelled: AtomicBool,
71}
72
73impl JobContext {
74    fn new(job: AsyncJob) -> Self {
75        JobContext {
76            job,
77            pending_progress: Mutex::new(None),
78            cancelled: AtomicBool::new(false),
79        }
80    }
81
82    /// Records progress to send with the next lease renewal. Only the latest
83    /// message is kept.
84    pub async fn report(&self, progress: impl Into<String>) {
85        *self.pending_progress.lock().await = Some(progress.into());
86    }
87
88    /// True once the job is no longer this worker's to run: the client
89    /// cancelled it, or the lease lapsed and another worker took it. A
90    /// handler should check this between steps and stop; whatever it returns
91    /// afterwards is dropped.
92    pub fn is_cancelled(&self) -> bool {
93        self.cancelled.load(Ordering::Relaxed)
94    }
95
96    async fn take_progress(&self) -> Option<String> {
97        self.pending_progress.lock().await.take()
98    }
99}
100
101pub type HandlerFuture<'a, T> =
102    Pin<Box<dyn Future<Output = Result<T, OperationOutcomeError>> + Send + 'a>>;
103
104/// Runs jobs of one kind.
105pub trait JobHandler: Send + Sync {
106    /// The `kind` of the jobs this handler runs.
107    fn kind(&self) -> &'static str;
108
109    /// Runs one job to its result. An `Err` fails the job with that outcome;
110    /// a panic fails it too.
111    fn run<'a>(&'a self, context: &'a JobContext) -> HandlerFuture<'a, JobOutput>;
112
113    /// Called once for a job that has expired (or was cancelled) and is being
114    /// removed, so the handler can remove whatever it produced for it.
115    fn expire<'a>(&'a self, job: &'a AsyncJob) -> HandlerFuture<'a, ()> {
116        let _ = job;
117        Box::pin(async { Ok(()) })
118    }
119}
120
121#[derive(Clone, Debug)]
122pub struct AsyncJobWorkerConfig {
123    /// How long to wait before looking for work again after finding none.
124    pub poll_interval: Duration,
125    /// How long one claim or renewal holds a job. Renewals happen at a third
126    /// of this, so a worker has to miss two in a row before its job can be
127    /// taken over.
128    pub lease_duration: Duration,
129    /// Jobs this worker runs at once.
130    pub max_concurrent_jobs: usize,
131    /// How often the sweeps for abandoned and expired jobs run.
132    pub maintenance_interval: Duration,
133    /// How long a finished job and its result are kept.
134    pub retention: Duration,
135}
136
137impl Default for AsyncJobWorkerConfig {
138    fn default() -> Self {
139        Self {
140            poll_interval: Duration::from_secs(1),
141            lease_duration: Duration::from_secs(60),
142            max_concurrent_jobs: 4,
143            maintenance_interval: Duration::from_secs(60),
144            retention: Duration::from_hours(24),
145        }
146    }
147}
148
149/// Expired rows removed per sweep. Bounds one sweep's work, and the time it
150/// holds row locks; what is left waits for the next sweep.
151const EXPIRED_SWEEP_LIMIT: i64 = 100;
152
153/// Nothing to poll at all (the claim failed): no reason to look again soon.
154const ERROR_BACKOFF: Duration = Duration::from_secs(5);
155
156pub struct AsyncJobWorker<Repo> {
157    repo: Arc<Repo>,
158    handlers: Arc<HashMap<&'static str, Arc<dyn JobHandler>>>,
159    config: AsyncJobWorkerConfig,
160    /// Names this worker in `async_jobs.locked_by`.
161    worker_id: String,
162    running: Arc<Mutex<bool>>,
163}
164
165impl<Repo: AsyncJobProvider + Send + Sync + 'static> AsyncJobWorker<Repo> {
166    /// A worker that runs the kinds `handlers` cover, and nothing else.
167    #[must_use]
168    pub fn new(
169        repo: Arc<Repo>,
170        handlers: Vec<Arc<dyn JobHandler>>,
171        config: AsyncJobWorkerConfig,
172    ) -> Self {
173        let handlers = handlers
174            .into_iter()
175            .map(|handler| (handler.kind(), handler))
176            .collect();
177
178        let host = std::env::var("HOSTNAME").unwrap_or_else(|_| "worker".to_string());
179        let worker_id = format!("{host}-{}-{}", std::process::id(), generate_id(Some(6)));
180
181        Self {
182            repo,
183            handlers: Arc::new(handlers),
184            config,
185            worker_id,
186            running: Arc::new(Mutex::new(true)),
187        }
188    }
189
190    #[must_use]
191    pub fn worker_id(&self) -> &str {
192        &self.worker_id
193    }
194}
195
196/// The outcome recorded on a job whose worker stopped renewing and which has
197/// no attempts left.
198fn abandoned_outcome() -> serde_json::Value {
199    let error = OperationOutcomeError::fatal(
200        IssueType::exception(),
201        "The job was abandoned: the worker running it stopped responding, and it has no \
202         attempts left."
203            .to_string(),
204    );
205
206    serde_json::to_value(error.outcome()).unwrap_or(serde_json::Value::Null)
207}
208
209/// Failing a job after a panic: the panic is logged with the job, the client
210/// only sees that it failed.
211fn panicked_outcome() -> serde_json::Value {
212    let error = OperationOutcomeError::fatal(
213        IssueType::exception(),
214        "The job failed with an internal error.".to_string(),
215    );
216
217    serde_json::to_value(error.outcome()).unwrap_or(serde_json::Value::Null)
218}
219
220/// Runs one claimed job: the handler in its own task (so a panic is caught
221/// here and fails the job rather than the worker), renewing the lease
222/// alongside it, then records the result. `_permit` holds this worker's
223/// concurrency slot until the job is done.
224async fn run_job<Repo: AsyncJobProvider + Send + Sync + 'static>(
225    repo: Arc<Repo>,
226    handlers: Arc<HashMap<&'static str, Arc<dyn JobHandler>>>,
227    lease: JobLease,
228    retention: Duration,
229    job: AsyncJob,
230    _permit: tokio::sync::OwnedSemaphorePermit,
231) {
232    let job_id = job.id.clone();
233    let kind = job.kind.clone();
234    let tenant = job.tenant.clone();
235    let project = job.project.clone();
236    let attempt = job.attempts;
237    let started = Instant::now();
238
239    tracing::info!(
240        job = %job_id,
241        kind = %kind,
242        tenant = %tenant,
243        project = %project.as_ref(),
244        attempt,
245        "Running job"
246    );
247
248    let context = Arc::new(JobContext::new(job));
249
250    let renewals = tokio::spawn(renew_lease_until_lost(
251        repo.clone(),
252        lease.clone(),
253        context.clone(),
254    ));
255
256    let outcome = match handlers.get(kind.as_str()).cloned() {
257        Some(handler) => {
258            let handler_context = context.clone();
259            match tokio::spawn(async move { handler.run(&handler_context).await }).await {
260                Ok(Ok(output)) => output.into_result(),
261                Ok(Err(error)) => Ok(JobResult::Failed {
262                    outcome: serde_json::to_value(error.outcome())
263                        .unwrap_or(serde_json::Value::Null),
264                }),
265                Err(join_error) => {
266                    tracing::error!(job = %job_id, kind = %kind, "Job panicked: {join_error:?}");
267                    Ok(JobResult::Failed {
268                        outcome: panicked_outcome(),
269                    })
270                }
271            }
272        }
273        // Only registered kinds are claimed, so this cannot happen; if it
274        // does, the job must not sit in-progress until its lease lapses.
275        None => Ok(JobResult::Failed {
276            outcome: serde_json::to_value(
277                OperationOutcomeError::fatal(
278                    IssueType::not_supported(),
279                    format!("No handler for job kind '{kind}'."),
280                )
281                .outcome(),
282            )
283            .unwrap_or(serde_json::Value::Null),
284        }),
285    };
286
287    renewals.abort();
288
289    if context.is_cancelled() {
290        tracing::info!(
291            job = %job_id,
292            kind = %kind,
293            "Job was cancelled or taken over while running; dropping its result"
294        );
295        return;
296    }
297
298    let result = match outcome {
299        Ok(result) => result,
300        Err(error) => JobResult::Failed {
301            outcome: serde_json::to_value(error.outcome()).unwrap_or(serde_json::Value::Null),
302        },
303    };
304    let status = match &result {
305        JobResult::Completed { .. } => "completed",
306        JobResult::Failed { .. } => "failed",
307    };
308
309    match repo
310        .finish_job(&job_id, &lease.worker, result, retention)
311        .await
312    {
313        Ok(true) => tracing::info!(
314            job = %job_id,
315            kind = %kind,
316            status,
317            elapsed = ?started.elapsed(),
318            "Job finished"
319        ),
320        Ok(false) => tracing::warn!(
321            job = %job_id,
322            kind = %kind,
323            "Job was no longer this worker's to finish; dropping its result"
324        ),
325        Err(error) => tracing::error!(
326            job = %job_id,
327            kind = %kind,
328            "Failed to record the job's result: {error:?}"
329        ),
330    }
331}
332
333/// Renews the job's lease at a third of its length until the job is lost,
334/// carrying any progress the handler reported since the last renewal. A
335/// failed renewal (the database was unreachable for a moment) is retried at
336/// the next tick; only a definite `Lost` stops the handler.
337async fn renew_lease_until_lost<Repo: AsyncJobProvider + Send + Sync + 'static>(
338    repo: Arc<Repo>,
339    lease: JobLease,
340    context: Arc<JobContext>,
341) {
342    let interval = lease.duration / 3;
343
344    loop {
345        tokio::time::sleep(interval).await;
346
347        let progress = context.take_progress().await;
348        match repo
349            .renew_lease(&context.job.id, &lease, progress.as_deref())
350            .await
351        {
352            Ok(LeaseRenewal::Renewed) => {}
353            Ok(LeaseRenewal::Lost) => {
354                context.cancelled.store(true, Ordering::Relaxed);
355                return;
356            }
357            Err(error) => tracing::warn!(
358                job = %context.job.id,
359                "Failed to renew the job's lease, will retry: {error:?}"
360            ),
361        }
362    }
363}
364
365/// The sweeps: fail what was abandoned past its last attempt, and remove what
366/// has expired, giving each expired job to its handler first.
367async fn maintain<Repo: AsyncJobProvider + Send + Sync + 'static>(
368    repo: &Repo,
369    handlers: &HashMap<&'static str, Arc<dyn JobHandler>>,
370    retention: Duration,
371) {
372    match repo
373        .fail_abandoned_jobs(&abandoned_outcome(), retention)
374        .await
375    {
376        Ok(0) => {}
377        Ok(failed) => tracing::warn!("Failed {failed} abandoned job(s)"),
378        Err(error) => tracing::error!("Failed to sweep abandoned jobs: {error:?}"),
379    }
380
381    match repo.delete_expired_jobs(EXPIRED_SWEEP_LIMIT).await {
382        Ok(expired) => {
383            for job in expired {
384                tracing::debug!(job = %job.id, kind = %job.kind, "Removed expired job");
385                if let Some(handler) = handlers.get(job.kind.as_str())
386                    && let Err(error) = handler.expire(&job).await
387                {
388                    tracing::error!(
389                        job = %job.id,
390                        kind = %job.kind,
391                        "Handler failed to clean up an expired job: {error:?}"
392                    );
393                }
394            }
395        }
396        Err(error) => tracing::error!("Failed to sweep expired jobs: {error:?}"),
397    }
398}
399
400impl<Repo: AsyncJobProvider + Send + Sync + 'static> Worker for AsyncJobWorker<Repo> {
401    // Nothing here waits: the loop is spawned and its handle returned.
402    fn run(&self) -> impl Future<Output = Result<JoinHandle<()>, OperationOutcomeError>> {
403        let repo = self.repo.clone();
404        let handlers = self.handlers.clone();
405        let config = self.config.clone();
406        let running = self.running.clone();
407        let lease = JobLease {
408            worker: self.worker_id.clone(),
409            duration: config.lease_duration,
410        };
411        let kinds: Vec<&'static str> = handlers.keys().copied().collect();
412        let slots = Arc::new(Semaphore::new(config.max_concurrent_jobs.max(1)));
413
414        tracing::info!(
415            worker = %self.worker_id,
416            kinds = ?kinds,
417            "Starting async job worker..."
418        );
419
420        let spawned = tokio::spawn(async move {
421            // Sweep once at start, then on the interval.
422            let mut last_maintenance: Option<Instant> = None;
423
424            while *running.lock().await {
425                if last_maintenance.is_none_or(|at| at.elapsed() >= config.maintenance_interval) {
426                    maintain(repo.as_ref(), &handlers, config.retention).await;
427                    last_maintenance = Some(Instant::now());
428                }
429
430                // Hold a slot before claiming, so a claimed job is never left
431                // waiting for one while its lease runs down.
432                let Ok(permit) = slots.clone().acquire_owned().await else {
433                    break;
434                };
435
436                let backoff = match repo.claim_job(&kinds, &lease).await {
437                    Ok(Some(job)) => {
438                        tokio::spawn(run_job(
439                            repo.clone(),
440                            handlers.clone(),
441                            lease.clone(),
442                            config.retention,
443                            job,
444                            permit,
445                        ));
446                        None
447                    }
448                    Ok(None) => {
449                        drop(permit);
450                        Some(config.poll_interval)
451                    }
452                    Err(error) => {
453                        drop(permit);
454                        tracing::error!("Failed to claim a job: {error:?}");
455                        Some(ERROR_BACKOFF)
456                    }
457                };
458
459                if let Some(backoff) = backoff {
460                    tokio::time::sleep(backoff).await;
461                }
462            }
463        });
464
465        std::future::ready(Ok(spawned))
466    }
467
468    async fn stop(&mut self) -> Result<(), OperationOutcomeError> {
469        *self.running.lock().await = false;
470        Ok(())
471    }
472}
haste_worker/async_jobs/mod.rs

haste_worker/async_jobs/
mod.rs
