1use std::sync::RwLock;
4
5use async_trait::async_trait;
6use boson_core::{
7 Job, JobEnqueueDisposition, JobStatus, QueueBackend, Result, Run, RunStatus, TaskConfig,
8 TaskRunStats,
9};
10use chrono::{DateTime, Utc};
11
12use crate::enqueue_rate::EnqueueRateLimiter;
13use crate::store::{read, write, write_fallible, Inner};
14
15#[derive(Debug)]
30pub struct MemQueueBackend {
31 inner: RwLock<Inner>,
32 enqueue_rate: EnqueueRateLimiter,
33}
34
35impl Default for MemQueueBackend {
36 fn default() -> Self {
37 Self::new()
38 }
39}
40
41impl MemQueueBackend {
42 #[must_use]
62 pub fn new() -> Self {
63 Self {
64 inner: RwLock::new(Inner::new()),
65 enqueue_rate: EnqueueRateLimiter::new(),
66 }
67 }
68}
69
70#[async_trait]
71impl QueueBackend for MemQueueBackend {
72 async fn upsert_job(&self, job: &Job) -> Result<()> {
73 write(&self.inner, |inner| crate::jobs::upsert_job(inner, job))
74 }
75
76 async fn enqueue_with_policies(
77 &self,
78 job: Job,
79 task_config: &TaskConfig,
80 ) -> Result<(String, JobEnqueueDisposition)> {
81 write_fallible(&self.inner, |inner| {
82 crate::jobs::enqueue_with_policies(inner, &self.enqueue_rate, &job, task_config)
83 })
84 }
85
86 async fn get_job(&self, job_id: &str) -> Result<Option<Job>> {
87 read(&self.inner, |inner| crate::jobs::get_job(inner, job_id))
88 }
89
90 async fn list_jobs(
91 &self,
92 status_filter: Option<JobStatus>,
93 offset: usize,
94 limit: usize,
95 ) -> Result<Vec<Job>> {
96 read(&self.inner, |inner| {
97 crate::jobs::list_jobs(inner, status_filter, offset, limit)
98 })
99 }
100
101 async fn cancel_job_if_active(&self, job_id: &str) -> Result<()> {
102 write_fallible(&self.inner, |inner| {
103 crate::jobs::cancel_job_if_active(inner, job_id)
104 })
105 }
106
107 async fn try_claim_job(&self, job_id: &str) -> Result<Option<Job>> {
108 write(&self.inner, |inner| {
109 crate::jobs::try_claim_job(inner, job_id)
110 })
111 }
112
113 async fn revert_job_to_queued(&self, job_id: &str) -> Result<()> {
114 write(&self.inner, |inner| {
115 crate::jobs::revert_job_to_queued(inner, job_id);
116 })
117 }
118
119 async fn distinct_pools_queued(&self) -> Result<Vec<String>> {
120 read(&self.inner, crate::jobs::distinct_pools_queued)
121 }
122
123 async fn list_queued_for_pool_sorted(&self, pool: &str, limit: usize) -> Result<Vec<Job>> {
124 read(&self.inner, |inner| {
125 crate::jobs::list_queued_for_pool_sorted(inner, pool, limit)
126 })
127 }
128
129 async fn pop_claim_from_pool(&self, pool: &str) -> Result<Option<Job>> {
130 write(&self.inner, |inner| {
131 crate::jobs::pop_claim_from_pool(inner, pool)
132 })
133 }
134
135 async fn count_jobs(&self, status_filter: Option<JobStatus>) -> Result<u64> {
136 read(&self.inner, |inner| {
137 crate::jobs::count_jobs(inner, status_filter)
138 })
139 }
140
141 async fn count_jobs_for_task(&self, task_name: &str, status: Option<JobStatus>) -> Result<u64> {
142 read(&self.inner, |inner| {
143 crate::jobs::count_jobs_for_task(inner, task_name, status)
144 })
145 }
146
147 async fn count_active_jobs_for_task(&self, task_name: &str) -> Result<u32> {
148 read(&self.inner, |inner| {
149 crate::jobs::count_active_jobs_for_task(inner, task_name)
150 })
151 }
152
153 async fn find_nonterminal_by_idempotency_key(&self, key: &str) -> Result<Option<String>> {
154 read(&self.inner, |inner| {
155 crate::jobs::find_nonterminal_by_idempotency_key(inner, key)
156 })
157 }
158
159 async fn upsert_run(&self, run: &Run) -> Result<()> {
160 write(&self.inner, |inner| crate::runs::upsert_run(inner, run))
161 }
162
163 async fn get_run(&self, run_id: &str) -> Result<Option<Run>> {
164 read(&self.inner, |inner| crate::runs::get_run(inner, run_id))
165 }
166
167 async fn list_runs(
168 &self,
169 job_id_filter: Option<&str>,
170 offset: usize,
171 limit: usize,
172 ) -> Result<Vec<Run>> {
173 read(&self.inner, |inner| {
174 crate::runs::list_runs(inner, job_id_filter, offset, limit)
175 })
176 }
177
178 async fn finish_run(
179 &self,
180 run_id: &str,
181 status: RunStatus,
182 duration_ms: Option<i64>,
183 error_message: Option<String>,
184 ) -> Result<()> {
185 write(&self.inner, |inner| {
186 crate::runs::finish_run(inner, run_id, status, duration_ms, error_message);
187 })
188 }
189
190 async fn count_runs(&self, job_id_filter: Option<&str>) -> Result<u64> {
191 read(&self.inner, |inner| {
192 crate::runs::count_runs(inner, job_id_filter)
193 })
194 }
195
196 async fn count_runs_since(&self, since: DateTime<Utc>) -> Result<u64> {
197 read(&self.inner, |inner| {
198 crate::runs::count_runs_since(inner, since)
199 })
200 }
201
202 async fn task_run_stats(&self, task_name: &str) -> Result<TaskRunStats> {
203 read(&self.inner, |inner| {
204 crate::runs::task_run_stats(inner, task_name)
205 })
206 }
207
208 async fn get_task_config(&self, task_name: &str) -> Result<Option<TaskConfig>> {
209 read(&self.inner, |inner| {
210 crate::task_config::get_task_config(inner, task_name)
211 })
212 }
213
214 async fn upsert_task_config(&self, config: &TaskConfig) -> Result<()> {
215 write(&self.inner, |inner| {
216 crate::task_config::upsert_task_config(inner, config);
217 })
218 }
219
220 async fn try_claim_run_lease(
221 &self,
222 job_id: &str,
223 worker_id: &str,
224 ttl_secs: i64,
225 ) -> Result<Option<String>> {
226 write(&self.inner, |inner| {
227 crate::leases::try_claim_run_lease(inner, job_id, worker_id, ttl_secs)
228 })
229 }
230
231 async fn extend_lease(&self, lease_id: &str, ttl_secs: i64) -> Result<()> {
232 write(&self.inner, |inner| {
233 crate::leases::extend_lease(inner, lease_id, ttl_secs);
234 })
235 }
236
237 async fn release_lease(&self, lease_id: &str) -> Result<()> {
238 write(&self.inner, |inner| {
239 crate::leases::release_lease(inner, lease_id);
240 })
241 }
242
243 async fn expired_lease_job_pairs(&self) -> Result<Vec<(String, String)>> {
244 read(&self.inner, crate::leases::expired_lease_job_pairs)
245 }
246}