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

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()
     }