feat(tracing): self-register module log tags

Author: BohuTANG

Cargo.lock

@@ -549,6 +549,7 @@ cargo_metadata = "0.19"
 fast-float2 = "0.2.3"
 gix = "0.71.0"
 indent = "0.1.1"
+inventory = "0.3.15"
 logos = "0.12.1"
 nom = "7.1.1"
 nom-rule = "0.4"

Cargo.toml

@@ -20,6 +20,7 @@ defer = { workspace = true }
 fastrace = { workspace = true }
 fastrace-opentelemetry = { workspace = true }
 itertools = { workspace = true }
+inventory = { workspace = true }
 jiff = { workspace = true }
 libc = { workspace = true }
 log = { workspace = true }

src/common/tracing/Cargo.toml

@@ -22,6 +22,7 @@ mod crash_hook;
 mod filter;
 mod init;
 mod loggers;
+pub mod module_tag;
 mod panic_hook;
 mod remote_log;
 mod structlog;

src/common/tracing/src/lib.rs

@@ -36,6 +36,8 @@ use logforth::Layout;
 use serde_json::Map;
 use serde_json::Value as JsonValue;
 
+use crate::module_tag::label_for_module;
+
 const PRINTER: DateTimePrinter = DateTimePrinter::new().precision(Some(6));
 
 pub fn format_timestamp(zdt: &Zoned) -> String {
@@ -124,53 +126,20 @@ impl<const FIXED_TIME: bool> Layout for TextLayout<FIXED_TIME> {
         }
 
         let level = record.level();
-        let module = map_module(record.module_path().unwrap_or(""));
+        let module = record.module_path().map(label_for_module).unwrap_or("");
         let log_file = Path::new(record.file().unwrap_or_default())
             .file_name()
             .and_then(|name| name.to_str())
             .unwrap_or_default();
         let line = record.line().unwrap_or(0);
         let msg = record.args();
-        write!(buf, " {level:>5} {module}: {log_file}:{line} {msg}",)?;
+        write!(buf, " {level:>5} {module}: {log_file}:{line} {msg}")?;
         record.key_values().visit(&mut KvWriter(&mut buf))?;
 
         Ok(buf)
     }
 }
 
-fn map_module(full_name: &str) -> &str {
-    if full_name.starts_with("databend_query") {
-        if full_name.starts_with("databend_query::sessions::query_ctx") {
-            return "[QUERY-CTX]";
-        }
-        if full_name.starts_with("databend_query::sessions::queue_mgr") {
-            return "[QUERY-QUEUE]";
-        }
-        if full_name.starts_with("databend_query::interpreters::interpreter") {
-            return "[INTERPRETER]";
-        }
-        if full_name.starts_with("databend_query::servers::http::v1::query") {
-            return "[HTTP-QUERY]";
-        }
-        if full_name.starts_with("databend_query::pipelines::builders") {
-            return "[BUILD-PIPELINES]";
-        }
-        if full_name.starts_with("databend_query::pipelines::executor::query_pipeline_executor") {
-            return "[PIPELINE-EXECUTOR]";
-        }
-        if full_name.starts_with("databend_query::pipelines::processors::transforms") {
-            return "[TRANSFORMS]";
-        }
-    }
-    if full_name.starts_with("databend_common_pipeline_transforms") {
-        return "[TRANSFORMS]";
-    }
-    if full_name.starts_with("databend_common_sql::planner::planner") {
-        return "[SQL-PLANNER]";
-    }
-    full_name
-}
-
 #[derive(Debug)]
 pub struct JsonLayout<const FIXED_TIME: bool = false>;
 

src/common/tracing/src/loggers.rs

@@ -0,0 +1,89 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/// A mapping between a module path prefix and a short log label.
+#[derive(Debug)]
+pub struct ModuleTag {
+    pub prefix: &'static str,
+    pub label: &'static str,
+}
+
+inventory::collect!(ModuleTag);
+
+/// Resolve the best matching label for the provided module path.
+pub fn label_for_module<'a>(full_name: &'a str) -> &'a str {
+    inventory::iter::<ModuleTag>
+        .into_iter()
+        .filter_map(|entry| {
+            full_name
+                .starts_with(entry.prefix)
+                .then_some((entry.prefix.len(), entry.label))
+        })
+        .max_by_key(|(len, _)| *len)
+        .map(|(_, label)| label)
+        .unwrap_or(full_name)
+}
+
+#[macro_export]
+macro_rules! register_module_tag {
+    ($label:expr $(,)?) => {
+        inventory::submit! {
+            #![crate = databend_common_tracing]
+
+            databend_common_tracing::module_tag::ModuleTag {
+                prefix: module_path!(),
+                label: $label,
+            }
+        }
+    };
+    ($label:expr, $prefix:expr $(,)?) => {
+        inventory::submit! {
+            #![crate = databend_common_tracing]
+
+            databend_common_tracing::module_tag::ModuleTag {
+                prefix: $prefix,
+                label: $label,
+            }
+        }
+    };
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    inventory::submit! {
+        ModuleTag {
+            prefix: "module::submodule",
+            label: "[SUB]",
+        }
+    }
+
+    inventory::submit! {
+        ModuleTag {
+            prefix: "module",
+            label: "[MODULE]",
+        }
+    }
+
+    #[test]
+    fn chooses_longest_prefix() {
+        assert_eq!(label_for_module("module::submodule::leaf"), "[SUB]");
+    }
+
+    #[test]
+    fn falls_back_to_full_module_name() {
+        assert_eq!(label_for_module("unknown::module"), "unknown::module");
+    }
+}

src/common/tracing/src/module_tag.rs

@@ -21,6 +21,7 @@ databend-common-functions = { workspace = true }
 databend-common-license = { workspace = true }
 databend-common-pipeline-core = { workspace = true }
 databend-common-sql = { workspace = true }
+databend-common-tracing = { workspace = true }
 databend-storages-common-cache = { workspace = true }
 enum-as-inner = { workspace = true }
 fastrace = { workspace = true }

src/query/pipeline/transforms/Cargo.toml

@@ -21,5 +21,7 @@
 #![feature(assert_matches)]
 #![feature(debug_closure_helpers)]
 
+databend_common_tracing::register_module_tag!("[TRANSFORMS]");
+
 pub mod processors;
 pub use processors::*;

src/query/pipeline/transforms/src/lib.rs

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+databend_common_tracing::register_module_tag!("[INTERPRETER]");
+
 use std::collections::BTreeMap;
 use std::sync::Arc;
 use std::time::SystemTime;

src/query/service/src/interpreters/interpreter.rs

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+databend_common_tracing::register_module_tag!("[BUILD-PIPELINES]");
+
 mod builder_aggregate;
 mod builder_append_table;