Skip to main content

featherbit/plugins/native/
store_incr.rs

1//! `store-incr` — atomically increments a counter in a named store.
2//!
3//! The TTL is applied when the key is **created** and never refreshed. A
4//! counter that bounds retries must expire a fixed time after it first appears;
5//! refreshing on every increment means a client that keeps retrying keeps the
6//! counter alive and the bound never resets.
7
8use async_trait::async_trait;
9use std::collections::HashMap;
10use std::sync::Arc;
11
12use crate::context::Context;
13use crate::plugins::resources::PluginResources;
14use crate::plugins::util::store_kv::{self, StoreHandle};
15use crate::plugins::{Plugin, PluginResult};
16use crate::vars::template::Template;
17
18#[cfg(feature = "redis-store")]
19use crate::plugins::PluginOutput;
20
21/// INCRBY, then set the expiry only if the key does not already have one.
22/// `TTL` returns a negative value when the key has no expiry, so the guard also
23/// repairs a key that somehow lost one. One round trip, atomic.
24///
25/// Only ever loaded into a `redis::Script` on the redis-backed data path (see
26/// `StoreIncrPlugin::script`), so this is `#[cfg(feature = "redis-store")]`
27/// like the rest of that path -- without the feature there is no consumer.
28#[cfg(feature = "redis-store")]
29const INCR_SCRIPT: &str = r#"
30local n = redis.call('INCRBY', KEYS[1], ARGV[1])
31if tonumber(ARGV[2]) > 0 and (tonumber(ARGV[3]) == 1 or redis.call('TTL', KEYS[1]) < 0) then
32  redis.call('EXPIRE', KEYS[1], ARGV[2])
33end
34return n
35"#;
36
37pub struct StoreIncrPlugin {
38    store: StoreHandle,
39    key: Template,
40    by: i64,
41    ttl_seconds: Option<u64>,
42    refresh_ttl: bool,
43    name: String,
44    #[cfg(feature = "redis-store")]
45    script: redis::Script,
46}
47
48// `StoreHandle`, `Template` and `redis::Script` are all `Debug`, so this
49// could derive -- except `by`/`ttl_seconds`/`name` are only read inside the
50// `#[cfg(feature = "redis-store")]` `execute` body, which does not exist in
51// a headless (`--no-default-features`) build. A derived impl doesn't count
52// as a read for the dead-code pass, so deriving would leave those three
53// fields read nowhere there and fail `-D warnings`; `#[allow(dead_code)]` is
54// off the table. Reading them here, unconditionally, keeps the headless
55// build clean without the attribute.
56impl std::fmt::Debug for StoreIncrPlugin {
57    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
58        f.debug_struct("StoreIncrPlugin")
59            .field("store", &self.store)
60            .field("key", &self.key)
61            .field("by", &self.by)
62            .field("ttl_seconds", &self.ttl_seconds)
63            .field("refresh_ttl", &self.refresh_ttl)
64            .field("name", &self.name)
65            .finish()
66    }
67}
68
69/// Parses the optional `refresh_ttl` flag.
70///
71/// Defaults to `false`, which keeps the expiry pinned to the key's creation.
72/// That default is load-bearing: a retry bound whose expiry refreshed on every
73/// increment could be held open indefinitely by the very client it limits.
74/// Setting it to `true` is the deliberate opposite -- a sliding window of
75/// "N events within `ttl_seconds` of each other".
76fn parse_refresh_ttl(config: &HashMap<String, serde_json::Value>) -> Result<bool, String> {
77    match config.get("refresh_ttl") {
78        None | Some(serde_json::Value::Null) => Ok(false),
79        Some(v) => v
80            .as_bool()
81            .ok_or_else(|| "store-incr: 'refresh_ttl' must be a boolean".to_string()),
82    }
83}
84
85/// Parses the optional `by` amount, which defaults to 1 and may be negative.
86fn parse_by(config: &HashMap<String, serde_json::Value>) -> Result<i64, String> {
87    match config.get("by") {
88        None | Some(serde_json::Value::Null) => Ok(1),
89        Some(v) => v
90            .as_i64()
91            .ok_or_else(|| "store-incr: 'by' must be an integer".to_string()),
92    }
93}
94
95impl StoreIncrPlugin {
96    /// Accepted keys:
97    /// - `store` (string, required): a declared `stores:` entry.
98    /// - `key` (string, required, templated): the counter key.
99    /// - `by` (integer, default `1`): amount to add; may be negative.
100    /// - `ttl_seconds` (integer, optional): expiry applied **only when the key
101    ///   is created**.
102    /// - `name` (string, required): `context.message` key receiving the new value.
103    pub fn from_config(
104        config: &HashMap<String, serde_json::Value>,
105        resources: &Arc<PluginResources>,
106    ) -> Result<Self, String> {
107        let name = config
108            .get("name")
109            .and_then(|v| v.as_str())
110            .filter(|s| !s.is_empty())
111            .ok_or_else(|| {
112                "store-incr: 'name' is required (the context.message key receiving the new value)"
113                    .to_string()
114            })?
115            .to_string();
116        let key = store_kv::required_template(config, "key", "store-incr")?;
117        let by = parse_by(config)?;
118        let ttl_seconds = store_kv::optional_ttl(config, "store-incr")?;
119        let refresh_ttl = parse_refresh_ttl(config)?;
120        if refresh_ttl && ttl_seconds.is_none() {
121            return Err(
122                "store-incr: 'refresh_ttl' requires 'ttl_seconds'; there is no expiry to refresh"
123                    .to_string(),
124            );
125        }
126        Ok(Self {
127            store: store_kv::resolve(config, resources, "store-incr")?,
128            key,
129            by,
130            ttl_seconds,
131            refresh_ttl,
132            name,
133            #[cfg(feature = "redis-store")]
134            script: redis::Script::new(INCR_SCRIPT),
135        })
136    }
137}
138
139#[async_trait]
140impl Plugin for StoreIncrPlugin {
141    fn plugin_type(&self) -> &str {
142        "store-incr"
143    }
144
145    fn reads_response_body(&self) -> bool {
146        self.key.references_response_body()
147    }
148
149    #[cfg(feature = "redis-store")]
150    async fn execute(&self, mut ctx: Context) -> PluginResult {
151        let rendered = self.key.render(&ctx).to_string();
152        if rendered.is_empty() {
153            return Err(store_kv::key_invalid(
154                ctx,
155                "store-incr",
156                "INCRBY",
157                &self.store.name,
158            ));
159        }
160        let key = self.store.key_for(&rendered);
161
162        let mut conn = match self.store.conn().await {
163            Ok(c) => c,
164            Err(e) => {
165                return Err(store_kv::store_error(
166                    ctx,
167                    "store-incr",
168                    "INCRBY",
169                    &self.store.name,
170                    e,
171                ))
172            }
173        };
174
175        let n: i64 = match self
176            .script
177            .key(key.as_str())
178            .arg(self.by)
179            .arg(self.ttl_seconds.unwrap_or(0))
180            .arg(i64::from(self.refresh_ttl))
181            .invoke_async(&mut conn)
182            .await
183        {
184            Ok(n) => n,
185            Err(e) => {
186                // A counter key holding a non-numeric value, or one holding
187                // the wrong Redis type entirely (WRONGTYPE, e.g. a list), is
188                // a config/data problem, not an outage, and gets its own
189                // code. The rendered (un-namespaced) key is used in the
190                // message, matching `store-get`.
191                return Err(if store_kv::is_value_type_error(&e) {
192                    store_kv::value_invalid(
193                        ctx,
194                        "store-incr",
195                        "INCRBY",
196                        &self.store.name,
197                        format!("value at '{}' is not usable as an integer", rendered),
198                    )
199                } else {
200                    store_kv::store_error(
201                        ctx,
202                        "store-incr",
203                        "INCRBY",
204                        &self.store.name,
205                        e.to_string(),
206                    )
207                });
208            }
209        };
210
211        ctx.message
212            .insert(self.name.clone(), serde_json::Value::from(n));
213        Ok(PluginOutput::success(ctx))
214    }
215
216    #[cfg(not(feature = "redis-store"))]
217    async fn execute(&self, ctx: Context) -> PluginResult {
218        Err(store_kv::store_error(
219            ctx,
220            "store-incr",
221            "INCRBY",
222            &self.store.name,
223            "built without the redis-store feature".to_string(),
224        ))
225    }
226}
227
228#[cfg(test)]
229mod tests {
230    use super::*;
231    use crate::plugins::resources::PluginResources;
232    use std::collections::HashMap;
233
234    fn cfg(json: serde_json::Value) -> HashMap<String, serde_json::Value> {
235        serde_json::from_value(json).unwrap()
236    }
237
238    #[test]
239    fn test_requires_a_name() {
240        let r = PluginResources::empty();
241        let err =
242            StoreIncrPlugin::from_config(&cfg(serde_json::json!({ "store": "s", "key": "k" })), &r)
243                .unwrap_err();
244        assert!(err.contains("name"), "{err}");
245    }
246
247    #[test]
248    fn test_by_defaults_to_one_and_accepts_negatives() {
249        assert_eq!(parse_by(&cfg(serde_json::json!({}))).unwrap(), 1);
250        assert_eq!(parse_by(&cfg(serde_json::json!({ "by": -2 }))).unwrap(), -2);
251        assert!(parse_by(&cfg(serde_json::json!({ "by": "x" }))).is_err());
252    }
253
254    /// `refresh_ttl` only means something alongside an expiry; accepting it
255    /// silently would hide a config typo in exactly the throttling rules it
256    /// exists for.
257    #[test]
258    fn test_refresh_ttl_without_a_ttl_is_a_config_error() {
259        let r = PluginResources::empty();
260        let err = StoreIncrPlugin::from_config(
261            &cfg(serde_json::json!({
262                "store": "s", "key": "k", "name": "n", "refresh_ttl": true
263            })),
264            &r,
265        )
266        .unwrap_err();
267        assert!(err.contains("refresh_ttl"), "{err}");
268        assert!(err.contains("ttl_seconds"), "{err}");
269    }
270
271    /// The default must stay create-only, so a retry bound cannot silently
272    /// become refreshable -- a client that keeps retrying would then keep its
273    /// own counter alive and the bound would never reset.
274    #[test]
275    fn test_refresh_ttl_defaults_to_false_and_parses() {
276        assert!(!parse_refresh_ttl(&cfg(serde_json::json!({}))).unwrap());
277        assert!(parse_refresh_ttl(&cfg(serde_json::json!({ "refresh_ttl": true }))).unwrap());
278        assert!(parse_refresh_ttl(&cfg(serde_json::json!({ "refresh_ttl": "yes" }))).is_err());
279    }
280}