add error tracking to UDP flood counter and implement signal-based run method

ubw-sward/src/udp/counter.rs

@@ -4,12 +4,14 @@ use std::sync::atomic::AtomicU64;
 pub struct UdpFloodCounter {
     pub sent_bytes: u64,
     pub sent_packets: u64,
+    pub error: u64,
 }
 
 #[derive(Default)]
 pub struct AtomicUdpFloodCounter {
     pub sent_bytes: AtomicU64,
     pub sent_packets: AtomicU64,
+    pub error: AtomicU64,
 }
 
 impl AtomicUdpFloodCounter {
@@ -17,6 +19,7 @@ impl AtomicUdpFloodCounter {
         UdpFloodCounter {
             sent_bytes: self.sent_bytes.load(std::sync::atomic::Ordering::Relaxed),
             sent_packets: self.sent_packets.load(std::sync::atomic::Ordering::Relaxed),
+            error: self.error.load(std::sync::atomic::Ordering::Relaxed),
         }
     }
 }

ubw-sward/src/udp/integrated.rs

@@ -4,7 +4,7 @@ use crate::udp::random_flood::UdpSward;
 use crate::udp::{BoxedUdpRequest, SizedUdpRequest};
 use rand::Rng;
 use std::net::SocketAddr;
-use std::sync::{atomic, Arc};
+use std::sync::{Arc, atomic};
 use tokio::net::UdpSocket;
 use tower::{Service, ServiceExt};
 use wyrand::WyRand;
@@ -55,6 +55,24 @@ impl IntegratedUdpSward {
             .call(content)
             .await
     }
+    pub async fn run_with_signal<const SIZE: usize>(
+        mut self,
+        mut signal: tokio::sync::oneshot::Receiver<()>,
+    ) -> Result<(), tokio::sync::oneshot::error::RecvError> {
+        loop {
+            tokio::select! {
+                receive_result = &mut signal => {
+                    receive_result?;
+                    return Ok(())
+                },
+                oneshot_result = self.oneshot_array::<SIZE>() => {
+                    if oneshot_result.is_err() {
+                        self.counter.error.fetch_add(1, atomic::Ordering::Relaxed);
+                    }
+                }
+            }
+        }
+    }
     pub fn get_counter(&self) -> Arc<AtomicUdpFloodCounter> {
         self.counter.clone()
     }