featherbit/plugins/native/

limit_conn.rs

//! Concurrent-request limiting (`limit-conn`), ported from APISIX's
//! `apisix/plugins/limit-conn`.
//!
//! Concurrency is a property of the *span* between entering and leaving the
//! upstream call, which a single featherbit node cannot observe — a node runs
//! once, at one point in the graph. The behavior is therefore expressed as a
//! **pair of nodes** sharing one in-flight counter:
//!
//! - an **acquire** node placed *before* `upstream`, which increments the
//!   counter and rejects when the ceiling is reached, and
//! - a **release** node placed *after* `upstream` (on **both** the success and
//!   error paths), which decrements the counter.
//!
//! Both nodes are configured with the same `key` template (and identical
//! `conn`/`burst`), so [`crate::traffic::ConnRegistry`] hands them the same
//! [`std::sync::atomic::AtomicI64`]. This is the same "two phases, one shared
//! key" shape `proxy-rewrite` uses for request/response.
//!
//! # Wiring
//!
//! ```text
//!            ┌──────────────────┐        ┌──────────┐        ┌──────────────────┐
//!  listener →│ limit-conn       │success →│ upstream │success →│ limit-conn       │→ client
//!            │  (phase=acquire) │        │          │        │  (phase=release) │
//!            └──────────────────┘        └──────────┘        └──────────────────┘
//!                    │ limited                 │ error               ▲
//!                    ▼                         └─────────────────────┘
//!               client.in                 (release also runs on the error path
//!             (503 rejection)              so the counter always decrements)
//! ```
//!
//! The acquire node's `limited` port goes to `client.in`: an over-limit
//! request short-circuits straight to the client with the rejection
//! response. The release node must sit on *every* path out of `upstream`
//! (success and error) so a failed upstream call still frees the slot.

use async_trait::async_trait;
use bytes::Bytes;
use std::collections::HashMap;
use std::sync::atomic::Ordering;
use std::sync::Arc;

use crate::context::Context;
use crate::plugins::resources::PluginResources;
use crate::plugins::{Plugin, PluginOutput, PluginResult};
use crate::vars::template::Template;

/// Which half of the pair this node is.
#[derive(Debug, Clone, Copy, PartialEq)]
enum Role {
    /// Runs before `upstream`: increments the counter, rejects at the ceiling.
    Acquire,
    /// Runs after `upstream`: decrements the counter (floored at zero).
    Release,
}

/// One node of a `limit-conn` acquire/release pair.
///
/// Holds a handle to the process-wide [`crate::traffic::ConnRegistry`]; the
/// per-key counter is resolved at request time from the interpolated `key` so
/// concurrency is measured per client (or per whatever the key selects), shared
/// across both nodes of the pair.
pub struct LimitConnPlugin {
    role: Role,
    /// Maximum sustained concurrent requests.
    conn: i64,
    /// Extra concurrent requests tolerated above `conn`.
    burst: i64,
    /// Key template (interpolated per request unless `key_type` is `constant`).
    key: String,
    /// Pre-parsed `key`: supports `{{namespace.path}}` references and legacy
    /// `$var` interpolation (see [`Template::render_with_legacy`]). Only used
    /// when `key_constant` is false.
    key_tpl: Template,
    /// When `true`, `key` is used verbatim rather than interpolated.
    key_constant: bool,
    /// Status returned when the limit is exceeded.
    rejected_code: u16,
    /// Optional human-readable rejection message (JSON `{"error_msg": ...}`).
    /// Supports `{{namespace.path}}` references (no legacy `$var`
    /// interpolation — this field never supported it, so this sweep must not
    /// start).
    rejected_msg: Option<Template>,
    resources: Arc<PluginResources>,
}

