From 19277831c822f232b43a4e6aaf60c857843d065f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Joakim=20Frosteg=C3=A5rd?= <joakim.frostegard@gmail.com>
Date: Tue, 7 Apr 2020 13:22:28 +0200
Subject: [PATCH] aquatic: count received requests / sent responses, print on
 interval

---
 TODO.md                    |  3 +--
 aquatic/src/lib/common.rs  | 11 +++++++++++
 aquatic/src/lib/config.rs  |  2 ++
 aquatic/src/lib/lib.rs     | 24 ++++++++++++++++++++++++
 aquatic/src/lib/network.rs | 14 ++++++++++++++
 5 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/TODO.md b/TODO.md
index cd38c12..42d059d 100644
--- a/TODO.md
+++ b/TODO.md
@@ -15,8 +15,7 @@
 * Performance
     * cpu-target=native good?
     * mialloc good?
-* Track received requests and sent responses, as well as bytes transferred,
-  in atomic u64 / usize (repurposed to hold f64 for bytes)?
+* Track bytes transferred?
 
 ## Don't do
 
diff --git a/aquatic/src/lib/common.rs b/aquatic/src/lib/common.rs
index 6be43e5..0c953fb 100644
--- a/aquatic/src/lib/common.rs
+++ b/aquatic/src/lib/common.rs
@@ -117,10 +117,20 @@ impl Default for TorrentData {
 pub type TorrentMap = DashMap<InfoHash, TorrentData>;
 
 
+#[derive(Default)]
+pub struct Statistics {
+    pub requests_received: AtomicUsize,
+    pub responses_sent: AtomicUsize,
+}
+
+
+
+
 #[derive(Clone)]
 pub struct State {
     pub connections: Arc<ConnectionMap>,
     pub torrents: Arc<TorrentMap>,
+    pub statistics: Arc<Statistics>,
 }
 
 impl State {
@@ -128,6 +138,7 @@ impl State {
         Self {
             connections: Arc::new(DashMap::new()),
             torrents: Arc::new(DashMap::new()),
+            statistics: Arc::new(Statistics::default()),
         }
     }
 }
diff --git a/aquatic/src/lib/config.rs b/aquatic/src/lib/config.rs
index da3677f..6465694 100644
--- a/aquatic/src/lib/config.rs
+++ b/aquatic/src/lib/config.rs
@@ -7,6 +7,7 @@ pub struct Config {
     pub recv_buffer_size: usize,
     pub max_scrape_torrents: u8,
     pub max_response_peers: usize,
+    pub statistics_interval: u64,
 }
 
 
@@ -17,6 +18,7 @@ impl Default for Config {
             recv_buffer_size: 4096 * 16,
             max_scrape_torrents: 255,
             max_response_peers: 255,
+            statistics_interval: 5,
         }
     }
 }
\ No newline at end of file
diff --git a/aquatic/src/lib/lib.rs b/aquatic/src/lib/lib.rs
index 0e7e1d6..9aa8f60 100644
--- a/aquatic/src/lib/lib.rs
+++ b/aquatic/src/lib/lib.rs
@@ -1,3 +1,4 @@
+use std::sync::atomic::Ordering;
 use std::time::Duration;
 
 pub mod common;
@@ -25,6 +26,29 @@ pub fn run(){
         });
     }
 
+    {
+        let state = state.clone();
+
+        ::std::thread::spawn(move || {
+            let interval = config.statistics_interval;
+
+            loop {
+                ::std::thread::sleep(Duration::from_secs(interval));
+
+                let requests_per_second: f64 = state.statistics.requests_received
+                    .fetch_and(0, Ordering::SeqCst) as f64 / interval as f64;
+                let responses_per_second: f64 = state.statistics.responses_sent
+                    .fetch_and(0, Ordering::SeqCst) as f64 / interval as f64;
+
+                println!(
+                    "stats: {} requests/second, {} responses/second",
+                    requests_per_second,
+                    responses_per_second
+                );
+            }
+        });
+    }
+
     loop {
         ::std::thread::sleep(Duration::from_secs(30));
 
diff --git a/aquatic/src/lib/network.rs b/aquatic/src/lib/network.rs
index 85f37b9..0860840 100644
--- a/aquatic/src/lib/network.rs
+++ b/aquatic/src/lib/network.rs
@@ -1,3 +1,4 @@
+use std::sync::atomic::Ordering;
 use std::net::SocketAddr;
 use std::io::ErrorKind;
 
@@ -109,6 +110,9 @@ fn handle_readable_socket(
     announce_requests: &mut Vec<(AnnounceRequest, SocketAddr)>,
     scrape_requests: &mut Vec<(ScrapeRequest, SocketAddr)>,
 ){
+    let mut requests_received: usize = 0;
+    let mut responses_sent: usize = 0;
+
     loop {
         match socket.recv_from(buffer) {
             Ok((amt, src)) => {
@@ -117,6 +121,10 @@ fn handle_readable_socket(
                     config.max_scrape_torrents
                 );
 
+                if request.is_ok(){
+                    requests_received += 1;
+                }
+
                 match request {
                     Ok(Request::Connect(r)) => {
                         connect_requests.push((r, src)); 
@@ -175,6 +183,7 @@ fn handle_readable_socket(
 
         match socket.send_to(&bytes[..], src){
             Ok(_bytes_sent) => {
+                responses_sent += 1;
             },
             Err(err) => {
                 match err.kind(){
@@ -188,4 +197,9 @@ fn handle_readable_socket(
             }
         }
     }
+
+    state.statistics.requests_received
+        .fetch_add(requests_received, Ordering::SeqCst);
+    state.statistics.responses_sent
+        .fetch_add(responses_sent, Ordering::SeqCst);
 }
\ No newline at end of file