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