1
use crate::db::DBError;
2
use sqlx::types::Uuid;
3
use sqlx::{Pool, Postgres};
4
use std::collections::HashMap;
5
use std::sync::{Arc, Mutex, MutexGuard};
6
use std::time::{Duration, Instant};
7
use tokio::sync::{OnceCell, OwnedSemaphorePermit, Semaphore};
8

            
9
#[cfg(test)]
10
mod tests;
11

            
12
/// Connections a single tenant pool may open.
13
pub(super) const TENANT_MAX_CONNECTIONS: u32 = 2;
14

            
15
/// Tenant pools resident in one process.
16
pub(super) const MAX_RESIDENT_POOLS: usize = 8;
17

            
18
/// How long a newcomer waits for a residency permit before giving up. Matches
19
/// the pool's own `acquire_timeout`, so a caller cannot wait appreciably longer
20
/// for a slot than for a connection once it has one.
21
const RESERVE_DEADLINE: Duration = Duration::from_secs(10);
22

            
23
/// How long each wait attempt runs before looking for a victim again.
24
const RESERVE_RETRY: Duration = Duration::from_millis(250);
25

            
26
/// A tenant pool and the residency permit that pays for it.
27
///
28
/// The permit lives *here*, not in the registry entry, and that placement is
29
/// what makes the budget hold under failure. The pool and its permit are
30
/// created and dropped as one value, so every path that discards a pool —
31
/// eviction, expiry, a failed build, a cancelled request — returns the permit
32
/// automatically. Nothing has to remember to.
33
///
34
/// It also decouples eviction from use: a caller holds an `Arc` of this while
35
/// working, so evicting the entry merely un-registers the pool. It stays alive,
36
/// and its permit stays spent, until the last user lets go.
37
pub(crate) struct TenantPool {
38
    pool: Pool<Postgres>,
39
    _permit: OwnedSemaphorePermit,
40
}
41

            
42
impl TenantPool {
43
143
    pub(crate) const fn new(pool: Pool<Postgres>, permit: OwnedSemaphorePermit) -> Self {
44
143
        Self {
45
143
            pool,
46
143
            _permit: permit,
47
143
        }
48
143
    }
49

            
50
235
    pub(crate) const fn pool(&self) -> &Pool<Postgres> {
51
235
        &self.pool
52
235
    }
53
}
54

            
55
/// One user's slot. The cell is empty while a build is in flight; `OnceCell`
56
/// serializes creation per user, so concurrent first requests await one
57
/// initialization instead of each building a pool and overwriting the rest.
58
pub(super) type Slot = Arc<OnceCell<Arc<TenantPool>>>;
59

            
60
struct Entry {
61
    cell: Slot,
62
    last_used: Instant,
63
    /// When this tenant's quarantine status was last read from the directory.
64
    checked_at: Instant,
65
}
66

            
67
pub(super) struct PoolRegistry {
68
    pools: Arc<Mutex<HashMap<Uuid, Entry>>>,
69
    residency: Arc<Semaphore>,
70
    expiration: Duration,
71
    reserve_deadline: Duration,
72
}
73

            
74
/// Removes pools idle for longer than `expiration`. Dropping the entry drops
75
/// the pool and returns its permit; nothing else is required.
76
2033
fn sweep_expired(pools: &mut HashMap<Uuid, Entry>, expiration: Duration) {
77
    // `last_used` is stamped on acquire, so a connection held longer than the
78
    // idle timeout looks idle to the clock while it is very much in use: only
79
    // drop what nobody holds. An ABANDONED slot — never built, nobody
80
    // building — is exactly such a thing, and requiring an initialised pool to
81
    // sweep left one behind for every request cancelled mid-build, forever.
82
2033
    pools.retain(|_, e| e.last_used.elapsed() < expiration || !unheld(e));
83
2033
}
84

            
85
/// Whether nobody holds this entry: no caller has the slot, and no connection
86
/// is out on the pool if one was ever built.
87
32
fn unheld(entry: &Entry) -> bool {
88
32
    Arc::strong_count(&entry.cell) == 1
89
31
        && entry
90
31
            .cell
91
31
            .get()
92
31
            .is_none_or(|pool| Arc::strong_count(pool) == 1)
93
32
}
94

            
95
/// Removes the least-recently-used slot so a newcomer can take its permit.
96
/// Returns whether anything was actually evicted — a registry holding only
97
/// in-flight builds has no permit to give back yet.
98
10
fn evict_lru(pools: &mut HashMap<Uuid, Entry>) -> bool {
99
10
    let Some(victim) = pools
100
10
        .iter()
101
47
        .filter(|(_, e)| is_idle(e))
102
10
        .min_by_key(|(_, e)| e.last_used)
103
10
        .map(|(id, _)| *id)
104
    else {
105
6
        return false;
106
    };
107

            
108
4
    pools.remove(&victim).is_some()
109
10
}
110

            
111
/// Whether a pool can be dropped right now: nobody is using it, and nobody is
112
/// on their way to using it.
113
///
114
/// The permit travels with the pool and comes back only when its last
115
/// [`TenantConnection`] drops, so dropping a BUSY pool frees nothing — the
116
/// newcomer waits out the whole reserve deadline and fails, while the victim
117
/// rebuilds a second pool beside its still-live one and takes the tenant over
118
/// its own connection cap.
119
///
120
/// A slot still building is never idle: its permit is
121
/// taken inside the build, so dropping the slot frees nothing and only strands
122
/// the builder.
123
///
124
/// BOTH counts matter. A connection anchors the pool itself, but a caller that
125
/// has taken the slot and not yet reached the pool — it may be in `routing()`,
126
/// an admin-database round trip — holds only the slot, and evicting under it
127
/// puts that caller on an unregistered pool while its replacement is built.
128
47
fn is_idle(entry: &Entry) -> bool {
129
47
    entry.cell.get().is_some() && unheld(entry)
130
47
}
131

            
132
impl PoolRegistry {
133
176
    pub(super) fn new(expiration: Duration) -> Self {
134
176
        Self {
135
176
            pools: Arc::new(Mutex::new(HashMap::new())),
136
176
            residency: Arc::new(Semaphore::new(MAX_RESIDENT_POOLS)),
137
176
            expiration,
138
176
            reserve_deadline: RESERVE_DEADLINE,
139
176
        }
140
176
    }
141

            
142
118
    pub(super) const fn idle_timeout(&self) -> Duration {
143
118
        self.expiration
144
118
    }
145

            
146
7827
    fn lock(&self) -> Result<MutexGuard<'_, HashMap<Uuid, Entry>>, DBError> {
147
7827
        self.pools
148
7827
            .lock()
149
7827
            .map_err(|_| DBError::Config(crate::config::ConfigError::DB))
150
7827
    }
