Remove ttl and add main loop

1 files changed, 99 insertions, 76 deletions

diff --git a/src/main.rs b/src/main.rs
index a5996d4..675634d 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -17,6 +17,87 @@ use print::{print_legende_line, print_ping_line, print_stats};
 use state::State;
 use stats::compute_stats;
 
+fn main_loop(args : &Args, states: &mut Vec<State>, target: &IpAddr, tx: &mpsc::Sender<PingResult>, rx: &mut mpsc::Receiver<PingResult>, packet_id: &mut u32) {
+
+    let current_time = State::current_time_ms();
+    for state in states.iter_mut() {
+        if let State::NotYetReceived { time_sent } = state {
+            // TODO: 2000ms Should be a configurable variable
+            if current_time - *time_sent > 2000 {
+                // 2s = 2000ms
+                *state = State::Lost;
+            }
+        }
+    }
+
+    // Send the ping
+    states.push(State::NotYetReceived {
+        time_sent: current_time,
+    });
+    spawn_ping(*target, tx.clone(), *packet_id, args.ttl.unwrap_or(64));
+    *packet_id += 1;
+
+    // The update state loop
+    loop {
+        match rx.try_recv() {
+            // Update the state of the received packet
+            Ok(result) => {
+                if args.beep {
+                    print!("\x07"); // Rust doesn't support \a ?!
+                }
+                match result {
+                    PingResult::Ping { rtt, id } => {
+                        if (id as usize) < states.len() {
+                            states[id as usize] = State::Received { rtt };
+                        }
+                    }
+                    PingResult::NoPingErr { id } => {
+                        states[id as usize] = State::SendError;
+                    }
+                }
+            }
+            Err(TryRecvError::Disconnected) => {
+                eprintln!("Critical error, pinger thread killed");
+                std::process::exit(2);
+            }
+            Err(TryRecvError::Empty) => {
+                // No more ping to read
+                break;
+            }
+        }
+    }
+
+    // Put the cursor back up
+    if !std::process::Command::new("tput")
+        .arg("rc")
+            .status()
+            .map(|s| s.success())
+            .unwrap_or(false)
+    {
+        println!("Process failed too");
+        std::process::exit(1);
+    }
+
+    // Prints lines based on the state of each ping
+    print_ping_line(&states);
+
+    // Print stats
+    if args.stats {
+        if let Some(stats) = compute_stats(&states) {
+            print_stats(&stats);
+        }
+        // Print stats of the last 30 pings
+        let first_index = if (states.len() as i32) - 30 > 0 {
+            states.len() - 30
+        } else {
+            0
+        };
+        if let Some(stats) = compute_stats(&states[first_index..]) {
+            print_stats(&stats);
+        }
+    }
+}
+
 #[tokio::main]
 async fn main() {
     let args = Args::parse();
@@ -36,7 +117,14 @@ async fn main() {
     let mut states: Vec<State> = Vec::new();
     let mut packet_id: u32 = 0;
 
-    print_legende_line();
+    if args.legende {
+        print_legende_line();
+    }
+
+    if args.remainder {
+        // TODO: Not yet implemented
+        // print_remainder_line
+    }
 
     // Save the current cursor position
     if !std::process::Command::new("tput")
@@ -49,83 +137,18 @@ async fn main() {
         std::process::exit(1);
     }
 
-    // Should be quit with Ctrl + C
-    loop {
-        let current_time = State::current_time_ms();
-        for state in states.iter_mut() {
-            if let State::NotYetReceived { time_sent } = state {
-                // TODO: 2000ms Should be a configurable variable
-                if current_time - *time_sent > 2000 {
-                    // 2s = 2000ms
-                    *state = State::Lost;
-                }
-            }
-        }
-
-        // Send the ping
-        states.push(State::NotYetReceived {
-            time_sent: current_time,
-        });
-        spawn_ping(target, tx.clone(), packet_id);
-        packet_id += 1;
 
-        // The update state loop
-        loop {
-            match rx.try_recv() {
-                // Update the state of the received packet
-                Ok(result) => match result {
-                    PingResult::Ping { rtt, id } => {
-                        if (id as usize) < states.len() {
-                            states[id as usize] = State::Received { rtt };
-                        }
-                    }
-                    PingResult::NoPingErr { id } => {
-                        states[id as usize] = State::SendError;
-                    }
-                },
-                Err(TryRecvError::Disconnected) => {
-                    eprintln!("Critical error, pinger thread killed");
-                    std::process::exit(2);
-                }
-                Err(TryRecvError::Empty) => {
-                    // No more ping to read
-                    break;
-                }
-            }
-        }
-
-        // Put the cursor back up
-        if !std::process::Command::new("tput")
-            .arg("rc")
-            .status()
-            .map(|s| s.success())
-            .unwrap_or(false)
-        {
-            println!("Process failed too");
-            std::process::exit(1);
-        }
-
-        // Prints lines based on the state of each ping
-        print_ping_line(&states);
-
-        // Print stats
-        if args.stats {
-            if let Some(stats) = compute_stats(&states) {
-                print_stats(&stats);
-            }
-            // Print stats of the last 30 pings
-            let first_index = if (states.len() as i32) - 30 > 0 {
-                states.len() - 30
-            } else {
-                0
-            };
-            if let Some(stats) = compute_stats(&states[first_index..]) {
-                print_stats(&stats);
+    let wait_time = args.interval.unwrap_or(1000);
+    match args.count {
+        Some(c) => {
+            for _ in 1..c {
+                main_loop(&args, &mut states, &target, &tx, &mut rx, &mut packet_id);
+                sleep(tokio::time::Duration::from_millis(wait_time as u64)).await;
             }
+        },
+        None => loop {
+            main_loop(&args, &mut states, &target, &tx, &mut rx, &mut packet_id);
+            sleep(tokio::time::Duration::from_millis(wait_time as u64)).await;
         }
-
-        // Sleep for 1 second before the next ping
-        // FEAT: Should be configurable via CLI
-        sleep(tokio::time::Duration::from_millis(1000)).await;
     }
 }