Skip to main content

statsig_rust/specs_adapter/
statsig_bootstrap_specs_adapter.rs

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