1use crate::{
2 client::{
3 AggregateQuery, AggregateResponse, BreakdownQuery, BreakdownResponse, CreateSiteRequest,
4 RealtimeVisitorsResponse, ResetSiteStatsRequest, SiteSummary, TimeseriesQuery,
5 TimeseriesResponse, UpdateSiteRequest,
6 },
7 rate_limit::{RateLimitError, RateLimiter, RateStatus},
8};
9use async_trait::async_trait;
10use std::num::NonZeroU32;
11use std::sync::atomic::{AtomicU64, Ordering};
12use std::sync::Arc;
13use std::time::Duration as StdDuration;
14use time::{Duration as TimeDuration, OffsetDateTime};
15use tokio::sync::{broadcast, mpsc, oneshot, Mutex, Notify};
16use tokio::time::sleep;
17
18const DEFAULT_QUEUE_CAPACITY: usize = 128;
19pub const DEFAULT_MAX_RETRIES: u32 = 3;
20const DEFAULT_RETRY_BASE_MS: u64 = 1_000;
21const MAX_RETRY_DELAY_MS: u64 = 30_000;
22
23pub type JobId = u64;
24pub type JobResult = Result<JobResponse, WorkerError>;
25
26#[derive(Debug, Clone)]
27pub struct JobRequest {
28 pub account: String,
29 pub kind: JobKind,
30 pub max_retries: u32,
31}
32
33impl JobRequest {
34 pub fn description(&self) -> String {
35 match &self.kind {
36 JobKind::ListSites => format!("List sites for {}", self.account),
37 JobKind::StatsAggregate { query } => {
38 format!("Aggregate stats for {} ({:?})", self.account, query.metrics)
39 }
40 JobKind::StatsTimeseries { query } => format!(
41 "Timeseries stats for {} ({:?})",
42 self.account, query.metrics
43 ),
44 JobKind::StatsBreakdown { query } => {
45 format!("Breakdown stats for {} ({:?})", self.account, query.metrics)
46 }
47 JobKind::SiteCreate { request } => {
48 format!("Create site {}", request.domain)
49 }
50 JobKind::SiteUpdate { site_id, .. } => {
51 format!("Update site {}", site_id)
52 }
53 JobKind::SiteReset { site_id, .. } => {
54 format!("Reset stats for {}", site_id)
55 }
56 JobKind::SiteDelete { site_id } => format!("Delete site {}", site_id),
57 JobKind::StatsRealtime { site_id } => {
58 format!("Realtime stats for {}", site_id)
59 }
60 JobKind::EventSend { .. } => format!("Send event for {}", self.account),
61 JobKind::EventsImport { events } => {
62 format!("Import {} events for {}", events.len(), self.account)
63 }
64 JobKind::Custom { label } => label.clone(),
65 }
66 }
67}
68
69#[derive(Debug, Clone)]
70pub enum JobKind {
71 ListSites,
72 StatsAggregate {
73 query: Box<AggregateQuery>,
74 },
75 StatsTimeseries {
76 query: Box<TimeseriesQuery>,
77 },
78 StatsBreakdown {
79 query: Box<BreakdownQuery>,
80 },
81 SiteCreate {
82 request: Box<CreateSiteRequest>,
83 },
84 SiteUpdate {
85 site_id: String,
86 request: Box<UpdateSiteRequest>,
87 },
88 SiteReset {
89 site_id: String,
90 request: Box<ResetSiteStatsRequest>,
91 },
92 SiteDelete {
93 site_id: String,
94 },
95 StatsRealtime {
96 site_id: String,
97 },
98 EventSend {
99 event: serde_json::Value,
100 },
101 EventsImport {
102 events: Vec<serde_json::Value>,
103 },
104 Custom {
105 label: String,
106 },
107}
108
109#[derive(Debug, Clone)]
110pub enum JobResponse {
111 Sites(Vec<SiteSummary>),
112 StatsAggregate(AggregateResponse),
113 StatsTimeseries(TimeseriesResponse),
114 StatsBreakdown(BreakdownResponse),
115 SiteCreated(SiteSummary),
116 SiteUpdated(SiteSummary),
117 SiteReset,
118 SiteDeleted,
119 StatsRealtime(RealtimeVisitorsResponse),
120 EventAck,
121 EventsProcessed { processed: usize },
122 Acknowledged,
123 Custom(serde_json::Value),
124}
125
126#[derive(Debug, Clone)]
127pub struct QueueSnapshot {
128 pub id: JobId,
129 pub account: String,
130 pub description: String,
131 pub state: QueueJobState,
132 pub enqueued_at: OffsetDateTime,
133 pub started_at: Option<OffsetDateTime>,
134 pub attempt: u32,
135 pub max_retries: u32,
136 pub last_error: Option<String>,
137 pub next_retry_at: Option<OffsetDateTime>,
138}
139
140#[derive(Debug, Clone, Copy, PartialEq, Eq)]
141pub enum QueueJobState {
142 Pending,
143 InFlight,
144}
145
146#[derive(Debug, Clone)]
147struct JobStateEntry {
148 id: JobId,
149 account: String,
150 description: String,
151 enqueued_at: OffsetDateTime,
152 started_at: Option<OffsetDateTime>,
153 attempt: u32,
154 max_retries: u32,
155 last_error: Option<String>,
156 next_retry_at: Option<OffsetDateTime>,
157}
158
159#[derive(Debug, Clone)]
160pub struct TelemetryEvent {
161 pub job_id: JobId,
162 pub account: String,
163 pub description: String,
164 pub kind: TelemetryKind,
165 pub timestamp: OffsetDateTime,
166 pub status: Option<RateStatus>,
167 pub attempt: u32,
168 pub max_retries: u32,
169 pub error: Option<String>,
170 pub next_retry_at: Option<OffsetDateTime>,
171}
172
173#[derive(Debug, Clone, Copy, PartialEq, Eq)]
174pub enum TelemetryKind {
175 Enqueued,
176 Started,
177 Succeeded,
178 Failed,
179}
180
181#[derive(Clone)]
182pub struct QueueHandle {
183 sender: mpsc::Sender<JobMessage>,
184 telemetry: broadcast::Sender<TelemetryEvent>,
185 sequence: Arc<AtomicU64>,
186 state: Arc<Mutex<Vec<JobStateEntry>>>,
187 idle_notify: Arc<Notify>,
188}
189
190impl QueueHandle {
191 pub async fn submit(
192 &self,
193 request: JobRequest,
194 weight: NonZeroU32,
195 ) -> Result<JobTicket, WorkerError> {
196 let id = self.sequence.fetch_add(1, Ordering::SeqCst) + 1;
197 let description = request.description();
198 let account = request.account.clone();
199 let enqueued_at = OffsetDateTime::now_utc();
200 let (tx, rx) = oneshot::channel();
201 let message = JobMessage {
202 id,
203 request: request.clone(),
204 weight,
205 responder: Some(tx),
206 enqueued_at,
207 attempt: 0,
208 max_retries: request.max_retries,
209 };
210 {
211 let mut state = self.state.lock().await;
212 state.push(JobStateEntry {
213 id,
214 account: account.clone(),
215 description: description.clone(),
216 enqueued_at,
217 started_at: None,
218 attempt: 0,
219 max_retries: request.max_retries,
220 last_error: None,
221 next_retry_at: None,
222 });
223 }
224 self.telemetry
225 .send(TelemetryEvent {
226 job_id: id,
227 account: account.clone(),
228 description: description.clone(),
229 kind: TelemetryKind::Enqueued,
230 timestamp: message.enqueued_at,
231 status: None,
232 attempt: 0,
233 max_retries: request.max_retries,
234 error: None,
235 next_retry_at: None,
236 })
237 .ok();
238 if self.sender.send(message).await.is_err() {
239 let mut state = self.state.lock().await;
240 state.retain(|entry| entry.id != id);
241 if state.is_empty() {
242 self.idle_notify.notify_waiters();
243 }
244 return Err(WorkerError::QueueClosed);
245 }
246 Ok(JobTicket { id, receiver: rx })
247 }
248
249 pub fn subscribe(&self) -> broadcast::Receiver<TelemetryEvent> {
250 self.telemetry.subscribe()
251 }
252
253 pub async fn snapshot(&self) -> Vec<QueueSnapshot> {
254 let state = self.state.lock().await;
255 state
256 .iter()
257 .map(|entry| QueueSnapshot {
258 id: entry.id,
259 account: entry.account.clone(),
260 description: entry.description.clone(),
261 state: if entry.started_at.is_some() {
262 QueueJobState::InFlight
263 } else {
264 QueueJobState::Pending
265 },
266 enqueued_at: entry.enqueued_at,
267 started_at: entry.started_at,
268 attempt: entry.attempt,
269 max_retries: entry.max_retries,
270 last_error: entry.last_error.clone(),
271 next_retry_at: entry.next_retry_at,
272 })
273 .collect()
274 }
275
276 pub async fn wait_idle(&self) {
277 loop {
278 if self.is_idle().await {
279 return;
280 }
281 let notified = self.idle_notify.notified();
282 if self.is_idle().await {
283 return;
284 }
285 notified.await;
286 }
287 }
288
289 pub async fn is_idle(&self) -> bool {
290 let state = self.state.lock().await;
291 state.is_empty()
292 }
293}
294
295pub struct JobTicket {
296 id: JobId,
297 receiver: oneshot::Receiver<JobResult>,
298}
299
300impl JobTicket {
301 pub fn id(&self) -> JobId {
302 self.id
303 }
304
305 pub async fn await_result(self) -> JobResult {
306 self.receiver.await.unwrap_or(Err(WorkerError::QueueClosed))
307 }
308}
309
310struct JobMessage {
311 id: JobId,
312 request: JobRequest,
313 weight: NonZeroU32,
314 responder: Option<oneshot::Sender<JobResult>>,
315 enqueued_at: OffsetDateTime,
316 attempt: u32,
317 max_retries: u32,
318}
319
320#[async_trait]
321pub trait JobExecutor: Send + Sync + 'static {
322 async fn execute(&self, request: JobRequest) -> JobResult;
323}
324
325pub struct Worker<E: JobExecutor> {
326 queue: mpsc::Receiver<JobMessage>,
327 sender: mpsc::Sender<JobMessage>,
328 telemetry: broadcast::Sender<TelemetryEvent>,
329 executor: Arc<E>,
330 rate_limiter: RateLimiter,
331 state: Arc<Mutex<Vec<JobStateEntry>>>,
332 idle_notify: Arc<Notify>,
333}
334
335impl<E: JobExecutor> Worker<E> {
336 pub fn spawn(
337 executor: Arc<E>,
338 rate_limiter: RateLimiter,
339 capacity: Option<usize>,
340 ) -> QueueHandle {
341 let (tx, rx) = mpsc::channel(capacity.unwrap_or(DEFAULT_QUEUE_CAPACITY));
342 let (telemetry_tx, _) = broadcast::channel(256);
343 let state = Arc::new(Mutex::new(Vec::new()));
344 let idle_notify = Arc::new(Notify::new());
345 let worker = Self {
346 queue: rx,
347 sender: tx.clone(),
348 telemetry: telemetry_tx.clone(),
349 executor,
350 rate_limiter: rate_limiter.clone(),
351 state: state.clone(),
352 idle_notify: idle_notify.clone(),
353 };
354 tokio::spawn(worker.run());
355 QueueHandle {
356 sender: tx,
357 telemetry: telemetry_tx,
358 sequence: Arc::new(AtomicU64::new(0)),
359 state,
360 idle_notify,
361 }
362 }
363
364 async fn run(mut self) {
365 while let Some(mut message) = self.queue.recv().await {
366 let id = message.id;
367 let attempt_number = message.attempt + 1;
368 let max_retries = message.max_retries;
369 let description = message.request.description();
370 let account = message.request.account.clone();
371 let weight = message.weight;
372
373 let start = OffsetDateTime::now_utc();
374 self.update_state_on_start(id, start, attempt_number).await;
375 self.telemetry
376 .send(TelemetryEvent {
377 job_id: id,
378 account: account.clone(),
379 description: description.clone(),
380 kind: TelemetryKind::Started,
381 timestamp: start,
382 status: None,
383 attempt: attempt_number,
384 max_retries,
385 error: None,
386 next_retry_at: None,
387 })
388 .ok();
389
390 self.rate_limiter.acquire(weight).await;
391 let result = self.executor.execute(message.request.clone()).await;
392
393 match result {
394 Ok(response) => {
395 let status = self
396 .rate_limiter
397 .record_success(weight.get(), OffsetDateTime::now_utc())
398 .await
399 .ok();
400 self.telemetry
401 .send(TelemetryEvent {
402 job_id: id,
403 account: account.clone(),
404 description: description.clone(),
405 kind: TelemetryKind::Succeeded,
406 timestamp: OffsetDateTime::now_utc(),
407 status,
408 attempt: attempt_number,
409 max_retries,
410 error: None,
411 next_retry_at: None,
412 })
413 .ok();
414 if let Some(tx) = message.responder.take() {
415 let _ = tx.send(Ok(response));
416 }
417 self.finish_job(id).await;
418 }
419 Err(err) => {
420 let error_msg = err.to_string();
421 let now = OffsetDateTime::now_utc();
422 let next_attempt = message.attempt + 1;
423 if next_attempt > max_retries {
424 self.telemetry
425 .send(TelemetryEvent {
426 job_id: id,
427 account: account.clone(),
428 description: description.clone(),
429 kind: TelemetryKind::Failed,
430 timestamp: now,
431 status: None,
432 attempt: attempt_number,
433 max_retries,
434 error: Some(error_msg.clone()),
435 next_retry_at: None,
436 })
437 .ok();
438 if let Some(tx) = message.responder.take() {
439 let _ = tx.send(Err(err));
440 }
441 self.finish_job(id).await;
442 } else {
443 let delay = retry_delay(attempt_number);
444 let next_retry_at = now.checked_add(to_time_duration(delay));
445 self.update_state_on_retry(
446 id,
447 next_attempt,
448 max_retries,
449 error_msg.clone(),
450 next_retry_at,
451 )
452 .await;
453 self.telemetry
454 .send(TelemetryEvent {
455 job_id: id,
456 account: account.clone(),
457 description: description.clone(),
458 kind: TelemetryKind::Failed,
459 timestamp: now,
460 status: None,
461 attempt: attempt_number,
462 max_retries,
463 error: Some(error_msg.clone()),
464 next_retry_at,
465 })
466 .ok();
467 message.attempt += 1;
468 let sender = self.sender.clone();
469 tokio::spawn(async move {
470 sleep(delay).await;
471 if let Err(err) = sender.send(message).await {
472 let mut failed_message = err.0;
473 if let Some(tx) = failed_message.responder.take() {
474 let _ = tx.send(Err(WorkerError::QueueClosed));
475 }
476 }
477 });
478 }
479 }
480 }
481 }
482 }
483
484 async fn finish_job(&self, id: JobId) {
485 let mut state = self.state.lock().await;
486 state.retain(|entry| entry.id != id);
487 if state.is_empty() {
488 self.idle_notify.notify_waiters();
489 }
490 }
491
492 async fn update_state_on_start(&self, id: JobId, started: OffsetDateTime, attempt: u32) {
493 let mut state = self.state.lock().await;
494 if let Some(entry) = state.iter_mut().find(|entry| entry.id == id) {
495 entry.started_at = Some(started);
496 entry.attempt = attempt;
497 entry.last_error = None;
498 entry.next_retry_at = None;
499 }
500 }
501
502 async fn update_state_on_retry(
503 &self,
504 id: JobId,
505 attempt: u32,
506 max_retries: u32,
507 error: String,
508 next_retry_at: Option<OffsetDateTime>,
509 ) {
510 let mut state = self.state.lock().await;
511 if let Some(entry) = state.iter_mut().find(|entry| entry.id == id) {
512 entry.started_at = None;
513 entry.attempt = attempt;
514 entry.max_retries = max_retries;
515 entry.last_error = Some(error);
516 entry.next_retry_at = next_retry_at;
517 }
518 }
519}
520
521fn retry_delay(attempt_number: u32) -> StdDuration {
522 let exponent = attempt_number.saturating_sub(1).min(8);
523 let factor = 1u64 << exponent;
524 let millis = DEFAULT_RETRY_BASE_MS
525 .saturating_mul(factor)
526 .min(MAX_RETRY_DELAY_MS);
527 StdDuration::from_millis(millis)
528}
529
530fn to_time_duration(delay: StdDuration) -> TimeDuration {
531 let millis = delay.as_millis();
532 if millis > i64::MAX as u128 {
533 TimeDuration::milliseconds(i64::MAX)
534 } else {
535 TimeDuration::milliseconds(millis as i64)
536 }
537}
538
539#[derive(thiserror::Error, Debug)]
540pub enum WorkerError {
541 #[error("queue closed")]
542 QueueClosed,
543 #[error("rate limiter error: {0}")]
544 RateLimit(#[from] RateLimitError),
545 #[error("execution error: {0}")]
546 Execution(String),
547}
548
549#[cfg(test)]
550mod tests {
551 use super::*;
552 use crate::config::ConfigPaths;
553 use crate::rate_limit::RateLimitConfig;
554 use std::num::NonZeroU32;
555 use std::sync::Arc;
556 use std::sync::Mutex as StdMutex;
557 use tempfile::tempdir;
558 use tokio::sync::Notify;
559
560 struct TestExecutor {
561 calls: StdMutex<Vec<JobKind>>,
562 responses: StdMutex<Vec<JobResponse>>,
563 }
564
565 #[async_trait]
566 impl JobExecutor for TestExecutor {
567 async fn execute(&self, request: JobRequest) -> JobResult {
568 self.calls.lock().unwrap().push(request.kind.clone());
569 let response = self
570 .responses
571 .lock()
572 .unwrap()
573 .pop()
574 .unwrap_or(JobResponse::Acknowledged);
575 Ok(response)
576 }
577 }
578
579 async fn setup_rate_limiter() -> RateLimiter {
580 let tmp = tempdir().expect("tmpdir");
581 let paths = ConfigPaths::from_base_dir(tmp.path());
582 let config = RateLimitConfig::new(NonZeroU32::new(100).unwrap());
583 RateLimiter::new(paths, "test", config)
584 .await
585 .expect("limiter")
586 }
587
588 struct BlockingExecutor {
589 resume: Arc<Notify>,
590 started: Arc<Notify>,
591 }
592
593 impl BlockingExecutor {
594 fn new() -> Self {
595 Self {
596 resume: Arc::new(Notify::new()),
597 started: Arc::new(Notify::new()),
598 }
599 }
600
601 fn resume_notifier(&self) -> Arc<Notify> {
602 self.resume.clone()
603 }
604
605 fn started_notifier(&self) -> Arc<Notify> {
606 self.started.clone()
607 }
608 }
609
610 #[async_trait]
611 impl JobExecutor for BlockingExecutor {
612 async fn execute(&self, _request: JobRequest) -> JobResult {
613 let resume_wait = self.resume.notified();
614 self.started.notify_waiters();
615 resume_wait.await;
616 Ok(JobResponse::Acknowledged)
617 }
618 }
619
620 struct RetryingExecutor {
621 fail_until: usize,
622 calls: StdMutex<usize>,
623 }
624
625 #[async_trait]
626 impl JobExecutor for RetryingExecutor {
627 async fn execute(&self, _request: JobRequest) -> JobResult {
628 let mut guard = self.calls.lock().unwrap();
629 if *guard < self.fail_until {
630 *guard += 1;
631 Err(WorkerError::Execution("forced failure".into()))
632 } else {
633 Ok(JobResponse::Acknowledged)
634 }
635 }
636 }
637
638 #[tokio::test]
639 async fn worker_processes_jobs_and_respects_order() {
640 let executor = Arc::new(TestExecutor {
641 calls: StdMutex::new(Vec::new()),
642 responses: StdMutex::new(vec![JobResponse::Acknowledged]),
643 });
644 let rate_limiter = setup_rate_limiter().await;
645 let handle = Worker::spawn(executor.clone(), rate_limiter.clone(), Some(4));
646
647 let request = JobRequest {
648 account: "test".into(),
649 kind: JobKind::Custom {
650 label: "custom".into(),
651 },
652 max_retries: 0,
653 };
654 let ticket = handle
655 .submit(request, NonZeroU32::new(1).unwrap())
656 .await
657 .expect("enqueue");
658 let outcome = ticket.await_result().await.expect("result");
659 matches!(outcome, JobResponse::Acknowledged);
660 }
661
662 #[tokio::test]
663 async fn telemetry_receives_events() {
664 let executor = Arc::new(TestExecutor {
665 calls: StdMutex::new(Vec::new()),
666 responses: StdMutex::new(vec![JobResponse::Acknowledged]),
667 });
668 let rate_limiter = setup_rate_limiter().await;
669 let handle = Worker::spawn(executor, rate_limiter, Some(4));
670 let mut telemetry = handle.subscribe();
671
672 let request = JobRequest {
673 account: "acct".into(),
674 kind: JobKind::ListSites,
675 max_retries: 0,
676 };
677 let ticket = handle
678 .submit(request, NonZeroU32::new(1).unwrap())
679 .await
680 .expect("enqueue");
681
682 for _ in 0..3 {
683 let _event = telemetry.recv().await.expect("tele");
684 }
685
686 let _ = ticket.await_result().await.expect("result");
687 }
688
689 #[tokio::test]
690 async fn snapshot_reflects_inflight_jobs() {
691 let executor = Arc::new(BlockingExecutor::new());
692 let resume = executor.resume_notifier();
693 let started = executor.started_notifier();
694 let rate_limiter = setup_rate_limiter().await;
695 let handle = Worker::spawn(executor, rate_limiter, Some(4));
696
697 let ticket = handle
698 .submit(
699 JobRequest {
700 account: "acct".into(),
701 kind: JobKind::Custom {
702 label: "blocking".into(),
703 },
704 max_retries: 0,
705 },
706 NonZeroU32::new(1).unwrap(),
707 )
708 .await
709 .expect("enqueue");
710
711 started.notified().await;
712 let snapshot = handle.snapshot().await;
713 assert_eq!(snapshot.len(), 1);
714 assert_eq!(snapshot[0].account, "acct");
715 assert_eq!(snapshot[0].state, QueueJobState::InFlight);
716 assert!(snapshot[0].started_at.is_some());
717
718 resume.notify_waiters();
719 ticket.await_result().await.expect("result");
720 handle.wait_idle().await;
721 assert!(handle.snapshot().await.is_empty());
722 }
723
724 #[tokio::test]
725 async fn retries_until_success() {
726 let executor = Arc::new(RetryingExecutor {
727 fail_until: 1,
728 calls: StdMutex::new(0),
729 });
730 let rate_limiter = setup_rate_limiter().await;
731 let handle = Worker::spawn(executor.clone(), rate_limiter, Some(4));
732
733 let ticket = handle
734 .submit(
735 JobRequest {
736 account: "acct".into(),
737 kind: JobKind::Custom {
738 label: "retry".into(),
739 },
740 max_retries: 3,
741 },
742 NonZeroU32::new(1).unwrap(),
743 )
744 .await
745 .expect("enqueue");
746
747 let result = ticket.await_result().await;
748 assert!(result.is_ok());
749 handle.wait_idle().await;
750 assert_eq!(*executor.calls.lock().unwrap(), 1);
751 }
752
753 #[tokio::test]
754 async fn returns_error_after_max_retries() {
755 let executor = Arc::new(RetryingExecutor {
756 fail_until: 10,
757 calls: StdMutex::new(0),
758 });
759 let rate_limiter = setup_rate_limiter().await;
760 let handle = Worker::spawn(executor, rate_limiter, Some(4));
761
762 let ticket = handle
763 .submit(
764 JobRequest {
765 account: "acct".into(),
766 kind: JobKind::Custom {
767 label: "retry-fail".into(),
768 },
769 max_retries: 1,
770 },
771 NonZeroU32::new(1).unwrap(),
772 )
773 .await
774 .expect("enqueue");
775
776 let result = ticket.await_result().await;
777 assert!(matches!(result, Err(WorkerError::Execution(_))));
778 handle.wait_idle().await;
779 }
780}