From 8e4b6be991365cc22e8ae0e9c45c27c90ac69024 Mon Sep 17 00:00:00 2001
From: Nikaidou Haruki <nharuki@pm.me>
Date: Sat, 6 Sep 2025 19:07:28 +0900
Subject: [PATCH] add error tracking to UDP flood counter and implement
 signal-based run method

---
 ubw-sward/src/udp/counter.rs    |  3 +++
 ubw-sward/src/udp/integrated.rs | 20 +++++++++++++++++++-
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/ubw-sward/src/udp/counter.rs b/ubw-sward/src/udp/counter.rs
index fd61ee7..1678fed 100644
--- a/ubw-sward/src/udp/counter.rs
+++ b/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),
         }
     }
 }
diff --git a/ubw-sward/src/udp/integrated.rs b/ubw-sward/src/udp/integrated.rs
index 65ab1a7..2033b1c 100644
--- a/ubw-sward/src/udp/integrated.rs
+++ b/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()
     }