g3tiles: use common stream task wrapper stats

g3tiles/src/module/stream/mod.rs

@@ -18,8 +18,8 @@ use tokio::io::{AsyncRead, AsyncWrite};
 
 mod stats;
 pub(crate) use stats::{
-    StreamBackendDurationRecorder, StreamBackendDurationStats, StreamBackendStats,
-    StreamServerStats,
+    StreamAcceptTaskCltWrapperStats, StreamBackendDurationRecorder, StreamBackendDurationStats,
+    StreamBackendStats, StreamRelayTaskCltWrapperStats, StreamServerStats,
 };
 
 mod error;

g3tiles/src/module/stream/stats/mod.rs

@@ -17,6 +17,9 @@
 mod server;
 pub(crate) use server::StreamServerStats;
 
+mod task;
+pub(crate) use task::{StreamAcceptTaskCltWrapperStats, StreamRelayTaskCltWrapperStats};
+
 mod backend;
 pub(crate) use backend::{
     StreamBackendDurationRecorder, StreamBackendDurationStats, StreamBackendStats,

g3tiles/src/serve/openssl_proxy/task/accept/stats.rs → g3tiles/src/module/stream/stats/task.rs

@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 ByteDance and/or its affiliates.
+ * Copyright 2024 ByteDance and/or its affiliates.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,41 +16,72 @@
 
 use std::sync::Arc;
 
-use g3_daemon::stat::task::TcpStreamConnectionStats;
+use g3_daemon::stat::task::{TcpStreamConnectionStats, TcpStreamTaskStats};
 use g3_io_ext::{LimitedReaderStats, LimitedWriterStats};
 
 use crate::module::stream::StreamServerStats;
 
 #[derive(Clone)]
-pub(crate) struct OpensslAcceptTaskCltWrapperStats {
+pub(crate) struct StreamAcceptTaskCltWrapperStats {
     server: Arc<StreamServerStats>,
     conn: Arc<TcpStreamConnectionStats>,
 }
 
-impl OpensslAcceptTaskCltWrapperStats {
+impl StreamAcceptTaskCltWrapperStats {
     pub(crate) fn new(
         server: &Arc<StreamServerStats>,
         conn: &Arc<TcpStreamConnectionStats>,
     ) -> Self {
-        OpensslAcceptTaskCltWrapperStats {
+        StreamAcceptTaskCltWrapperStats {
             server: Arc::clone(server),
             conn: Arc::clone(conn),
         }
     }
 }
 
-impl LimitedReaderStats for OpensslAcceptTaskCltWrapperStats {
+impl LimitedReaderStats for StreamAcceptTaskCltWrapperStats {
     fn add_read_bytes(&self, size: usize) {
         let size = size as u64;
         self.conn.read.add_bytes(size);
         self.server.add_read(size);
     }
 }
 
-impl LimitedWriterStats for OpensslAcceptTaskCltWrapperStats {
+impl LimitedWriterStats for StreamAcceptTaskCltWrapperStats {
     fn add_write_bytes(&self, size: usize) {
         let size = size as u64;
         self.conn.write.add_bytes(size);
         self.server.add_write(size);
     }
 }
+
+#[derive(Clone)]
+pub(crate) struct StreamRelayTaskCltWrapperStats {
+    server: Arc<StreamServerStats>,
+    task: Arc<TcpStreamTaskStats>,
+}
+
+impl StreamRelayTaskCltWrapperStats {
+    pub(crate) fn new(server: &Arc<StreamServerStats>, task: &Arc<TcpStreamTaskStats>) -> Self {
+        StreamRelayTaskCltWrapperStats {
+            server: Arc::clone(server),
+            task: Arc::clone(task),
+        }
+    }
+}
+
+impl LimitedReaderStats for StreamRelayTaskCltWrapperStats {
+    fn add_read_bytes(&self, size: usize) {
+        let size = size as u64;
+        self.task.clt.read.add_bytes(size);
+        self.server.add_read(size);
+    }
+}
+
+impl LimitedWriterStats for StreamRelayTaskCltWrapperStats {
+    fn add_write_bytes(&self, size: usize) {
+        let size = size as u64;
+        self.task.clt.write.add_bytes(size);
+        self.server.add_write(size);
+    }
+}

g3tiles/src/serve/openssl_proxy/task/accept/mod.rs → g3tiles/src/serve/openssl_proxy/task/accept.rs

@@ -30,11 +30,9 @@ use g3_types::limit::GaugeSemaphorePermit;
 use g3_types::route::HostMatch;
 
 use super::{CommonTaskContext, OpensslRelayTask};
+use crate::module::stream::StreamAcceptTaskCltWrapperStats;
 use crate::serve::openssl_proxy::OpensslHost;
 
-mod stats;
-use stats::OpensslAcceptTaskCltWrapperStats;
-
 pub(crate) struct OpensslAcceptTask {
     ctx: CommonTaskContext,
     hosts: Arc<HostMatch<Arc<OpensslHost>>>,
@@ -55,7 +53,7 @@ impl OpensslAcceptTask {
 
         let pre_handshake_stats = Arc::new(TcpStreamConnectionStats::default());
         let wrapper_stats =
-            OpensslAcceptTaskCltWrapperStats::new(&self.ctx.server_stats, &pre_handshake_stats);
+            StreamAcceptTaskCltWrapperStats::new(&self.ctx.server_stats, &pre_handshake_stats);
 
         let limit_config = self.ctx.server_config.tcp_sock_speed_limit;
         let stream = LimitedStream::new(

g3tiles/src/serve/openssl_proxy/task/relay/task.rs → g3tiles/src/serve/openssl_proxy/task/relay.rs

@@ -26,10 +26,11 @@ use g3_io_ext::{LimitedCopy, LimitedCopyConfig, LimitedCopyError, LimitedStream}
 use g3_openssl::SslStream;
 use g3_types::limit::GaugeSemaphorePermit;
 
-use super::{CommonTaskContext, OpensslRelayTaskCltWrapperStats};
+use super::CommonTaskContext;
 use crate::backend::ArcBackend;
 use crate::config::server::ServerConfig;
 use crate::log::task::tcp_connect::TaskLogForTcpConnect;
+use crate::module::stream::StreamRelayTaskCltWrapperStats;
 use crate::serve::openssl_proxy::OpensslHost;
 use crate::serve::{ServerTaskError, ServerTaskNotes, ServerTaskResult, ServerTaskStage};
 
@@ -229,7 +230,7 @@ impl OpensslRelayTask {
         // reset io stats
         // TODO add host level stats
         let clt_wrapper_stats =
-            OpensslRelayTaskCltWrapperStats::new(&self.ctx.server_stats, &self.task_stats);
+            StreamRelayTaskCltWrapperStats::new(&self.ctx.server_stats, &self.task_stats);
         ssl_stream
             .get_mut()
             .reset_stats(Arc::new(clt_wrapper_stats));

g3tiles/src/serve/rustls_proxy/task/accept/mod.rs → g3tiles/src/serve/rustls_proxy/task/accept.rs

@@ -31,11 +31,9 @@ use g3_types::net::Host;
 use g3_types::route::HostMatch;
 
 use super::{CommonTaskContext, RustlsRelayTask};
+use crate::module::stream::StreamAcceptTaskCltWrapperStats;
 use crate::serve::rustls_proxy::RustlsHost;
 
-mod stats;
-use stats::RustlsAcceptTaskCltWrapperStats;
-
 pub(crate) struct RustlsAcceptTask {
     ctx: CommonTaskContext,
     alive_permit: Option<GaugeSemaphorePermit>,
@@ -58,7 +56,7 @@ impl RustlsAcceptTask {
 
         let pre_handshake_stats = Arc::new(TcpStreamConnectionStats::default());
         let wrapper_stats =
-            RustlsAcceptTaskCltWrapperStats::new(&self.ctx.server_stats, &pre_handshake_stats);
+            StreamAcceptTaskCltWrapperStats::new(&self.ctx.server_stats, &pre_handshake_stats);
 
         let limit_config = self.ctx.server_config.tcp_sock_speed_limit;
         let stream = LimitedStream::new(

g3tiles/src/serve/rustls_proxy/task/relay/task.rs → g3tiles/src/serve/rustls_proxy/task/relay.rs

@@ -26,10 +26,11 @@ use g3_daemon::stat::task::{TcpStreamConnectionStats, TcpStreamTaskStats};
 use g3_io_ext::{LimitedCopy, LimitedCopyConfig, LimitedCopyError, LimitedStream};
 use g3_types::limit::GaugeSemaphorePermit;
 
-use super::{CommonTaskContext, RustlsRelayTaskCltWrapperStats};
+use super::CommonTaskContext;
 use crate::backend::ArcBackend;
 use crate::config::server::ServerConfig;
 use crate::log::task::tcp_connect::TaskLogForTcpConnect;
+use crate::module::stream::StreamRelayTaskCltWrapperStats;
 use crate::serve::rustls_proxy::RustlsHost;
 use crate::serve::{ServerTaskError, ServerTaskNotes, ServerTaskResult, ServerTaskStage};
 
@@ -231,7 +232,7 @@ impl RustlsRelayTask {
         // reset io stats
         // TODO add host level stats
         let clt_wrapper_stats =
-            RustlsRelayTaskCltWrapperStats::new(&self.ctx.server_stats, &self.task_stats);
+            StreamRelayTaskCltWrapperStats::new(&self.ctx.server_stats, &self.task_stats);
         tls_stream
             .get_mut()
             .0