1use crate::{
2 pool_cleared_error, pool_clearing_error, pool_invalid_config_error, PoolError, PoolResult,
3};
4use notify_future::Notify;
5use std::collections::{HashMap, VecDeque};
6use std::hash::Hash;
7use std::ops::{Deref, DerefMut};
8use std::sync::{Arc, Mutex};
9use std::time::{Duration, Instant};
10
11pub trait WorkerClassification: Send + 'static + Clone + Hash + Eq + PartialEq {}
12
13impl<T: Send + 'static + Clone + Hash + Eq + PartialEq> WorkerClassification for T {}
14
15#[derive(Debug, Clone, Default)]
16pub struct ClassifiedWorkerPoolConfig {
17 pub idle_timeout: Option<Duration>,
18}
19
20#[async_trait::async_trait]
21pub trait ClassifiedWorker<C: WorkerClassification>: Send + 'static {
22 fn is_work(&self) -> bool;
23 fn is_valid(&self, c: C) -> bool;
26 fn classification(&self) -> C;
28}
29
30pub struct ClassifiedWorkerGuard<
31 C: WorkerClassification,
32 W: ClassifiedWorker<C>,
33 F: ClassifiedWorkerFactory<C, W>,
34> {
35 pool_ref: ClassifiedWorkerPoolRef<C, W, F>,
36 worker: Option<W>,
37}
38
39impl<C: WorkerClassification, W: ClassifiedWorker<C>, F: ClassifiedWorkerFactory<C, W>>
40 ClassifiedWorkerGuard<C, W, F>
41{
42 fn new(worker: W, pool_ref: ClassifiedWorkerPoolRef<C, W, F>) -> Self {
43 ClassifiedWorkerGuard {
44 pool_ref,
45 worker: Some(worker),
46 }
47 }
48}
49
50impl<C: WorkerClassification, W: ClassifiedWorker<C>, F: ClassifiedWorkerFactory<C, W>> Deref
51 for ClassifiedWorkerGuard<C, W, F>
52{
53 type Target = W;
54
55 fn deref(&self) -> &Self::Target {
56 self.worker.as_ref().unwrap()
57 }
58}
59
60impl<C: WorkerClassification, W: ClassifiedWorker<C>, F: ClassifiedWorkerFactory<C, W>> DerefMut
61 for ClassifiedWorkerGuard<C, W, F>
62{
63 fn deref_mut(&mut self) -> &mut Self::Target {
64 self.worker.as_mut().unwrap()
65 }
66}
67
68impl<C: WorkerClassification, W: ClassifiedWorker<C>, F: ClassifiedWorkerFactory<C, W>> Drop
69 for ClassifiedWorkerGuard<C, W, F>
70{
71 fn drop(&mut self) {
72 if let Some(worker) = self.worker.take() {
73 self.pool_ref.release(worker);
74 }
75 }
76}
77
78#[async_trait::async_trait]
79pub trait ClassifiedWorkerFactory<C: WorkerClassification, W: ClassifiedWorker<C>>:
80 Send + Sync + 'static
81{
82 async fn create(&self, c: Option<C>) -> PoolResult<W>;
83}
84
85struct WaitingItem<
86 C: WorkerClassification,
87 W: ClassifiedWorker<C>,
88 F: ClassifiedWorkerFactory<C, W>,
89> {
90 future: Notify<ClassifiedWorkerWaitResult<C, W, F>>,
91 condition: Option<C>,
92}
93
94struct IdleWorker<W> {
95 worker: W,
96 idle_since: Instant,
97}
98
99enum ClassifiedWorkerWaitResult<
100 C: WorkerClassification,
101 W: ClassifiedWorker<C>,
102 F: ClassifiedWorkerFactory<C, W>,
103> {
104 Worker(ClassifiedWorkerGuard<C, W, F>),
105 Retry,
106 Error(PoolError),
107}
108
109struct WorkerPoolState<
110 C: WorkerClassification,
111 W: ClassifiedWorker<C>,
112 F: ClassifiedWorkerFactory<C, W>,
113> {
114 current_count: u16,
115 classified_count_map: HashMap<C, u16>,
116 pending_classified_count_map: HashMap<C, u16>,
117 worker_list: VecDeque<IdleWorker<W>>,
118 waiting_list: Vec<WaitingItem<C, W, F>>,
119 clearing: bool,
120 clear_waiting_list: Vec<Notify<()>>,
121}
122
123impl<C: WorkerClassification, W: ClassifiedWorker<C>, F: ClassifiedWorkerFactory<C, W>>
124 WorkerPoolState<C, W, F>
125{
126 fn inc_classified_count(&mut self, c: C) {
127 let count = self.classified_count_map.entry(c).or_insert(0);
128 *count += 1;
129 }
130
131 fn dec_classified_count(&mut self, c: C) {
132 let mut should_remove = false;
133 if let Some(count) = self.classified_count_map.get_mut(&c) {
134 debug_assert!(*count > 0);
135 *count -= 1;
136 should_remove = *count == 0;
137 }
138 if should_remove {
139 self.classified_count_map.remove(&c);
140 }
141 }
142
143 fn inc_pending_classified_count(&mut self, c: C) {
144 let count = self.pending_classified_count_map.entry(c).or_insert(0);
145 *count += 1;
146 }
147
148 fn dec_pending_classified_count(&mut self, c: C) {
149 let mut should_remove = false;
150 if let Some(count) = self.pending_classified_count_map.get_mut(&c) {
151 debug_assert!(*count > 0);
152 *count -= 1;
153 should_remove = *count == 0;
154 }
155 if should_remove {
156 self.pending_classified_count_map.remove(&c);
157 }
158 }
159
160 fn reserved_classified_count(&self, c: &C) -> u16 {
161 self.classified_count_map.get(c).copied().unwrap_or(0)
162 + self
163 .pending_classified_count_map
164 .get(c)
165 .copied()
166 .unwrap_or(0)
167 }
168
169 fn take_clear_waiters_if_done(&mut self) -> Vec<Notify<()>> {
170 if self.clearing && self.current_count == 0 {
171 self.clearing = false;
172 self.clear_waiting_list.drain(..).collect()
173 } else {
174 Vec::new()
175 }
176 }
177
178 fn find_matching_waiter_index_for_worker(&self, worker: &W) -> Option<usize> {
179 self.waiting_list.iter().position(|waiting| {
180 if waiting.future.is_canceled() {
181 return false;
182 }
183 waiting
184 .condition
185 .as_ref()
186 .map(|condition| worker.is_valid(condition.clone()))
187 .unwrap_or(true)
188 })
189 }
190
191 fn find_any_classified_waiter_index(&self) -> Option<usize> {
192 self.waiting_list
193 .iter()
194 .position(|waiting| waiting.condition.is_some() && !waiting.future.is_canceled())
195 }
196
197 fn drain_waiters(&mut self) -> Vec<Notify<ClassifiedWorkerWaitResult<C, W, F>>> {
198 self.waiting_list
199 .drain(..)
200 .map(|waiting| waiting.future)
201 .collect()
202 }
203}
204
205pub struct ClassifiedWorkerPool<
206 C: WorkerClassification,
207 W: ClassifiedWorker<C>,
208 F: ClassifiedWorkerFactory<C, W>,
209> {
210 factory: Arc<F>,
211 max_count: u16,
212 config: ClassifiedWorkerPoolConfig,
213 state: Mutex<WorkerPoolState<C, W, F>>,
214}
215pub type ClassifiedWorkerPoolRef<C, W, F> = Arc<ClassifiedWorkerPool<C, W, F>>;
216
217impl<C: WorkerClassification, W: ClassifiedWorker<C>, F: ClassifiedWorkerFactory<C, W>>
218 ClassifiedWorkerPool<C, W, F>
219{
220 fn validate_created_worker(requested_classification: Option<&C>, worker: &W) -> PoolResult<()> {
221 let worker_classification = worker.classification();
222 if !worker.is_valid(worker_classification.clone()) {
223 return Err(pool_invalid_config_error(
224 "worker primary classification is not valid for itself",
225 ));
226 }
227 if let Some(classification) = requested_classification {
228 if worker_classification != classification.clone() {
229 return Err(pool_invalid_config_error(
230 "factory returned worker with mismatched classification",
231 ));
232 }
233 }
234 Ok(())
235 }
236
237 pub fn new(max_count: u16, factory: F) -> ClassifiedWorkerPoolRef<C, W, F> {
238 Self::new_with_config(max_count, factory, ClassifiedWorkerPoolConfig::default())
239 }
240
241 pub fn new_with_config(
242 max_count: u16,
243 factory: F,
244 config: ClassifiedWorkerPoolConfig,
245 ) -> ClassifiedWorkerPoolRef<C, W, F> {
246 Arc::new(ClassifiedWorkerPool {
247 factory: Arc::new(factory),
248 max_count,
249 config,
250 state: Mutex::new(WorkerPoolState {
251 current_count: 0,
252 classified_count_map: HashMap::new(),
253 pending_classified_count_map: HashMap::new(),
254 worker_list: VecDeque::with_capacity(max_count as usize),
255 waiting_list: Vec::new(),
256 clearing: false,
257 clear_waiting_list: Vec::new(),
258 }),
259 })
260 }
261
262 fn remove_expired_idle_workers(
263 state: &mut WorkerPoolState<C, W, F>,
264 idle_timeout: Option<std::time::Duration>,
265 ) -> u16 {
266 let Some(idle_timeout) = idle_timeout else {
267 return 0;
268 };
269 let mut removed_count = 0;
270 let now = Instant::now();
271 while state
272 .worker_list
273 .front()
274 .map(|idle_worker| now.duration_since(idle_worker.idle_since) >= idle_timeout)
275 .unwrap_or(false)
276 {
277 let idle_worker = state.worker_list.pop_front().unwrap();
278 state.current_count -= 1;
279 state.dec_classified_count(idle_worker.worker.classification());
280 removed_count += 1;
281 }
282 removed_count
283 }
284
285 pub fn cleanup_idle_worker(&self) -> u16 {
286 let (removed_count, clear_waiters) = {
287 let mut state = self.state.lock().unwrap();
288 let removed_count =
289 Self::remove_expired_idle_workers(&mut state, self.config.idle_timeout);
290 let clear_waiters = state.take_clear_waiters_if_done();
291 (removed_count, clear_waiters)
292 };
293 for waiter in clear_waiters {
294 waiter.notify(());
295 }
296 removed_count
297 }
298
299 pub async fn get_worker(
300 self: &ClassifiedWorkerPoolRef<C, W, F>,
301 ) -> PoolResult<ClassifiedWorkerGuard<C, W, F>> {
302 loop {
303 if self.max_count == 0 {
304 return Err(pool_invalid_config_error("pool max_count is zero"));
305 }
306
307 let wait = {
308 let mut state = self.state.lock().unwrap();
309 if state.clearing {
310 return Err(pool_clearing_error());
311 }
312
313 Self::remove_expired_idle_workers(&mut state, self.config.idle_timeout);
314
315 while let Some(idle_worker) = state.worker_list.pop_back() {
316 let worker = idle_worker.worker;
317 if !worker.is_work() {
318 state.current_count -= 1;
319 state.dec_classified_count(worker.classification());
320 continue;
321 }
322 return Ok(ClassifiedWorkerGuard::new(worker, self.clone()));
323 }
324
325 if state.current_count < self.max_count {
326 state.current_count += 1;
327 None
328 } else {
329 let (notify, waiter) = Notify::new();
330 state.waiting_list.push(WaitingItem {
331 future: notify,
332 condition: None,
333 });
334 Some(waiter)
335 }
336 };
337
338 if let Some(wait) = wait {
339 match wait.await {
340 ClassifiedWorkerWaitResult::Worker(worker) => return Ok(worker),
341 ClassifiedWorkerWaitResult::Retry => continue,
342 ClassifiedWorkerWaitResult::Error(err) => return Err(err),
343 }
344 }
345
346 let worker = match self.factory.create(None).await {
347 Ok(worker) => {
348 if let Err(err) = Self::validate_created_worker(None, &worker) {
349 let (retry_waiters, clear_waiters) = {
350 let mut state = self.state.lock().unwrap();
351 state.current_count -= 1;
352 let retry_waiters = state.drain_waiters();
353 let clear_waiters = state.take_clear_waiters_if_done();
354 (retry_waiters, clear_waiters)
355 };
356 Self::notify_retry_waiters(retry_waiters);
357 for waiter in clear_waiters {
358 waiter.notify(());
359 }
360 return Err(err);
361 }
362 worker
363 }
364 Err(err) => {
365 let (retry_waiters, clear_waiters) = {
366 let mut state = self.state.lock().unwrap();
367 state.current_count -= 1;
368 let retry_waiters = state.drain_waiters();
369 let clear_waiters = state.take_clear_waiters_if_done();
370 (retry_waiters, clear_waiters)
371 };
372 Self::notify_retry_waiters(retry_waiters);
373 for waiter in clear_waiters {
374 waiter.notify(());
375 }
376 return Err(err);
377 }
378 };
379 let (clearing, clear_waiters) = {
380 let mut state = self.state.lock().unwrap();
381 if state.clearing {
382 state.current_count -= 1;
383 (true, state.take_clear_waiters_if_done())
384 } else {
385 state.inc_classified_count(worker.classification());
386 (false, Vec::new())
387 }
388 };
389 for waiter in clear_waiters {
390 waiter.notify(());
391 }
392 if clearing {
393 return Err(pool_cleared_error());
394 }
395 return Ok(ClassifiedWorkerGuard::new(worker, self.clone()));
396 }
397 }
398
399 pub async fn get_classified_worker(
400 self: &ClassifiedWorkerPoolRef<C, W, F>,
401 classification: C,
402 ) -> PoolResult<ClassifiedWorkerGuard<C, W, F>> {
403 loop {
404 if self.max_count == 0 {
405 return Err(pool_invalid_config_error("pool max_count is zero"));
406 }
407
408 let wait =
409 {
410 let mut state = self.state.lock().unwrap();
411 if state.clearing {
412 return Err(pool_clearing_error());
413 }
414
415 Self::remove_expired_idle_workers(&mut state, self.config.idle_timeout);
416
417 let old_count = state.worker_list.len() as u16;
418 let unwork_classification = state
419 .worker_list
420 .iter()
421 .filter(|idle_worker| !idle_worker.worker.is_work())
422 .map(|idle_worker| idle_worker.worker.classification())
423 .collect::<Vec<C>>();
424 for classification in unwork_classification.iter() {
425 state.dec_classified_count(classification.clone());
426 }
427 state
428 .worker_list
429 .retain(|idle_worker| idle_worker.worker.is_work());
430 state.current_count -= old_count - state.worker_list.len() as u16;
431 if let Some(index) = state.worker_list.iter().rposition(|idle_worker| {
432 idle_worker.worker.is_valid(classification.clone())
433 }) {
434 let idle_worker = state.worker_list.remove(index).unwrap();
435 return Ok(ClassifiedWorkerGuard::new(idle_worker.worker, self.clone()));
436 }
437
438 if state.current_count < self.max_count {
439 state.current_count += 1;
440 state.inc_pending_classified_count(classification.clone());
441 None
442 } else if let Some(idle_worker) = state.worker_list.pop_front() {
443 state.dec_classified_count(idle_worker.worker.classification());
444 state.inc_pending_classified_count(classification.clone());
445 None
446 } else if state.reserved_classified_count(&classification) == 0 {
447 if state.current_count == u16::MAX {
448 return Err(pool_invalid_config_error(
449 "pool current_count reached u16 max",
450 ));
451 }
452 state.current_count += 1;
453 state.inc_pending_classified_count(classification.clone());
454 None
455 } else {
456 let (notify, waiter) = Notify::new();
457 state.waiting_list.push(WaitingItem {
458 future: notify,
459 condition: Some(classification.clone()),
460 });
461 Some(waiter)
462 }
463 };
464
465 if let Some(wait) = wait {
466 match wait.await {
467 ClassifiedWorkerWaitResult::Worker(worker) => return Ok(worker),
468 ClassifiedWorkerWaitResult::Retry => continue,
469 ClassifiedWorkerWaitResult::Error(err) => return Err(err),
470 }
471 }
472
473 let worker = match self.factory.create(Some(classification.clone())).await {
474 Ok(worker) => {
475 if let Err(err) = Self::validate_created_worker(Some(&classification), &worker)
476 {
477 let (retry_waiters, clear_waiters) = {
478 let mut state = self.state.lock().unwrap();
479 state.current_count -= 1;
480 state.dec_pending_classified_count(classification.clone());
481 let retry_waiters = state.drain_waiters();
482 let clear_waiters = state.take_clear_waiters_if_done();
483 (retry_waiters, clear_waiters)
484 };
485 Self::notify_retry_waiters(retry_waiters);
486 for waiter in clear_waiters {
487 waiter.notify(());
488 }
489 return Err(err);
490 }
491 worker
492 }
493 Err(err) => {
494 let (retry_waiters, clear_waiters) = {
495 let mut state = self.state.lock().unwrap();
496 state.current_count -= 1;
497 state.dec_pending_classified_count(classification.clone());
498 let retry_waiters = state.drain_waiters();
499 let clear_waiters = state.take_clear_waiters_if_done();
500 (retry_waiters, clear_waiters)
501 };
502 Self::notify_retry_waiters(retry_waiters);
503 for waiter in clear_waiters {
504 waiter.notify(());
505 }
506 return Err(err);
507 }
508 };
509 let (clearing, clear_waiters) = {
510 let mut state = self.state.lock().unwrap();
511 state.dec_pending_classified_count(classification.clone());
512 if state.clearing {
513 state.current_count -= 1;
514 (true, state.take_clear_waiters_if_done())
515 } else {
516 state.inc_classified_count(worker.classification());
517 (false, Vec::new())
518 }
519 };
520 for waiter in clear_waiters {
521 waiter.notify(());
522 }
523 if clearing {
524 return Err(pool_cleared_error());
525 }
526 return Ok(ClassifiedWorkerGuard::new(worker, self.clone()));
527 }
528 }
529
530 pub async fn clear_all_worker(&self) {
531 let (waiter, waiting_list, clear_waiters) = {
532 let mut state = self.state.lock().unwrap();
533 if !state.clearing {
534 state.clearing = true;
535 let idle_classifications = state
536 .worker_list
537 .iter()
538 .map(|idle_worker| idle_worker.worker.classification())
539 .collect::<Vec<_>>();
540 let cur_worker_count = idle_classifications.len();
541 state.worker_list.clear();
542 state.current_count -= cur_worker_count as u16;
543 for classification in idle_classifications {
544 state.dec_classified_count(classification);
545 }
546 }
547
548 let waiting_list = state.waiting_list.drain(..).collect::<Vec<_>>();
549 if state.current_count == 0 {
550 let clear_waiters = state.take_clear_waiters_if_done();
551 (None, waiting_list, clear_waiters)
552 } else {
553 let (notify, waiter) = Notify::new();
554 state.clear_waiting_list.push(notify);
555 (Some(waiter), waiting_list, Vec::new())
556 }
557 };
558 for waiting in waiting_list {
559 waiting
560 .future
561 .notify(ClassifiedWorkerWaitResult::Error(pool_cleared_error()));
562 }
563 for waiter in clear_waiters {
564 waiter.notify(());
565 }
566 if let Some(waiter) = waiter {
567 waiter.await;
568 }
569 }
570
571 fn notify_retry_waiters(waiters: Vec<Notify<ClassifiedWorkerWaitResult<C, W, F>>>) {
572 for waiter in waiters {
573 waiter.notify(ClassifiedWorkerWaitResult::Retry);
574 }
575 }
576
577 fn release(self: &ClassifiedWorkerPoolRef<C, W, F>, work: W) {
578 enum ReleaseAction<
579 C: WorkerClassification,
580 W: ClassifiedWorker<C>,
581 F: ClassifiedWorkerFactory<C, W>,
582 > {
583 None,
584 Notify(
585 Notify<ClassifiedWorkerWaitResult<C, W, F>>,
586 ClassifiedWorkerGuard<C, W, F>,
587 ),
588 Retry(Vec<Notify<ClassifiedWorkerWaitResult<C, W, F>>>),
589 }
590
591 let mut clear_waiters = Vec::new();
592 let action = {
593 let mut state = self.state.lock().unwrap();
594 if state.clearing {
595 state.current_count -= 1;
596 let classification = work.classification();
597 state.dec_classified_count(classification);
598 clear_waiters = state.take_clear_waiters_if_done();
599 ReleaseAction::None
600 } else if work.is_work() {
601 if let Some(index) = state.find_matching_waiter_index_for_worker(&work) {
602 let waiting_item = state.waiting_list.remove(index);
603 ReleaseAction::Notify(
604 waiting_item.future,
605 ClassifiedWorkerGuard::new(work, self.clone()),
606 )
607 } else if let Some(index) = state.find_any_classified_waiter_index() {
608 let classification = work.classification();
609 state.current_count -= 1;
610 state.dec_classified_count(classification);
611 let mut waiters = state.drain_waiters();
612 if index < waiters.len() {
613 waiters.swap(0, index);
614 }
615 ReleaseAction::Retry(waiters)
616 } else if state.current_count > self.max_count {
617 state.current_count -= 1;
618 let classification = work.classification();
619 state.dec_classified_count(classification);
620 clear_waiters = state.take_clear_waiters_if_done();
621 ReleaseAction::None
622 } else {
623 state.worker_list.push_back(IdleWorker {
624 worker: work,
625 idle_since: Instant::now(),
626 });
627 ReleaseAction::None
628 }
629 } else {
630 let classification = work.classification();
631 state.dec_classified_count(classification);
632 state.current_count -= 1;
633 let waiters = state.drain_waiters();
634 if !waiters.is_empty() {
635 ReleaseAction::Retry(waiters)
636 } else {
637 clear_waiters = state.take_clear_waiters_if_done();
638 ReleaseAction::None
639 }
640 }
641 };
642
643 for waiter in clear_waiters {
644 waiter.notify(());
645 }
646
647 match action {
648 ReleaseAction::None => {}
649 ReleaseAction::Notify(waiting, worker) => {
650 waiting.notify(ClassifiedWorkerWaitResult::Worker(worker));
651 }
652 ReleaseAction::Retry(waiters) => {
653 Self::notify_retry_waiters(waiters);
654 }
655 }
656 }
657}
658
659#[tokio::test]
660async fn test_pool() {
661 struct TestWorker {
662 work: bool,
663 classification: TestWorkerClassification,
664 }
665
666 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
667 enum TestWorkerClassification {
668 A,
669 B,
670 }
671 #[async_trait::async_trait]
672 impl ClassifiedWorker<TestWorkerClassification> for TestWorker {
673 fn is_work(&self) -> bool {
674 self.work
675 }
676
677 fn is_valid(&self, c: TestWorkerClassification) -> bool {
678 self.classification == c
679 }
680
681 fn classification(&self) -> TestWorkerClassification {
682 self.classification.clone()
683 }
684 }
685
686 struct TestWorkerFactory;
687
688 #[async_trait::async_trait]
689 impl ClassifiedWorkerFactory<TestWorkerClassification, TestWorker> for TestWorkerFactory {
690 async fn create(
691 &self,
692 classification: Option<TestWorkerClassification>,
693 ) -> PoolResult<TestWorker> {
694 if let Some(classification) = classification {
695 Ok(TestWorker {
696 work: true,
697 classification,
698 })
699 } else {
700 Ok(TestWorker {
701 work: true,
702 classification: TestWorkerClassification::A,
703 })
704 }
705 }
706 }
707
708 let pool = ClassifiedWorkerPool::new(3, TestWorkerFactory);
709
710 let worker_a1 = pool.get_worker().await.unwrap();
711 let worker_a2 = pool.get_worker().await.unwrap();
712 let worker_b = pool
713 .get_classified_worker(TestWorkerClassification::B)
714 .await
715 .unwrap();
716
717 let pool_ref = pool.clone();
718 let classified_waiter = tokio::spawn(async move {
719 pool_ref
720 .get_classified_worker(TestWorkerClassification::B)
721 .await
722 });
723 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
724 assert!(!classified_waiter.is_finished());
725
726 drop(worker_b);
727 let worker_b = tokio::time::timeout(std::time::Duration::from_secs(1), classified_waiter)
728 .await
729 .unwrap()
730 .unwrap()
731 .unwrap();
732 drop(worker_a1);
733 drop(worker_a2);
734 drop(worker_b);
735
736 let worker3 = pool
737 .get_classified_worker(TestWorkerClassification::B)
738 .await
739 .unwrap();
740 let worker1 = pool.get_worker().await.unwrap();
741 let worker2 = pool.get_worker().await.unwrap();
742
743 let pool_ref = pool.clone();
744 let generic_waiter = tokio::spawn(async move { pool_ref.get_worker().await });
745 let pool_ref = pool.clone();
746 let classified_waiter = tokio::spawn(async move {
747 pool_ref
748 .get_classified_worker(TestWorkerClassification::B)
749 .await
750 });
751 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
752 assert!(!generic_waiter.is_finished());
753 assert!(!classified_waiter.is_finished());
754
755 let pool_ref = pool.clone();
756 let clear_task = tokio::spawn(async move {
757 pool_ref.clear_all_worker().await;
758 });
759 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
760
761 assert!(generic_waiter.await.unwrap().is_err());
762 assert!(classified_waiter.await.unwrap().is_err());
763
764 drop(worker1);
765 drop(worker2);
766 drop(worker3);
767
768 tokio::time::timeout(std::time::Duration::from_secs(1), clear_task)
769 .await
770 .unwrap()
771 .unwrap();
772}
773
774#[tokio::test]
775async fn test_clear_all_worker_waits_for_inflight_create() {
776 use std::sync::atomic::{AtomicUsize, Ordering};
777 use std::sync::Arc;
778
779 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
780 enum TestWorkerClassification {
781 A,
782 }
783
784 struct TestWorker {
785 classification: TestWorkerClassification,
786 }
787
788 #[async_trait::async_trait]
789 impl ClassifiedWorker<TestWorkerClassification> for TestWorker {
790 fn is_work(&self) -> bool {
791 true
792 }
793
794 fn is_valid(&self, c: TestWorkerClassification) -> bool {
795 self.classification == c
796 }
797
798 fn classification(&self) -> TestWorkerClassification {
799 self.classification.clone()
800 }
801 }
802
803 struct TestWorkerFactory {
804 create_count: Arc<AtomicUsize>,
805 }
806
807 #[async_trait::async_trait]
808 impl ClassifiedWorkerFactory<TestWorkerClassification, TestWorker> for TestWorkerFactory {
809 async fn create(
810 &self,
811 classification: Option<TestWorkerClassification>,
812 ) -> PoolResult<TestWorker> {
813 self.create_count.fetch_add(1, Ordering::SeqCst);
814 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
815 Ok(TestWorker {
816 classification: classification.unwrap_or(TestWorkerClassification::A),
817 })
818 }
819 }
820
821 let create_count = Arc::new(AtomicUsize::new(0));
822 let pool = ClassifiedWorkerPool::new(
823 1,
824 TestWorkerFactory {
825 create_count: create_count.clone(),
826 },
827 );
828
829 let pool_ref = pool.clone();
830 let worker_task = tokio::spawn(async move { pool_ref.get_worker().await });
831 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
832
833 pool.clear_all_worker().await;
834
835 let worker = worker_task.await.unwrap();
836 assert!(worker.is_err());
837 assert_eq!(create_count.load(Ordering::SeqCst), 1);
838}
839
840#[tokio::test]
841async fn test_concurrent_clear_all_worker() {
842 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
843 enum TestWorkerClassification {
844 A,
845 }
846
847 struct TestWorker {
848 classification: TestWorkerClassification,
849 }
850
851 #[async_trait::async_trait]
852 impl ClassifiedWorker<TestWorkerClassification> for TestWorker {
853 fn is_work(&self) -> bool {
854 true
855 }
856
857 fn is_valid(&self, c: TestWorkerClassification) -> bool {
858 self.classification == c
859 }
860
861 fn classification(&self) -> TestWorkerClassification {
862 self.classification.clone()
863 }
864 }
865
866 struct TestWorkerFactory;
867
868 #[async_trait::async_trait]
869 impl ClassifiedWorkerFactory<TestWorkerClassification, TestWorker> for TestWorkerFactory {
870 async fn create(
871 &self,
872 classification: Option<TestWorkerClassification>,
873 ) -> PoolResult<TestWorker> {
874 Ok(TestWorker {
875 classification: classification.unwrap_or(TestWorkerClassification::A),
876 })
877 }
878 }
879
880 let pool = ClassifiedWorkerPool::new(1, TestWorkerFactory);
881 let worker = pool.get_worker().await.unwrap();
882
883 let pool_ref = pool.clone();
884 let clear_task1 = tokio::spawn(async move {
885 pool_ref.clear_all_worker().await;
886 });
887
888 let pool_ref = pool.clone();
889 let clear_task2 = tokio::spawn(async move {
890 pool_ref.clear_all_worker().await;
891 });
892
893 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
894 drop(worker);
895
896 tokio::time::timeout(std::time::Duration::from_secs(1), async {
897 clear_task1.await.unwrap();
898 clear_task2.await.unwrap();
899 })
900 .await
901 .unwrap();
902}
903
904#[tokio::test]
905async fn test_zero_max_count_returns_error() {
906 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
907 enum TestWorkerClassification {
908 A,
909 }
910
911 struct TestWorker {
912 classification: TestWorkerClassification,
913 }
914
915 #[async_trait::async_trait]
916 impl ClassifiedWorker<TestWorkerClassification> for TestWorker {
917 fn is_work(&self) -> bool {
918 true
919 }
920
921 fn is_valid(&self, c: TestWorkerClassification) -> bool {
922 self.classification == c
923 }
924
925 fn classification(&self) -> TestWorkerClassification {
926 self.classification.clone()
927 }
928 }
929
930 struct TestWorkerFactory;
931
932 #[async_trait::async_trait]
933 impl ClassifiedWorkerFactory<TestWorkerClassification, TestWorker> for TestWorkerFactory {
934 async fn create(
935 &self,
936 classification: Option<TestWorkerClassification>,
937 ) -> PoolResult<TestWorker> {
938 Ok(TestWorker {
939 classification: classification.unwrap_or(TestWorkerClassification::A),
940 })
941 }
942 }
943
944 let pool = ClassifiedWorkerPool::new(0, TestWorkerFactory);
945 let worker = pool.get_worker().await;
946 assert!(worker.is_err());
947 assert_eq!(
948 worker.err().unwrap().code(),
949 crate::PoolErrorCode::InvalidConfig
950 );
951}
952
953#[tokio::test]
954async fn test_classified_pool_waits_when_classification_already_has_worker() {
955 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
956 enum TestWorkerClassification {
957 A,
958 B,
959 }
960
961 struct TestWorker {
962 classification: TestWorkerClassification,
963 }
964
965 #[async_trait::async_trait]
966 impl ClassifiedWorker<TestWorkerClassification> for TestWorker {
967 fn is_work(&self) -> bool {
968 true
969 }
970
971 fn is_valid(&self, c: TestWorkerClassification) -> bool {
972 self.classification == c
973 }
974
975 fn classification(&self) -> TestWorkerClassification {
976 self.classification.clone()
977 }
978 }
979
980 struct TestWorkerFactory;
981
982 #[async_trait::async_trait]
983 impl ClassifiedWorkerFactory<TestWorkerClassification, TestWorker> for TestWorkerFactory {
984 async fn create(
985 &self,
986 classification: Option<TestWorkerClassification>,
987 ) -> PoolResult<TestWorker> {
988 Ok(TestWorker {
989 classification: classification.unwrap_or(TestWorkerClassification::A),
990 })
991 }
992 }
993
994 let pool = ClassifiedWorkerPool::new(1, TestWorkerFactory);
995 let _worker = pool
996 .get_classified_worker(TestWorkerClassification::B)
997 .await
998 .unwrap();
999
1000 let pool_ref = pool.clone();
1001 let result = tokio::time::timeout(std::time::Duration::from_millis(100), async move {
1002 pool_ref
1003 .get_classified_worker(TestWorkerClassification::B)
1004 .await
1005 })
1006 .await;
1007
1008 assert!(result.is_err());
1009}
1010
1011#[tokio::test]
1012async fn test_missing_classification_can_exceed_max_count_once() {
1013 use std::sync::atomic::{AtomicUsize, Ordering};
1014 use std::sync::Arc;
1015
1016 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
1017 enum TestWorkerClassification {
1018 A,
1019 B,
1020 }
1021
1022 struct TestWorker {
1023 id: usize,
1024 classification: TestWorkerClassification,
1025 }
1026
1027 #[async_trait::async_trait]
1028 impl ClassifiedWorker<TestWorkerClassification> for TestWorker {
1029 fn is_work(&self) -> bool {
1030 true
1031 }
1032
1033 fn is_valid(&self, c: TestWorkerClassification) -> bool {
1034 self.classification == c
1035 }
1036
1037 fn classification(&self) -> TestWorkerClassification {
1038 self.classification.clone()
1039 }
1040 }
1041
1042 struct TestWorkerFactory {
1043 create_count: Arc<AtomicUsize>,
1044 }
1045
1046 #[async_trait::async_trait]
1047 impl ClassifiedWorkerFactory<TestWorkerClassification, TestWorker> for TestWorkerFactory {
1048 async fn create(
1049 &self,
1050 classification: Option<TestWorkerClassification>,
1051 ) -> PoolResult<TestWorker> {
1052 let id = self.create_count.fetch_add(1, Ordering::SeqCst);
1053 Ok(TestWorker {
1054 id,
1055 classification: classification.unwrap_or(TestWorkerClassification::A),
1056 })
1057 }
1058 }
1059
1060 let create_count = Arc::new(AtomicUsize::new(0));
1061 let pool = ClassifiedWorkerPool::new(
1062 1,
1063 TestWorkerFactory {
1064 create_count: create_count.clone(),
1065 },
1066 );
1067
1068 let worker_a = pool
1069 .get_classified_worker(TestWorkerClassification::A)
1070 .await
1071 .unwrap();
1072 let worker_b = pool
1073 .get_classified_worker(TestWorkerClassification::B)
1074 .await
1075 .unwrap();
1076
1077 assert_eq!(worker_a.id, 0);
1078 assert_eq!(worker_b.id, 1);
1079 assert_eq!(worker_b.classification(), TestWorkerClassification::B);
1080 assert_eq!(create_count.load(Ordering::SeqCst), 2);
1081}
1082
1083#[tokio::test]
1084async fn test_classified_create_failure_fails_same_classification_waiters() {
1085 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
1086 enum TestWorkerClassification {
1087 B,
1088 }
1089
1090 struct TestWorker {
1091 classification: TestWorkerClassification,
1092 }
1093
1094 #[async_trait::async_trait]
1095 impl ClassifiedWorker<TestWorkerClassification> for TestWorker {
1096 fn is_work(&self) -> bool {
1097 true
1098 }
1099
1100 fn is_valid(&self, c: TestWorkerClassification) -> bool {
1101 self.classification == c
1102 }
1103
1104 fn classification(&self) -> TestWorkerClassification {
1105 self.classification.clone()
1106 }
1107 }
1108
1109 struct TestWorkerFactory;
1110
1111 #[async_trait::async_trait]
1112 impl ClassifiedWorkerFactory<TestWorkerClassification, TestWorker> for TestWorkerFactory {
1113 async fn create(
1114 &self,
1115 _classification: Option<TestWorkerClassification>,
1116 ) -> PoolResult<TestWorker> {
1117 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1118 Err(crate::pool_invalid_config_error("create failed"))
1119 }
1120 }
1121
1122 let pool = ClassifiedWorkerPool::new(1, TestWorkerFactory);
1123
1124 let pool_ref = pool.clone();
1125 let worker1 = tokio::spawn(async move {
1126 pool_ref
1127 .get_classified_worker(TestWorkerClassification::B)
1128 .await
1129 });
1130 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
1131
1132 let pool_ref = pool.clone();
1133 let worker2 = tokio::spawn(async move {
1134 pool_ref
1135 .get_classified_worker(TestWorkerClassification::B)
1136 .await
1137 });
1138
1139 let (worker1, worker2) = tokio::time::timeout(std::time::Duration::from_secs(1), async {
1140 (worker1.await.unwrap(), worker2.await.unwrap())
1141 })
1142 .await
1143 .unwrap();
1144
1145 assert_eq!(
1146 worker1.err().unwrap().code(),
1147 crate::PoolErrorCode::InvalidConfig
1148 );
1149 assert_eq!(
1150 worker2.err().unwrap().code(),
1151 crate::PoolErrorCode::InvalidConfig
1152 );
1153}
1154
1155#[tokio::test]
1156async fn test_classified_create_failure_wakes_generic_waiter_to_create() {
1157 use std::sync::atomic::{AtomicUsize, Ordering};
1158 use std::sync::Arc;
1159
1160 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
1161 enum TestWorkerClassification {
1162 A,
1163 }
1164
1165 struct TestWorker {
1166 id: usize,
1167 classification: TestWorkerClassification,
1168 }
1169
1170 #[async_trait::async_trait]
1171 impl ClassifiedWorker<TestWorkerClassification> for TestWorker {
1172 fn is_work(&self) -> bool {
1173 true
1174 }
1175
1176 fn is_valid(&self, c: TestWorkerClassification) -> bool {
1177 self.classification == c
1178 }
1179
1180 fn classification(&self) -> TestWorkerClassification {
1181 self.classification.clone()
1182 }
1183 }
1184
1185 struct TestWorkerFactory {
1186 create_count: Arc<AtomicUsize>,
1187 }
1188
1189 #[async_trait::async_trait]
1190 impl ClassifiedWorkerFactory<TestWorkerClassification, TestWorker> for TestWorkerFactory {
1191 async fn create(
1192 &self,
1193 classification: Option<TestWorkerClassification>,
1194 ) -> PoolResult<TestWorker> {
1195 let id = self.create_count.fetch_add(1, Ordering::SeqCst);
1196 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1197 if id == 0 && classification == Some(TestWorkerClassification::A) {
1198 Err(crate::pool_invalid_config_error("create failed"))
1199 } else {
1200 Ok(TestWorker {
1201 id,
1202 classification: classification.unwrap_or(TestWorkerClassification::A),
1203 })
1204 }
1205 }
1206 }
1207
1208 let create_count = Arc::new(AtomicUsize::new(0));
1209 let pool = ClassifiedWorkerPool::new(
1210 1,
1211 TestWorkerFactory {
1212 create_count: create_count.clone(),
1213 },
1214 );
1215
1216 let pool_ref = pool.clone();
1217 let classified = tokio::spawn(async move {
1218 pool_ref
1219 .get_classified_worker(TestWorkerClassification::A)
1220 .await
1221 });
1222 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
1223
1224 let pool_ref = pool.clone();
1225 let generic = tokio::spawn(async move { pool_ref.get_worker().await });
1226
1227 let (classified, generic) = tokio::time::timeout(std::time::Duration::from_secs(1), async {
1228 (classified.await.unwrap(), generic.await.unwrap())
1229 })
1230 .await
1231 .unwrap();
1232
1233 assert_eq!(
1234 classified.err().unwrap().code(),
1235 crate::PoolErrorCode::InvalidConfig
1236 );
1237 let generic = generic.unwrap();
1238 assert_eq!(generic.id, 1);
1239 assert_eq!(create_count.load(Ordering::SeqCst), 2);
1240}
1241
1242#[tokio::test]
1243async fn test_classified_retry_notification_skips_canceled_waiter() {
1244 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
1245 enum TestWorkerClassification {
1246 A,
1247 }
1248
1249 struct TestWorker {
1250 classification: TestWorkerClassification,
1251 }
1252
1253 #[async_trait::async_trait]
1254 impl ClassifiedWorker<TestWorkerClassification> for TestWorker {
1255 fn is_work(&self) -> bool {
1256 true
1257 }
1258
1259 fn is_valid(&self, c: TestWorkerClassification) -> bool {
1260 self.classification == c
1261 }
1262
1263 fn classification(&self) -> TestWorkerClassification {
1264 self.classification.clone()
1265 }
1266 }
1267
1268 struct TestWorkerFactory;
1269
1270 #[async_trait::async_trait]
1271 impl ClassifiedWorkerFactory<TestWorkerClassification, TestWorker> for TestWorkerFactory {
1272 async fn create(
1273 &self,
1274 classification: Option<TestWorkerClassification>,
1275 ) -> PoolResult<TestWorker> {
1276 Ok(TestWorker {
1277 classification: classification.unwrap_or(TestWorkerClassification::A),
1278 })
1279 }
1280 }
1281
1282 let (canceled_notify, canceled_waiter) = Notify::new();
1283 drop(canceled_waiter);
1284 let (notify, waiter) = Notify::new();
1285
1286 ClassifiedWorkerPool::<TestWorkerClassification, TestWorker, TestWorkerFactory>::notify_retry_waiters(
1287 vec![canceled_notify, notify],
1288 );
1289
1290 let result = tokio::time::timeout(std::time::Duration::from_secs(1), waiter)
1291 .await
1292 .unwrap();
1293 assert!(matches!(result, ClassifiedWorkerWaitResult::Retry));
1294}
1295
1296#[tokio::test]
1297async fn test_classified_request_replaces_non_matching_idle_worker() {
1298 use std::sync::atomic::{AtomicUsize, Ordering};
1299 use std::sync::Arc;
1300
1301 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
1302 enum TestWorkerClassification {
1303 A,
1304 B,
1305 }
1306
1307 struct TestWorker {
1308 id: usize,
1309 classification: TestWorkerClassification,
1310 }
1311
1312 #[async_trait::async_trait]
1313 impl ClassifiedWorker<TestWorkerClassification> for TestWorker {
1314 fn is_work(&self) -> bool {
1315 true
1316 }
1317
1318 fn is_valid(&self, c: TestWorkerClassification) -> bool {
1319 self.classification == c
1320 }
1321
1322 fn classification(&self) -> TestWorkerClassification {
1323 self.classification.clone()
1324 }
1325 }
1326
1327 struct TestWorkerFactory {
1328 create_count: Arc<AtomicUsize>,
1329 }
1330
1331 #[async_trait::async_trait]
1332 impl ClassifiedWorkerFactory<TestWorkerClassification, TestWorker> for TestWorkerFactory {
1333 async fn create(
1334 &self,
1335 classification: Option<TestWorkerClassification>,
1336 ) -> PoolResult<TestWorker> {
1337 let id = self.create_count.fetch_add(1, Ordering::SeqCst);
1338 Ok(TestWorker {
1339 id,
1340 classification: classification.unwrap_or(TestWorkerClassification::A),
1341 })
1342 }
1343 }
1344
1345 let create_count = Arc::new(AtomicUsize::new(0));
1346 let pool = ClassifiedWorkerPool::new(
1347 1,
1348 TestWorkerFactory {
1349 create_count: create_count.clone(),
1350 },
1351 );
1352
1353 {
1354 let worker = pool
1355 .get_classified_worker(TestWorkerClassification::A)
1356 .await
1357 .unwrap();
1358 assert_eq!(worker.id, 0);
1359 }
1360
1361 let worker = pool
1362 .get_classified_worker(TestWorkerClassification::B)
1363 .await
1364 .unwrap();
1365 assert_eq!(worker.id, 1);
1366 assert_eq!(worker.classification(), TestWorkerClassification::B);
1367 assert_eq!(create_count.load(Ordering::SeqCst), 2);
1368}
1369
1370#[tokio::test]
1371async fn test_classified_waiter_replaces_returned_non_matching_worker() {
1372 use std::sync::atomic::{AtomicUsize, Ordering};
1373 use std::sync::Arc;
1374
1375 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
1376 enum TestWorkerClassification {
1377 A,
1378 B,
1379 }
1380
1381 struct TestWorker {
1382 id: usize,
1383 classification: TestWorkerClassification,
1384 }
1385
1386 #[async_trait::async_trait]
1387 impl ClassifiedWorker<TestWorkerClassification> for TestWorker {
1388 fn is_work(&self) -> bool {
1389 true
1390 }
1391
1392 fn is_valid(&self, c: TestWorkerClassification) -> bool {
1393 self.classification == c
1394 }
1395
1396 fn classification(&self) -> TestWorkerClassification {
1397 self.classification.clone()
1398 }
1399 }
1400
1401 struct TestWorkerFactory {
1402 create_count: Arc<AtomicUsize>,
1403 }
1404
1405 #[async_trait::async_trait]
1406 impl ClassifiedWorkerFactory<TestWorkerClassification, TestWorker> for TestWorkerFactory {
1407 async fn create(
1408 &self,
1409 classification: Option<TestWorkerClassification>,
1410 ) -> PoolResult<TestWorker> {
1411 let id = self.create_count.fetch_add(1, Ordering::SeqCst);
1412 Ok(TestWorker {
1413 id,
1414 classification: classification.unwrap_or(TestWorkerClassification::A),
1415 })
1416 }
1417 }
1418
1419 let create_count = Arc::new(AtomicUsize::new(0));
1420 let pool = ClassifiedWorkerPool::new(
1421 2,
1422 TestWorkerFactory {
1423 create_count: create_count.clone(),
1424 },
1425 );
1426 let worker_a = pool
1427 .get_classified_worker(TestWorkerClassification::A)
1428 .await
1429 .unwrap();
1430 let _worker_b = pool
1431 .get_classified_worker(TestWorkerClassification::B)
1432 .await
1433 .unwrap();
1434
1435 let pool_ref = pool.clone();
1436 let waiter = tokio::spawn(async move {
1437 pool_ref
1438 .get_classified_worker(TestWorkerClassification::B)
1439 .await
1440 });
1441 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
1442 assert!(!waiter.is_finished());
1443
1444 drop(worker_a);
1445 let worker = tokio::time::timeout(std::time::Duration::from_secs(1), waiter)
1446 .await
1447 .unwrap()
1448 .unwrap()
1449 .unwrap();
1450 assert_eq!(worker.id, 2);
1451 assert_eq!(worker.classification(), TestWorkerClassification::B);
1452 assert_eq!(create_count.load(Ordering::SeqCst), 3);
1453}
1454
1455#[tokio::test]
1456async fn test_classified_waiter_replaces_unwork_non_matching_worker() {
1457 use std::sync::atomic::{AtomicUsize, Ordering};
1458 use std::sync::Arc;
1459
1460 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
1461 enum TestWorkerClassification {
1462 A,
1463 B,
1464 }
1465
1466 struct TestWorker {
1467 id: usize,
1468 work: bool,
1469 classification: TestWorkerClassification,
1470 }
1471
1472 #[async_trait::async_trait]
1473 impl ClassifiedWorker<TestWorkerClassification> for TestWorker {
1474 fn is_work(&self) -> bool {
1475 self.work
1476 }
1477
1478 fn is_valid(&self, c: TestWorkerClassification) -> bool {
1479 self.classification == c
1480 }
1481
1482 fn classification(&self) -> TestWorkerClassification {
1483 self.classification.clone()
1484 }
1485 }
1486
1487 struct TestWorkerFactory {
1488 create_count: Arc<AtomicUsize>,
1489 }
1490
1491 #[async_trait::async_trait]
1492 impl ClassifiedWorkerFactory<TestWorkerClassification, TestWorker> for TestWorkerFactory {
1493 async fn create(
1494 &self,
1495 classification: Option<TestWorkerClassification>,
1496 ) -> PoolResult<TestWorker> {
1497 let id = self.create_count.fetch_add(1, Ordering::SeqCst);
1498 Ok(TestWorker {
1499 id,
1500 work: true,
1501 classification: classification.unwrap_or(TestWorkerClassification::A),
1502 })
1503 }
1504 }
1505
1506 let create_count = Arc::new(AtomicUsize::new(0));
1507 let pool = ClassifiedWorkerPool::new(
1508 2,
1509 TestWorkerFactory {
1510 create_count: create_count.clone(),
1511 },
1512 );
1513 let mut worker_a = pool
1514 .get_classified_worker(TestWorkerClassification::A)
1515 .await
1516 .unwrap();
1517 let _worker_b = pool
1518 .get_classified_worker(TestWorkerClassification::B)
1519 .await
1520 .unwrap();
1521
1522 let pool_ref = pool.clone();
1523 let waiter = tokio::spawn(async move {
1524 pool_ref
1525 .get_classified_worker(TestWorkerClassification::B)
1526 .await
1527 });
1528 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
1529 assert!(!waiter.is_finished());
1530
1531 worker_a.work = false;
1532 drop(worker_a);
1533
1534 let worker = tokio::time::timeout(std::time::Duration::from_secs(1), waiter)
1535 .await
1536 .unwrap()
1537 .unwrap()
1538 .unwrap();
1539 assert_eq!(worker.id, 2);
1540 assert_eq!(worker.classification(), TestWorkerClassification::B);
1541 assert_eq!(create_count.load(Ordering::SeqCst), 3);
1542}
1543
1544#[tokio::test]
1545async fn test_factory_must_return_matching_classification() {
1546 use std::sync::atomic::{AtomicUsize, Ordering};
1547 use std::sync::Arc;
1548
1549 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
1550 enum TestWorkerClassification {
1551 A,
1552 B,
1553 }
1554
1555 struct TestWorker {
1556 classification: TestWorkerClassification,
1557 }
1558
1559 #[async_trait::async_trait]
1560 impl ClassifiedWorker<TestWorkerClassification> for TestWorker {
1561 fn is_work(&self) -> bool {
1562 true
1563 }
1564
1565 fn is_valid(&self, c: TestWorkerClassification) -> bool {
1566 self.classification == c
1567 }
1568
1569 fn classification(&self) -> TestWorkerClassification {
1570 self.classification.clone()
1571 }
1572 }
1573
1574 struct TestWorkerFactory {
1575 create_count: Arc<AtomicUsize>,
1576 }
1577
1578 #[async_trait::async_trait]
1579 impl ClassifiedWorkerFactory<TestWorkerClassification, TestWorker> for TestWorkerFactory {
1580 async fn create(
1581 &self,
1582 classification: Option<TestWorkerClassification>,
1583 ) -> PoolResult<TestWorker> {
1584 let count = self.create_count.fetch_add(1, Ordering::SeqCst);
1585 let classification = if count == 0 {
1586 TestWorkerClassification::A
1587 } else {
1588 classification.unwrap_or(TestWorkerClassification::A)
1589 };
1590 Ok(TestWorker { classification })
1591 }
1592 }
1593
1594 let create_count = Arc::new(AtomicUsize::new(0));
1595 let pool = ClassifiedWorkerPool::new(
1596 1,
1597 TestWorkerFactory {
1598 create_count: create_count.clone(),
1599 },
1600 );
1601 let worker = pool
1602 .get_classified_worker(TestWorkerClassification::B)
1603 .await;
1604 assert!(worker.is_err());
1605 assert_eq!(
1606 worker.err().unwrap().code(),
1607 crate::PoolErrorCode::InvalidConfig
1608 );
1609
1610 let worker = pool
1611 .get_classified_worker(TestWorkerClassification::B)
1612 .await;
1613 assert!(worker.is_ok());
1614 assert_eq!(create_count.load(Ordering::SeqCst), 2);
1615}
1616
1617#[tokio::test(flavor = "multi_thread")]
1618async fn test_classified_waiter_keeps_queue_priority_over_later_generic_waiter() {
1619 use std::sync::mpsc;
1620
1621 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
1622 enum TestWorkerClassification {
1623 B,
1624 }
1625
1626 struct TestWorker {
1627 classification: TestWorkerClassification,
1628 }
1629
1630 #[async_trait::async_trait]
1631 impl ClassifiedWorker<TestWorkerClassification> for TestWorker {
1632 fn is_work(&self) -> bool {
1633 true
1634 }
1635
1636 fn is_valid(&self, c: TestWorkerClassification) -> bool {
1637 self.classification == c
1638 }
1639
1640 fn classification(&self) -> TestWorkerClassification {
1641 self.classification.clone()
1642 }
1643 }
1644
1645 struct TestWorkerFactory;
1646
1647 #[async_trait::async_trait]
1648 impl ClassifiedWorkerFactory<TestWorkerClassification, TestWorker> for TestWorkerFactory {
1649 async fn create(
1650 &self,
1651 classification: Option<TestWorkerClassification>,
1652 ) -> PoolResult<TestWorker> {
1653 Ok(TestWorker {
1654 classification: classification.unwrap_or(TestWorkerClassification::B),
1655 })
1656 }
1657 }
1658
1659 let pool = ClassifiedWorkerPool::new(1, TestWorkerFactory);
1660 let worker = pool
1661 .get_classified_worker(TestWorkerClassification::B)
1662 .await
1663 .unwrap();
1664
1665 let (tx, rx) = mpsc::channel();
1666
1667 let pool_ref = pool.clone();
1668 let tx_classified = tx.clone();
1669 let classified_task = tokio::spawn(async move {
1670 let _worker = pool_ref
1671 .get_classified_worker(TestWorkerClassification::B)
1672 .await
1673 .unwrap();
1674 tx_classified.send("classified").unwrap();
1675 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1676 });
1677
1678 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
1679
1680 let pool_ref = pool.clone();
1681 let generic_task = tokio::spawn(async move {
1682 let _worker = pool_ref.get_worker().await.unwrap();
1683 tx.send("generic").unwrap();
1684 });
1685
1686 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
1687 drop(worker);
1688
1689 let first = rx.recv_timeout(std::time::Duration::from_secs(2)).unwrap();
1690 assert_eq!(first, "classified");
1691
1692 classified_task.await.unwrap();
1693 generic_task.await.unwrap();
1694}
1695
1696#[tokio::test]
1697async fn test_generic_factory_worker_must_be_valid_for_its_primary_classification() {
1698 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
1699 enum TestWorkerClassification {
1700 A,
1701 B,
1702 }
1703
1704 struct TestWorker {
1705 classification: TestWorkerClassification,
1706 }
1707
1708 #[async_trait::async_trait]
1709 impl ClassifiedWorker<TestWorkerClassification> for TestWorker {
1710 fn is_work(&self) -> bool {
1711 true
1712 }
1713
1714 fn is_valid(&self, c: TestWorkerClassification) -> bool {
1715 c == TestWorkerClassification::B
1716 }
1717
1718 fn classification(&self) -> TestWorkerClassification {
1719 self.classification.clone()
1720 }
1721 }
1722
1723 struct TestWorkerFactory;
1724
1725 #[async_trait::async_trait]
1726 impl ClassifiedWorkerFactory<TestWorkerClassification, TestWorker> for TestWorkerFactory {
1727 async fn create(
1728 &self,
1729 _classification: Option<TestWorkerClassification>,
1730 ) -> PoolResult<TestWorker> {
1731 Ok(TestWorker {
1732 classification: TestWorkerClassification::A,
1733 })
1734 }
1735 }
1736
1737 let pool = ClassifiedWorkerPool::new(1, TestWorkerFactory);
1738 let worker = pool.get_worker().await;
1739 assert!(worker.is_err());
1740 assert_eq!(
1741 worker.err().unwrap().code(),
1742 crate::PoolErrorCode::InvalidConfig
1743 );
1744}
1745
1746#[tokio::test]
1747async fn test_classified_idle_worker_timeout_releases_worker() {
1748 use std::sync::atomic::{AtomicUsize, Ordering};
1749 use std::sync::Arc;
1750
1751 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
1752 enum TestWorkerClassification {
1753 A,
1754 B,
1755 }
1756
1757 struct TestWorker {
1758 id: usize,
1759 classification: TestWorkerClassification,
1760 }
1761
1762 #[async_trait::async_trait]
1763 impl ClassifiedWorker<TestWorkerClassification> for TestWorker {
1764 fn is_work(&self) -> bool {
1765 true
1766 }
1767
1768 fn is_valid(&self, c: TestWorkerClassification) -> bool {
1769 self.classification == c
1770 }
1771
1772 fn classification(&self) -> TestWorkerClassification {
1773 self.classification.clone()
1774 }
1775 }
1776
1777 struct TestWorkerFactory {
1778 create_count: Arc<AtomicUsize>,
1779 }
1780
1781 #[async_trait::async_trait]
1782 impl ClassifiedWorkerFactory<TestWorkerClassification, TestWorker> for TestWorkerFactory {
1783 async fn create(
1784 &self,
1785 classification: Option<TestWorkerClassification>,
1786 ) -> PoolResult<TestWorker> {
1787 let id = self.create_count.fetch_add(1, Ordering::SeqCst);
1788 Ok(TestWorker {
1789 id,
1790 classification: classification.unwrap_or(TestWorkerClassification::A),
1791 })
1792 }
1793 }
1794
1795 let create_count = Arc::new(AtomicUsize::new(0));
1796 let pool = ClassifiedWorkerPool::new_with_config(
1797 1,
1798 TestWorkerFactory {
1799 create_count: create_count.clone(),
1800 },
1801 ClassifiedWorkerPoolConfig {
1802 idle_timeout: Some(std::time::Duration::from_millis(30)),
1803 ..ClassifiedWorkerPoolConfig::default()
1804 },
1805 );
1806
1807 {
1808 let worker = pool
1809 .get_classified_worker(TestWorkerClassification::B)
1810 .await
1811 .unwrap();
1812 assert_eq!(worker.id, 0);
1813 assert_eq!(worker.classification(), TestWorkerClassification::B);
1814 }
1815
1816 tokio::time::sleep(std::time::Duration::from_millis(80)).await;
1817
1818 let worker = pool
1819 .get_classified_worker(TestWorkerClassification::A)
1820 .await
1821 .unwrap();
1822 assert_eq!(worker.id, 1);
1823 assert_eq!(worker.classification(), TestWorkerClassification::A);
1824 assert_eq!(create_count.load(Ordering::SeqCst), 2);
1825}
1826
1827#[tokio::test]
1828async fn test_get_classified_worker_uses_most_recent_matching_idle_worker() {
1829 use std::sync::atomic::{AtomicUsize, Ordering};
1830 use std::sync::Arc;
1831
1832 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
1833 enum TestWorkerClassification {
1834 A,
1835 }
1836
1837 struct TestWorker {
1838 id: usize,
1839 classification: TestWorkerClassification,
1840 }
1841
1842 #[async_trait::async_trait]
1843 impl ClassifiedWorker<TestWorkerClassification> for TestWorker {
1844 fn is_work(&self) -> bool {
1845 true
1846 }
1847
1848 fn is_valid(&self, c: TestWorkerClassification) -> bool {
1849 self.classification == c
1850 }
1851
1852 fn classification(&self) -> TestWorkerClassification {
1853 self.classification.clone()
1854 }
1855 }
1856
1857 struct TestWorkerFactory {
1858 create_count: Arc<AtomicUsize>,
1859 }
1860
1861 #[async_trait::async_trait]
1862 impl ClassifiedWorkerFactory<TestWorkerClassification, TestWorker> for TestWorkerFactory {
1863 async fn create(
1864 &self,
1865 classification: Option<TestWorkerClassification>,
1866 ) -> PoolResult<TestWorker> {
1867 let id = self.create_count.fetch_add(1, Ordering::SeqCst);
1868 Ok(TestWorker {
1869 id,
1870 classification: classification.unwrap_or(TestWorkerClassification::A),
1871 })
1872 }
1873 }
1874
1875 let create_count = Arc::new(AtomicUsize::new(0));
1876 let pool = ClassifiedWorkerPool::new(
1877 2,
1878 TestWorkerFactory {
1879 create_count: create_count.clone(),
1880 },
1881 );
1882
1883 let worker1 = pool
1884 .get_classified_worker(TestWorkerClassification::A)
1885 .await
1886 .unwrap();
1887 let worker2 = pool
1888 .get_classified_worker(TestWorkerClassification::A)
1889 .await
1890 .unwrap();
1891 assert_eq!(worker1.id, 0);
1892 assert_eq!(worker2.id, 1);
1893
1894 drop(worker1);
1895 drop(worker2);
1896
1897 let worker = pool
1898 .get_classified_worker(TestWorkerClassification::A)
1899 .await
1900 .unwrap();
1901 assert_eq!(worker.id, 1);
1902 assert_eq!(create_count.load(Ordering::SeqCst), 2);
1903}
1904
1905#[tokio::test]
1906async fn test_classified_cleanup_idle_worker_can_be_triggered_externally() {
1907 use std::sync::atomic::{AtomicUsize, Ordering};
1908 use std::sync::Arc;
1909
1910 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
1911 enum TestWorkerClassification {
1912 A,
1913 B,
1914 }
1915
1916 struct TestWorker {
1917 id: usize,
1918 classification: TestWorkerClassification,
1919 }
1920
1921 #[async_trait::async_trait]
1922 impl ClassifiedWorker<TestWorkerClassification> for TestWorker {
1923 fn is_work(&self) -> bool {
1924 true
1925 }
1926
1927 fn is_valid(&self, c: TestWorkerClassification) -> bool {
1928 self.classification == c
1929 }
1930
1931 fn classification(&self) -> TestWorkerClassification {
1932 self.classification.clone()
1933 }
1934 }
1935
1936 struct TestWorkerFactory {
1937 create_count: Arc<AtomicUsize>,
1938 }
1939
1940 #[async_trait::async_trait]
1941 impl ClassifiedWorkerFactory<TestWorkerClassification, TestWorker> for TestWorkerFactory {
1942 async fn create(
1943 &self,
1944 classification: Option<TestWorkerClassification>,
1945 ) -> PoolResult<TestWorker> {
1946 let id = self.create_count.fetch_add(1, Ordering::SeqCst);
1947 Ok(TestWorker {
1948 id,
1949 classification: classification.unwrap_or(TestWorkerClassification::A),
1950 })
1951 }
1952 }
1953
1954 let create_count = Arc::new(AtomicUsize::new(0));
1955 let pool = ClassifiedWorkerPool::new_with_config(
1956 1,
1957 TestWorkerFactory {
1958 create_count: create_count.clone(),
1959 },
1960 ClassifiedWorkerPoolConfig {
1961 idle_timeout: Some(std::time::Duration::from_millis(30)),
1962 ..ClassifiedWorkerPoolConfig::default()
1963 },
1964 );
1965
1966 {
1967 let worker = pool
1968 .get_classified_worker(TestWorkerClassification::B)
1969 .await
1970 .unwrap();
1971 assert_eq!(worker.id, 0);
1972 assert_eq!(worker.classification(), TestWorkerClassification::B);
1973 }
1974
1975 tokio::time::sleep(std::time::Duration::from_millis(80)).await;
1976
1977 assert_eq!(pool.cleanup_idle_worker(), 1);
1978
1979 let worker = pool
1980 .get_classified_worker(TestWorkerClassification::A)
1981 .await
1982 .unwrap();
1983 assert_eq!(worker.id, 1);
1984 assert_eq!(worker.classification(), TestWorkerClassification::A);
1985 assert_eq!(create_count.load(Ordering::SeqCst), 2);
1986}