refactor: almost everything

fix: return traffic result when io error occurs
feat: allow copy from file

Author: cxw620

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "tokio-splice2"
-version = "0.3.0-alpha.10"
+version = "0.3.0-alpha.11"
 edition = "2021"
 rust-version = "1.70.0"
 
@@ -13,17 +13,23 @@ readme = "README.md"
 repository = "https://github.com/hanyu-dev/tokio-splice2"
 
 [dependencies]
-pin-project-lite = "0.2.16"
 rustix = { version = "1.0.0", features = ["pipe"] }
-tokio = { version = "1.39.0", features = ["fs", "net"] }
+tokio = { version = "1.46.0", features = ["fs", "net"] }
 tracing = { version = "0.1.41", optional = true }
 
 [dev-dependencies]
-tokio = { version = "1.39.0", features = ["net", "macros", "rt", "signal", "io-util", "time"] }
+human-format-next = "0.2.2"
+rand = "0.9.1"
+tokio = { version = "1.46.0", features = ["net", "macros", "rt", "signal", "io-util", "time"] }
+
 
 [features]
-# Enable tracing support
+default = ["feat-tracing"]
+
+# Enable crate `tracing` support
 feat-tracing = ["dep:tracing"]
+# Enable crate `tracing` support and enable Level::TRACE log in release profile build.
+feat-tracing-trace = ["feat-tracing"]
 
 # Enable nightly Rust feature
 feat-nightly = []

README.md

@@ -13,11 +13,14 @@ See [examples](./examples/).
 ## Changelog
 
 - 0.3.0:
+
   - [BREAKING] use `rustix` instead of `libc`
-  - [BREAKING] Do not shutdown stream after transfer finished (let user controls).
   - [BREAKING] MSRV is changed to 1.70.0
-  - Add unidirectional splicing.
-  - Accept `tokio::fs::File` as input.
+  - Add tracing log support.
+  - Add unidirectional copy.
+  - Add blocking unidirectional copy.
+  - (Experimental) Add `std::fs::File`/ `tokio::fs::File` support to splice from (like `sendfile`) / to (not fully tested).
+  - Returns `TrafficResult` instead of `io::Result<T>` to have traffic transferred returned when error occurs (e.g. the sender force closes the stream).
 
 - 0.2.1:
   - Fix the maximum value of the `size_t` type. Closes: [https://github.com/Hanaasagi/tokio-splice/issues/2](https://github.com/Hanaasagi/tokio-splice/issues/2).
@@ -28,7 +31,7 @@ See [BENCHMARK](./BENCHMARK.md).
 
 ## MSRV
 
-1.70.0
+1.70.0 (For running the examples, we recommend using the latest stable Rust version)
 
 ## LICENSE
 

examples/go.mod

@@ -0,0 +1,3 @@
+module proxy
+
+go 1.21

examples/proxy.go

@@ -0,0 +1,152 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"log"
+	"net"
+	"os"
+	"os/signal"
+	"sync"
+	"syscall"
+	"time"
+)
+
+func main() {
+	fmt.Printf("PID is %d\n", os.Getpid())
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	sigChan := make(chan os.Signal, 1)
+	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
+
+	go func() {
+		if err := serve(ctx); err != nil {
+			log.Printf("Serve failed: %v", err)
+		}
+	}()
+
+	<-sigChan
+	fmt.Println("Received Ctrl + C, shutting down")
+	cancel()
+
+	time.Sleep(100 * time.Millisecond)
+}
+
+func serve(ctx context.Context) error {
+	listenAddr := os.Getenv("EXAMPLE_LISTEN_ADDR")
+	if listenAddr == "" {
+		listenAddr = "0.0.0.0:5201"
+	}
+
+	listener, err := net.Listen("tcp", listenAddr)
+	if err != nil {
+		return fmt.Errorf("failed to listen on %s: %w", listenAddr, err)
+	}
+	defer listener.Close()
+
+	fmt.Printf("Listening on %s\n", listenAddr)
+
+	for {
+		select {
+		case <-ctx.Done():
+			return nil
+		default:
+		}
+
+		if tcpListener, ok := listener.(*net.TCPListener); ok {
+			tcpListener.SetDeadline(time.Now().Add(100 * time.Millisecond))
+		}
+
+		conn, err := listener.Accept()
+		if err != nil {
+			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
+				continue
+			}
+			if netErr, ok := err.(net.Error); ok && netErr.Temporary() {
+				log.Printf("Temporary accept error: %v", err)
+				continue
+			}
+			return fmt.Errorf("failed to accept: %w", err)
+		}
+
+		remoteAddr := conn.RemoteAddr()
+		fmt.Printf("Process incoming connection from %s\n", remoteAddr)
+
+		go forwarding(conn)
+	}
+}
+
+func forwarding(stream1 net.Conn) error {
+	defer stream1.Close()
+
+	remoteAddr := os.Getenv("EXAMPLE_REMOTE_ADDR")
+	if remoteAddr == "" {
+		remoteAddr = "127.0.0.1:5202"
+	}
+
+	stream2, err := net.Dial("tcp", remoteAddr)
+	if err != nil {
+		log.Printf("Failed to connect to remote server: %v", err)
+		return err
+	}
+	defer stream2.Close()
+
+	result, err := copyBidirectional(stream1, stream2)
+
+	fmt.Printf("Forwarded traffic: %+v\n", result)
+
+	if err != nil {
+		log.Printf("Failed to copy data: %v", err)
+		return err
+	}
+
+	return nil
+}
+
+type TrafficStats struct {
+	BytesForward uint64
+	BytesReverse uint64
+}
+
+func (t TrafficStats) String() string {
+	return fmt.Sprintf("TrafficStats { bytes_forward: %d, bytes_reverse: %d }",
+		t.BytesForward, t.BytesReverse)
+}
+
+func copyBidirectional(conn1, conn2 net.Conn) (*TrafficStats, error) {
+	var stats TrafficStats
+	var wg sync.WaitGroup
+	var err1, err2 error
+
+	wg.Add(2)
+
+	go func() {
+		defer wg.Done()
+		defer conn2.Close()
+		n, err := io.Copy(conn2, conn1)
+		stats.BytesForward = uint64(n)
+		err1 = err
+	}()
+
+	go func() {
+		defer wg.Done()
+		defer conn1.Close()
+		n, err := io.Copy(conn1, conn2)
+		stats.BytesReverse = uint64(n)
+		err2 = err
+	}()
+
+	wg.Wait()
+
+	if err1 != nil && err1 != io.EOF {
+		return &stats, err1
+	}
+	if err2 != nil && err2 != io.EOF {
+		return &stats, err2
+	}
+
+	return &stats, nil
+}

