1use crate::executor::{DispatchedJob, JobExecutor};
2use crate::runtime::InFlightMap;
3use crate::storage::RuntimeStorage;
4use awa_model::JobRow;
5use sqlx::PgPool;
6use std::collections::HashMap;
7use std::sync::atomic::{AtomicBool, Ordering};
8use std::sync::{Arc, Mutex as StdMutex, OnceLock};
9use std::time::{Duration, Instant};
10use tokio::sync::{Mutex, Notify, Semaphore};
11use tokio::task::JoinSet;
12use tokio_util::sync::CancellationToken;
13use tracing::{debug, error, info, warn};
14use uuid::Uuid;
15
16const DEFAULT_CLAIM_BATCH_LIMIT: usize = 512;
17const MAX_CLAIMERS_PER_QUEUE: i16 = 4;
18const CLAIMER_LEASE_TTL: Duration = Duration::from_secs(3);
19const CLAIMER_IDLE_THRESHOLD: Duration = Duration::from_millis(500);
20
21fn max_claimers_per_queue() -> i16 {
22 static MAX_CLAIMERS: OnceLock<i16> = OnceLock::new();
23 *MAX_CLAIMERS.get_or_init(|| {
24 let Ok(raw) = std::env::var("AWA_MAX_CLAIMERS_PER_QUEUE") else {
25 return MAX_CLAIMERS_PER_QUEUE;
26 };
27
28 match raw.parse::<i16>() {
29 Ok(value) if value > 0 => value,
30 Ok(value) => {
31 warn!(
32 value,
33 default = MAX_CLAIMERS_PER_QUEUE,
34 "AWA_MAX_CLAIMERS_PER_QUEUE must be positive; using default"
35 );
36 MAX_CLAIMERS_PER_QUEUE
37 }
38 Err(error) => {
39 warn!(
40 raw = %raw,
41 %error,
42 default = MAX_CLAIMERS_PER_QUEUE,
43 "Failed to parse AWA_MAX_CLAIMERS_PER_QUEUE; using default"
44 );
45 MAX_CLAIMERS_PER_QUEUE
46 }
47 }
48 })
49}
50
51#[derive(Debug, Clone, Copy)]
52enum WakeReason {
53 Notify,
54 Capacity,
55 Poll,
56}
57
58impl WakeReason {
59 fn as_str(self) -> &'static str {
60 match self {
61 WakeReason::Notify => "notify",
62 WakeReason::Capacity => "capacity",
63 WakeReason::Poll => "poll",
64 }
65 }
66}
67
68#[derive(Debug, Clone)]
70pub struct RateLimit {
71 pub max_rate: f64,
73 pub burst: u32,
75}
76
77pub(crate) struct TokenBucket {
79 tokens: f64,
80 max_tokens: f64,
81 refill_rate: f64,
82 last_refill: Instant,
83}
84
85impl TokenBucket {
86 fn new(rate_limit: &RateLimit) -> Self {
87 let burst = if rate_limit.burst == 0 {
88 (rate_limit.max_rate.ceil() as u32).max(1)
89 } else {
90 rate_limit.burst
91 };
92 Self {
93 tokens: burst as f64,
94 max_tokens: burst as f64,
95 refill_rate: rate_limit.max_rate,
96 last_refill: Instant::now(),
97 }
98 }
99
100 fn available(&mut self) -> u32 {
102 let now = Instant::now();
103 let elapsed = now.duration_since(self.last_refill).as_secs_f64();
104 self.tokens = (self.tokens + elapsed * self.refill_rate).min(self.max_tokens);
105 self.last_refill = now;
106 self.tokens.floor() as u32
107 }
108
109 fn consume(&mut self, n: u32) {
111 self.tokens -= n as f64;
112 }
113
114 fn refund(&mut self, n: u32) {
115 self.tokens = (self.tokens + n as f64).min(self.max_tokens);
116 }
117}
118
119pub(crate) fn shared_rate_limiter(config: &QueueConfig) -> Option<Arc<StdMutex<TokenBucket>>> {
120 config
121 .rate_limit
122 .as_ref()
123 .map(TokenBucket::new)
124 .map(|bucket| Arc::new(StdMutex::new(bucket)))
125}
126
127#[derive(Debug, Clone)]
129pub struct QueueConfig {
130 pub max_workers: u32,
131 pub poll_interval: Duration,
132 pub deadline_duration: Duration,
133 pub priority_aging_interval: Duration,
134 pub rate_limit: Option<RateLimit>,
136 pub min_workers: u32,
138 pub weight: u32,
140 pub claimers: u16,
146 pub claim_batch_size: usize,
148}
149
150impl Default for QueueConfig {
151 fn default() -> Self {
152 Self {
153 max_workers: 50,
154 poll_interval: Duration::from_millis(200),
155 deadline_duration: Duration::from_secs(300), priority_aging_interval: Duration::from_secs(60),
157 rate_limit: None,
158 min_workers: 0,
159 weight: 1,
160 claimers: 1,
161 claim_batch_size: DEFAULT_CLAIM_BATCH_LIMIT,
162 }
163 }
164}
165
166#[allow(dead_code)]
169pub(crate) enum DispatchPermit {
170 Hard(tokio::sync::OwnedSemaphorePermit),
172 Local(tokio::sync::OwnedSemaphorePermit),
174 Overflow {
176 pool: Arc<OverflowPool>,
177 queue: String,
178 },
179}
180
181impl Drop for DispatchPermit {
182 fn drop(&mut self) {
183 if let DispatchPermit::Overflow { pool, queue } = self {
184 pool.release(queue, 1);
185 }
186 }
188}
189
190pub(crate) enum ConcurrencyMode {
192 HardReserved { semaphore: Arc<Semaphore> },
194 Weighted {
196 local_semaphore: Arc<Semaphore>,
197 overflow_pool: Arc<OverflowPool>,
198 queue_name: String,
199 },
200}
201
202pub(crate) struct OverflowPool {
205 total: u32,
206 state: std::sync::Mutex<OverflowState>,
207}
208
209struct OverflowState {
210 held: HashMap<String, u32>,
212 demand: HashMap<String, u32>,
214 weights: HashMap<String, u32>,
216}
217
218impl OverflowPool {
219 pub fn new(total: u32, weights: HashMap<String, u32>) -> Self {
220 Self {
221 total,
222 state: std::sync::Mutex::new(OverflowState {
223 held: HashMap::new(),
224 demand: HashMap::new(),
225 weights,
226 }),
227 }
228 }
229
230 pub fn try_acquire(&self, queue: &str, wanted: u32) -> u32 {
235 let mut state = self.state.lock().unwrap();
236
237 state.demand.insert(queue.to_string(), wanted);
239
240 if wanted == 0 {
241 return 0;
242 }
243
244 let currently_used: u32 = state.held.values().sum();
245 let available = self.total.saturating_sub(currently_used);
246 if available == 0 {
247 return 0;
248 }
249
250 let my_weight = state.weights.get(queue).copied().unwrap_or(1);
251
252 let contending_weight: u32 = state
254 .weights
255 .iter()
256 .filter(|(q, _)| {
257 state.demand.get(q.as_str()).copied().unwrap_or(0) > 0
258 || state.held.get(q.as_str()).copied().unwrap_or(0) > 0
259 })
260 .map(|(_, w)| *w)
261 .sum();
262
263 if contending_weight == 0 {
264 return 0;
265 }
266
267 let my_fair_share =
269 ((self.total as f64) * (my_weight as f64 / contending_weight as f64)).ceil() as u32;
270 let my_held = state.held.get(queue).copied().unwrap_or(0);
271 let room = my_fair_share.saturating_sub(my_held);
272
273 let granted = wanted.min(available).min(room);
274 if granted > 0 {
275 *state.held.entry(queue.to_string()).or_insert(0) += granted;
276 }
277 granted
278 }
279
280 pub fn release(&self, queue: &str, n: u32) {
282 let mut state = self.state.lock().unwrap();
283 if let Some(held) = state.held.get_mut(queue) {
284 *held = held.saturating_sub(n);
285 }
286 }
287
288 pub fn held(&self, queue: &str) -> u32 {
290 let state = self.state.lock().unwrap();
291 state.held.get(queue).copied().unwrap_or(0)
292 }
293}
294
295pub struct Dispatcher {
297 queue: String,
298 runtime_instance_id: Uuid,
299 config: QueueConfig,
300 pool: PgPool,
301 executor: Arc<JobExecutor>,
302 metrics: crate::metrics::AwaMetrics,
303 _in_flight: InFlightMap,
304 concurrency: ConcurrencyMode,
305 alive: Arc<AtomicBool>,
306 cancel: CancellationToken,
307 job_set: Arc<Mutex<JoinSet<()>>>,
308 rate_limiter: Option<Arc<StdMutex<TokenBucket>>>,
309 storage: RuntimeStorage,
310 capacity_wake: Arc<Notify>,
311 claimer_owner_id: Uuid,
312}
313
314impl Dispatcher {
315 #[allow(clippy::too_many_arguments)]
317 pub(crate) fn with_concurrency(
318 queue: String,
319 runtime_instance_id: Uuid,
320 config: QueueConfig,
321 pool: PgPool,
322 executor: Arc<JobExecutor>,
323 metrics: crate::metrics::AwaMetrics,
324 in_flight: InFlightMap,
325 alive: Arc<AtomicBool>,
326 cancel: CancellationToken,
327 job_set: Arc<Mutex<JoinSet<()>>>,
328 concurrency: ConcurrencyMode,
329 rate_limiter: Option<Arc<StdMutex<TokenBucket>>>,
330 capacity_wake: Arc<Notify>,
331 claimer_owner_id: Uuid,
332 storage: RuntimeStorage,
333 ) -> Self {
334 Self {
335 queue,
336 runtime_instance_id,
337 config,
338 pool,
339 executor,
340 metrics,
341 _in_flight: in_flight,
342 concurrency,
343 alive,
344 cancel,
345 job_set,
346 rate_limiter,
347 storage,
348 capacity_wake,
349 claimer_owner_id,
350 }
351 }
352
353 #[tracing::instrument(skip(self), fields(queue = %self.queue))]
355 pub async fn run(mut self) {
356 self.alive.store(true, Ordering::SeqCst);
357 info!(
358 queue = %self.queue,
359 runtime_instance_id = %self.runtime_instance_id,
360 claimer_owner_id = %self.claimer_owner_id,
361 poll_interval_ms = self.config.poll_interval.as_millis(),
362 "Dispatcher started"
363 );
364
365 let notify_channel = format!("awa:{}", self.queue);
367 let mut listener = match sqlx::postgres::PgListener::connect_with(&self.pool).await {
368 Ok(listener) => listener,
369 Err(err) => {
370 error!(error = %err, "Failed to create PG listener, falling back to polling only");
371 self.poll_loop_only().await;
373 self.alive.store(false, Ordering::SeqCst);
374 return;
375 }
376 };
377
378 if let Err(err) = listener.listen(¬ify_channel).await {
379 warn!(error = %err, channel = %notify_channel, "Failed to LISTEN, falling back to polling");
380 self.poll_loop_only().await;
381 self.alive.store(false, Ordering::SeqCst);
382 return;
383 }
384
385 debug!(channel = %notify_channel, "Listening for job notifications");
386
387 loop {
388 tokio::select! {
389 _ = self.cancel.cancelled() => {
390 debug!(queue = %self.queue, "Dispatcher shutting down");
391 break;
392 }
393 notification = listener.recv() => {
395 match notification {
396 Ok(_) => {
397 debug!(queue = %self.queue, "Woken by NOTIFY");
398 self.drain_ready(WakeReason::Notify, Instant::now()).await;
399 }
400 Err(err) => {
401 warn!(error = %err, "PG listener error, will retry");
402 tokio::time::sleep(Duration::from_secs(1)).await;
403 }
404 }
405 }
406 _ = self.capacity_wake.notified() => {
407 self.drain_ready(WakeReason::Capacity, Instant::now()).await;
408 }
409 _ = tokio::time::sleep(self.config.poll_interval) => {
410 self.drain_ready(WakeReason::Poll, Instant::now()).await;
411 }
412 }
413 }
414
415 self.alive.store(false, Ordering::SeqCst);
416 }
417
418 async fn poll_loop_only(&mut self) {
420 loop {
421 tokio::select! {
422 _ = self.cancel.cancelled() => {
423 debug!(queue = %self.queue, "Dispatcher (poll-only) shutting down");
424 break;
425 }
426 _ = self.capacity_wake.notified() => {
427 self.drain_ready(WakeReason::Capacity, Instant::now()).await;
428 }
429 _ = tokio::time::sleep(self.config.poll_interval) => {
430 self.drain_ready(WakeReason::Poll, Instant::now()).await;
431 }
432 }
433 }
434 }
435
436 fn acquire_permits(&mut self) -> Vec<DispatchPermit> {
438 let mut permits = Vec::new();
439 match &self.concurrency {
440 ConcurrencyMode::HardReserved { semaphore } => {
441 for _ in 0..self.config.claim_batch_size {
442 match semaphore.clone().try_acquire_owned() {
443 Ok(p) => permits.push(DispatchPermit::Hard(p)),
444 Err(_) => break,
445 }
446 }
447 }
448 ConcurrencyMode::Weighted {
449 local_semaphore,
450 overflow_pool,
451 queue_name,
452 } => {
453 for _ in 0..self.config.claim_batch_size {
455 match local_semaphore.clone().try_acquire_owned() {
456 Ok(p) => permits.push(DispatchPermit::Local(p)),
457 Err(_) => break,
458 }
459 }
460 let overflow_wanted =
462 (self.config.claim_batch_size.saturating_sub(permits.len())) as u32;
463 let granted = overflow_pool.try_acquire(queue_name, overflow_wanted);
464 for _ in 0..granted {
465 permits.push(DispatchPermit::Overflow {
466 pool: overflow_pool.clone(),
467 queue: queue_name.clone(),
468 });
469 }
470 }
471 }
472 permits
473 }
474
475 fn refund_rate_limit(&self, n: usize) {
476 if n == 0 {
477 return;
478 }
479 if let Some(rate_limiter) = &self.rate_limiter {
480 let mut rate_limiter = rate_limiter.lock().expect("rate limiter lock poisoned");
481 rate_limiter.refund(n as u32);
482 }
483 }
484
485 async fn drain_ready(&mut self, wake_reason: WakeReason, woke_at: Instant) {
488 self.metrics
489 .record_dispatch_wake(&self.queue, wake_reason.as_str());
490 let mut first_iteration = true;
491 while !self.cancel.is_cancelled() {
492 let wake_context = first_iteration.then_some((wake_reason, woke_at));
493 if !self.poll_once(wake_context).await {
494 break;
495 }
496 first_iteration = false;
497 }
498 }
499
500 #[tracing::instrument(skip(self), fields(queue = %self.queue))]
502 async fn poll_once(&mut self, wake_context: Option<(WakeReason, Instant)>) -> bool {
503 let mut permits = self.acquire_permits();
505 if permits.is_empty() {
506 return false;
507 }
508 if let Some((reason, woke_at)) = wake_context {
509 self.metrics.record_dispatch_wake_to_claim(
510 &self.queue,
511 reason.as_str(),
512 woke_at.elapsed(),
513 );
514 self.metrics.record_dispatch_capacity_available(
515 &self.queue,
516 reason.as_str(),
517 permits.len() as u64,
518 );
519 }
520
521 let batch_size = if let Some(rate_limiter) = &self.rate_limiter {
523 let mut rate_limiter = rate_limiter.lock().expect("rate limiter lock poisoned");
524 let rate_available = rate_limiter.available() as usize;
525 let batch_size = permits
526 .len()
527 .min(rate_available)
528 .min(self.config.claim_batch_size);
529 if batch_size > 0 {
530 rate_limiter.consume(batch_size as u32);
531 }
532 batch_size
533 } else {
534 permits.len().min(self.config.claim_batch_size)
535 };
536 if batch_size == 0 {
537 if let Some((reason, _)) = wake_context {
539 self.metrics
540 .record_dispatch_rate_limited(&self.queue, reason.as_str());
541 }
542 return false;
543 }
544 while permits.len() > batch_size {
546 permits.pop(); }
548
549 let deadline_secs = self.config.deadline_duration.as_secs_f64();
561 let claim_start = Instant::now();
562
563 let jobs: Vec<DispatchedJob> = match &self.storage {
564 RuntimeStorage::Canonical => match sqlx::query_as::<_, JobRow>(
565 r#"
566 WITH claimed AS (
567 SELECT id
568 FROM awa.jobs_hot
569 WHERE state = 'available'
570 AND queue = $1
571 AND run_at <= now()
572 AND NOT EXISTS (
573 SELECT 1 FROM awa.queue_meta
574 WHERE queue = $1 AND paused = TRUE
575 )
576 ORDER BY priority ASC, run_at ASC, id ASC
577 LIMIT $2
578 FOR UPDATE SKIP LOCKED
579 )
580 UPDATE awa.jobs_hot
581 SET state = 'running',
582 attempt = attempt + 1,
583 run_lease = run_lease + 1,
584 attempted_at = now(),
585 heartbeat_at = now(),
586 deadline_at = now() + make_interval(secs => $3)
587 FROM claimed
588 WHERE awa.jobs_hot.id = claimed.id
589 AND awa.jobs_hot.state = 'available'
590 RETURNING awa.jobs_hot.*
591 "#,
592 )
593 .bind(&self.queue)
594 .bind(batch_size as i32)
595 .bind(deadline_secs)
596 .fetch_all(&self.pool)
597 .await
598 {
599 Ok(jobs) => jobs
600 .into_iter()
601 .map(|job| DispatchedJob {
602 job,
603 queue_storage_claim: None,
604 queue_storage_unique_states: None,
605 })
606 .collect(),
607 Err(err) => {
608 warn!(queue = %self.queue, error = %err, "Failed to claim jobs");
609 self.refund_rate_limit(batch_size);
610 return false;
611 }
612 },
613 RuntimeStorage::QueueStorage(runtime) => match runtime
614 .store
615 .claim_runtime_batch_with_aging_for_instance(
616 &self.pool,
617 &self.queue,
618 batch_size as i64,
619 self.config.deadline_duration,
620 self.config.priority_aging_interval,
621 self.claimer_owner_id,
622 max_claimers_per_queue().max(self.config.claimers as i16),
623 CLAIMER_LEASE_TTL,
624 CLAIMER_IDLE_THRESHOLD,
625 )
626 .await
627 {
628 Ok(jobs) => jobs
629 .into_iter()
630 .map(|claimed| DispatchedJob {
631 job: claimed.job,
632 queue_storage_claim: Some(claimed.claim),
633 queue_storage_unique_states: claimed.unique_states,
634 })
635 .collect(),
636 Err(err) => {
637 warn!(
638 queue = %self.queue,
639 error = %err,
640 "Failed to claim queue storage jobs"
641 );
642 self.refund_rate_limit(batch_size);
643 return false;
644 }
645 },
646 };
647 self.metrics
648 .record_claim_batch(&self.queue, jobs.len() as u64, claim_start.elapsed());
649 if !jobs.is_empty() {
650 let mut by_shard: std::collections::BTreeMap<i16, u64> =
656 std::collections::BTreeMap::new();
657 let mut canonical_count: u64 = 0;
658 for dispatched in &jobs {
659 match &dispatched.queue_storage_claim {
660 Some(claim) => *by_shard.entry(claim.enqueue_shard).or_default() += 1,
661 None => canonical_count += 1,
662 }
663 }
664 for (shard, count) in by_shard {
665 self.metrics
666 .record_job_claimed_by_shard(&self.queue, shard, count);
667 }
668 if canonical_count > 0 {
669 self.metrics
670 .record_job_claimed(&self.queue, canonical_count);
671 }
672 let now = chrono::Utc::now();
685 for job in &jobs {
686 let claim_at = job.job.attempted_at.unwrap_or(now);
687 let wait_secs = (claim_at - job.job.created_at).num_milliseconds() as f64 / 1000.0;
688 if wait_secs >= 0.0 {
689 self.metrics.record_wait_duration(&self.queue, wait_secs);
690 }
691 }
692 }
693 if jobs.len() < batch_size {
694 self.refund_rate_limit(batch_size - jobs.len());
695 }
696
697 let unused_permits = permits.len().saturating_sub(jobs.len());
699 while permits.len() > jobs.len() {
700 permits.pop();
701 }
702 if unused_permits > 0 {
703 self.metrics
704 .record_dispatch_unused_permits(&self.queue, unused_permits as u64);
705 }
706
707 if jobs.is_empty() {
709 if let Some((reason, _)) = wake_context {
710 self.metrics
711 .record_dispatch_empty_claim(&self.queue, reason.as_str());
712 }
713 if let ConcurrencyMode::Weighted {
714 overflow_pool,
715 queue_name,
716 ..
717 } = &self.concurrency
718 {
719 overflow_pool.try_acquire(queue_name, 0);
720 }
721 return false;
722 }
723
724 debug!(queue = %self.queue, count = jobs.len(), "Claimed jobs");
725
726 let mut set = self.job_set.lock().await;
728 while set.try_join_next().is_some() {}
734 for (job, permit) in jobs.into_iter().zip(permits) {
735 let cancel_flag = Arc::new(AtomicBool::new(false));
736 let task = self.executor.execute_task(job, cancel_flag);
737 let capacity_wake = self.capacity_wake.clone();
738 set.spawn(async move {
739 task.await;
740 drop(permit);
741 capacity_wake.notify_one();
742 });
743 }
744
745 true
746 }
747}