151

            
152
    /// Marks the user's slot as used now and hands it back, if resident.
153
2033
    pub(super) fn touch(&self, user_id: Uuid) -> Result<Option<Slot>, DBError> {
154
2033
        let mut pools = self.lock()?;
155
2033
        sweep_expired(&mut pools, self.expiration);
156

            
157
2033
        Ok(pools.get_mut(&user_id).map(|entry| {
158
118
            entry.last_used = Instant::now();
159
118
            Arc::clone(&entry.cell)
160
118
        }))
161
2033
    }
162

            
163
    /// Reserves a residency permit, evicting the least-recently-used pool when
164
    /// the process is already at its ceiling.
165
    ///
166
    /// Retried on every pass, which is safe only because `is_idle` refuses to
167
    /// take a pool anybody holds: a newcomer can therefore reclaim a resident
168
    /// that fell idle while it waited, instead of giving up on the one victim
169
    /// it found at the start and timing out beside an idle pool.
170
1948
    pub(super) async fn reserve(&self) -> Result<OwnedSemaphorePermit, DBError> {
171
1948
        let deadline = Instant::now() + self.reserve_deadline;
172

            
173
        loop {
174
1950
            if let Ok(permit) = Arc::clone(&self.residency).try_acquire_owned() {
175
1945
                return Ok(permit);
176
5
            }
177

            
178
            {
179
5
                let mut pools = self.lock()?;
180
5
                evict_lru(&mut pools);
181
            }
182

            
183
5
            let remaining = deadline.saturating_duration_since(Instant::now());
184
5
            if remaining.is_zero() {
185
2
                return Err(DBError::PoolBudgetExhausted);
186
3
            }
187

            
188
3
            let wait = Arc::clone(&self.residency).acquire_owned();
189
3
            if let Ok(Ok(permit)) = tokio::time::timeout(remaining.min(RESERVE_RETRY), wait).await {
190
1
                return Ok(permit);
191
2
            }
192
        }
193
1948
    }