examples/proxy.rs

@@ -23,7 +23,10 @@ async fn main() -> io::Result<()> {
 }
 
 async fn serve() -> io::Result<()> {
-    let listener = TcpListener::bind("127.0.0.1:8989").await?;
+    let listener = TcpListener::bind(
+        env::var("EXAMPLE_LISTEN_ADDR").unwrap_or_else(|_| "0.0.0.0:5201".to_string()),
+    )
+    .await?;
 
     loop {
         let (incoming, remote_addr) = match listener.accept().await {
@@ -33,7 +36,7 @@ async fn serve() -> io::Result<()> {
                 continue;
             }
             Err(e) => {
-                eprintln!("Failed to accept: {:#?}", e);
+                eprintln!("Failed to accept: {e:#?}");
                 break Err(e);
             }
         };
@@ -45,23 +48,46 @@ async fn serve() -> io::Result<()> {
 }
 
 async fn forwarding(mut stream1: TcpStream) -> io::Result<()> {
-    let mut stream2 = TcpStream::connect(
-        env::var("EXAMPLE_REMOTE_ADDR").unwrap_or_else(|_| "4.ipw.cn:80".to_string()),
+    let stream2 = TcpStream::connect(
+        env::var("EXAMPLE_REMOTE_ADDR").unwrap_or_else(|_| "127.0.0.1:5202".to_string()),
     )
-    .await
-    .inspect_err(|e| {
-        eprintln!("Failed to connect to remote server: {e}");
-    })?;
+    .await;
+
+    let mut stream2 = match stream2 {
+        Ok(s) => s,
+        Err(e) => {
+            eprintln!("Failed to connect to remote server: {e}");
+            return Err(e);
+        }
+    };
+
+    // let result = tokio_splice::zero_copy_bidirectional(&mut stream1, &mut
+    // stream2).await;
+    let instant = std::time::Instant::now();
+    let result = tokio_splice2::copy_bidirectional(&mut stream1, &mut stream2).await;
+    // let result = realm_io::bidi_zero_copy(&mut stream1, &mut stream2).await;
 
-    tokio_splice2::copy_bidirectional(&mut stream1, &mut stream2)
-        .await
-        .inspect(|(r, w)| {
+    match result {
+        Ok(traffic) => {
+            let total = traffic.sum();
+            let cost = instant.elapsed();
             println!(
-                "Forwarded {r} bytes from stream1 to stream2, {w} bytes from stream2 to stream1"
+                "Forwarded traffic: {traffic:?}, total: {}, time: {:.2}s, avg: {}",
+                human_format_next::Formatter::BINARY
+                    .with_custom_unit("B")
+                    .with_decimals::<4>()
+                    .format(total as f64),
+                cost.as_secs_f64(),
+                human_format_next::Formatter::BINARY
+                    .with_custom_unit("B/s")
+                    .with_decimals::<4>()
+                    .format(total as f64 / cost.as_secs_f64())
             );
-        })
-        .inspect_err(|e| {
+            Ok(())
+        }
+        Err(e) => {
             eprintln!("Failed to copy data: {e}");
-        })?;
-    Ok(())
+            Err(e)
+        }
+    }
 }