Skip to main content

ff_script/functions/
quota.rs

1//! Typed FCALL wrapper for quota admission function (lua/quota.lua).
2
3use ff_core::contracts::*;
4use crate::error::ScriptError;
5use ff_core::keys::QuotaKeyContext;
6
7use crate::result::{FcallResult, FromFcallResult};
8
9/// Key context for quota admission check on {q:K}.
10pub struct QuotaOpKeys<'a> {
11    pub ctx: &'a QuotaKeyContext,
12    pub dimension: &'a str,
13    pub execution_id: &'a ff_core::types::ExecutionId,
14}
15
16// ─── ff_create_quota_policy ───────────────────────────────────────────
17//
18// Lua KEYS (5): quota_def, quota_window_zset, quota_concurrency_counter,
19//               admitted_set, quota_policies_index
20// Lua ARGV (5): quota_policy_id, window_seconds, max_requests_per_window,
21//               max_concurrent, now_ms
22
23ff_function! {
24    pub ff_create_quota_policy(args: CreateQuotaPolicyArgs) -> CreateQuotaPolicyResult {
25        keys(k: &QuotaOpKeys<'_>) {
26            k.ctx.definition(),
27            k.ctx.window(k.dimension),
28            k.ctx.concurrency(),
29            k.ctx.admitted_set(),
30            ff_core::keys::quota_policies_index(k.ctx.hash_tag()),
31        }
32        argv {
33            args.quota_policy_id.to_string(),
34            args.window_seconds.to_string(),
35            args.max_requests_per_window.to_string(),
36            args.max_concurrent.to_string(),
37            args.now.to_string(),
38        }
39    }
40}
41
42impl FromFcallResult for CreateQuotaPolicyResult {
43    fn from_fcall_result(raw: &ferriskey::Value) -> Result<Self, ScriptError> {
44        let r = FcallResult::parse(raw)?.into_success()?;
45        let id_str = r.field_str(0);
46        let qid = ff_core::types::QuotaPolicyId::parse(&id_str)
47            .map_err(|e| ScriptError::Parse {
48                fcall: "ff_create_quota_policy".into(),
49                execution_id: None,
50                message: format!("invalid quota_policy_id: {e}"),
51            })?;
52        match r.status.as_str() {
53            "OK" => Ok(CreateQuotaPolicyResult::Created { quota_policy_id: qid }),
54            "ALREADY_SATISFIED" => Ok(CreateQuotaPolicyResult::AlreadySatisfied { quota_policy_id: qid }),
55            _ => Err(ScriptError::Parse {
56                fcall: "ff_create_quota_policy".into(),
57                execution_id: None,
58                message: format!("unexpected status: {}", r.status),
59            }),
60        }
61    }
62}
63
64// ─── ff_check_admission_and_record ────────────────────────────────────
65//
66// Lua KEYS (5): window_zset, concurrency_counter, quota_def, admitted_guard,
67//               admitted_set
68// Lua ARGV (6): now_ms, window_seconds, rate_limit, concurrency_cap,
69//               execution_id, jitter_ms
70
71ff_function! {
72    pub ff_check_admission_and_record(args: CheckAdmissionArgs) -> CheckAdmissionResult {
73        keys(k: &QuotaOpKeys<'_>) {
74            k.ctx.window(k.dimension),
75            k.ctx.concurrency(),
76            k.ctx.definition(),
77            k.ctx.admitted(k.execution_id),
78            k.ctx.admitted_set(),
79        }
80        argv {
81            args.now.to_string(),
82            args.window_seconds.to_string(),
83            args.rate_limit.to_string(),
84            args.concurrency_cap.to_string(),
85            args.execution_id.to_string(),
86            args.jitter_ms.unwrap_or(0).to_string(),
87        }
88    }
89}
90
91impl FromFcallResult for CheckAdmissionResult {
92    fn from_fcall_result(raw: &ferriskey::Value) -> Result<Self, ScriptError> {
93        // Domain-specific return: {"ADMITTED"}, {"ALREADY_ADMITTED"},
94        // {"RATE_EXCEEDED", retry_after_ms}, {"CONCURRENCY_EXCEEDED"}
95        let arr = match raw {
96            ferriskey::Value::Array(arr) => arr,
97            _ => return Err(ScriptError::Parse {
98                fcall: "ff_check_admission".into(),
99                execution_id: None,
100                message: "expected Array".into(),
101            }),
102        };
103        let status = match arr.first() {
104            Some(Ok(ferriskey::Value::BulkString(b))) => String::from_utf8_lossy(b).into_owned(),
105            Some(Ok(ferriskey::Value::SimpleString(s))) => s.clone(),
106            _ => return Err(ScriptError::Parse {
107                fcall: "ff_check_admission".into(),
108                execution_id: None,
109                message: "expected status string".into(),
110            }),
111        };
112        match status.as_str() {
113            "ADMITTED" => Ok(CheckAdmissionResult::Admitted),
114            "ALREADY_ADMITTED" => Ok(CheckAdmissionResult::AlreadyAdmitted),
115            "RATE_EXCEEDED" => {
116                let retry_str = match arr.get(1) {
117                    Some(Ok(ferriskey::Value::BulkString(b))) => {
118                        String::from_utf8_lossy(b).into_owned()
119                    }
120                    Some(Ok(ferriskey::Value::Int(n))) => n.to_string(),
121                    _ => "0".to_string(),
122                };
123                let retry_after: u64 = retry_str.parse().unwrap_or(0);
124                Ok(CheckAdmissionResult::RateExceeded {
125                    retry_after_ms: retry_after,
126                })
127            }
128            "CONCURRENCY_EXCEEDED" => Ok(CheckAdmissionResult::ConcurrencyExceeded),
129            _ => Err(ScriptError::Parse {
130                fcall: "ff_check_admission".into(),
131                execution_id: None,
132                message: format!(
133                "unknown admission status: {status}"
134            ),
135            }),
136        }
137    }
138}
139
140// ─── ff_release_admission ────────────────────────────────────────────
141//
142// Lua KEYS (3): admitted_guard_key, admitted_set, concurrency_counter
143// Lua ARGV (1): execution_id
144
145ff_function! {
146    pub ff_release_admission(args: ReleaseAdmissionArgs) -> ReleaseAdmissionResult {
147        keys(k: &QuotaOpKeys<'_>) {
148            k.ctx.admitted(k.execution_id),
149            k.ctx.admitted_set(),
150            k.ctx.concurrency(),
151        }
152        argv {
153            args.execution_id.to_string(),
154        }
155    }
156}
157
158impl FromFcallResult for ReleaseAdmissionResult {
159    fn from_fcall_result(raw: &ferriskey::Value) -> Result<Self, ScriptError> {
160        let _r = FcallResult::parse(raw)?.into_success()?;
161        Ok(ReleaseAdmissionResult::Released)
162    }
163}