parent
a128fdcf5a
commit
dbc89831b1
|
|
@ -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(¶ms)
|
.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
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -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;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue