ff_script/functions/
quota.rs1use ff_core::contracts::*;
4use crate::error::ScriptError;
5use ff_core::keys::QuotaKeyContext;
6
7use crate::result::{FcallResult, FromFcallResult};
8
9pub struct QuotaOpKeys<'a> {
11 pub ctx: &'a QuotaKeyContext,
12 pub dimension: &'a str,
13 pub execution_id: &'a ff_core::types::ExecutionId,
14}
15
16ff_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
64ff_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 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
140ff_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}