Skip to main content

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}