feat: add serve health heartbeat (#343)

* add serve health tracking

* fix lint

* fix
release-0.0
Meng Zhang 2023-08-09 16:08:42 +08:00 committed by GitHub
parent a128fdcf5a
commit dbc89831b1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 41 additions and 20 deletions

View File

@ -1,10 +1,8 @@
use std::{ use std::fs;
collections::HashMap,
fs::{self},
};
use lazy_static::lazy_static; use lazy_static::lazy_static;
use reqwest::Client; use reqwest::Client;
use serde::Serialize;
use uuid::Uuid; use uuid::Uuid;
use crate::path::usage_id_file; use crate::path::usage_id_file;
@ -34,12 +32,19 @@ impl UsageTracker {
Self { id, client } 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 { 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 client
.post(USAGE_API_ENDPOINT) .post(USAGE_API_ENDPOINT)
.json(&params) .json(&payload)
.send() .send()
.await .await
.ok(); .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! { lazy_static! {
static ref TRACKER: UsageTracker = UsageTracker::new(); static ref TRACKER: UsageTracker = UsageTracker::new();
} }
pub async fn capture(event: &str) { pub async fn capture<T>(event: &str, properties: T)
TRACKER.capture(event).await where
} T: Serialize,
{
#[cfg(test)] TRACKER.capture(event, properties).await
mod tests {
use super::capture;
#[tokio::test]
async fn it_fire_event() {
capture("UsageTest").await
}
} }

View File

@ -5,17 +5,20 @@ mod health;
use std::{ use std::{
net::{Ipv4Addr, SocketAddr}, net::{Ipv4Addr, SocketAddr},
sync::Arc, sync::Arc,
time::Duration,
}; };
use axum::{routing, Router, Server}; use axum::{routing, Router, Server};
use axum_tracing_opentelemetry::opentelemetry_tracing_layer; use axum_tracing_opentelemetry::opentelemetry_tracing_layer;
use clap::Args; use clap::Args;
use tabby_common::config::Config; use tabby_common::{config::Config, usage};
use tokio::time::sleep;
use tower_http::cors::CorsLayer; use tower_http::cors::CorsLayer;
use tracing::info; use tracing::info;
use utoipa::OpenApi; use utoipa::OpenApi;
use utoipa_swagger_ui::SwaggerUi; use utoipa_swagger_ui::SwaggerUi;
use self::health::HealthState;
use crate::fatal; use crate::fatal;
#[derive(OpenApi)] #[derive(OpenApi)]
@ -129,6 +132,8 @@ pub async fn main(config: &Config, args: &ServeArgs) {
let address = SocketAddr::from((Ipv4Addr::UNSPECIFIED, args.port)); let address = SocketAddr::from((Ipv4Addr::UNSPECIFIED, args.port));
info!("Listening at {}", address); info!("Listening at {}", address);
start_heartbeat(args);
Server::bind(&address) Server::bind(&address)
.serve(app.into_make_service()) .serve(app.into_make_service())
.await .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;
}
});
}