1- use tracing:: { debug, level_filters:: LevelFilter } ;
1+ use std:: {
2+ io:: Write ,
3+ sync:: atomic:: { AtomicBool , Ordering } ,
4+ } ;
5+
6+ use once_cell:: sync:: Lazy ;
7+ use tracing:: level_filters:: LevelFilter ;
28use tracing_subscriber:: {
3- fmt:: format:: { Format , Writer } ,
9+ fmt:: {
10+ format:: { Format , Writer } ,
11+ MakeWriter ,
12+ } ,
413 EnvFilter ,
514} ;
615
16+ static SHUTDOWN_IN_PROGRESS : Lazy < AtomicBool > = Lazy :: new ( || AtomicBool :: new ( false ) ) ;
17+
18+ struct ShutdownAwareWriter {
19+ buffer : std:: io:: BufWriter < std:: io:: Stdout > ,
20+ }
21+
22+ impl ShutdownAwareWriter {
23+ fn new ( ) -> Self {
24+ Self { buffer : std:: io:: BufWriter :: new ( std:: io:: stdout ( ) ) }
25+ }
26+ }
27+
28+ impl Write for ShutdownAwareWriter {
29+ fn write ( & mut self , buf : & [ u8 ] ) -> std:: io:: Result < usize > {
30+ if SHUTDOWN_IN_PROGRESS . load ( Ordering :: Relaxed ) {
31+ // During shutdown, write directly to stdout
32+ let stdout = std:: io:: stdout ( ) ;
33+ let mut handle = stdout. lock ( ) ;
34+ handle. write ( buf)
35+ } else {
36+ self . buffer . write ( buf)
37+ }
38+ }
39+
40+ fn flush ( & mut self ) -> std:: io:: Result < ( ) > {
41+ if SHUTDOWN_IN_PROGRESS . load ( Ordering :: Relaxed ) {
42+ let stdout = std:: io:: stdout ( ) ;
43+ let mut handle = stdout. lock ( ) ;
44+ handle. flush ( )
45+ } else {
46+ self . buffer . flush ( )
47+ }
48+ }
49+ }
50+
51+ struct ShutdownAwareWriterMaker ;
52+
53+ impl < ' a > MakeWriter < ' a > for ShutdownAwareWriterMaker {
54+ type Writer = ShutdownAwareWriter ;
55+
56+ fn make_writer ( & ' a self ) -> Self :: Writer {
57+ ShutdownAwareWriter :: new ( )
58+ }
59+ }
60+
761struct CustomTimer ;
862
963impl tracing_subscriber:: fmt:: time:: FormatTime for CustomTimer {
1064 fn format_time ( & self , writer : & mut Writer < ' _ > ) -> std:: fmt:: Result {
11- let now = chrono:: Local :: now ( ) ;
12- write ! ( writer, "{} - {}" , now. format( "%d %B" ) , now. format( "%H:%M:%S%.6f" ) )
65+ // Use a simpler time format during shutdown
66+ if SHUTDOWN_IN_PROGRESS . load ( Ordering :: Relaxed ) {
67+ let now = chrono:: Local :: now ( ) ;
68+ write ! ( writer, "{}" , now. format( "%H:%M:%S" ) )
69+ } else {
70+ let now = chrono:: Local :: now ( ) ;
71+ write ! ( writer, "{} - {}" , now. format( "%d %B" ) , now. format( "%H:%M:%S%.6f" ) )
72+ }
1373 }
1474}
1575
@@ -18,22 +78,32 @@ pub fn setup_logger(log_level: LevelFilter) {
1878
1979 let format = Format :: default ( ) . with_timer ( CustomTimer ) . with_level ( true ) . with_target ( false ) ;
2080
21- let subscriber =
22- tracing_subscriber:: fmt ( ) . with_env_filter ( filter) . event_format ( format) . finish ( ) ;
81+ let subscriber = tracing_subscriber:: fmt ( )
82+ . with_writer ( ShutdownAwareWriterMaker )
83+ . with_env_filter ( filter)
84+ . event_format ( format)
85+ . finish ( ) ;
2386
2487 if tracing:: subscriber:: set_global_default ( subscriber) . is_err ( ) {
25- debug ! ( "Logger has already been set up, continuing..." ) ;
88+ // Use println! here since logging might not be set up yet
89+ println ! ( "Logger has already been set up, continuing..." ) ;
2690 }
2791}
2892
2993pub fn setup_info_logger ( ) {
3094 setup_logger ( LevelFilter :: INFO ) ;
3195}
3296
33- // pub fn set_no_op_logger() -> DefaultGuard {
34- // let no_op_subscriber = FmtSubscriber::builder().with_writer(|| NullWriter).finish();
35- //
36- // let no_op_dispatch = Dispatch::new(no_op_subscriber);
37- //
38- // tracing::dispatcher::set_default(&no_op_dispatch)
39- // }
97+ // Call this when starting shutdown
98+ pub fn mark_shutdown_started ( ) {
99+ SHUTDOWN_IN_PROGRESS . store ( true , Ordering :: Relaxed ) ;
100+ }
101+
102+ // Optional guard for temporary logger suppression
103+ pub struct LoggerGuard ;
104+
105+ impl Drop for LoggerGuard {
106+ fn drop ( & mut self ) {
107+ SHUTDOWN_IN_PROGRESS . store ( false , Ordering :: Relaxed ) ;
108+ }
109+ }
0 commit comments