context: add io counter, use millis timestamp

Signed-off-by: Bearice Ren <[email protected]>
MengJiangProject · Sep 3, 2021 · b18be04 · b18be04
1 parent 0edbfcc
commit b18be04
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 6 deletions.
diff --git a/src/common/copy.rs b/src/common/copy.rs
@@ -1,3 +1,5 @@
+use std::sync::{atomic::AtomicUsize, Arc};
+
 use crate::context::ContextRef;
 use easy_error::{Error, ResultExt};
 use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadHalf, WriteHalf};
@@ -7,6 +9,7 @@ pub async fn copy_stream<T: AsyncRead + AsyncWrite>(
     rn: &str,
     mut w: WriteHalf<T>,
     wn: &str,
+    cnt: Arc<AtomicUsize>,
 ) -> Result<(), Error> {
     let mut buf = [0u8; 65536];
     loop {
@@ -26,6 +29,7 @@ pub async fn copy_stream<T: AsyncRead + AsyncWrite>(
             w.flush()
                 .await
                 .with_context(|| format!("flush {} buffer", wn))?;
+            cnt.fetch_add(len, std::sync::atomic::Ordering::Relaxed);
         } else {
             break;
         }
@@ -41,8 +45,20 @@ pub async fn copy_bidi(ctx: ContextRef) -> Result<(), Error> {
     let mut server = ctx.get_server_stream().await;
     let (cread, cwrite) = tokio::io::split(client.get_mut());
     let (sread, swrite) = tokio::io::split(server.get_mut());
-    let copy_a_to_b = copy_stream(cread, "client", swrite, "server");
-    let copy_b_to_a = copy_stream(sread, "server", cwrite, "client");
+    let copy_a_to_b = copy_stream(
+        cread,
+        "client",
+        swrite,
+        "server",
+        ctx.props().client_sent.clone(),
+    );
+    let copy_b_to_a = copy_stream(
+        sread,
+        "server",
+        cwrite,
+        "client",
+        ctx.props().server_sent.clone(),
+    );
     tokio::try_join!(copy_a_to_b, copy_b_to_a)?;
     Ok(())
 }
diff --git a/src/context.rs b/src/context.rs
@@ -10,7 +10,7 @@ use std::{
     pin::Pin,
     str::FromStr,
     sync::{
-        atomic::{AtomicU64, Ordering},
+        atomic::{AtomicU64, AtomicUsize, Ordering},
         Arc, Weak,
     },
     time::{Duration, SystemTime},
@@ -149,7 +149,7 @@ impl Serialize for ContextStatusLog {
                 .time
                 .duration_since(SystemTime::UNIX_EPOCH)
                 .unwrap()
-                .as_secs(),
+                .as_millis(),
         )?;
         st.end()
     }
@@ -162,13 +162,27 @@ impl From<(ContextStatus, SystemTime)> for ContextStatusLog {
 }
 
 // this is the value object of Context, chould be used in filter evaluation or stored after Context is terminated, for statistics.
-#[derive(Debug, Hash, Eq, Clone, PartialEq, Serialize)]
+#[derive(Debug, Clone, Serialize)]
 pub struct ContextProps {
     pub id: u64,
     pub status: Vec<ContextStatusLog>,
     pub listener: String,
     pub source: SocketAddr,
     pub target: TargetAddress,
+    pub client_sent: Arc<AtomicUsize>,
+    pub server_sent: Arc<AtomicUsize>,
+}
+
+impl std::hash::Hash for ContextProps {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.id.hash(state);
+    }
+}
+
+impl PartialEq for ContextProps {
+    fn eq(&self, other: &Self) -> bool {
+        self.id == other.id
+    }
 }
 
 impl Default for ContextProps {
@@ -179,6 +193,8 @@ impl Default for ContextProps {
             listener: Default::default(),
             source: ([0, 0, 0, 0], 0).into(),
             target: TargetAddress::Unknown,
+            client_sent: Default::default(),
+            server_sent: Default::default(),
         }
     }
 }
@@ -207,7 +223,7 @@ impl GlobalState {
             listener,
             source,
             status: vec![(ContextStatus::ClientConnected, SystemTime::now()).into()],
-            target: TargetAddress::Unknown,
+            ..Default::default()
         });
         let ret = Arc::new(RwLock::new(Context {
             props,