From dbc89831b1790df9e6806efe2e3915ee321c811a Mon Sep 17 00:00:00 2001
From: Meng Zhang <meng@tabbyml.com>
Date: Wed, 9 Aug 2023 16:08:42 +0800
Subject: [PATCH] feat: add serve health heartbeat (#343)

* add serve health tracking

* fix lint

* fix
---
 crates/tabby-common/src/usage.rs | 44 ++++++++++++++++++--------------
 crates/tabby/src/serve/mod.rs    | 17 +++++++++++-
 2 files changed, 41 insertions(+), 20 deletions(-)

diff --git a/crates/tabby-common/src/usage.rs b/crates/tabby-common/src/usage.rs
index ab74b03..1392cc0 100644
--- a/crates/tabby-common/src/usage.rs
+++ b/crates/tabby-common/src/usage.rs
@@ -1,10 +1,8 @@
-use std::{
-    collections::HashMap,
-    fs::{self},
-};
+use std::fs;
 
 use lazy_static::lazy_static;
 use reqwest::Client;
+use serde::Serialize;
 use uuid::Uuid;
 
 use crate::path::usage_id_file;
@@ -34,12 +32,19 @@ impl UsageTracker {
         Self { id, client }
     }
 
-    async fn capture(&self, event: &str) {
+    async fn capture<T>(&self, event: &str, properties: T)
+    where
+        T: Serialize,
+    {
         if let Some(client) = &self.client {
-            let params = HashMap::from([("distinctId", self.id.as_ref()), ("event", event)]);
+            let payload = Payload {
+                distinct_id: self.id.as_ref(),
+                event,
+                properties,
+            };
             client
                 .post(USAGE_API_ENDPOINT)
-                .json(&params)
+                .json(&payload)
                 .send()
                 .await
                 .ok();
@@ -47,20 +52,21 @@ impl UsageTracker {
     }
 }
 
+#[derive(Serialize)]
+#[serde(rename_all = "camelCase")]
+struct Payload<'a, T> {
+    distinct_id: &'a str,
+    event: &'a str,
+    properties: T,
+}
+
 lazy_static! {
     static ref TRACKER: UsageTracker = UsageTracker::new();
 }
 
-pub async fn capture(event: &str) {
-    TRACKER.capture(event).await
-}
-
-#[cfg(test)]
-mod tests {
-    use super::capture;
-
-    #[tokio::test]
-    async fn it_fire_event() {
-        capture("UsageTest").await
-    }
+pub async fn capture<T>(event: &str, properties: T)
+where
+    T: Serialize,
+{
+    TRACKER.capture(event, properties).await
 }
diff --git a/crates/tabby/src/serve/mod.rs b/crates/tabby/src/serve/mod.rs
index 03fff5c..43de7ce 100644
--- a/crates/tabby/src/serve/mod.rs
+++ b/crates/tabby/src/serve/mod.rs
@@ -5,17 +5,20 @@ mod health;
 use std::{
     net::{Ipv4Addr, SocketAddr},
     sync::Arc,
+    time::Duration,
 };
 
 use axum::{routing, Router, Server};
 use axum_tracing_opentelemetry::opentelemetry_tracing_layer;
 use clap::Args;
-use tabby_common::config::Config;
+use tabby_common::{config::Config, usage};
+use tokio::time::sleep;
 use tower_http::cors::CorsLayer;
 use tracing::info;
 use utoipa::OpenApi;
 use utoipa_swagger_ui::SwaggerUi;
 
+use self::health::HealthState;
 use crate::fatal;
 
 #[derive(OpenApi)]
@@ -129,6 +132,8 @@ pub async fn main(config: &Config, args: &ServeArgs) {
 
     let address = SocketAddr::from((Ipv4Addr::UNSPECIFIED, args.port));
     info!("Listening at {}", address);
+
+    start_heartbeat(args);
     Server::bind(&address)
         .serve(app.into_make_service())
         .await
@@ -172,3 +177,13 @@ fn valid_args(args: &ServeArgs) {
         }
     }
 }
+
+fn start_heartbeat(args: &ServeArgs) {
+    let state = HealthState::new(args);
+    tokio::spawn(async move {
+        loop {
+            usage::capture("ServeHealth", &state).await;
+            sleep(Duration::from_secs(60)).await;
+        }
+    });
+}