11use binaryninja:: headless:: Session ;
2+ use binaryninja:: tracing:: TracingLogListener ;
23use clap:: Parser ;
3- use indicatif:: { ProgressBar , ProgressStyle } ;
44use std:: collections:: HashMap ;
55use std:: path:: PathBuf ;
6+ use std:: sync:: atomic:: { AtomicBool , Ordering } ;
67use std:: sync:: Arc ;
78use std:: time:: Duration ;
9+ use tracing_indicatif:: span_ext:: IndicatifSpanExt ;
10+ use tracing_indicatif:: style:: ProgressStyle ;
11+ use tracing_indicatif:: IndicatifLayer ;
12+ use tracing_subscriber:: layer:: SubscriberExt ;
13+ use tracing_subscriber:: util:: SubscriberInitExt ;
814use warp_ninja:: processor:: {
915 CompressionTypeField , ProcessingFileState , ProcessingState , WarpFileProcessor ,
1016} ;
@@ -50,9 +56,15 @@ struct Args {
5056}
5157
5258fn main ( ) {
59+ let indicatif_layer = IndicatifLayer :: new ( ) ;
60+ tracing_subscriber:: registry ( )
61+ . with ( tracing_subscriber:: fmt:: layer ( ) . with_writer ( indicatif_layer. get_stderr_writer ( ) ) )
62+ . with ( indicatif_layer)
63+ . init ( ) ;
64+ let _listener = TracingLogListener :: new ( ) . register ( ) ;
65+
5366 let _session = Session :: new ( ) . expect ( "Failed to create session" ) ;
5467 let args = Args :: parse ( ) ;
55- env_logger:: init ( ) ;
5668
5769 let compression_ty = match args. compressed {
5870 true => CompressionTypeField :: Zstd ,
@@ -72,22 +84,27 @@ fn main() {
7284 } )
7385 . expect ( "Error setting Ctrl-C handler" ) ;
7486
75- let progress_guard = run_progress_bar ( processor. state ( ) ) ;
87+ let finished_signal = Arc :: new ( AtomicBool :: new ( false ) ) ;
88+ // Report progress to the terminal.
89+ let progress_state = processor. state ( ) ;
90+ let progress_finished_signal = finished_signal. clone ( ) ;
91+ let progress_thread =
92+ std:: thread:: spawn ( || run_progress_bar ( progress_state, progress_finished_signal) ) ;
7693
7794 let outputs: HashMap < PathBuf , WarpFile < ' static > > = args
7895 . input
7996 . into_iter ( )
8097 . filter_map ( |i| match processor. process ( i. clone ( ) ) {
8198 Ok ( o) => Some ( ( i, o) ) ,
8299 Err ( err) => {
83- log :: error!( "{}" , err) ;
100+ tracing :: error!( "{}" , err) ;
84101 None
85102 }
86103 } )
87104 . collect ( ) ;
88105
89- // Stop progress UI
90- progress_guard . finish ( ) ;
106+ finished_signal . store ( true , Ordering :: Relaxed ) ;
107+ progress_thread . join ( ) . expect ( "Progress thread exited" ) ;
91108
92109 match args. output . is_dir ( ) {
93110 true => {
@@ -98,85 +115,56 @@ fn main() {
98115 None => input. components ( ) . last ( ) . unwrap ( ) . as_os_str ( ) ,
99116 } ;
100117 let output_path = args. output . join ( output_name) . with_extension ( "warp" ) ;
101- log :: info!( "Writing to {:?}" , output_path) ;
118+ tracing :: info!( "Writing to {:?}" , output_path) ;
102119 std:: fs:: write ( output_path, output. to_bytes ( ) ) . unwrap ( ) ;
103120 }
104121 }
105122 false => {
106123 // Given a non-existing directory, merge all outputs and place at the output path.
107124 match processor. merge_files ( outputs. values ( ) . cloned ( ) . collect ( ) ) {
108125 Ok ( output) => {
109- log :: info!( "Writing to {:?}" , args. output) ;
126+ tracing :: info!( "Writing to {:?}" , args. output) ;
110127 std:: fs:: write ( args. output , output. to_bytes ( ) ) . unwrap ( ) ;
111128 }
112129 Err ( err) => {
113- log :: error!( "{}" , err) ;
130+ tracing :: error!( "{}" , err) ;
114131 }
115132 }
116133 }
117134 }
118135}
119136
120137// TODO: Also poll for background tasks and display them as independent progress bars.
121- fn run_progress_bar ( state : Arc < ProcessingState > ) -> ProgressGuard {
122- let pb = ProgressBar :: new ( 0 ) ;
123- pb. set_style (
124- ProgressStyle :: with_template (
125- "{spinner} {bar:40.cyan/blue} {pos}/{len} ({percent}%) [{elapsed_precise}] {msg}" ,
126- )
127- . unwrap ( )
128- . progress_chars ( "=>-" ) ,
129- ) ;
130-
131- let pb_clone = pb. clone ( ) ;
132- let state_clone = state. clone ( ) ;
133- let handle = std:: thread:: spawn ( move || loop {
134- let total = state_clone. total_files ( ) as u64 ;
135- let done = state_clone. files_with_state ( ProcessingFileState :: Processed ) as u64 ;
136- let unprocessed = state_clone. files_with_state ( ProcessingFileState :: Unprocessed ) ;
137- let analyzing = state_clone. files_with_state ( ProcessingFileState :: Analyzing ) ;
138- let processing = state_clone. files_with_state ( ProcessingFileState :: Processing ) ;
139-
140- if pb_clone. length ( ) . unwrap_or ( 0 ) != total {
141- pb_clone. set_length ( total) ;
142- }
143- pb_clone. set_position ( done. min ( total) ) ;
144- pb_clone. set_message ( format ! (
138+ fn run_progress_bar ( state : Arc < ProcessingState > , finished : Arc < AtomicBool > ) {
139+ let progress_span = tracing:: info_span!( "loading" ) ;
140+ let progress_style = ProgressStyle :: with_template ( "{msg} {elapsed} {wide_bar}" ) . unwrap ( ) ;
141+ progress_span. pb_set_style ( & progress_style) ;
142+ progress_span. pb_set_message ( "Processing files" ) ;
143+ progress_span. pb_set_finish_message ( "Files processed" ) ;
144+
145+ let elapsed = std:: time:: Instant :: now ( ) ;
146+ progress_span. in_scope ( || loop {
147+ let total = state. total_files ( ) as u64 ;
148+ let done = state. files_with_state ( ProcessingFileState :: Processed ) as u64 ;
149+ let unprocessed = state. files_with_state ( ProcessingFileState :: Unprocessed ) ;
150+ let analyzing = state. files_with_state ( ProcessingFileState :: Analyzing ) ;
151+ let processing = state. files_with_state ( ProcessingFileState :: Processing ) ;
152+
153+ progress_span. pb_set_length ( total) ;
154+ progress_span. pb_set_position ( done. min ( total) ) ;
155+ progress_span. pb_set_message ( & format ! (
145156 "{{u:{}|a:{}|p:{}|d:{}}}" ,
146157 unprocessed, analyzing, processing, done
147158 ) ) ;
148159
149- if pb_clone . is_finished ( ) {
160+ if state . is_cancelled ( ) || finished . load ( Ordering :: Relaxed ) {
150161 break ;
151162 }
152163 std:: thread:: sleep ( Duration :: from_millis ( 100 ) ) ;
153164 } ) ;
154165
155- ProgressGuard {
156- pb,
157- handle : Some ( handle) ,
158- }
159- }
160-
161- struct ProgressGuard {
162- pb : ProgressBar ,
163- handle : Option < std:: thread:: JoinHandle < ( ) > > ,
164- }
165-
166- impl ProgressGuard {
167- fn finish ( mut self ) {
168- self . pb . finish_and_clear ( ) ;
169- if let Some ( h) = self . handle . take ( ) {
170- let _ = h. join ( ) ;
171- }
172- }
173- }
174-
175- impl Drop for ProgressGuard {
176- fn drop ( & mut self ) {
177- self . pb . finish_and_clear ( ) ;
178- if let Some ( h) = self . handle . take ( ) {
179- let _ = h. join ( ) ;
180- }
181- }
166+ tracing:: info!(
167+ "Finished processing in {:.2} seconds" ,
168+ elapsed. elapsed( ) . as_secs_f32( )
169+ ) ;
182170}
0 commit comments