krabka_raft/
connection_limiter.rs

1//! Live-connection accounting for Kafka's `ConnectionQuotas`: the
2//! `max.connections` caps, broker-wide and per listener, the
3//! `max.connections.per.ip` cap and its per-host overrides, and the KIP-612
4//! `max.connection.creation.rate` gates, together with the RAII guard that
5//! releases a slot when a connection ends. It sits apart from the accept
6//! loops that consult it: one limiter serves every broker listener of a node,
7//! and the controller listener has one of its own, as Kafka's
8//! `ControllerServer` builds its own `SocketServer`.
9//!
10//! Every quota is resolved per accept from the metadata image
11//! ([`ConnectionQuotaSettings::resolve`]), so a dynamic change applies to the
12//! next connection without a restart.
13
14mod overrides;
15mod rate;
16mod settings;
17
18use {
19    crate::types::NodeId,
20    krabka_metadata::MetadataImage,
21    overrides::OverrideAddresses,
22    rate::CreationRates,
23    std::{
24        collections::HashMap,
25        hash,
26        net::IpAddr,
27        sync::{Arc, Mutex, MutexGuard, PoisonError},
28        time::{Duration, Instant},
29    },
30};
31pub use {
32    overrides::IpOverrides,
33    settings::{
34        ConnectionLimits, ConnectionQuotaSettings, MAX_CONNECTION_CREATION_RATE, MAX_CONNECTIONS,
35        MAX_CONNECTIONS_PER_IP, MAX_CONNECTIONS_PER_IP_OVERRIDES,
36    },
37};
38
39/// The quotas one accept is checked against: the settings in force and the
40/// resolved override addresses.
41#[derive(Debug, Clone)]
42pub struct ConnectionQuotas {
43    pub settings: ConnectionQuotaSettings,
44    pub per_ip_overrides: IpOverrides,
45}
46
47impl ConnectionQuotas {
48    /// The cap for `ip`: its override, else `max.connections.per.ip`.
49    #[must_use]
50    pub fn max_connections_for(&self, ip: IpAddr) -> usize {
51        self.per_ip_overrides
52            .get(&ip.to_canonical())
53            .copied()
54            .unwrap_or(self.settings.max_connections_per_ip)
55    }
56}
57
58/// The listener an accept arrived on.
59#[derive(Debug, Clone, Copy, PartialEq, Eq)]
60pub struct AcceptingListener<'a> {
61    /// The listener's name, matched without regard to case as Kafka's
62    /// `ListenerName.normalised` matches it.
63    pub name: &'a str,
64    /// Whether the listener is exempt from the broker-wide cap and rate:
65    /// Kafka's `protectedListener`, the inter-broker listener of a node with
66    /// more than one listener.
67    pub protected: bool,
68}
69
70/// Live connection counts, under one lock as Kafka keeps them under `counts`.
71#[derive(Debug, Default)]
72struct Counts {
73    total: usize,
74    per_ip: HashMap<IpAddr, usize>,
75    per_listener: HashMap<String, usize>,
76}
77
78/// Live-connection accounting for one node's listeners. A clone shares the
79/// same counters, so every accept loop and every [`ConnectionGuard`] account
80/// against one set of totals.
81#[derive(Clone)]
82pub struct ConnectionLimiter {
83    configured: ConnectionLimits,
84    counts: Arc<Mutex<Counts>>,
85    rates: Arc<Mutex<CreationRates>>,
86    overrides: Arc<OverrideAddresses>,
87    /// The settings last resolved for each listener, which a refused value
88    /// leaves in force.
89    applied: Arc<Mutex<HashMap<String, ConnectionQuotaSettings>>>,
90}
91
92/// The connection ceiling that refused an inbound connection.
93#[derive(Debug, Clone, Copy, PartialEq, Eq)]
94pub enum ConnectionLimit {
95    /// `max.connections`, broker-wide.
96    Global,
97    /// `listener.name.<listener>.max.connections`.
98    Listener,
99    /// `max.connections.per.ip` or the address's override.
100    PerIp,
101}
102
103/// Lock `mutex`, recovering the data of a poisoned lock: a count is still
104/// right after a panic elsewhere, because every update is one statement.
105fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
106    mutex.lock().unwrap_or_else(PoisonError::into_inner)
107}
108
109impl ConnectionLimiter {
110    /// A limiter for a node started with `max_connections` and
111    /// `max_connections_per_ip`.
112    #[must_use]
113    pub fn new(max_connections: usize, max_connections_per_ip: usize) -> Self {
114        Self {
115            configured: ConnectionLimits {
116                max_connections,
117                max_connections_per_ip,
118            },
119            counts: Arc::default(),
120            rates: Arc::default(),
121            overrides: Arc::default(),
122            applied: Arc::default(),
123        }
124    }
125
126    /// The caps this limiter was started with.
127    #[must_use]
128    pub const fn limits(&self) -> ConnectionLimits { self.configured }
129
130    /// The quotas in force on `listener` of `node_id` as of `image`. A value
131    /// Kafka would refuse leaves the listener's previous value in force. A new
132    /// `max.connections.per.ip.overrides` set that names a host is resolved in
133    /// the background, so this never waits on DNS. It must be called within a
134    /// Tokio runtime.
135    #[must_use]
136    pub fn quotas(
137        &self,
138        image: &MetadataImage,
139        node_id: NodeId,
140        listener: &str,
141    ) -> ConnectionQuotas {
142        let name = listener.to_ascii_uppercase();
143        let mut applied = lock(&self.applied);
144        let settings = ConnectionQuotaSettings::resolve(
145            image,
146            node_id,
147            listener,
148            self.configured,
149            applied.get(&name),
150        );
151        applied.insert(name, settings.clone());
152        drop(applied);
153        let per_ip_overrides = self.overrides.addresses(&settings.per_ip_overrides);
154        ConnectionQuotas {
155            settings,
156            per_ip_overrides,
157        }
158    }
159
160    /// Record a new connection on `listener` at `now` against the creation
161    /// rates of `quotas`, and return how long the acceptor waits before it
162    /// takes the connection, at most `max_delay`
163    /// (`quota.window.size.seconds`). Kafka's acceptor waits the throttle out
164    /// and then accepts.
165    #[must_use]
166    pub fn creation_delay(
167        &self,
168        listener: AcceptingListener<'_>,
169        quotas: &ConnectionQuotas,
170        now: Instant,
171        max_delay: Duration,
172    ) -> Duration {
173        lock(&self.rates).admit(
174            listener.name,
175            listener.protected,
176            (
177                quotas.settings.max_connection_creation_rate,
178                quotas.settings.listener_max_connection_creation_rate,
179            ),
180            now,
181            max_delay,
182        )
183    }
184
185    /// Take back a connection [`Self::creation_delay`] recorded, for one the
186    /// peer's `ip` quota throttles, as Kafka's `updateListenerMetrics`
187    /// un-records it from the broker-wide and listener rates.
188    pub fn refund_creation(&self, listener: AcceptingListener<'_>, quotas: &ConnectionQuotas) {
189        lock(&self.rates).refund(
190            listener.name,
191            listener.protected,
192            (
193                quotas.settings.max_connection_creation_rate,
194                quotas.settings.listener_max_connection_creation_rate,
195            ),
196        );
197    }
198
199    /// Try to reserve a slot for a connection from `ip` on `listener` under
200    /// `quotas`, which the caller resolves per accept. On success returns a
201    /// [`ConnectionGuard`] that releases the slot on drop. Lowering a cap
202    /// closes no connection that is already counted, as in Kafka's
203    /// `ConnectionQuotas`.
204    ///
205    /// The checks run in Kafka's order: the listener's cap, then, unless the
206    /// listener is protected, the broker-wide cap, then the cap of the
207    /// address.
208    ///
209    /// # Errors
210    /// Returns the [`ConnectionLimit`] that refused the connection, and
211    /// reserves nothing.
212    pub fn try_acquire_within(
213        &self,
214        ip: IpAddr,
215        listener: AcceptingListener<'_>,
216        quotas: &ConnectionQuotas,
217    ) -> Result<ConnectionGuard, ConnectionLimit> {
218        let name = listener.name.to_ascii_uppercase();
219        let mut counts = lock(&self.counts);
220        let on_listener = counts.per_listener.get(&name).copied().unwrap_or(0);
221        if on_listener >= quotas.settings.listener_max_connections {
222            return Err(ConnectionLimit::Listener);
223        }
224        if !listener.protected && counts.total >= quotas.settings.max_connections {
225            return Err(ConnectionLimit::Global);
226        }
227        let from_ip = counts.per_ip.get(&ip).copied().unwrap_or(0);
228        if from_ip >= quotas.max_connections_for(ip) {
229            return Err(ConnectionLimit::PerIp);
230        }
231        counts.total += 1;
232        *counts.per_ip.entry(ip).or_default() += 1;
233        *counts.per_listener.entry(name.clone()).or_default() += 1;
234        drop(counts);
235        Ok(ConnectionGuard {
236            limiter: self.clone(),
237            ip,
238            listener: name,
239        })
240    }
241
242    /// Diagnostic accessor: current live-connection count across listeners.
243    #[must_use]
244    pub fn total(&self) -> usize { lock(&self.counts).total }
245}
246
247/// RAII release for one accepted connection.
248///
249/// Moved into the spawned per-connection task, so it fires however the
250/// connection ends: clean close, error, panic, or task abort. On drop it
251/// decrements the total, the address and the listener counts, and removes a
252/// count that reaches 0 so the maps do not grow without bound.
253pub struct ConnectionGuard {
254    limiter: ConnectionLimiter,
255    ip: IpAddr,
256    listener: String,
257}
258
259/// Decrement `key`'s count in `map`, removing it at 0.
260fn release<K: Eq + hash::Hash>(map: &mut HashMap<K, usize>, key: &K) {
261    if let Some(count) = map.get_mut(key) {
262        *count -= 1;
263        if *count == 0 {
264            map.remove(key);
265        }
266    }
267}
268
269impl Drop for ConnectionGuard {
270    fn drop(&mut self) {
271        let mut counts = lock(&self.limiter.counts);
272        counts.total -= 1;
273        release(&mut counts.per_ip, &self.ip);
274        release(&mut counts.per_listener, &self.listener);
275    }
276}
277
278#[cfg(test)]
279mod tests {
280    use {super::*, assert2::check, std::collections::BTreeMap};
281
282    const PLAINTEXT: AcceptingListener<'static> = AcceptingListener {
283        name: "PLAINTEXT",
284        protected: false,
285    };
286
287    const A: &str = "10.0.0.1";
288    const B: &str = "10.0.0.2";
289
290    fn ip(address: &str) -> IpAddr { address.parse().unwrap() }
291
292    /// Quotas with every cap at `usize::MAX` and no rate, before `customise`.
293    fn quotas(customise: impl FnOnce(&mut ConnectionQuotas)) -> ConnectionQuotas {
294        let mut quotas = ConnectionQuotas {
295            settings: ConnectionQuotaSettings {
296                max_connections: usize::MAX,
297                max_connections_per_ip: usize::MAX,
298                per_ip_overrides: BTreeMap::new(),
299                listener_max_connections: usize::MAX,
300                max_connection_creation_rate: None,
301                listener_max_connection_creation_rate: None,
302            },
303            per_ip_overrides: Arc::default(),
304        };
305        customise(&mut quotas);
306        quotas
307    }
308
309    impl ConnectionLimiter {
310        fn per_ip_count(&self, ip: IpAddr) -> usize {
311            lock(&self.counts).per_ip.get(&ip).copied().unwrap_or(0)
312        }
313
314        fn listener_count(&self, name: &str) -> usize {
315            lock(&self.counts)
316                .per_listener
317                .get(name)
318                .copied()
319                .unwrap_or(0)
320        }
321    }
322
323    #[test]
324    fn a_guard_counts_and_releases_the_total_the_address_and_the_listener() {
325        let limiter = ConnectionLimiter::new(usize::MAX, usize::MAX);
326        let unlimited = quotas(|_| {});
327        let first = limiter
328            .try_acquire_within(ip(A), PLAINTEXT, &unlimited)
329            .unwrap();
330        let second = limiter
331            .try_acquire_within(ip(A), PLAINTEXT, &unlimited)
332            .unwrap();
333        check!(
334            (
335                limiter.total(),
336                limiter.per_ip_count(ip(A)),
337                limiter.listener_count("PLAINTEXT")
338            ) == (2, 2, 2)
339        );
340        drop((first, second));
341        check!(limiter.total() == 0);
342        check!(lock(&limiter.counts).per_ip.is_empty());
343        check!(lock(&limiter.counts).per_listener.is_empty());
344    }
345
346    /// With one connection from `A` held on `PLAINTEXT`, which listener and
347    /// address the next connection may come from.
348    #[yare::parameterized(
349        global_cap = { |q: &mut ConnectionQuotas| q.settings.max_connections = 1, B, PLAINTEXT, Some(ConnectionLimit::Global) },
350        protected_listener_skips_global_cap = {
351            |q: &mut ConnectionQuotas| q.settings.max_connections = 1,
352            B,
353            AcceptingListener { name: "REPLICATION", protected: true },
354            None
355        },
356        listener_cap = { |q: &mut ConnectionQuotas| q.settings.listener_max_connections = 1, B, PLAINTEXT, Some(ConnectionLimit::Listener) },
357        listener_cap_ignores_case = {
358            |q: &mut ConnectionQuotas| q.settings.listener_max_connections = 1,
359            B,
360            AcceptingListener { name: "plaintext", protected: false },
361            Some(ConnectionLimit::Listener)
362        },
363        listener_cap_counts_its_own_listener = {
364            |q: &mut ConnectionQuotas| q.settings.listener_max_connections = 1,
365            B,
366            AcceptingListener { name: "EXTERNAL", protected: false },
367            None
368        },
369        per_ip_cap = { |q: &mut ConnectionQuotas| q.settings.max_connections_per_ip = 1, A, PLAINTEXT, Some(ConnectionLimit::PerIp) },
370        per_ip_cap_other_address = { |q: &mut ConnectionQuotas| q.settings.max_connections_per_ip = 1, B, PLAINTEXT, None },
371        override_raises_the_cap = {
372            |q: &mut ConnectionQuotas| {
373                q.settings.max_connections_per_ip = 1;
374                q.per_ip_overrides = Arc::new(HashMap::from([(ip(A), 2)]));
375            },
376            A,
377            PLAINTEXT,
378            None
379        },
380        override_of_zero_refuses_the_address = {
381            |q: &mut ConnectionQuotas| q.per_ip_overrides = Arc::new(HashMap::from([(ip(B), 0)])),
382            B,
383            PLAINTEXT,
384            Some(ConnectionLimit::PerIp)
385        },
386    )]
387    fn the_next_connection_meets_each_cap(
388        customise: fn(&mut ConnectionQuotas),
389        from: &str,
390        listener: AcceptingListener<'static>,
391        refused: Option<ConnectionLimit>,
392    ) {
393        let limiter = ConnectionLimiter::new(usize::MAX, usize::MAX);
394        let _held = limiter
395            .try_acquire_within(ip(A), PLAINTEXT, &quotas(|_| {}))
396            .unwrap();
397        let next = limiter.try_acquire_within(ip(from), listener, &quotas(customise));
398        assert2::assert!(next.as_ref().err().copied() == refused);
399        // A refusal reserves nothing.
400        assert2::assert!(limiter.total() == 1 + usize::from(refused.is_none()));
401    }
402
403    #[test]
404    fn a_v4_mapped_peer_meets_its_v4_override() {
405        let quotas = quotas(|q| q.per_ip_overrides = Arc::new(HashMap::from([(ip(A), 0)])));
406        assert2::assert!(quotas.max_connections_for(ip("::ffff:10.0.0.1")) == 0);
407    }
408}
krabka_raft/connection_limiter.rs

krabka_raft/
connection_limiter.rs