194

            
195
    /// Inserts an empty slot for `user_id`, or yields the existing one if
196
    /// another task won the race.
197
1953
    pub(super) fn install(&self, user_id: Uuid) -> Result<Slot, DBError> {
198
1953
        let mut pools = self.lock()?;
199

            
200
1953
        if let Some(existing) = pools.get_mut(&user_id) {
201
1
            existing.last_used = Instant::now();
202
1
            return Ok(Arc::clone(&existing.cell));
203
1952
        }
204

            
205
1952
        let cell = Slot::default();
206
1952
        pools.insert(
207
1952
            user_id,
208
1952
            Entry {
209
1952
                cell: Arc::clone(&cell),
210
1952
                last_used: Instant::now(),
211
1952
                checked_at: Instant::now(),
212
1952
            },
213
1952
        );
214

            
215
1952
        Ok(cell)
216
1953
    }
217

            
218
    /// Claims the right to re-read this tenant's quarantine status.
219
    ///
220
    /// Returns true to exactly one caller per interval. A plain "is it stale?"
221
    /// check would let every concurrent request for a tenant observe the same
222
    /// stale timestamp and each issue its own directory query — a stampede
223
    /// against an admin pool whose entire budget is 5. Claiming advances the
224
    /// clock under the lock, so the losers proceed on the previous answer.
225
    ///
226
    /// An absent entry answers true: that tenant is about to be built, and the
227
    /// build reads the directory anyway.
228
2038
    pub(super) fn claim_recheck(&self, user_id: Uuid, interval: Duration) -> Result<bool, DBError> {
229
2038
        let mut pools = self.lock()?;
230

            
231
2038
        let Some(entry) = pools.get_mut(&user_id) else {
232
1
            return Ok(true);
233
        };
234

            
235
2037
        if entry.checked_at.elapsed() < interval {
236
2034
            return Ok(false);
237
3
        }
238

            
239
3
        entry.checked_at = Instant::now();
240
3
        Ok(true)
241
2038
    }
242

            
243
    /// Rolls back a claimed recheck whose read failed, so the next caller
244
    /// retries instead of waiting out the interval on an unknown status.
245
2
    pub(super) fn release_recheck(&self, user_id: Uuid, interval: Duration) -> Result<(), DBError> {
246
2
        if let Some(entry) = self.lock()?.get_mut(&user_id) {
247
1
            entry.checked_at = Instant::now()
248
1
                .checked_sub(interval)
249
1
                .unwrap_or_else(Instant::now);
250
1
        }
251

            
252
2
        Ok(())
253
2
    }
254

            
255
    /// Drops a slot whose build failed — but only if it is still *this* slot.
256
    ///
257
    /// Without the identity check a caller whose build failed would evict
258
    /// whatever entry happens to sit under that user id, including a fresh one
259
    /// another task installed and populated in the meantime.
260
1796
    pub(super) fn forget(&self, user_id: Uuid, slot: &Slot) -> Result<(), DBError> {
261
1796
        let mut pools = self.lock()?;
262

            
263
1796
        if pools
264
1796
            .get(&user_id)
265
1796
            .is_some_and(|entry| Arc::ptr_eq(&entry.cell, slot))
266
1795
        {
267
1795
            pools.remove(&user_id);
268
1795
        }
269

            
270
1796
        Ok(())
271
1796
    }
272
}