Skip to main content

pylon_plugin/builtin/
feature_flags.rs

1use std::collections::HashMap;
2use std::sync::Mutex;
3
4use crate::Plugin;
5use pylon_auth::AuthContext;
6
7/// A feature flag rule.
8#[derive(Debug, Clone)]
9pub enum FlagRule {
10    /// Always on or off.
11    Boolean(bool),
12    /// On for specific user IDs.
13    UserList(Vec<String>),
14    /// On for a percentage of users (0-100).
15    Percentage(u8),
16}
17
18/// A feature flag definition.
19#[derive(Debug, Clone)]
20pub struct FeatureFlag {
21    pub name: String,
22    pub description: String,
23    pub rule: FlagRule,
24    pub enabled: bool,
25}
26
27/// Feature flags plugin. Toggle features per user/percentage.
28pub struct FeatureFlagsPlugin {
29    flags: Mutex<HashMap<String, FeatureFlag>>,
30}
31
32impl FeatureFlagsPlugin {
33    pub fn new() -> Self {
34        Self {
35            flags: Mutex::new(HashMap::new()),
36        }
37    }
38
39    /// Define a flag that's globally on or off.
40    pub fn add_boolean(&self, name: &str, description: &str, enabled: bool) {
41        self.flags.lock().unwrap().insert(
42            name.to_string(),
43            FeatureFlag {
44                name: name.to_string(),
45                description: description.to_string(),
46                rule: FlagRule::Boolean(enabled),
47                enabled,
48            },
49        );
50    }
51
52    /// Define a flag that's on for specific users.
53    pub fn add_user_list(&self, name: &str, description: &str, users: Vec<String>) {
54        self.flags.lock().unwrap().insert(
55            name.to_string(),
56            FeatureFlag {
57                name: name.to_string(),
58                description: description.to_string(),
59                rule: FlagRule::UserList(users),
60                enabled: true,
61            },
62        );
63    }
64
65    /// Define a flag that's on for a percentage of users.
66    pub fn add_percentage(&self, name: &str, description: &str, percent: u8) {
67        self.flags.lock().unwrap().insert(
68            name.to_string(),
69            FeatureFlag {
70                name: name.to_string(),
71                description: description.to_string(),
72                rule: FlagRule::Percentage(percent.min(100)),
73                enabled: true,
74            },
75        );
76    }
77
78    /// Check if a flag is enabled for a given auth context.
79    pub fn is_enabled(&self, flag_name: &str, auth: &AuthContext) -> bool {
80        let flags = self.flags.lock().unwrap();
81        let flag = match flags.get(flag_name) {
82            Some(f) => f,
83            None => return false, // unknown flag = off
84        };
85
86        if !flag.enabled {
87            return false;
88        }
89
90        match &flag.rule {
91            FlagRule::Boolean(on) => *on,
92            FlagRule::UserList(users) => auth
93                .user_id
94                .as_ref()
95                .map(|id| users.contains(id))
96                .unwrap_or(false),
97            FlagRule::Percentage(pct) => {
98                let hash = auth
99                    .user_id
100                    .as_ref()
101                    .map(|id| {
102                        id.bytes()
103                            .fold(0u64, |acc, b| acc.wrapping_mul(31).wrapping_add(b as u64))
104                    })
105                    .unwrap_or(0);
106                (hash % 100) < (*pct as u64)
107            }
108        }
109    }
110
111    /// Toggle a flag on or off.
112    pub fn set_enabled(&self, flag_name: &str, enabled: bool) -> bool {
113        let mut flags = self.flags.lock().unwrap();
114        if let Some(flag) = flags.get_mut(flag_name) {
115            flag.enabled = enabled;
116            true
117        } else {
118            false
119        }
120    }
121
122    /// List all flags.
123    pub fn list_flags(&self) -> Vec<FeatureFlag> {
124        self.flags.lock().unwrap().values().cloned().collect()
125    }
126
127    /// Remove a flag.
128    pub fn remove(&self, flag_name: &str) -> bool {
129        self.flags.lock().unwrap().remove(flag_name).is_some()
130    }
131}
132
133impl Plugin for FeatureFlagsPlugin {
134    fn name(&self) -> &str {
135        "feature-flags"
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142
143    #[test]
144    fn boolean_flag() {
145        let plugin = FeatureFlagsPlugin::new();
146        plugin.add_boolean("dark-mode", "Enable dark mode", true);
147
148        assert!(plugin.is_enabled("dark-mode", &AuthContext::anonymous()));
149
150        plugin.set_enabled("dark-mode", false);
151        assert!(!plugin.is_enabled("dark-mode", &AuthContext::anonymous()));
152    }
153
154    #[test]
155    fn user_list_flag() {
156        let plugin = FeatureFlagsPlugin::new();
157        plugin.add_user_list(
158            "beta",
159            "Beta features",
160            vec!["user-1".into(), "user-2".into()],
161        );
162
163        assert!(plugin.is_enabled("beta", &AuthContext::authenticated("user-1".into())));
164        assert!(plugin.is_enabled("beta", &AuthContext::authenticated("user-2".into())));
165        assert!(!plugin.is_enabled("beta", &AuthContext::authenticated("user-3".into())));
166        assert!(!plugin.is_enabled("beta", &AuthContext::anonymous()));
167    }
168
169    #[test]
170    fn percentage_flag() {
171        let plugin = FeatureFlagsPlugin::new();
172        plugin.add_percentage("new-ui", "New UI experiment", 50);
173
174        // Check that it's deterministic for the same user.
175        let auth = AuthContext::authenticated("test-user".into());
176        let result1 = plugin.is_enabled("new-ui", &auth);
177        let result2 = plugin.is_enabled("new-ui", &auth);
178        assert_eq!(result1, result2);
179    }
180
181    #[test]
182    fn percentage_zero_always_off() {
183        let plugin = FeatureFlagsPlugin::new();
184        plugin.add_percentage("disabled", "Always off", 0);
185
186        assert!(!plugin.is_enabled("disabled", &AuthContext::authenticated("user-1".into())));
187    }
188
189    #[test]
190    fn percentage_100_always_on() {
191        let plugin = FeatureFlagsPlugin::new();
192        plugin.add_percentage("enabled", "Always on", 100);
193
194        assert!(plugin.is_enabled("enabled", &AuthContext::authenticated("user-1".into())));
195        assert!(plugin.is_enabled("enabled", &AuthContext::authenticated("user-2".into())));
196    }
197
198    #[test]
199    fn unknown_flag_returns_false() {
200        let plugin = FeatureFlagsPlugin::new();
201        assert!(!plugin.is_enabled("nonexistent", &AuthContext::anonymous()));
202    }
203
204    #[test]
205    fn remove_flag() {
206        let plugin = FeatureFlagsPlugin::new();
207        plugin.add_boolean("test", "Test", true);
208        assert!(plugin.remove("test"));
209        assert!(!plugin.is_enabled("test", &AuthContext::anonymous()));
210    }
211
212    #[test]
213    fn list_flags() {
214        let plugin = FeatureFlagsPlugin::new();
215        plugin.add_boolean("a", "Flag A", true);
216        plugin.add_boolean("b", "Flag B", false);
217
218        let flags = plugin.list_flags();
219        assert_eq!(flags.len(), 2);
220    }
221
222    #[test]
223    fn disabled_flag_ignores_rules() {
224        let plugin = FeatureFlagsPlugin::new();
225        plugin.add_user_list("beta", "Beta", vec!["user-1".into()]);
226        plugin.set_enabled("beta", false);
227
228        assert!(!plugin.is_enabled("beta", &AuthContext::authenticated("user-1".into())));
229    }
230}