rocketmq_common/
log.rs

1/*
2 * Licensed to the Apache Software Foundation (ASF) under one or more
3 * contributor license agreements.  See the NOTICE file distributed with
4 * this work for additional information regarding copyright ownership.
5 * The ASF licenses this file to You under the Apache License, Version 2.0
6 * (the "License"); you may not use this file except in compliance with
7 * the License.  You may obtain a copy of the License at
8 *
9 *     http://www.apache.org/licenses/LICENSE-2.0
10 *
11 * Unless required by applicable law or agreed to in writing, software
12 * distributed under the License is distributed on an "AS IS" BASIS,
13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 * See the License for the specific language governing permissions and
15 * limitations under the License.
16 */
17use std::fmt;
18use std::path::Path;
19use std::str::FromStr;
20use std::sync::OnceLock;
21
22use rocketmq_error::RocketMQError;
23use rocketmq_error::RocketMQResult;
24use tracing_appender::non_blocking::WorkerGuard;
25use tracing_appender::rolling;
26use tracing_subscriber::fmt::writer::MakeWriterExt;
27use tracing_subscriber::layer::SubscriberExt;
28use tracing_subscriber::util::SubscriberInitExt;
29use tracing_subscriber::Layer;
30
31/// Static storage for the worker guard to prevent premature log flushing.
32/// This ensures logs are properly written before the program exits.
33static WORKER_GUARD: OnceLock<WorkerGuard> = OnceLock::new();
34
35/// Initializes the logger with the specified configuration.
36///
37/// This function sets up the logger using the `tracing_subscriber` crate.
38/// It reads the log level from the `RUST_LOG` environment variable, defaulting to "INFO" if not
39/// set. The logger is configured to include thread names, log levels, line numbers, and thread IDs
40/// in the log output.
41///
42/// # Returns
43///
44/// Returns `Ok(())` on success, or a `RocketMQError` if initialization fails.
45pub fn init_logger() -> RocketMQResult<()> {
46    let info_level = std::env::var("RUST_LOG").unwrap_or_else(|_| String::from("INFO"));
47    let level = tracing::Level::from_str(&info_level).map_err(|_| {
48        RocketMQError::illegal_argument(format!("Invalid log level: {}", info_level))
49    })?;
50
51    tracing_subscriber::fmt()
52        .with_thread_names(true)
53        .with_level(true)
54        .with_line_number(true)
55        .with_thread_ids(true)
56        .with_max_level(level)
57        .try_init()
58        .map_err(|e| RocketMQError::Internal(format!("Failed to initialize logger: {}", e)))?;
59
60    Ok(())
61}
62
63/// Initializes the logger with the specified log level.
64///
65/// This function sets up the logger using the `tracing_subscriber` crate.
66/// It configures the logger to include thread names, log levels, line numbers, and thread IDs
67/// in the log output. The maximum log level is set based on the provided `level`.
68///
69/// # Arguments
70///
71/// * `level` - A `Level` representing the desired log level.
72///
73/// # Returns
74///
75/// Returns `Ok(())` on success, or a `RocketMQError` if initialization fails.
76pub fn init_logger_with_level(level: Level) -> RocketMQResult<()> {
77    let tracing_level = level.to_tracing_level()?;
78
79    tracing_subscriber::fmt()
80        .with_thread_names(true)
81        .with_level(true)
82        .with_line_number(true)
83        .with_thread_ids(true)
84        .with_max_level(tracing_level)
85        .try_init()
86        .map_err(|e| RocketMQError::Internal(format!("Failed to initialize logger: {}", e)))?;
87
88    Ok(())
89}
90
91/// Initializes the logger with both file and console output.
92///
93/// This function configures logging to output to both a file and the console simultaneously.
94/// The file logger uses daily rotation to create new log files each day.
95///
96/// # Arguments
97///
98/// * `level` - The logging level to use for both outputs.
99/// * `directory` - The directory where log files will be stored.
100/// * `file_name_prefix` - The prefix to use for log filenames.
101///
102/// # Returns
103///
104/// Returns `Ok(())` on success, or a `RocketMQError` if initialization fails.
105///
106/// # Notes
107///
108/// This function creates a non-blocking file writer to improve performance.
109/// The worker guard is stored in a static `OnceLock` to prevent premature flushing.
110/// This ensures logs are properly written before the program exits.
111pub fn init_logger_with_file(
112    level: Level,
113    directory: impl AsRef<Path>,
114    file_name_prefix: impl AsRef<Path>,
115) -> RocketMQResult<()> {
116    // Convert custom Level to tracing LevelFilter
117    let max_level = tracing_subscriber::filter::LevelFilter::from_str(level.as_str())
118        .map_err(|_| RocketMQError::illegal_argument(format!("Invalid log level: {}", level)))?;
119
120    // log file output (daily rolling)
121    let file_appender = rolling::daily(directory, file_name_prefix);
122    let (file_writer, guard) = tracing_appender::non_blocking(file_appender);
123
124    // console layer (colorful output)
125    let console_layer = tracing_subscriber::fmt::layer()
126        .with_writer(std::io::stdout)
127        .with_thread_names(true)
128        .with_thread_ids(true)
129        .with_line_number(true)
130        .with_level(true)
131        .with_ansi(true)
132        .with_filter(max_level);
133
134    // file layer (non-color output)
135    let file_layer = tracing_subscriber::fmt::layer()
136        .with_writer(file_writer)
137        .with_thread_names(true)
138        .with_thread_ids(true)
139        .with_line_number(true)
140        .with_level(true)
141        .with_ansi(false)
142        .with_filter(max_level);
143
144    // Register two layers
145    tracing_subscriber::registry()
146        .with(console_layer)
147        .with(file_layer)
148        .try_init()
149        .map_err(|e| RocketMQError::Internal(format!("Failed to initialize logger: {}", e)))?;
150
151    // Store the guard in a static variable to prevent it from being dropped
152    // Using OnceLock instead of Box::leak to avoid memory leak
153    WORKER_GUARD
154        .set(guard)
155        .map_err(|_| RocketMQError::Internal("Logger already initialized".to_string()))?;
156
157    Ok(())
158}
159
160/// Custom log level type that wraps a static string.
161///
162/// This type provides a type-safe way to represent log levels while maintaining
163/// compatibility with the tracing crate's log levels.
164#[derive(Clone, Debug, PartialEq, Eq, Hash)]
165pub struct Level(&'static str);
166
167impl Level {
168    /// Constant representing the ERROR log level.
169    pub const ERROR: Level = Level("ERROR");
170
171    /// Constant representing the WARN log level.
172    pub const WARN: Level = Level("WARN");
173
174    /// Constant representing the INFO log level.
175    pub const INFO: Level = Level("INFO");
176
177    /// Constant representing the DEBUG log level.
178    pub const DEBUG: Level = Level("DEBUG");
179
180    /// Constant representing the TRACE log level.
181    pub const TRACE: Level = Level("TRACE");
182
183    /// Returns the string representation of the log level.
184    #[inline]
185    pub fn as_str(&self) -> &'static str {
186        self.0
187    }
188
189    /// Converts this Level to a tracing::Level.
190    ///
191    /// # Returns
192    ///
193    /// Returns `Ok(tracing::Level)` if the conversion is successful,
194    /// or a `RocketMQError` if the level string is invalid.
195    pub fn to_tracing_level(&self) -> RocketMQResult<tracing::Level> {
196        tracing::Level::from_str(self.0)
197            .map_err(|_| RocketMQError::illegal_argument(format!("Invalid log level: {}", self.0)))
198    }
199}
200
201impl From<&'static str> for Level {
202    fn from(level: &'static str) -> Self {
203        match level {
204            "ERROR" | "WARN" | "INFO" | "DEBUG" | "TRACE" => Level(level),
205            _ => panic!("Invalid log level: {level}"),
206        }
207    }
208}
209
210impl FromStr for Level {
211    type Err = String;
212
213    fn from_str(s: &str) -> Result<Self, Self::Err> {
214        match s {
215            "ERROR" => Ok(Level::ERROR),
216            "WARN" => Ok(Level::WARN),
217            "INFO" => Ok(Level::INFO),
218            "DEBUG" => Ok(Level::DEBUG),
219            "TRACE" => Ok(Level::TRACE),
220            _ => Err(format!("Invalid log level: {s}")),
221        }
222    }
223}
224
225impl fmt::Display for Level {
226    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
227        f.pad(self.0)
228    }
229}
230
231#[cfg(test)]
232mod tests {
233    use super::*;
234
235    #[test]
236    fn level_as_str_returns_correct_value() {
237        assert_eq!(Level::ERROR.as_str(), "ERROR");
238        assert_eq!(Level::WARN.as_str(), "WARN");
239        assert_eq!(Level::INFO.as_str(), "INFO");
240        assert_eq!(Level::DEBUG.as_str(), "DEBUG");
241        assert_eq!(Level::TRACE.as_str(), "TRACE");
242    }
243
244    #[test]
245    fn level_from_str_creates_correct_level() {
246        assert_eq!(Level::from("ERROR"), Level::ERROR);
247        assert_eq!(Level::from("WARN"), Level::WARN);
248        assert_eq!(Level::from("INFO"), Level::INFO);
249        assert_eq!(Level::from("DEBUG"), Level::DEBUG);
250        assert_eq!(Level::from("TRACE"), Level::TRACE);
251    }
252
253    #[test]
254    fn level_from_str_trait_works() {
255        assert_eq!("ERROR".parse::<Level>().unwrap(), Level::ERROR);
256        assert_eq!("INFO".parse::<Level>().unwrap(), Level::INFO);
257        assert!("invalid".parse::<Level>().is_err());
258    }
259
260    #[test]
261    fn level_display_formats_correctly() {
262        assert_eq!(format!("{}", Level::ERROR), "ERROR");
263        assert_eq!(format!("{}", Level::WARN), "WARN");
264        assert_eq!(format!("{}", Level::INFO), "INFO");
265        assert_eq!(format!("{}", Level::DEBUG), "DEBUG");
266        assert_eq!(format!("{}", Level::TRACE), "TRACE");
267    }
268
269    #[test]
270    fn level_to_tracing_level_success() {
271        assert!(Level::ERROR.to_tracing_level().is_ok());
272        assert!(Level::WARN.to_tracing_level().is_ok());
273        assert!(Level::INFO.to_tracing_level().is_ok());
274        assert!(Level::DEBUG.to_tracing_level().is_ok());
275        assert!(Level::TRACE.to_tracing_level().is_ok());
276    }
277
278    #[test]
279    fn level_clone_works() {
280        let level = Level::INFO;
281        let cloned = level.clone();
282        assert_eq!(level, cloned);
283    }
284
285    #[test]
286    fn level_hash_works() {
287        use std::collections::HashSet;
288        let mut set = HashSet::new();
289        set.insert(Level::ERROR);
290        set.insert(Level::INFO);
291        assert!(set.contains(&Level::ERROR));
292        assert!(!set.contains(&Level::DEBUG));
293    }
294}
rocketmq_common/log.rs

rocketmq_common/
log.rs
