notify/raw_tcp_benchmark_test.go

package notify

import (
	"errors"
	"io"
	"net"
	"testing"
	"time"
)

func BenchmarkRawTCPLocalhostThroughput(b *testing.B) {
	cases := []struct {
		name        string
		payloadSize int
	}{
		{
			name:        "raw_64KiB",
			payloadSize: 64 * 1024,
		},
		{
			name:        "raw_256KiB",
			payloadSize: 256 * 1024,
		},
		{
			name:        "raw_512KiB",
			payloadSize: 512 * 1024,
		},
		{
			name:        "raw_1MiB",
			payloadSize: 1024 * 1024,
		},
	}

	for _, tc := range cases {
		b.Run(tc.name, func(b *testing.B) {
			benchmarkRawTCPLocalhostThroughput(b, tc.payloadSize)
		})
	}
}

func benchmarkRawTCPLocalhostThroughput(b *testing.B, payloadSize int) {
	b.Helper()

	listener, err := net.Listen("tcp", benchmarkTCPListenAddr(b))
	if err != nil {
		b.Fatalf("net.Listen failed: %v", err)
	}
	b.Cleanup(func() {
		_ = listener.Close()
	})

	acceptCh := make(chan net.Conn, 1)
	acceptErrCh := make(chan error, 1)
	go func() {
		conn, err := listener.Accept()
		if err != nil {
			acceptErrCh <- err
			return
		}
		acceptCh <- conn
	}()

	clientConn, err := net.Dial("tcp", benchmarkTCPDialAddr(b, listener.Addr().String()))
	if err != nil {
		b.Fatalf("net.Dial failed: %v", err)
	}
	b.Cleanup(func() {
		_ = clientConn.Close()
	})
	if tcpConn, ok := clientConn.(*net.TCPConn); ok {
		_ = tcpConn.SetNoDelay(true)
	}

	var serverConn net.Conn
	select {
	case conn := <-acceptCh:
		serverConn = conn
	case err := <-acceptErrCh:
		b.Fatalf("Accept failed: %v", err)
	case <-time.After(5 * time.Second):
		b.Fatal("timed out waiting for accept")
	}
	b.Cleanup(func() {
		if serverConn != nil {
			_ = serverConn.Close()
		}
	})

	drainDone := make(chan error, 1)
	go func() {
		_, err := io.Copy(io.Discard, serverConn)
		if err != nil && !errors.Is(err, io.EOF) {
			drainDone <- err
			return
		}
		drainDone <- nil
	}()

	payload := make([]byte, payloadSize)
	for i := range payload {
		payload[i] = byte(i)
	}

	b.ReportAllocs()
	b.SetBytes(int64(payloadSize))
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if err := benchmarkRawTCPWriteFull(clientConn, payload); err != nil {
			b.Fatalf("raw tcp write failed at iter %d: %v", i, err)
		}
	}
	b.StopTimer()

	if tcpConn, ok := clientConn.(*net.TCPConn); ok {
		_ = tcpConn.CloseWrite()
	} else {
		_ = clientConn.Close()
	}

	select {
	case err := <-drainDone:
		if err != nil {
			b.Fatalf("server drain failed: %v", err)
		}
	case <-time.After(10 * time.Second):
		b.Fatal("timed out waiting for server drain")
	}
}

func benchmarkRawTCPWriteFull(conn net.Conn, payload []byte) error {
	for len(payload) > 0 {
		n, err := conn.Write(payload)
		if n > 0 {
			payload = payload[n:]
		}
		if err != nil {
			return err
		}
		if n == 0 {
			return io.ErrNoProgress
		}
	}
	return nil
}