impl LimitConnPlugin {
    /// Builds one node of the pair from node config.
    ///
    /// Accepted keys:
    /// - `phase` / `role` (string, **required**): `acquire` (before upstream)
    ///   or `release` (after upstream). Any other value is a config error.
    /// - `conn` (integer, default `20`): maximum sustained concurrent requests.
    ///   Must be `> 0`.
    /// - `burst` (integer, default `0`): extra concurrent requests tolerated
    ///   above `conn`. Must be `>= 0`. The hard ceiling is `conn + burst`.
    /// - `key` (string template, default `"$remote_addr"`): the concurrency
    ///   key, rendered per request — supports `{{namespace.path}}` references
    ///   plus legacy `$var` interpolation (see
    ///   [`crate::vars::template::Template::render_with_legacy`]) unless
    ///   `key_type: constant`. Both nodes of a pair **must** use the same
    ///   `key` and `conn`/`burst` so they share one counter.
    /// - `key_type` (string, default `var`): `constant` uses `key` verbatim;
    ///   any other value (`var`, `var_combination`) interpolates it.
    /// - `rejected_code` (integer, default `503`): status for over-limit
    ///   requests.
    /// - `rejected_msg` (string, optional): message returned as a JSON body
    ///   `{"error_msg": "..."}` on rejection. Supports `{{namespace.path}}`
    ///   references.
    /// - `default_conn_delay` (number, optional): accepted for APISIX config
    ///   compatibility but ignored — featherbit rejects rather than delays.
    ///
    /// ```yaml
    /// # before upstream
    /// type: limit-conn
    /// config:
    ///   phase: acquire
    ///   conn: 100
    ///   burst: 50
    ///   key: $remote_addr
    ///   rejected_code: 503
    /// ---
    /// # after upstream (on both success and error paths)
    /// type: limit-conn
    /// config:
    ///   phase: release
    ///   conn: 100
    ///   burst: 50
    ///   key: $remote_addr
    /// ```
    pub fn from_config(
        config: &HashMap<String, serde_json::Value>,
        resources: &Arc<PluginResources>,
    ) -> Result<Self, String> {
        let role = match config
            .get("phase")
            .or_else(|| config.get("role"))
            .and_then(|v| v.as_str())
        {
            Some("acquire") => Role::Acquire,
            Some("release") => Role::Release,
            Some(other) => {
                return Err(format!(
                    "limit-conn: unknown phase/role '{}' (expected 'acquire' or 'release')",
                    other
                ))
            }
            None => {
                return Err(
                    "limit-conn: 'phase' (or 'role') is required: 'acquire' or 'release'"
                        .to_string(),
                )
            }
        };

        let conn = config.get("conn").and_then(|v| v.as_i64()).unwrap_or(20);
        if conn <= 0 {
            return Err(format!(
                "limit-conn: 'conn' must be a positive integer, got {}",
                conn
            ));
        }

        let burst = config.get("burst").and_then(|v| v.as_i64()).unwrap_or(0);
        if burst < 0 {
            return Err(format!(
                "limit-conn: 'burst' must be non-negative, got {}",
                burst
            ));
        }

        let key = config
            .get("key")
            .and_then(|v| v.as_str())
            .unwrap_or("$remote_addr")
            .to_string();
        // Discard warnings here — the compile-time walk (a later task)
        // reports well-formed-but-unknown references; execution must not.
        let key_tpl = Template::parse(&key).0;

        let key_constant = matches!(
            config.get("key_type").and_then(|v| v.as_str()),
            Some("constant")
        );

        let rejected_code = config
            .get("rejected_code")
            .and_then(|v| v.as_u64())
            .map(|c| c as u16)
            .unwrap_or(503);

        let rejected_msg = config
            .get("rejected_msg")
            .and_then(|v| v.as_str())
            // Discard warnings here — the compile-time walk (a later task)
            // reports well-formed-but-unknown references; execution must not.
            .map(|s| Template::parse(s).0);

        Ok(Self {
            role,
            conn,
            burst,
            key,
            key_tpl,
            key_constant,
            rejected_code,
            rejected_msg,
            resources: resources.clone(),
        })
    }

    /// Resolves the concurrency key for this request.
    fn resolve_key(&self, ctx: &Context) -> String {
        if self.key_constant {
            self.key.clone()
        } else {
            self.key_tpl.render_with_legacy(ctx)
        }
    }
}

#[async_trait]
impl Plugin for LimitConnPlugin {
    fn plugin_type(&self) -> &str {
        "limit-conn"
    }

