1use std::{
2 collections::HashMap,
3 net::SocketAddr,
4 num::NonZeroU16,
5 panic::AssertUnwindSafe,
6 sync::{
7 Arc,
8 atomic::{AtomicU16, Ordering},
9 },
10 time::Duration,
11};
12
13use futures::FutureExt;
14use pgtask_core::{
15 HandlerVersion, LeaseRenewal, QueueName, ScheduleConfig, Task, TaskId, TaskName, TaskState, WorkerId,
16};
17use pgtask_postgres::{PostgresError, ReadyListener, Store};
18use serde_json::json;
19use thiserror::Error;
20use tokio::{
21 sync::{Mutex, Notify},
22 task::{JoinError, JoinSet},
23 time::{Instant, MissedTickBehavior},
24};
25use tokio_util::sync::CancellationToken;
26use tracing::{Instrument, info_span, warn};
27
28use crate::{
29 HandlerRegistry,
30 health::{Health, Supervisor},
31 registry::RegisteredHandler,
32};
33
34#[derive(Clone, Debug)]
35pub struct WorkerConfig {
36 pub queues: Vec<QueueName>,
38 pub concurrency: NonZeroU16,
39 pub claim_batch_size: NonZeroU16,
40 pub lease_duration: Duration,
41 pub poll_interval: Duration,
42 pub shutdown_grace: Duration,
43 pub worker_heartbeat_interval: Duration,
44 pub worker_ttl: Duration,
45 pub scheduler_enabled: bool,
46 pub schedule_batch_size: NonZeroU16,
47 pub wait_batch_size: NonZeroU16,
48 pub schedule_reconciliation_interval: Duration,
49 pub retention_enabled: bool,
50 pub retention_batch_size: NonZeroU16,
51 pub retention_interval: Duration,
52 pub declared_schedules: Vec<ScheduleConfig>,
53 pub health_address: Option<SocketAddr>,
54 pub supervisor_interval: Duration,
55 pub overload_protection: OverloadProtectionConfig,
56}
57
58#[derive(Clone, Debug)]
59pub struct OverloadProtectionConfig {
60 pub enabled: bool,
61 pub enforce: bool,
62 pub event_loop_lag_threshold: Duration,
63 pub sustained_samples: NonZeroU16,
64 pub recovery_samples: NonZeroU16,
65 pub minimum_concurrency: NonZeroU16,
66}
67
68impl Default for OverloadProtectionConfig {
69 fn default() -> Self {
70 Self {
71 enabled: true,
72 enforce: false,
73 event_loop_lag_threshold: Duration::from_millis(250),
74 sustained_samples: NonZeroU16::new(3).expect("3 is nonzero"),
75 recovery_samples: NonZeroU16::new(5).expect("5 is nonzero"),
76 minimum_concurrency: NonZeroU16::MIN,
77 }
78 }
79}
80
81impl WorkerConfig {
82 pub fn new(queue_name: QueueName) -> Self {
83 Self::with_queues(vec![queue_name])
84 }
85
86 pub fn with_queues(queues: Vec<QueueName>) -> Self {
87 Self {
88 queues,
89 concurrency: NonZeroU16::new(10).expect("10 is nonzero"),
90 claim_batch_size: NonZeroU16::new(10).expect("10 is nonzero"),
91 lease_duration: Duration::from_secs(30),
92 poll_interval: Duration::from_secs(30),
93 shutdown_grace: Duration::from_secs(30),
94 worker_heartbeat_interval: Duration::from_secs(10),
95 worker_ttl: Duration::from_secs(30),
96 scheduler_enabled: true,
97 schedule_batch_size: NonZeroU16::new(100).expect("100 is nonzero"),
98 wait_batch_size: NonZeroU16::new(100).expect("100 is nonzero"),
99 schedule_reconciliation_interval: Duration::from_secs(30),
100 retention_enabled: true,
101 retention_batch_size: NonZeroU16::new(100).expect("100 is nonzero"),
102 retention_interval: Duration::from_mins(1),
103 declared_schedules: Vec::new(),
104 health_address: None,
105 supervisor_interval: Duration::from_secs(1),
106 overload_protection: OverloadProtectionConfig::default(),
107 }
108 }
109}
110
111#[derive(Debug, Error)]
112pub enum WorkerError {
113 #[error(transparent)]
114 Postgres(#[from] PostgresError),
115 #[error("lease duration must be at least three milliseconds")]
116 InvalidLeaseDuration,
117 #[error("poll interval must be greater than zero")]
118 InvalidPollInterval,
119 #[error("worker heartbeat interval must be nonzero and shorter than its time to live")]
120 InvalidWorkerHeartbeat,
121 #[error("schedule reconciliation interval must be greater than zero")]
122 InvalidScheduleReconciliationInterval,
123 #[error("retention interval must be greater than zero")]
124 InvalidRetentionInterval,
125 #[error("supervisor interval must be greater than zero")]
126 InvalidSupervisorInterval,
127 #[error("overload protection minimum concurrency exceeds configured concurrency")]
128 InvalidMinimumConcurrency,
129 #[error("worker supervisor failed: {0}")]
130 Supervisor(#[source] std::io::Error),
131 #[error(
132 "database storage protocols {database_minimum}..={database_maximum} are incompatible with worker protocols {worker_minimum}..={worker_maximum}"
133 )]
134 IncompatibleStorageProtocol {
135 database_minimum: u32,
136 database_maximum: u32,
137 worker_minimum: u32,
138 worker_maximum: u32,
139 },
140 #[error("worker has no registered handlers")]
141 MissingHandlers,
142 #[error("worker has no queues")]
143 MissingQueues,
144 #[error("worker queue list contains duplicates")]
145 DuplicateQueues,
146 #[error("effective concurrency {requested} exceeds configured concurrency {configured}")]
147 AdmissionLimitExceedsConfigured { requested: u16, configured: u16 },
148 #[error("declared schedule {0} targets another queue or an unregistered handler")]
149 InvalidDeclaredSchedule(String),
150 #[error("claimed task {0} has no lease token")]
151 MissingLeaseToken(pgtask_core::TaskId),
152 #[error("claimed task has no registered handler")]
153 MissingHandler,
154}
155
156pub struct Worker {
157 store: Store,
158 registry: Arc<HandlerRegistry>,
159 config: WorkerConfig,
160 control: WorkerControl,
161 health: Health,
162 id: WorkerId,
163}
164
165#[derive(Clone)]
166pub struct WorkerControl {
167 configured: NonZeroU16,
168 effective: Arc<AtomicU16>,
169 proposed: Arc<AtomicU16>,
170 changed: Arc<Notify>,
171 queue_name: QueueName,
172}
173
174impl WorkerControl {
175 pub fn configured_concurrency(&self) -> NonZeroU16 {
176 self.configured
177 }
178
179 pub fn effective_concurrency(&self) -> NonZeroU16 {
180 NonZeroU16::new(self.effective.load(Ordering::Acquire)).expect("the admission limit is always nonzero")
181 }
182
183 pub fn proposed_concurrency(&self) -> NonZeroU16 {
184 NonZeroU16::new(self.proposed.load(Ordering::Acquire)).expect("the proposed admission limit is always nonzero")
185 }
186
187 pub fn set_effective_concurrency(&self, limit: NonZeroU16) -> Result<(), WorkerError> {
188 self.apply_effective_concurrency(limit, "manual")
189 }
190
191 pub(crate) fn apply_effective_concurrency(
192 &self,
193 limit: NonZeroU16,
194 reason: &'static str,
195 ) -> Result<(), WorkerError> {
196 if limit > self.configured {
197 return Err(WorkerError::AdmissionLimitExceedsConfigured {
198 requested: limit.get(),
199 configured: self.configured.get(),
200 });
201 }
202 let previous = self.effective.swap(limit.get(), Ordering::AcqRel);
203 if previous != limit.get() {
204 pgtask_otel::record_worker_admission_limit(self.queue_name.as_str(), "applied", reason, limit.get());
205 }
206 self.changed.notify_waiters();
207 Ok(())
208 }
209
210 pub(crate) fn record_proposed_concurrency(&self, limit: NonZeroU16, reason: &'static str) {
211 let previous = self.proposed.swap(limit.get(), Ordering::AcqRel);
212 if previous != limit.get() {
213 pgtask_otel::record_worker_admission_limit(self.queue_name.as_str(), "proposed", reason, limit.get());
214 }
215 }
216}
217
218type ActiveLeases = Arc<Mutex<HashMap<TaskId, ActiveLease>>>;
219
220#[derive(Clone)]
221struct ActiveLease {
222 renewal: LeaseRenewal,
223 queue_name: QueueName,
224 task_name: TaskName,
225 lost: CancellationToken,
226 last_renewed: Instant,
227}
228
229struct HeartbeatConfig {
230 worker_id: WorkerId,
231 queue_name: QueueName,
232 capabilities: Vec<(TaskName, HandlerVersion)>,
233 interval: Duration,
234 ttl: Duration,
235}
236
237impl Worker {
238 pub fn new(store: Store, registry: HandlerRegistry, config: WorkerConfig) -> Result<Self, WorkerError> {
239 if config.lease_duration < Duration::from_millis(3) {
240 return Err(WorkerError::InvalidLeaseDuration);
241 }
242 if config.poll_interval.is_zero() {
243 return Err(WorkerError::InvalidPollInterval);
244 }
245 if config.worker_heartbeat_interval.is_zero() || config.worker_heartbeat_interval >= config.worker_ttl {
246 return Err(WorkerError::InvalidWorkerHeartbeat);
247 }
248 if config.schedule_reconciliation_interval.is_zero() {
249 return Err(WorkerError::InvalidScheduleReconciliationInterval);
250 }
251 if config.retention_interval.is_zero() {
252 return Err(WorkerError::InvalidRetentionInterval);
253 }
254 if config.supervisor_interval.is_zero() {
255 return Err(WorkerError::InvalidSupervisorInterval);
256 }
257 if config.overload_protection.minimum_concurrency > config.concurrency {
258 return Err(WorkerError::InvalidMinimumConcurrency);
259 }
260 if registry.capabilities().is_empty() {
261 return Err(WorkerError::MissingHandlers);
262 }
263 if config.queues.is_empty() {
264 return Err(WorkerError::MissingQueues);
265 }
266 if config
267 .queues
268 .iter()
269 .enumerate()
270 .any(|(index, queue)| config.queues[..index].contains(queue))
271 {
272 return Err(WorkerError::DuplicateQueues);
273 }
274 if let Some(schedule) = config.declared_schedules.iter().find(|schedule| {
275 !config.queues.contains(&schedule.task.queue_name)
276 || registry
277 .get(&schedule.task.task_name, schedule.task.handler_version)
278 .is_none()
279 }) {
280 return Err(WorkerError::InvalidDeclaredSchedule(schedule.name.to_string()));
281 }
282 let control = WorkerControl {
283 configured: config.concurrency,
284 effective: Arc::new(AtomicU16::new(config.concurrency.get())),
285 proposed: Arc::new(AtomicU16::new(config.concurrency.get())),
286 changed: Arc::new(Notify::new()),
287 queue_name: config.queues[0].clone(),
288 };
289 Ok(Self {
290 store,
291 registry: Arc::new(registry),
292 config,
293 control,
294 health: Health::new(),
295 id: WorkerId::new(),
296 })
297 }
298
299 pub fn control(&self) -> WorkerControl {
300 self.control.clone()
301 }
302
303 pub async fn run(self, shutdown: CancellationToken) -> Result<(), WorkerError> {
304 let database_protocol = self.store.storage_protocol_range().await?;
305 if !database_protocol.overlaps(pgtask_core::STORAGE_PROTOCOL_RANGE) {
306 return Err(WorkerError::IncompatibleStorageProtocol {
307 database_minimum: database_protocol.minimum,
308 database_maximum: database_protocol.maximum,
309 worker_minimum: pgtask_core::STORAGE_PROTOCOL_MIN_VERSION,
310 worker_maximum: pgtask_core::STORAGE_PROTOCOL_MAX_VERSION,
311 });
312 }
313 let _supervisor = Supervisor::start(
314 self.health.clone(),
315 self.config.queues[0].clone(),
316 self.config.supervisor_interval,
317 self.config.health_address,
318 self.control.clone(),
319 self.config.overload_protection.clone(),
320 self.config.lease_duration * 2 / 3,
321 )
322 .map_err(WorkerError::Supervisor)?;
323 let active_leases = Arc::new(Mutex::new(HashMap::new()));
324 let task_wakeup = Arc::new(Notify::new());
325 let schedule_wakeup = Arc::new(Notify::new());
326 let runtime_shutdown = CancellationToken::new();
327 let capabilities = self.registry.capabilities();
328 let registrations = self.registry.registrations();
329 let ready_listener = self.store.ready_listener_for(&self.config.queues).await?;
330 self.health.set_listener(true);
331 for schedule in &self.config.declared_schedules {
332 self.store.put_schedule(schedule).await?;
333 }
334 self.store
335 .register_worker(
336 self.id,
337 &self.config.queues[0],
338 env!("CARGO_PKG_VERSION"),
339 ®istrations,
340 self.config.worker_ttl,
341 )
342 .await?;
343 self.health.set_database(true);
344 self.health.set_admission(true);
345 let renewer = renew_leases(
346 self.store.clone(),
347 Arc::clone(&active_leases),
348 self.health.clone(),
349 self.config.lease_duration,
350 runtime_shutdown.clone(),
351 );
352 let listener = listen_for_ready(
353 self.store.clone(),
354 self.config.queues.clone(),
355 Arc::clone(&task_wakeup),
356 Arc::clone(&schedule_wakeup),
357 runtime_shutdown.clone(),
358 ready_listener,
359 self.health.clone(),
360 );
361 let scheduler = materialize_schedules(
362 self.store.clone(),
363 self.config.scheduler_enabled,
364 self.config.schedule_batch_size,
365 self.config.wait_batch_size,
366 self.config.schedule_reconciliation_interval,
367 schedule_wakeup,
368 runtime_shutdown.clone(),
369 );
370 let retention = delete_expired_terminal(
371 self.store.clone(),
372 self.config.queues.clone(),
373 self.config.retention_enabled,
374 self.config.retention_batch_size,
375 self.config.retention_interval,
376 runtime_shutdown.clone(),
377 );
378 let heartbeat = heartbeat_worker(
379 self.store.clone(),
380 HeartbeatConfig {
381 worker_id: self.id,
382 queue_name: self.config.queues[0].clone(),
383 capabilities: capabilities.clone(),
384 interval: self.config.worker_heartbeat_interval,
385 ttl: self.config.worker_ttl,
386 },
387 runtime_shutdown.clone(),
388 self.health.clone(),
389 );
390 let handlers = async {
391 let result = self
392 .run_handlers(shutdown, Arc::clone(&active_leases), task_wakeup)
393 .await;
394 runtime_shutdown.cancel();
395 self.health.set_admission(false);
396 result
397 };
398 let ((), (), (), (), (), result) = tokio::join!(renewer, listener, scheduler, retention, heartbeat, handlers);
399 result
400 }
401
402 async fn run_handlers(
403 &self,
404 shutdown: CancellationToken,
405 active_leases: ActiveLeases,
406 wakeup: Arc<Notify>,
407 ) -> Result<(), WorkerError> {
408 let mut handlers = JoinSet::new();
409 let capabilities = self.registry.capabilities();
410 loop {
411 self.health.record_runtime_progress();
412 while let Some(result) = handlers.try_join_next() {
413 handle_handler_result(result);
414 }
415 if shutdown.is_cancelled() {
416 break;
417 }
418
419 let Some((limit, tasks)) = self
420 .claim_tasks(&shutdown, &wakeup, handlers.len(), &capabilities)
421 .await
422 else {
423 continue;
424 };
425 let claimed_any = !tasks.is_empty();
426 for task in tasks {
427 self.spawn_task(&mut handlers, &active_leases, task).await?;
428 }
429
430 if !claimed_any {
431 let deadline_delay = if limit == 0 {
432 self.config.poll_interval
433 } else {
434 self.next_task_delay(&capabilities).await
435 };
436 if handlers.is_empty() {
437 tokio::select! {
438 () = shutdown.cancelled() => break,
439 () = self.control.changed.notified() => {}
440 () = wakeup.notified() => {}
441 () = tokio::time::sleep(deadline_delay) => {}
442 }
443 } else {
444 tokio::select! {
445 () = shutdown.cancelled() => break,
446 () = self.control.changed.notified() => {}
447 result = handlers.join_next() => handle_handler_result(
448 result.expect("a nonempty handler set returns one task"),
449 ),
450 () = wakeup.notified() => {}
451 () = tokio::time::sleep(deadline_delay) => {}
452 }
453 }
454 }
455 }
456
457 let deadline = Instant::now() + self.config.shutdown_grace;
458 while !handlers.is_empty() {
459 tokio::select! {
460 result = handlers.join_next() => handle_handler_result(
461 result.expect("a nonempty handler set returns one task"),
462 ),
463 () = tokio::time::sleep_until(deadline) => {
464 handlers.abort_all();
465 break;
466 }
467 }
468 }
469 active_leases.lock().await.clear();
470 Ok(())
471 }
472
473 async fn next_task_delay(&self, capabilities: &[(TaskName, HandlerVersion)]) -> Duration {
474 let mut delay = self.config.poll_interval;
475 for queue_name in &self.config.queues {
476 match self.store.next_task_delay(queue_name, capabilities).await {
477 Ok(Some(queue_delay)) => delay = delay.min(queue_delay),
478 Ok(None) => {}
479 Err(error) => {
480 self.health.set_database(false);
481 warn!(%error, "could not read the next task deadline");
482 return Duration::from_secs(1).min(self.config.poll_interval);
483 }
484 }
485 }
486 delay.max(Duration::from_millis(1))
487 }
488
489 async fn claim_tasks(
490 &self,
491 shutdown: &CancellationToken,
492 wakeup: &Notify,
493 active_handlers: usize,
494 capabilities: &[(TaskName, HandlerVersion)],
495 ) -> Option<(usize, Vec<Task>)> {
496 let effective_concurrency = self.control.effective_concurrency().get();
497 pgtask_otel::record_worker_capacity(
498 self.config.queues[0].as_str(),
499 self.config.concurrency.get(),
500 effective_concurrency,
501 active_handlers,
502 );
503 for queue_name in &self.config.queues {
504 if let Err(error) = self
505 .store
506 .recover_expired(queue_name, self.config.claim_batch_size.get())
507 .await
508 {
509 self.health.set_database(false);
510 warn!(%error, "could not recover expired task leases");
511 wait_after_database_error(shutdown, wakeup).await;
512 return None;
513 }
514 }
515 self.health.set_database(true);
516 let available = usize::from(effective_concurrency).saturating_sub(active_handlers);
517 let limit = available.min(usize::from(self.config.claim_batch_size.get()));
518 if limit == 0 {
519 return Some((limit, Vec::new()));
520 }
521 let mut tasks = Vec::new();
522 for queue_name in &self.config.queues {
523 let remaining = limit - tasks.len();
524 if remaining == 0 {
525 break;
526 }
527 match self
528 .store
529 .claim(
530 queue_name,
531 self.id,
532 capabilities,
533 u16::try_from(remaining).expect("limit is bounded by a u16 configuration value"),
534 self.config.lease_duration,
535 )
536 .await
537 {
538 Ok(claimed) => {
539 self.health.set_database(true);
540 tasks.extend(claimed);
541 }
542 Err(error) => {
543 self.health.set_database(false);
544 warn!(%error, "could not claim tasks");
545 wait_after_database_error(shutdown, wakeup).await;
546 return None;
547 }
548 }
549 }
550 Some((limit, tasks))
551 }
552
553 async fn spawn_task(
554 &self,
555 handlers: &mut JoinSet<Result<(), PostgresError>>,
556 active_leases: &ActiveLeases,
557 task: Task,
558 ) -> Result<(), WorkerError> {
559 let lease_token = task.lease_token.ok_or(WorkerError::MissingLeaseToken(task.id))?;
560 let handler = self
561 .registry
562 .get(&task.task_name, task.handler_version)
563 .ok_or(WorkerError::MissingHandler)?
564 .clone();
565 let lost = CancellationToken::new();
566 active_leases.lock().await.insert(
567 task.id,
568 ActiveLease {
569 renewal: LeaseRenewal {
570 task_id: task.id,
571 attempt: task.attempt,
572 lease_token,
573 },
574 queue_name: task.queue_name.clone(),
575 task_name: task.task_name.clone(),
576 lost: lost.clone(),
577 last_renewed: Instant::now(),
578 },
579 );
580 self.health.set_active_leases(true);
581 let span = info_span!(
582 "pgtask.execute",
583 otel.kind = "consumer",
584 pgtask.task.id = %task.id,
585 pgtask.task.name = %task.task_name,
586 pgtask.task.attempt = task.attempt,
587 pgtask.queue.name = %task.queue_name,
588 );
589 pgtask_otel::set_parent_from_headers(&span, &task.headers)
590 .unwrap_or_else(|error| warn!(%error, "could not attach the producer trace context"));
591 let active_leases = Arc::clone(active_leases);
592 let store = self.store.clone();
593 let health = self.health.clone();
594 handlers.spawn(
595 async move {
596 let task_id = task.id;
597 let result = execute(store, handler, task, lease_token, lost).await;
598 let mut leases = active_leases.lock().await;
599 leases.remove(&task_id);
600 health.set_active_leases(!leases.is_empty());
601 result
602 }
603 .instrument(span),
604 );
605 Ok(())
606 }
607}
608
609fn handle_handler_result(result: Result<Result<(), PostgresError>, JoinError>) {
610 if let Err(error) = result.expect("engine execution tasks do not panic") {
611 warn!(%error, "task state transition failed; its lease will be recovered");
612 }
613}
614
615async fn wait_after_database_error(shutdown: &CancellationToken, wakeup: &Notify) {
616 tokio::select! {
617 () = shutdown.cancelled() => {}
618 () = wakeup.notified() => {}
619 () = tokio::time::sleep(Duration::from_secs(1)) => {}
620 }
621}
622
623async fn execute(
624 store: Store,
625 handler: RegisteredHandler,
626 task: Task,
627 lease_token: pgtask_core::LeaseToken,
628 lease_lost: CancellationToken,
629) -> Result<(), PostgresError> {
630 let queue_latency = task
631 .updated_at
632 .signed_duration_since(task.created_at)
633 .to_std()
634 .unwrap_or_default();
635 pgtask_otel::record_queue_latency(task.queue_name.as_str(), task.task_name.as_str(), queue_latency);
636 let started_at = std::time::Instant::now();
637 let context = crate::TaskContext::new(store.clone(), &task, lease_token, lease_lost.clone());
638 let handler_future = AssertUnwindSafe((handler.function)(task.clone(), context)).catch_unwind();
639 tokio::pin!(handler_future);
640
641 tokio::select! {
642 result = &mut handler_future => {
643 match result {
644 Ok(Ok(result)) => {
645 if store.complete(task.id, task.attempt, lease_token, Some(&result)).await? {
646 pgtask_otel::record_succeeded(task.queue_name.as_str(), task.task_name.as_str());
647 pgtask_otel::record_execution(
648 task.queue_name.as_str(),
649 task.task_name.as_str(),
650 "succeeded",
651 started_at.elapsed(),
652 );
653 } else {
654 pgtask_otel::record_lease_lost(task.queue_name.as_str(), task.task_name.as_str());
655 warn!("task completion lost its lease");
656 }
657 }
658 Ok(Err(error)) => {
659 if error.is_suspended() {
660 pgtask_otel::record_execution(
661 task.queue_name.as_str(),
662 task.task_name.as_str(),
663 "suspended",
664 started_at.elapsed(),
665 );
666 return Ok(());
667 }
668 let retry_after = if error.retryable {
669 task.retry_policy.unwrap_or(handler.retry_policy).delay_for(task.attempt)
670 } else {
671 None
672 };
673 let state = store.fail(task.id, task.attempt, lease_token, &error.error, retry_after).await?;
674 if state.is_none() {
675 pgtask_otel::record_lease_lost(task.queue_name.as_str(), task.task_name.as_str());
676 warn!("task failure lost its lease");
677 } else if state == Some(TaskState::Pending) {
678 tracing::debug!("task scheduled for retry");
679 }
680 record_failure_state(&task, state);
681 pgtask_otel::record_execution(
682 task.queue_name.as_str(),
683 task.task_name.as_str(),
684 if state == Some(TaskState::Pending) { "retry" } else { "failed" },
685 started_at.elapsed(),
686 );
687 }
688 Err(_) => {
689 let error = json!({"type": "handler_panic"});
690 let state = store
691 .fail(
692 task.id,
693 task.attempt,
694 lease_token,
695 &error,
696 task.retry_policy.unwrap_or(handler.retry_policy).delay_for(task.attempt),
697 )
698 .await?;
699 if state.is_none() {
700 pgtask_otel::record_lease_lost(task.queue_name.as_str(), task.task_name.as_str());
701 warn!("panicked task lost its lease");
702 }
703 record_failure_state(&task, state);
704 pgtask_otel::record_execution(
705 task.queue_name.as_str(),
706 task.task_name.as_str(),
707 "panic",
708 started_at.elapsed(),
709 );
710 }
711 }
712 }
713 () = lease_lost.cancelled() => {
714 pgtask_otel::record_execution(
715 task.queue_name.as_str(),
716 task.task_name.as_str(),
717 "lease_lost",
718 started_at.elapsed(),
719 );
720 warn!("task lost its lease during execution");
721 }
722 }
723 Ok(())
724}
725
726fn record_failure_state(task: &Task, state: Option<TaskState>) {
727 match state {
728 Some(TaskState::Pending) => pgtask_otel::record_retried(task.queue_name.as_str(), task.task_name.as_str()),
729 Some(_) => pgtask_otel::record_failed(task.queue_name.as_str(), task.task_name.as_str()),
730 None => {}
731 }
732}
733
734async fn renew_leases(
735 store: Store,
736 active: ActiveLeases,
737 health: Health,
738 lease_duration: Duration,
739 shutdown: CancellationToken,
740) {
741 let renewal_interval = lease_duration / 3;
742 let mut interval = tokio::time::interval_at(Instant::now() + renewal_interval, renewal_interval);
743 interval.set_missed_tick_behavior(MissedTickBehavior::Delay);
744 loop {
745 tokio::select! {
746 () = shutdown.cancelled() => break,
747 _ = interval.tick() => {
748 let leases: Vec<_> = active.lock().await.values().cloned().collect();
749 if leases.is_empty() {
750 health.set_active_leases(false);
751 continue;
752 }
753 match store.renew_leases(
754 &leases.iter().map(|lease| lease.renewal).collect::<Vec<_>>(),
755 lease_duration,
756 ).await {
757 Ok(renewed) => {
758 health.set_database(true);
759 health.record_lease_renewal(renewed.len() == leases.len());
760 update_renewed_leases(&active, &leases, &renewed).await;
761 }
762 Err(error) => {
763 health.set_database(false);
764 health.record_lease_renewal(false);
765 warn!(%error, "could not renew active task leases");
766 cancel_uncertain_leases(&active, &leases, lease_duration).await;
767 }
768 }
769 }
770 }
771 }
772}
773
774async fn update_renewed_leases(active: &ActiveLeases, leases: &[ActiveLease], renewed: &[TaskId]) {
775 let now = Instant::now();
776 let mut active = active.lock().await;
777 for lease in leases {
778 let was_renewed = renewed.contains(&lease.renewal.task_id);
779 pgtask_otel::record_renewed(lease.queue_name.as_str(), lease.task_name.as_str(), was_renewed);
780 if let Some(current) = active.get_mut(&lease.renewal.task_id)
781 && current.renewal == lease.renewal
782 {
783 if was_renewed {
784 current.last_renewed = now;
785 } else {
786 current.lost.cancel();
787 pgtask_otel::record_lease_lost(lease.queue_name.as_str(), lease.task_name.as_str());
788 }
789 }
790 }
791}
792
793async fn cancel_uncertain_leases(active: &ActiveLeases, leases: &[ActiveLease], lease_duration: Duration) {
794 let mut active = active.lock().await;
795 for lease in leases {
796 if lease.last_renewed.elapsed() >= lease_duration * 2 / 3
797 && let Some(current) = active.get_mut(&lease.renewal.task_id)
798 && current.renewal == lease.renewal
799 {
800 current.lost.cancel();
801 pgtask_otel::record_lease_lost(lease.queue_name.as_str(), lease.task_name.as_str());
802 }
803 }
804}
805
806async fn listen_for_ready(
807 store: Store,
808 queues: Vec<QueueName>,
809 task_wakeup: Arc<Notify>,
810 schedule_wakeup: Arc<Notify>,
811 shutdown: CancellationToken,
812 mut listener: ReadyListener,
813 health: Health,
814) {
815 let mut retry_delay = Duration::from_millis(100);
816 loop {
817 loop {
818 let notification = tokio::select! {
819 () = shutdown.cancelled() => return,
820 result = listener.recv() => result,
821 };
822 match notification {
823 Ok(notification)
824 if notification.channel().starts_with("pgtask_ready_")
825 && queues.iter().any(|queue| notification.payload() == queue.as_str()) =>
826 {
827 task_wakeup.notify_one();
828 }
829 Ok(notification) if matches!(notification.channel(), "pgtask_schedule" | "pgtask_wait") => {
830 schedule_wakeup.notify_one();
831 }
832 Ok(_) => {}
833 Err(error) => {
834 health.set_listener(false);
835 warn!(%error, "task notification listener disconnected");
836 break;
837 }
838 }
839 }
840 loop {
841 let reconnected = tokio::select! {
842 () = shutdown.cancelled() => return,
843 result = store.ready_listener_for(&queues) => result,
844 };
845 match reconnected {
846 Ok(reconnected) => {
847 listener = reconnected;
848 health.set_database(true);
849 health.set_listener(true);
850 retry_delay = Duration::from_millis(100);
851 task_wakeup.notify_one();
852 schedule_wakeup.notify_one();
853 break;
854 }
855 Err(error) => {
856 health.set_database(false);
857 warn!(%error, "could not reconnect the task notification listener");
858 tokio::select! {
859 () = shutdown.cancelled() => return,
860 () = tokio::time::sleep(retry_delay) => {}
861 }
862 retry_delay = (retry_delay * 2).min(Duration::from_secs(5));
863 }
864 }
865 }
866 }
867}
868
869async fn materialize_schedules(
870 store: Store,
871 enabled: bool,
872 schedule_batch_size: NonZeroU16,
873 wait_batch_size: NonZeroU16,
874 reconciliation_interval: Duration,
875 wakeup: Arc<Notify>,
876 shutdown: CancellationToken,
877) {
878 loop {
879 if enabled && let Err(error) = store.materialize_due_schedules(schedule_batch_size.get()).await {
880 warn!(%error, "could not materialize due schedules");
881 }
882 if let Err(error) = store.recover_wait_timeouts(wait_batch_size.get()).await {
883 warn!(%error, "could not recover signal wait timeouts");
884 }
885 if let Err(error) = store.recover_result_wait_timeouts(wait_batch_size.get()).await {
886 warn!(%error, "could not recover result wait timeouts");
887 }
888 let mut delay = reconciliation_interval;
889 if enabled {
890 match store.next_schedule_delay().await {
891 Ok(schedule_delay) => {
892 if let Some(schedule_delay) = schedule_delay {
893 delay = delay.min(schedule_delay);
894 }
895 }
896 Err(error) => warn!(%error, "could not read the next schedule deadline"),
897 }
898 }
899 match store.next_wait_delay().await {
900 Ok(wait_delay) => {
901 if let Some(wait_delay) = wait_delay {
902 delay = delay.min(wait_delay);
903 }
904 }
905 Err(error) => warn!(%error, "could not read the next wait deadline"),
906 }
907 tokio::select! {
908 () = shutdown.cancelled() => return,
909 () = wakeup.notified() => {}
910 () = tokio::time::sleep(delay) => {}
911 }
912 }
913}
914
915async fn delete_expired_terminal(
916 store: Store,
917 queues: Vec<QueueName>,
918 enabled: bool,
919 batch_size: NonZeroU16,
920 retention_interval: Duration,
921 shutdown: CancellationToken,
922) {
923 let mut interval = tokio::time::interval(retention_interval);
924 interval.set_missed_tick_behavior(MissedTickBehavior::Delay);
925 loop {
926 tokio::select! {
927 () = shutdown.cancelled() => return,
928 _ = interval.tick() => {
929 for queue_name in &queues {
930 if enabled
931 && let Err(error) = store.delete_expired_terminal(queue_name, batch_size.get()).await
932 {
933 warn!(%error, "could not delete expired terminal tasks");
934 }
935 if enabled
936 && let Err(error) = store.delete_expired_idempotency_keys(queue_name, batch_size.get()).await
937 {
938 warn!(%error, "could not delete expired idempotency keys");
939 }
940 }
941 }
942 }
943 }
944}
945
946async fn heartbeat_worker(store: Store, config: HeartbeatConfig, shutdown: CancellationToken, health: Health) {
947 let mut interval = tokio::time::interval(config.interval);
948 interval.set_missed_tick_behavior(MissedTickBehavior::Delay);
949 interval.tick().await;
950 loop {
951 tokio::select! {
952 () = shutdown.cancelled() => {
953 if let Err(error) = store.heartbeat_worker(config.worker_id, Duration::from_millis(1), true).await {
954 warn!(%error, "could not mark worker as stopped");
955 }
956 break;
957 }
958 _ = interval.tick() => {
959 match store.heartbeat_worker(config.worker_id, config.ttl, false).await {
960 Ok(true) => {
961 health.set_database(true);
962 pgtask_otel::record_heartbeat(config.queue_name.as_str(), "ok");
963 }
964 Ok(false) => {
965 health.set_database(false);
966 pgtask_otel::record_heartbeat(config.queue_name.as_str(), "missing");
967 warn!("worker registration disappeared");
968 }
969 Err(error) => {
970 health.set_database(false);
971 pgtask_otel::record_heartbeat(config.queue_name.as_str(), "error");
972 warn!(%error, "could not update worker heartbeat");
973 }
974 }
975 match store.live_worker_count(&config.queue_name).await {
976 Ok(live) => pgtask_otel::record_live_workers(config.queue_name.as_str(), live),
977 Err(error) => warn!(%error, "could not read the live worker count"),
978 }
979 match store.queue_demand(&config.queue_name, &config.capabilities).await {
980 Ok(demand) => pgtask_otel::record_queue_demand(
981 config.queue_name.as_str(),
982 demand.capable_tasks,
983 demand.unroutable_tasks,
984 ),
985 Err(error) => {
986 health.set_database(false);
987 warn!(%error, "could not read queue demand");
988 }
989 }
990 }
991 }
992 }
993}