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::{
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
}

View File

@ -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;
}
});
}