    async fn execute(&self, mut ctx: Context) -> PluginResult {
        let key = self.resolve_key(&ctx);
        let counter = self.resources.traffic.conn.counter(&key);

        match self.role {
            Role::Acquire => {
                // fetch_add returns the value *before* the increment, i.e. the
                // number of requests already in flight.
                let in_flight = counter.fetch_add(1, Ordering::SeqCst);
                if in_flight >= self.conn + self.burst {
                    // Over the ceiling — undo our increment and reject.
                    counter.fetch_sub(1, Ordering::SeqCst);

                    ctx.response.status_code = self.rejected_code;
                    let body = match &self.rejected_msg {
                        Some(msg) => {
                            let rendered = msg.render(&ctx).into_owned();
                            format!(r#"{{"error_msg":{}}}"#, json_string(&rendered))
                        }
                        None => r#"{"error":"limit_conn_exceeded"}"#.to_string(),
                    };
                    ctx.response.body = Bytes::from(body);
                    ctx.response.headers.insert(
                        "content-type".to_string(),
                        vec!["application/json".to_string()],
                    );

                    return Ok(PluginOutput::on_port(ctx, "limited"));
                }
                Ok(PluginOutput::success(ctx))
            }
            Role::Release => {
                // Decrement, floored at zero so a stray release (e.g. an
                // acquire that rejected but was still wired to a release) can
                // never drive the counter negative.
                let mut cur = counter.load(Ordering::SeqCst);
                while cur > 0 {
                    match counter.compare_exchange_weak(
                        cur,
                        cur - 1,
                        Ordering::SeqCst,
                        Ordering::SeqCst,
                    ) {
                        Ok(_) => break,
                        Err(actual) => cur = actual,
                    }
                }
                Ok(PluginOutput::success(ctx))
            }
        }
    }
}

/// Minimal JSON string escaping for the rejection message.
fn json_string(s: &str) -> String {
    let mut out = String::with_capacity(s.len() + 2);
    out.push('"');
    for c in s.chars() {
        match c {
            '"' => out.push_str("\\\""),
            '\\' => out.push_str("\\\\"),
            '\n' => out.push_str("\\n"),
            '\r' => out.push_str("\\r"),
            '\t' => out.push_str("\\t"),
            c => out.push(c),
        }
    }
    out.push('"');
    out
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::context::{GatewayRequest, GatewayResponse, Protocol};

    fn ctx() -> Context {
        Context {
            request: GatewayRequest {
                method: "GET".to_string(),
                path: "/".to_string(),
                host: "localhost".to_string(),
                scheme: "http".to_string(),
                headers: HashMap::new(),
                query_params: HashMap::new(),
                body: Bytes::new(),
                remote_addr: "10.0.0.1:5000".to_string(),
                protocol: Protocol::Http1,
            },
            response: GatewayResponse {
                status_code: 0,
                headers: HashMap::new(),
                body: Bytes::new(),
                stream: None,
            },
            message: HashMap::new(),
            errors: Vec::new(),
        }
    }

    fn cfg(pairs: &[(&str, serde_json::Value)]) -> HashMap<String, serde_json::Value> {
        pairs
            .iter()
            .map(|(k, v)| (k.to_string(), v.clone()))
            .collect()
    }

    #[test]
    fn test_unknown_role_and_bad_thresholds_fail() {
        let r = PluginResources::empty();
        assert!(
            LimitConnPlugin::from_config(&cfg(&[("phase", serde_json::json!("nope"))]), &r)
                .is_err()
        );
        assert!(
            LimitConnPlugin::from_config(&cfg(&[]), &r).is_err(),
            "missing phase must fail"
        );
        assert!(LimitConnPlugin::from_config(
            &cfg(&[
                ("phase", serde_json::json!("acquire")),
                ("conn", serde_json::json!(0))
            ]),
            &r
        )
        .is_err());
        assert!(LimitConnPlugin::from_config(
            &cfg(&[
                ("phase", serde_json::json!("acquire")),
                ("burst", serde_json::json!(-1))
            ]),
            &r
        )
        .is_err());
    }

    #[tokio::test]
    async fn test_rejected_msg_renders_template() {
        let r = PluginResources::empty();
        let acquire = LimitConnPlugin::from_config(
            &cfg(&[
                ("phase", serde_json::json!("acquire")),
                ("conn", serde_json::json!(1)),
                ("key", serde_json::json!("const-conn")),
                ("key_type", serde_json::json!("constant")),
                (
                    "rejected_msg",
                    serde_json::json!("too many from {{client.ip}}"),
                ),
            ]),
            &r,
        )
        .unwrap();
        acquire.execute(ctx()).await.unwrap();
        let out = acquire.execute(ctx()).await.unwrap();
        assert_eq!(out.port, Some("limited"));
        let body = String::from_utf8(out.context.response.body.to_vec()).unwrap();
        assert!(body.contains("too many from 10.0.0.1"), "{body}");
    }

    #[tokio::test]
    async fn test_acquire_increments_rejects_at_limit_release_decrements() {
        // Shared resources → shared counter registry across the pair.
        let r = PluginResources::empty();
        let acquire = LimitConnPlugin::from_config(
            &cfg(&[
                ("phase", serde_json::json!("acquire")),
                ("conn", serde_json::json!(3)),
                ("key", serde_json::json!("$remote_addr")),
                ("rejected_code", serde_json::json!(503)),
            ]),
            &r,
        )
        .unwrap();
        let release = LimitConnPlugin::from_config(
            &cfg(&[
                ("phase", serde_json::json!("release")),
                ("conn", serde_json::json!(3)),
                ("key", serde_json::json!("$remote_addr")),
            ]),
            &r,
        )
        .unwrap();

        // conn=3, burst=0 → 3 concurrent allowed, the 4th rejected.
        for i in 0..3 {
            let out = acquire.execute(ctx()).await.unwrap();
            assert!(out.port.is_none(), "acquire #{} should be allowed", i + 1);
        }
        let rejected = acquire
            .execute(ctx())
            .await
            .expect("4th concurrent acquire completes with a `limited` outcome");
        assert_eq!(rejected.port, Some("limited"));
        assert_eq!(rejected.context.response.status_code, 503);

        // Release one slot → the next acquire succeeds again.
        release.execute(ctx()).await.unwrap();
        let out = acquire.execute(ctx()).await.unwrap();
        assert!(
            out.port.is_none(),
            "acquire should succeed after a release freed a slot"
        );
    }

    #[tokio::test]
    async fn test_release_floors_at_zero() {
        let r = PluginResources::empty();
        let release = LimitConnPlugin::from_config(
            &cfg(&[
                ("phase", serde_json::json!("release")),
                ("conn", serde_json::json!(5)),
                ("key", serde_json::json!("const")),
                ("key_type", serde_json::json!("constant")),
            ]),
            &r,
        )
        .unwrap();
        // Extra releases must not drive the counter negative.
        release.execute(ctx()).await.unwrap();
        release.execute(ctx()).await.unwrap();
        assert_eq!(r.traffic.conn.counter("const").load(Ordering::SeqCst), 0);
    }
}