amplitude_telemetry.rs

  1use crate::http::HttpClient;
  2use db::Db;
  3use gpui::{
  4    executor::Background,
  5    serde_json::{self, value::Map, Value},
  6    AppContext, Task,
  7};
  8use isahc::Request;
  9use lazy_static::lazy_static;
 10use parking_lot::Mutex;
 11use serde::Serialize;
 12use serde_json::json;
 13use std::{
 14    io::Write,
 15    mem,
 16    path::PathBuf,
 17    sync::Arc,
 18    time::{Duration, SystemTime, UNIX_EPOCH},
 19};
 20use tempfile::NamedTempFile;
 21use util::{post_inc, ResultExt, TryFutureExt};
 22use uuid::Uuid;
 23
 24pub struct AmplitudeTelemetry {
 25    http_client: Arc<dyn HttpClient>,
 26    executor: Arc<Background>,
 27    session_id: u128,
 28    state: Mutex<AmplitudeTelemetryState>,
 29}
 30
 31#[derive(Default)]
 32struct AmplitudeTelemetryState {
 33    metrics_id: Option<Arc<str>>,
 34    device_id: Option<Arc<str>>,
 35    app_version: Option<Arc<str>>,
 36    os_version: Option<Arc<str>>,
 37    os_name: &'static str,
 38    queue: Vec<AmplitudeEvent>,
 39    next_event_id: usize,
 40    flush_task: Option<Task<()>>,
 41    log_file: Option<NamedTempFile>,
 42}
 43
 44const AMPLITUDE_EVENTS_URL: &'static str = "https://api2.amplitude.com/batch";
 45
 46lazy_static! {
 47    static ref AMPLITUDE_API_KEY: Option<String> = std::env::var("ZED_AMPLITUDE_API_KEY")
 48        .ok()
 49        .or_else(|| option_env!("ZED_AMPLITUDE_API_KEY").map(|key| key.to_string()));
 50}
 51
 52#[derive(Serialize)]
 53struct AmplitudeEventBatch {
 54    api_key: &'static str,
 55    events: Vec<AmplitudeEvent>,
 56}
 57
 58#[derive(Serialize)]
 59struct AmplitudeEvent {
 60    #[serde(skip_serializing_if = "Option::is_none")]
 61    user_id: Option<Arc<str>>,
 62    device_id: Option<Arc<str>>,
 63    event_type: String,
 64    #[serde(skip_serializing_if = "Option::is_none")]
 65    event_properties: Option<Map<String, Value>>,
 66    #[serde(skip_serializing_if = "Option::is_none")]
 67    user_properties: Option<Map<String, Value>>,
 68    os_name: &'static str,
 69    os_version: Option<Arc<str>>,
 70    app_version: Option<Arc<str>>,
 71    #[serde(rename = "App")]
 72    app: &'static str,
 73    event_id: usize,
 74    session_id: u128,
 75    time: u128,
 76}
 77
 78#[cfg(debug_assertions)]
 79const MAX_QUEUE_LEN: usize = 1;
 80
 81#[cfg(not(debug_assertions))]
 82const MAX_QUEUE_LEN: usize = 10;
 83
 84#[cfg(debug_assertions)]
 85const DEBOUNCE_INTERVAL: Duration = Duration::from_secs(1);
 86
 87#[cfg(not(debug_assertions))]
 88const DEBOUNCE_INTERVAL: Duration = Duration::from_secs(30);
 89
 90impl AmplitudeTelemetry {
 91    pub fn new(client: Arc<dyn HttpClient>, cx: &AppContext) -> Arc<Self> {
 92        let platform = cx.platform();
 93        let this = Arc::new(Self {
 94            http_client: client,
 95            executor: cx.background().clone(),
 96            session_id: SystemTime::now()
 97                .duration_since(UNIX_EPOCH)
 98                .unwrap()
 99                .as_millis(),
100            state: Mutex::new(AmplitudeTelemetryState {
101                os_version: platform
102                    .os_version()
103                    .log_err()
104                    .map(|v| v.to_string().into()),
105                os_name: platform.os_name().into(),
106                app_version: platform
107                    .app_version()
108                    .log_err()
109                    .map(|v| v.to_string().into()),
110                device_id: None,
111                queue: Default::default(),
112                flush_task: Default::default(),
113                next_event_id: 0,
114                log_file: None,
115                metrics_id: None,
116            }),
117        });
118
119        if AMPLITUDE_API_KEY.is_some() {
120            this.executor
121                .spawn({
122                    let this = this.clone();
123                    async move {
124                        if let Some(tempfile) = NamedTempFile::new().log_err() {
125                            this.state.lock().log_file = Some(tempfile);
126                        }
127                    }
128                })
129                .detach();
130        }
131
132        this
133    }
134
135    pub fn log_file_path(&self) -> Option<PathBuf> {
136        Some(self.state.lock().log_file.as_ref()?.path().to_path_buf())
137    }
138
139    pub fn start(self: &Arc<Self>, db: Db) {
140        let this = self.clone();
141        self.executor
142            .spawn(
143                async move {
144                    let device_id = if let Ok(Some(device_id)) = db.read_kvp("device_id") {
145                        device_id
146                    } else {
147                        let device_id = Uuid::new_v4().to_string();
148                        db.write_kvp("device_id", &device_id)?;
149                        device_id
150                    };
151
152                    let device_id = Some(Arc::from(device_id));
153                    let mut state = this.state.lock();
154                    state.device_id = device_id.clone();
155                    for event in &mut state.queue {
156                        event.device_id = device_id.clone();
157                    }
158                    if !state.queue.is_empty() {
159                        drop(state);
160                        this.flush();
161                    }
162
163                    anyhow::Ok(())
164                }
165                .log_err(),
166            )
167            .detach();
168    }
169
170    pub fn set_authenticated_user_info(
171        self: &Arc<Self>,
172        metrics_id: Option<String>,
173        is_staff: bool,
174    ) {
175        let is_signed_in = metrics_id.is_some();
176        self.state.lock().metrics_id = metrics_id.map(|s| s.into());
177        if is_signed_in {
178            self.report_event_with_user_properties(
179                "$identify",
180                Default::default(),
181                json!({ "$set": { "staff": is_staff } }),
182            )
183        }
184    }
185
186    pub fn report_event(self: &Arc<Self>, kind: &str, properties: Value) {
187        self.report_event_with_user_properties(kind, properties, Default::default());
188    }
189
190    fn report_event_with_user_properties(
191        self: &Arc<Self>,
192        kind: &str,
193        properties: Value,
194        user_properties: Value,
195    ) {
196        if AMPLITUDE_API_KEY.is_none() {
197            return;
198        }
199
200        let mut state = self.state.lock();
201        let event = AmplitudeEvent {
202            event_type: kind.to_string(),
203            time: SystemTime::now()
204                .duration_since(UNIX_EPOCH)
205                .unwrap()
206                .as_millis(),
207            session_id: self.session_id,
208            event_properties: if let Value::Object(properties) = properties {
209                Some(properties)
210            } else {
211                None
212            },
213            user_properties: if let Value::Object(user_properties) = user_properties {
214                Some(user_properties)
215            } else {
216                None
217            },
218            user_id: state.metrics_id.clone(),
219            device_id: state.device_id.clone(),
220            os_name: state.os_name,
221            app: "Zed",
222            os_version: state.os_version.clone(),
223            app_version: state.app_version.clone(),
224            event_id: post_inc(&mut state.next_event_id),
225        };
226        state.queue.push(event);
227        if state.device_id.is_some() {
228            if state.queue.len() >= MAX_QUEUE_LEN {
229                drop(state);
230                self.flush();
231            } else {
232                let this = self.clone();
233                let executor = self.executor.clone();
234                state.flush_task = Some(self.executor.spawn(async move {
235                    executor.timer(DEBOUNCE_INTERVAL).await;
236                    this.flush();
237                }));
238            }
239        }
240    }
241
242    fn flush(self: &Arc<Self>) {
243        let mut state = self.state.lock();
244        let events = mem::take(&mut state.queue);
245        state.flush_task.take();
246        drop(state);
247
248        if let Some(api_key) = AMPLITUDE_API_KEY.as_ref() {
249            let this = self.clone();
250            self.executor
251                .spawn(
252                    async move {
253                        let mut json_bytes = Vec::new();
254
255                        if let Some(file) = &mut this.state.lock().log_file {
256                            let file = file.as_file_mut();
257                            for event in &events {
258                                json_bytes.clear();
259                                serde_json::to_writer(&mut json_bytes, event)?;
260                                file.write_all(&json_bytes)?;
261                                file.write(b"\n")?;
262                            }
263                        }
264
265                        let batch = AmplitudeEventBatch { api_key, events };
266                        json_bytes.clear();
267                        serde_json::to_writer(&mut json_bytes, &batch)?;
268                        let request =
269                            Request::post(AMPLITUDE_EVENTS_URL).body(json_bytes.into())?;
270                        this.http_client.send(request).await?;
271                        Ok(())
272                    }
273                    .log_err(),
274                )
275                .detach();
276        }
277    }
278}