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                }),
60                None => Err(StatsigErr::UnstartedAdapter("Listener not set".to_string())),
61            },
62            None => Err(StatsigErr::LockFailure(
63                "Failed to acquire read lock on listener".to_string(),
64            )),
65        }
66    }
67}
68
69#[async_trait]
70impl SpecsAdapter for StatsigBootstrapSpecsAdapter {
71    async fn start(
72        self: Arc<Self>,
73        _statsig_runtime: &Arc<StatsigRuntime>,
74    ) -> Result<(), StatsigErr> {
75        self.push_update()
76    }
77
78    fn initialize(&self, listener: Arc<dyn SpecsUpdateListener>) {
79        match self
80            .listener
81            .try_write_for(std::time::Duration::from_secs(5))
82        {
83            Some(mut lock) => *lock = Some(listener),
84            None => {
85                log_e!(TAG, "Failed to acquire write lock on listener");
86            }
87        }
88    }
89
90    async fn shutdown(
91        &self,
92        _timeout: Duration,
93        _statsig_runtime: &Arc<StatsigRuntime>,
94    ) -> Result<(), StatsigErr> {
95        Ok(())
96    }
97
98    async fn schedule_background_sync(
99        self: Arc<Self>,
100        _statsig_runtime: &Arc<StatsigRuntime>,
101    ) -> Result<(), StatsigErr> {
102        Ok(())
103    }
104
105    fn get_type_name(&self) -> String {
106        stringify!(StatsigBootstrapSpecsAdapter).to_string()
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use crate::SpecsInfo;
113
114    use super::*;
115    use std::sync::Arc;
116
117    struct TestListener {
118        received_update: RwLock<Option<SpecsUpdate>>,
119    }
120
121    impl TestListener {
122        fn new() -> Self {
123            Self {
124                received_update: RwLock::new(None),
125            }
126        }
127    }
128
129    #[async_trait]
130    impl SpecsUpdateListener for TestListener {
131        fn did_receive_specs_update(&self, update: SpecsUpdate) -> Result<(), StatsigErr> {
132            if let Some(mut lock) = self.received_update.try_write() {
133                *lock = Some(update);
134            }
135            Ok(())
136        }
137
138        fn get_current_specs_info(&self) -> SpecsInfo {
139            SpecsInfo::empty()
140        }
141    }
142
143    #[tokio::test]
144    async fn test_manually_sync_specs() {
145        let test_data = serde_json::json!({
146            "feature_gates": {},
147            "dynamic_configs": {},
148            "layer_configs": {},
149        })
150        .to_string();
151
152        let adapter = Arc::new(StatsigBootstrapSpecsAdapter::new(test_data.clone()));
153        let listener = Arc::new(TestListener::new());
154
155        let statsig_rt = StatsigRuntime::get_runtime();
156        adapter.initialize(listener.clone());
157        adapter.clone().start(&statsig_rt).await.unwrap();
158
159        let mut update = listener
160            .received_update
161            .try_write()
162            .unwrap()
163            .take()
164            .unwrap();
165        assert_eq!(update.source, SpecsSource::Bootstrap);
166        assert_eq!(update.data.read_to_string().unwrap(), test_data);
167    }
168
169    #[tokio::test]
170    async fn test_set_data() {
171        let statsig_rt = StatsigRuntime::get_runtime();
172
173        let adapter = Arc::new(StatsigBootstrapSpecsAdapter::new(String::new()));
174
175        let listener = Arc::new(TestListener::new());
176        adapter.initialize(listener.clone());
177        adapter.clone().start(&statsig_rt).await.unwrap();
178
179        let test_data = "{\"some\": \"value\"}".to_string();
180        let result = adapter.set_data(test_data.clone());
181        assert!(result.is_ok());
182
183        let mut update = listener
184            .received_update
185            .try_write()
186            .unwrap()
187            .take()
188            .unwrap();
189        assert_eq!(update.source, SpecsSource::Bootstrap);
190        assert_eq!(update.data.read_to_string().unwrap(), test_data);
191    }
192}