statsig_rust/specs_adapter/
statsig_bootstrap_specs_adapter.rs

1use crate::specs_adapter::{SpecsAdapter, SpecsSource, SpecsUpdate, SpecsUpdateListener};
2use crate::statsig_err::StatsigErr;
3use crate::{log_e, StatsigRuntime};
4use async_trait::async_trait;
5use chrono::Utc;
6
7use std::sync::{Arc, RwLock};
8use std::time::Duration;
9pub struct StatsigBootstrapSpecsAdapter {
10    data: RwLock<String>,
11    listener: RwLock<Option<Arc<dyn SpecsUpdateListener>>>,
12}
13const TAG: &str = stringify!(StatsigBootstrapSpecsAdapter);
14
15impl StatsigBootstrapSpecsAdapter {
16    #[must_use]
17    pub fn new(data: String) -> Self {
18        Self {
19            data: RwLock::new(data),
20            listener: RwLock::new(None),
21        }
22    }
23
24    pub fn set_data(&self, data: String) -> Result<(), StatsigErr> {
25        match self.data.write() {
26            Ok(mut lock) => *lock = data.clone(),
27            Err(e) => return Err(StatsigErr::LockFailure(e.to_string())),
28        };
29
30        self.push_update()
31    }
32
33    fn push_update(&self) -> Result<(), StatsigErr> {
34        let data = match self.data.read() {
35            Ok(lock) => lock.clone(),
36            Err(e) => return Err(StatsigErr::LockFailure(e.to_string())),
37        };
38
39        match &self.listener.read() {
40            Ok(lock) => match lock.as_ref() {
41                Some(listener) => listener.did_receive_specs_update(SpecsUpdate {
42                    data: data.into_bytes(),
43                    source: SpecsSource::Bootstrap,
44                    received_at: Utc::now().timestamp_millis() as u64,
45                }),
46                None => Err(StatsigErr::UnstartedAdapter("Listener not set".to_string())),
47            },
48            Err(e) => Err(StatsigErr::LockFailure(e.to_string())),
49        }
50    }
51}
52
53#[async_trait]
54impl SpecsAdapter for StatsigBootstrapSpecsAdapter {
55    async fn start(
56        self: Arc<Self>,
57        _statsig_runtime: &Arc<StatsigRuntime>,
58    ) -> Result<(), StatsigErr> {
59        self.push_update()
60    }
61
62    fn initialize(&self, listener: Arc<dyn SpecsUpdateListener>) {
63        match self.listener.write() {
64            Ok(mut lock) => *lock = Some(listener),
65            Err(e) => {
66                log_e!(TAG, "Failed to acquire write lock on listener: {}", e);
67            }
68        }
69    }
70
71    async fn shutdown(
72        &self,
73        _timeout: Duration,
74        _statsig_runtime: &Arc<StatsigRuntime>,
75    ) -> Result<(), StatsigErr> {
76        Ok(())
77    }
78
79    async fn schedule_background_sync(
80        self: Arc<Self>,
81        _statsig_runtime: &Arc<StatsigRuntime>,
82    ) -> Result<(), StatsigErr> {
83        Ok(())
84    }
85
86    fn get_type_name(&self) -> String {
87        stringify!(StatsigBootstrapSpecsAdapter).to_string()
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use crate::SpecsInfo;
94
95    use super::*;
96    use std::sync::Arc;
97
98    struct TestListener {
99        received_update: RwLock<Option<SpecsUpdate>>,
100    }
101
102    impl TestListener {
103        fn new() -> Self {
104            Self {
105                received_update: RwLock::new(None),
106            }
107        }
108    }
109
110    #[async_trait]
111    impl SpecsUpdateListener for TestListener {
112        fn did_receive_specs_update(&self, update: SpecsUpdate) -> Result<(), StatsigErr> {
113            if let Ok(mut lock) = self.received_update.write() {
114                *lock = Some(update);
115            }
116            Ok(())
117        }
118
119        fn get_current_specs_info(&self) -> SpecsInfo {
120            SpecsInfo::empty()
121        }
122    }
123
124    #[tokio::test]
125    async fn test_manually_sync_specs() {
126        let test_data = serde_json::json!({
127            "feature_gates": {},
128            "dynamic_configs": {},
129            "layer_configs": {},
130        })
131        .to_string();
132
133        let adapter = Arc::new(StatsigBootstrapSpecsAdapter::new(test_data.clone()));
134        let listener = Arc::new(TestListener::new());
135
136        let statsig_rt = StatsigRuntime::get_runtime();
137        adapter.initialize(listener.clone());
138        adapter.clone().start(&statsig_rt).await.unwrap();
139
140        if let Ok(lock) = listener.clone().received_update.read() {
141            let update = lock.as_ref().unwrap();
142            assert_eq!(update.source, SpecsSource::Bootstrap);
143            assert_eq!(update.data, test_data.into_bytes());
144        }
145    }
146
147    #[tokio::test]
148    async fn test_set_data() {
149        let statsig_rt = StatsigRuntime::get_runtime();
150
151        let adapter = Arc::new(StatsigBootstrapSpecsAdapter::new(String::new()));
152
153        let listener = Arc::new(TestListener::new());
154        adapter.initialize(listener.clone());
155        adapter.clone().start(&statsig_rt).await.unwrap();
156
157        let test_data = "{\"some\": \"value\"}".to_string();
158        let result = adapter.set_data(test_data.clone());
159        assert!(result.is_ok());
160
161        if let Ok(lock) = listener.clone().received_update.read() {
162            let update = lock.as_ref().unwrap();
163            assert_eq!(update.source, SpecsSource::Bootstrap);
164            assert_eq!(update.data, test_data.into_bytes());
165        }
166    }
167}