SQUASH-ME or DELETE-ME

bitcoin-dev-project · Oct 20, 2023 · acce1e1 · acce1e1
1 parent 255ea00
commit acce1e1
Showing 1 changed file with 19 additions and 19 deletions.
diff --git a/sim-lib/src/lib.rs b/sim-lib/src/lib.rs
@@ -893,49 +893,49 @@ async fn write_payment_results(
     print_batch_size: u32,
     no_results: bool,
 ) -> Result<(), SimulationError> {
-    let mut writer = if no_results {
-        WriterBuilder::new().from_path("/dev/null")?
-    } else {
-        WriterBuilder::new().from_path(format!(
+    let mut writer = if !no_results {
+        Some(WriterBuilder::new().from_path(format!(
             "simulation_{:?}.csv",
             SystemTime::now()
                 .duration_since(SystemTime::UNIX_EPOCH)
                 .unwrap()
                 .as_secs()
-        ))?
+        ))?)
+    } else {
+        None
     };
 
     let mut result_logger = PaymentResultLogger::new();
-
-    let mut counter = 1;
+    let mut counter = 0;
     loop {
         tokio::select! {
             biased;
             _ = listener.clone() => {
                 log::debug!("Simulation results consumer received shutdown signal.");
-                break writer.flush().map_err(|_| SimulationError::FileError)
+                return writer.map(|ref mut w| w.flush().map_err(|_| SimulationError::FileError)).unwrap_or(Ok(()));
+
             },
             payment_report = receiver.recv() => {
                 match payment_report {
                     Some((details, result)) => {
                         result_logger.report_result(&details, &result);
                         log::trace!("Resolved dispatched payment: {} with: {}.", details, result);
 
-                        writer.serialize((details, result)).map_err(|e| {
-                            let _ = writer.flush();
-                            SimulationError::CsvError(e)
-                        })?;
+                        if let Some(ref mut w) = writer {
+                            w.serialize((details, result)).map_err(|e| {
+                                let _ = w.flush();
+                                SimulationError::CsvError(e)
+                            })?;
 
-                        if print_batch_size == counter {
-                            writer.flush().map_err(|_| SimulationError::FileError)?;
-                            counter = 1;
-                        } else {
-                            counter += 1;
+                            counter = counter % print_batch_size + 1;
+                            if print_batch_size == counter {
+                                w.flush().map_err(|_| SimulationError::FileError)?;
+                            }
                         }
-                        continue;
                     },
                     None => {
-                        break writer.flush().map_err(|_| SimulationError::FileError)
+                        return writer.map(|ref mut w| w.flush().map_err(|_| SimulationError::FileError)).unwrap_or(Ok(()));
+
                     }
                 }
             }