statsig_rust/user/
statsig_user_loggable.rs

1use std::collections::HashMap;
2use std::sync::{Arc, RwLock, Weak};
3
4use serde::{Deserialize, Serialize};
5use serde_json::{json, Value};
6
7use crate::{log_e, user::StatsigUserInternal};
8
9const TAG: &str = "StatsigUserLoggable";
10
11lazy_static::lazy_static! {
12    static ref LOGGABLE_USER_STORE: RwLock<HashMap<String, Weak<UserLoggableData>>> =
13    RwLock::new(HashMap::new());
14}
15
16#[derive(Serialize, Deserialize)]
17pub struct UserLoggableData {
18    pub key: String,
19    pub value: Value,
20}
21
22#[derive(Clone)]
23pub struct StatsigUserLoggable {
24    pub data: Arc<UserLoggableData>,
25}
26
27impl Serialize for StatsigUserLoggable {
28    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
29    where
30        S: serde::Serializer,
31    {
32        self.data.value.serialize(serializer)
33    }
34}
35
36impl<'de> Deserialize<'de> for StatsigUserLoggable {
37    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
38    where
39        D: serde::Deserializer<'de>,
40    {
41        let value = Value::deserialize(deserializer)?;
42        Ok(StatsigUserLoggable {
43            data: Arc::new(UserLoggableData {
44                key: "".to_string(),
45                value,
46            }),
47        })
48    }
49}
50
51fn make_loggable(
52    full_user_key: String,
53    user_internal: &StatsigUserInternal,
54) -> Arc<UserLoggableData> {
55    let result = Arc::new(UserLoggableData {
56        key: full_user_key.clone(),
57        value: json!(user_internal),
58    });
59
60    let mut store = match LOGGABLE_USER_STORE.write() {
61        Ok(store) => store,
62        Err(e) => {
63            log_e!(TAG, "Error locking user loggable store: {:?}", e);
64            return result;
65        }
66    };
67
68    store.insert(full_user_key, Arc::downgrade(&result));
69
70    result
71}
72
73impl StatsigUserLoggable {
74    pub fn new(user_internal: &StatsigUserInternal) -> Self {
75        let user_key = user_internal.get_full_user_key();
76
77        let mut existing = None;
78
79        match LOGGABLE_USER_STORE.read() {
80            Ok(store) => existing = store.get(&user_key).map(|x| x.upgrade()),
81            Err(e) => {
82                log_e!(TAG, "Error locking user loggable store: {:?}", e);
83            }
84        };
85
86        let data = match existing {
87            Some(Some(x)) => x,
88            _ => make_loggable(user_key, user_internal),
89        };
90
91        Self { data }
92    }
93
94    pub fn null_user() -> Self {
95        Self {
96            data: Arc::new(UserLoggableData {
97                key: "".to_string(),
98                value: Value::Null,
99            }),
100        }
101    }
102
103    pub fn create_sampling_key(&self) -> String {
104        let user_data = &self.data.value;
105        let user_id = user_data
106            .get("userID")
107            .map(|x| x.as_str())
108            .unwrap_or_default()
109            .unwrap_or_default();
110
111        // done this way for perf reasons
112        let mut user_key = String::from("u:");
113        user_key += user_id;
114        user_key += ";";
115
116        let custom_ids = user_data
117            .get("customIDs")
118            .map(|x| x.as_object())
119            .unwrap_or_default();
120
121        if let Some(custom_ids) = custom_ids {
122            for (key, val) in custom_ids.iter() {
123                if let Some(string_value) = &val.as_str() {
124                    user_key += key;
125                    user_key += ":";
126                    user_key += string_value;
127                    user_key += ";";
128                }
129            }
130        };
131
132        user_key
133    }
134}
135
136impl StatsigUserInternal<'_, '_> {
137    pub fn to_loggable(&self) -> StatsigUserLoggable {
138        StatsigUserLoggable::new(self)
139    }
140}
141
142impl Drop for StatsigUserLoggable {
143    fn drop(&mut self) {
144        let strong_count = match LOGGABLE_USER_STORE.read() {
145            Ok(store) => match store.get(&self.data.key) {
146                Some(weak_ref) => weak_ref.strong_count(),
147                None => return,
148            },
149            Err(e) => {
150                log_e!(TAG, "Error locking user loggable store: {:?}", e);
151                return;
152            }
153        };
154
155        if strong_count > 1 {
156            return;
157        }
158
159        match LOGGABLE_USER_STORE.write() {
160            Ok(mut store) => {
161                store.remove(&self.data.key);
162            }
163            Err(e) => {
164                log_e!(TAG, "Error locking user loggable store: {:?}", e);
165            }
166        };
167    }
168}
169
170#[cfg(test)]
171mod tests {
172    use std::time::Duration;
173
174    use futures::future::join_all;
175
176    use crate::StatsigUser;
177
178    use super::*;
179
180    #[tokio::test]
181    async fn test_creating_many_loggable_users_and_map_growth() {
182        let mut handles = vec![];
183
184        for _ in 0..10 {
185            let handle = tokio::spawn(async move {
186                for i in 0..1000 {
187                    let user_data = StatsigUser::with_user_id(format!("user{}", i));
188                    let user_internal = StatsigUserInternal::new(&user_data, None);
189                    let loggable = user_internal.to_loggable();
190                    tokio::time::sleep(Duration::from_micros(1)).await;
191                    let _ = loggable; // held across the sleep
192                }
193            });
194
195            handles.push(handle);
196        }
197
198        join_all(handles).await;
199
200        assert_eq!(LOGGABLE_USER_STORE.read().unwrap().len(), 0);
201    }
202}
statsig_rust/user/statsig_user_loggable.rs

statsig_rust/user/
statsig_user_loggable.rs
