Add benchmark for tcp proxy use case

Signed-off-by: Fabian Ruff <[email protected]>

Author: databus23

protocol_test.go

@@ -798,6 +798,131 @@ func TestCopyFromWrappedConnectionToWrappedConnection(t *testing.T) {
 	}
 }
 
+func benchmarkTCPProxy(size int, b *testing.B) {
+	//create and start the echo backend
+	backend, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		b.Fatalf("err: %v", err)
+	}
+	defer backend.Close()
+	go func() {
+		for {
+			conn, err := backend.Accept()
+			if err != nil {
+				break
+			}
+			_, err = io.Copy(conn, conn)
+			conn.Close()
+			if err != nil {
+				b.Fatalf("Failed to read entire payload: %v", err)
+			}
+		}
+	}()
+
+	//start the proxyprotocol enabled tcp proxy
+	l, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		b.Fatalf("err: %v", err)
+	}
+	defer l.Close()
+	pl := &Listener{Listener: l}
+	go func() {
+		for {
+			conn, err := pl.Accept()
+			if err != nil {
+				break
+			}
+			bConn, err := net.Dial("tcp", backend.Addr().String())
+			if err != nil {
+				b.Fatalf("failed to dial backend: %v", err)
+			}
+			go func() {
+				_, err = io.Copy(bConn, conn)
+				if err != nil {
+					b.Fatalf("Failed to proxy incoming data to backend: %v", err)
+				}
+				bConn.(*net.TCPConn).CloseWrite()
+			}()
+			_, err = io.Copy(conn, bConn)
+			if err != nil {
+				b.Fatalf("Failed to proxy data from backend: %v", err)
+			}
+			conn.Close()
+			bConn.Close()
+		}
+	}()
+
+	data := make([]byte, size)
+
+	header := &Header{
+		Version:           2,
+		Command:           PROXY,
+		TransportProtocol: TCPv4,
+		SourceAddr: &net.TCPAddr{
+			IP:   net.ParseIP("10.1.1.1"),
+			Port: 1000,
+		},
+		DestinationAddr: &net.TCPAddr{
+			IP:   net.ParseIP("20.2.2.2"),
+			Port: 2000,
+		},
+	}
+
+	//now for the actual benchmark
+	b.ResetTimer()
+	for n := 0; n < b.N; n++ {
+		conn, err := net.Dial("tcp", pl.Addr().String())
+		if err != nil {
+			b.Fatalf("err: %v", err)
+		}
+		// Write out the header!
+		header.WriteTo(conn)
+		//send data
+		go func() {
+			_, err = conn.Write(data)
+			if err != nil {
+				b.Fatalf("Failed to write data: %v", err)
+			}
+			conn.(*net.TCPConn).CloseWrite()
+
+		}()
+		//receive data
+		n, err := io.Copy(ioutil.Discard, conn)
+		if n != int64(len(data)) {
+			b.Fatalf("Expected to receive %d bytes, got %d", len(data), n)
+		}
+		if err != nil {
+			b.Fatalf("Failed to read data: %v", err)
+		}
+		conn.Close()
+	}
+}
+
+func BenchmarkTCPProxy16KB(b *testing.B) {
+	benchmarkTCPProxy(16*1024, b)
+}
+func BenchmarkTCPProxy32KB(b *testing.B) {
+	benchmarkTCPProxy(32*1024, b)
+}
+func BenchmarkTCPProxy64KB(b *testing.B) {
+	benchmarkTCPProxy(64*1024, b)
+}
+func BenchmarkTCPProxy128KB(b *testing.B) {
+	benchmarkTCPProxy(128*1024, b)
+}
+func BenchmarkTCPProxy256KB(b *testing.B) {
+	benchmarkTCPProxy(256*1024, b)
+}
+func BenchmarkTCPProxy512KB(b *testing.B) {
+	benchmarkTCPProxy(512*1024, b)
+}
+func BenchmarkTCPProxy1024KB(b *testing.B) {
+	benchmarkTCPProxy(1024*1024, b)
+}
+func BenchmarkTCPProxy2048KB(b *testing.B) {
+	benchmarkTCPProxy(2048*1024, b)
+}
+
 // copied from src/net/http/internal/testcert.go
 
 // Copyright 2015 The Go Authors. All rights reserved.