This repository has no description
4.1 kB
135 lines
1use std::sync::{Arc, OnceLock};
2
3use tokio::sync::{OwnedSemaphorePermit, Semaphore};
4
5use crate::cpu::threads;
6
7const RESOLVE_SLOTS_PER_TRANSPORT: usize = 16;
8
9const SHARING_TRANSPORTS: usize = 2;
10
11const RESOLVE_SLOTS: usize = RESOLVE_SLOTS_PER_TRANSPORT * SHARING_TRANSPORTS;
12
13static PROCESS: OnceLock<Slots> = OnceLock::new();
14
15pub struct SlotPermit(#[allow(dead_code)] OwnedSemaphorePermit);
16
17macro_rules! slot_kind {
18 ($($name:ident),+ $(,)?) => {$(
19 #[derive(Clone)]
20 pub struct $name(Arc<Semaphore>);
21
22 impl $name {
23 pub fn new(permits: usize) -> Self {
24 Self(Arc::new(Semaphore::new(permits.max(1))))
25 }
26
27 pub async fn acquire(&self) -> SlotPermit {
28 SlotPermit(
29 Arc::clone(&self.0)
30 .acquire_owned()
31 .await
32 .expect("a slot budget closes only with the process that owns it"),
33 )
34 }
35
36 pub fn available(&self) -> usize {
37 self.0.available_permits()
38 }
39 }
40 )+};
41}
42
43slot_kind!(ResolveSlots, ReceiveSlots, PackSlots);
44
45impl ResolveSlots {
46 pub fn try_acquire(&self) -> Option<SlotPermit> {
47 Arc::clone(&self.0).try_acquire_owned().ok().map(SlotPermit)
48 }
49}
50
51#[derive(Clone)]
52pub struct Slots {
53 pub resolve: ResolveSlots,
54 pub receive: ReceiveSlots,
55 pub pack: PackSlots,
56}
57
58impl Slots {
59 pub fn for_machine() -> Self {
60 PROCESS
61 .get_or_init(|| Self {
62 resolve: ResolveSlots::new(RESOLVE_SLOTS),
63 receive: ReceiveSlots::new(threads().get()),
64 pack: PackSlots::new(threads().get()),
65 })
66 .clone()
67 }
68
69 pub fn testing(permits: usize) -> Self {
70 Self {
71 resolve: ResolveSlots::new(permits),
72 receive: ReceiveSlots::new(permits),
73 pack: PackSlots::new(permits),
74 }
75 }
76}
77
78#[cfg(test)]
79mod tests {
80 use super::*;
81
82 #[tokio::test]
83 async fn a_clone_shares_one_pool_per_kind_of_work() {
84 let slots = Slots::testing(1);
85 let other = slots.clone();
86 let held = slots.receive.acquire().await;
87 assert_eq!(
88 other.receive.available(),
89 0,
90 "a cloned budget mustn't grant a second permit for the one slot"
91 );
92 assert_eq!(
93 other.pack.available(),
94 1,
95 "spending a receive slot mustn't spend the pack budget"
96 );
97 let _resolving = slots.resolve.acquire().await;
98 assert!(
99 other.resolve.try_acquire().is_none(),
100 "a cosmetic lookup mustn't wait, \
101 or a push queues on it while its receive and pack slots stay spent"
102 );
103 drop(held);
104 assert_eq!(other.receive.available(), 1);
105 }
106
107 #[tokio::test]
108 async fn every_caller_of_for_machine_shares_one_budget_that_no_test_budget_touches() {
109 let ssh = Slots::for_machine();
110 let http = Slots::for_machine();
111 let isolated = Slots::testing(1);
112 let before = http.receive.available();
113 let _held = ssh.receive.acquire().await;
114 assert_eq!(
115 http.receive.available(),
116 before - 1,
117 "two transports asking the machine for a budget must get the same one, \
118 or the process grants twice the concurrency it was configured for"
119 );
120 let resolving: Vec<SlotPermit> = (0..RESOLVE_SLOTS_PER_TRANSPORT)
121 .filter_map(|_| ssh.resolve.try_acquire())
122 .collect();
123 assert_eq!(resolving.len(), RESOLVE_SLOTS_PER_TRANSPORT);
124 assert!(
125 http.resolve.try_acquire().is_some(),
126 "collapsing a per-transport pool into a process-wide one mustn't give a deployment \
127 running both transports less outbound resolution than either had on its own"
128 );
129 assert_eq!(
130 isolated.receive.available(),
131 1,
132 "whatever the process budget is doing mustn't spend a test budget"
133 );
134 }
135}