diff --git a/init.sql b/init.sql index 4a30e6a..3b04f1b 100644 --- a/init.sql +++ b/init.sql @@ -23,4 +23,7 @@ CREATE INDEX idx_logs_service_time ON app_logs(service_name, created_at); -- create_next_month.sql -- 这里的日期逻辑需要通过外部脚本(Python/Shell)动态生成 SQL 语句 CREATE TABLE IF NOT EXISTS app_logs_2026_01 PARTITION OF app_logs - FOR VALUES FROM ('2026-01-01') TO ('2026-02-01'); \ No newline at end of file + FOR VALUES FROM ('2026-01-01') TO ('2026-02-01'); + +ALTER TABLE app_logs ADD COLUMN trace_id VARCHAR(64); +CREATE INDEX idx_logs_trace_id ON app_logs(trace_id); -- 加上索引,查问题全靠它 \ No newline at end of file diff --git a/src/context.rs b/src/context.rs new file mode 100644 index 0000000..b080c13 --- /dev/null +++ b/src/context.rs @@ -0,0 +1,16 @@ +// src/context.rs +use tokio::task_local; + +// 定义一个 Tokio 任务局部变量 +task_local! { + pub static TRACE_ID: String; +} + +// 辅助函数:设置 TraceID 作用域 +// 在这个 scope 闭包内运行的所有代码,都能访问到这个 ID +pub async fn with_trace_id(id: String, f: F) -> R +where + F: std::future::Future, +{ + TRACE_ID.scope(id, f).await +} diff --git a/src/core.rs b/src/core.rs index 631b876..56ac9f3 100644 --- a/src/core.rs +++ b/src/core.rs @@ -1,4 +1,5 @@ // src/core.rs +use crate::context::TRACE_ID; use crate::model::{LogLevel, LogRecord}; use crate::outputs::LogOutput; use chrono::Utc; @@ -56,6 +57,8 @@ impl Logger { return; } + let trace_id = TRACE_ID.try_with(|id| id.clone()).ok(); + // 2. 构造记录 let record = LogRecord { service_name: self.service_name.clone(), @@ -63,6 +66,7 @@ impl Logger { level, message, module, + trace_id, }; // 3. 发送 (非阻塞,如果队列满则丢弃,防止拖死业务) diff --git a/src/lib.rs b/src/lib.rs index 29f77f7..5af9021 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,6 +4,7 @@ pub mod core; pub mod model; pub mod outputs; pub use cleaner::LogCleaner; +pub mod context; use crate::core::Logger; use once_cell::sync::OnceCell; diff --git a/src/model.rs b/src/model.rs index dc5ef2a..ce3b371 100644 --- a/src/model.rs +++ b/src/model.rs @@ -34,4 +34,5 @@ pub struct LogRecord { pub level: LogLevel, pub message: String, pub module: String, + pub trace_id: Option, } diff --git a/src/outputs/file.rs b/src/outputs/file.rs index 7688a2f..37be2af 100644 --- a/src/outputs/file.rs +++ b/src/outputs/file.rs @@ -110,10 +110,16 @@ impl LogOutput for FileOutput { return; } + let trace_part = match &record.trace_id { + Some(id) => format!("[{}] ", id), + None => "".to_string(), + }; + let log_line = format!( - "{} [{}] - {}\n", + "{} [{}] {}- {}\n", record.timestamp.format("%Y-%m-%d %H:%M:%S%.3f"), record.level, + trace_part, record.message ); diff --git a/src/outputs/postgres.rs b/src/outputs/postgres.rs index 07e92c0..1ec9b7c 100644 --- a/src/outputs/postgres.rs +++ b/src/outputs/postgres.rs @@ -85,8 +85,8 @@ impl LogOutput for PostgresOutput { self.ensure_partition_exists(&record.timestamp).await; let query = r#" - INSERT INTO app_logs (service_name, log_level, message, module, created_at) - VALUES ($1, $2, $3, $4, $5) + INSERT INTO app_logs (service_name, log_level, message, module, created_at, trace_id) + VALUES ($1, $2, $3, $4, $5, $6) "#; // 注意:这里的 write 是在后台任务中执行的,就算慢也不会阻塞主业务 @@ -98,6 +98,7 @@ impl LogOutput for PostgresOutput { .bind(&record.message) .bind(&record.module) .bind(record.timestamp) + .bind(&record.trace_id) .execute(&self.pool) .await {