perf(test): reduce 1GB echo allocations and random overhead

- Reuse buffers and use rng.Read to cut per-iteration allocations in random echo test
- Separate RNG for payload length to keep data stream deterministic
- Optimize UDP source validation by comparing UDPAddr directly to avoid per-packet String allocation

Author: xtaci

readloop.go

@@ -23,18 +23,34 @@
 package kcp
 
 import (
+	"net"
 	"sync/atomic"
 
 	"github.com/pkg/errors"
 )
 
+func sameUDPAddr(a, b *net.UDPAddr) bool {
+	if a == nil || b == nil {
+		return false
+	}
+	if a.Port != b.Port || a.Zone != b.Zone {
+		return false
+	}
+	return a.IP.Equal(b.IP)
+}
+
 // defaultReadLoop is the standard procedure for reading from a connection
 func (s *UDPSession) defaultReadLoop() {
 	buf := make([]byte, mtuLimit)
 
-	var src string
+	var src *net.UDPAddr
+	var srcStr string
 	if s.remote != nil {
-		src = s.remote.String()
+		if udp, ok := s.remote.(*net.UDPAddr); ok {
+			src = udp
+		} else {
+			srcStr = s.remote.String()
+		}
 	}
 	for {
 		n, addr, err := s.conn.ReadFrom(buf)
@@ -48,9 +64,19 @@ func (s *UDPSession) defaultReadLoop() {
 		}
 
 		// make sure the packet is from the same source
-		if src == "" { // set source address if nil
-			src = addr.String()
-		} else if addr.String() != src {
+		if src == nil && srcStr == "" { // set source address if nil
+			if udp, ok := addr.(*net.UDPAddr); ok {
+				src = udp
+			} else {
+				srcStr = addr.String()
+			}
+		} else if src != nil {
+			udp, ok := addr.(*net.UDPAddr)
+			if !ok || !sameUDPAddr(src, udp) {
+				atomic.AddUint64(&DefaultSnmp.InErrs, 1)
+				continue
+			}
+		} else if addr.String() != srcStr {
 			atomic.AddUint64(&DefaultSnmp.InErrs, 1)
 			continue
 		}

readloop_linux.go

@@ -25,6 +25,7 @@
 package kcp
 
 import (
+	"net"
 	"sync/atomic"
 
 	"github.com/pkg/errors"
@@ -44,9 +45,14 @@ func (s *UDPSession) readLoop() {
 	}
 
 	// x/net version
-	var src string
+	var src *net.UDPAddr
+	var srcStr string
 	if s.remote != nil {
-		src = s.remote.String()
+		if udp, ok := s.remote.(*net.UDPAddr); ok {
+			src = udp
+		} else {
+			srcStr = s.remote.String()
+		}
 	}
 	msgs := make([]ipv4.Message, batchSize)
 	for k := range msgs {
@@ -68,9 +74,19 @@ func (s *UDPSession) readLoop() {
 			msg := &msgs[i]
 
 			// make sure the packet is from the same source
-			if src == "" { // set source address if nil
-				src = msg.Addr.String()
-			} else if msg.Addr.String() != src {
+			if src == nil && srcStr == "" { // set source address if nil
+				if udp, ok := msg.Addr.(*net.UDPAddr); ok {
+					src = udp
+				} else {
+					srcStr = msg.Addr.String()
+				}
+			} else if src != nil {
+				udp, ok := msg.Addr.(*net.UDPAddr)
+				if !ok || !sameUDPAddr(src, udp) {
+					atomic.AddUint64(&DefaultSnmp.InErrs, 1)
+					continue
+				}
+			} else if msg.Addr.String() != srcStr {
 				atomic.AddUint64(&DefaultSnmp.InErrs, 1)
 				continue
 			}

sess_test.go

@@ -400,18 +400,21 @@ func randomEchoTest(t *testing.T, cli *UDPSession, N int64) {
 	// Writer goroutine
 	go func() {
 		r := mrand.New(writerSrc)
+		lenRand := mrand.New(mrand.NewSource(seed + 1))
 		lastPrint := int64(0)
+		sndbuf := make([]byte, 1<<20)
 		for bytesSent < N {
-			length := mrand.Intn(1<<20) + 1 // Random length between 1 and 1MB
+			length := lenRand.Intn(1<<20) + 1 // Random length between 1 and 1MB
 			if bytesSent+int64(length) > N {
 				length = int(N - bytesSent)
 			}
-			sndbuf := make([]byte, length)
-			for i := range sndbuf {
-				sndbuf[i] = byte(r.Int())
+			payload := sndbuf[:length]
+			if _, err := r.Read(payload); err != nil {
+				t.Errorf("Random fill error: %v", err)
+				return
 			}
 
-			n, err := cli.Write(sndbuf)
+			n, err := cli.Write(payload)
 			if err != nil {
 				t.Errorf("Write error: %v", err)
 				return
@@ -426,22 +429,31 @@ func randomEchoTest(t *testing.T, cli *UDPSession, N int64) {
 
 	// Reader goroutine
 	r := mrand.New(readerSrc)
+	lenRand := mrand.New(mrand.NewSource(seed + 2))
 	lastPrint := int64(0)
+	rcvbuf := make([]byte, 1<<20)
+	expbuf := make([]byte, 1<<20)
 	for bytesReceived < N {
-		length := mrand.Intn(1<<20) + 1 // Random length between 1 and 1MB
+		length := lenRand.Intn(1<<20) + 1 // Random length between 1 and 1MB
 		if bytesReceived+int64(length) > N {
 			length = int(N - bytesReceived)
 		}
-		rcvbuf := make([]byte, length)
-		n, err := cli.Read(rcvbuf)
+		buf := rcvbuf[:length]
+		n, err := cli.Read(buf)
 		if err != nil && err != io.EOF {
 			t.Fatalf("Read error: %v", err)
 		}
-		for i := range n {
-			expectedByte := byte(r.Int())
-			if rcvbuf[i] != expectedByte {
-				t.Fatalf("Data mismatch at byte %d: got %v, want %v", bytesReceived+int64(i), rcvbuf[i], expectedByte)
+		expected := expbuf[:n]
+		if _, err := r.Read(expected); err != nil {
+			t.Fatalf("Random fill error: %v", err)
+		}
+		if !bytes.Equal(buf[:n], expected) {
+			for i := 0; i < n; i++ {
+				if buf[i] != expected[i] {
+					t.Fatalf("Data mismatch at byte %d: got %v, want %v", bytesReceived+int64(i), buf[i], expected[i])
+				}
 			}
+			t.Fatalf("Data mismatch at byte %d", bytesReceived)
 		}
 		bytesReceived += int64(n)
 		if bytesReceived-lastPrint >= 1<<28 { // print every 256MB