polyc_runtime/admission.rs
1//! Shared admission control for edge webhook handlers (#795).
2//!
3//! Exactly one edge handler used to bound how many turns it dials
4//! concurrently (a `turn_limit` semaphore, acquired before spawning a
5//! turn and shed with a `503` when full). Every other edge — Slack,
6//! Telegram, Discord, email, A2A — accepted unbounded concurrent load,
7//! dialing the agent for every inbound webhook no matter how many were
8//! already in flight.
9//!
10//! [`AdmissionGate`] is that same bounded-concurrency pattern, extracted so
11//! every edge can front its agent dial with it. A caller that fails to
12//! [`AdmissionGate::try_admit`] must shed the request — respond with
13//! `polyc_proto::admission_shed_text` (or the edge's own transport-level
14//! equivalent, e.g. an HTTP `503`) — rather than dial the agent.
15
16use std::sync::Arc;
17
18use tokio::sync::{OwnedSemaphorePermit, Semaphore};
19
20/// Bounds how many turns an edge dials concurrently.
21///
22/// Cheap to clone — every handle shares the same underlying limit, the same
23/// shape as `AgentDialer`/`Client` elsewhere in an edge's `AppState`.
24#[derive(Clone)]
25pub struct AdmissionGate {
26 limit: Arc<Semaphore>,
27}
28
29impl AdmissionGate {
30 /// Builds a gate that admits at most `max_concurrent_turns` turns at
31 /// once. A caller beyond that bound must shed instead of dialing.
32 #[must_use]
33 pub fn new(max_concurrent_turns: usize) -> Self {
34 Self {
35 limit: Arc::new(Semaphore::new(max_concurrent_turns)),
36 }
37 }
38
39 /// Tries to admit one more turn. `None` means the gate is already at
40 /// capacity — the caller must shed rather than dial the agent.
41 ///
42 /// The returned [`AdmissionPermit`] holds the slot until dropped; hold
43 /// it for the lifetime of the spawned turn.
44 #[must_use]
45 pub fn try_admit(&self) -> Option<AdmissionPermit> {
46 Arc::clone(&self.limit)
47 .try_acquire_owned()
48 .ok()
49 .map(AdmissionPermit)
50 }
51}
52
53/// Held for the lifetime of one admitted turn; dropping it frees the slot
54/// back to the gate. Never read directly — its whole purpose is the `Drop`
55/// impl inherited from the wrapped [`OwnedSemaphorePermit`].
56#[must_use = "dropping this immediately releases the admission slot"]
57pub struct AdmissionPermit(
58 #[allow(dead_code, reason = "held only for its Drop")] OwnedSemaphorePermit,
59);
60
61#[cfg(test)]
62mod tests {
63 use super::AdmissionGate;
64
65 #[test]
66 fn admits_up_to_the_limit_then_sheds() {
67 let gate = AdmissionGate::new(2);
68 let first = gate.try_admit();
69 assert!(first.is_some());
70 let second = gate.try_admit();
71 assert!(second.is_some());
72 // The gate is now full — a third caller must shed, not dial.
73 assert!(gate.try_admit().is_none());
74 drop(first);
75 drop(second);
76 }
77
78 #[test]
79 fn dropping_a_permit_frees_the_slot() {
80 let gate = AdmissionGate::new(1);
81 let permit = gate.try_admit();
82 assert!(permit.is_some());
83 assert!(gate.try_admit().is_none(), "already at capacity");
84 drop(permit);
85 assert!(
86 gate.try_admit().is_some(),
87 "freed slot must be admittable again"
88 );
89 }
90
91 #[test]
92 fn clones_share_the_same_underlying_limit() {
93 let gate = AdmissionGate::new(1);
94 let clone = gate.clone();
95 let _held = gate.try_admit().expect("first admit");
96 // The clone shares the same semaphore, so it observes the same
97 // exhausted capacity — not an independent limit of its own.
98 assert!(clone.try_admit().is_none());
99 }
100}