Swap back to a blocking udp socket, test shutdown(2) (#1806)

Co-authored-by: Jack Doan <me@jackdoan.com>

udp/udp_linux_test.go

@@ -0,0 +1,179 @@
+//go:build linux && !android && !e2e_testing
+
+package udp
+
+import (
+	"errors"
+	"fmt"
+	"log/slog"
+	"net"
+	"net/netip"
+	"os"
+	"runtime"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"golang.org/x/sys/unix"
+)
+
+func testLogger() *slog.Logger {
+	return slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError}))
+}
+
+// TestShutdownWakesAfterRx_Mechanism exercises the kernel quirk our teardown
+// relies on: once a socket has received a packet, shutdown(2) wakes a blocked
+// recvmmsg with n>=1/Len==0 (not n==0). recvmmsg must turn that into net.ErrClosed
+// once Close set closed, so a parked reader exits instead of spinning.
+func TestShutdownWakesAfterRx_Mechanism(t *testing.T) {
+	c, err := NewListener(testLogger(), netip.MustParseAddr("127.0.0.1"), 0, true, 64)
+	if err != nil {
+		t.Fatalf("NewListener: %v", err)
+	}
+	sc := c.(*StdConn)
+	addr, err := sc.LocalAddr()
+	if err != nil {
+		t.Fatalf("LocalAddr: %v", err)
+	}
+	msgs, _, _ := sc.PrepareRawMessages(sc.batch)
+
+	// Receive a real packet so the socket has carried data.
+	send, err := net.Dial("udp", addr.String())
+	if err != nil {
+		t.Fatalf("dial: %v", err)
+	}
+	if _, err := send.Write([]byte("hello")); err != nil {
+		t.Fatalf("write: %v", err)
+	}
+	time.Sleep(50 * time.Millisecond)
+	n, err := sc.recvmmsg(msgs)
+	t.Logf("drain of real packet: n=%d err=%v msgs[0].Len=%d", n, err, msgs[0].Len)
+	_ = send.Close()
+
+	// Block a reader on the now-empty queue, then tear down as Close() does.
+	// recvmmsg must return net.ErrClosed (not hang, not spin) even post-rx.
+	done := make(chan error, 1)
+	go func() {
+		_, err := sc.recvmmsg(msgs)
+		done <- err
+	}()
+	time.Sleep(150 * time.Millisecond) // let it park in recvmmsg
+
+	sc.closed.Store(true)
+	if serr := unix.Shutdown(sc.sysFd, unix.SHUT_RDWR); serr != nil {
+		t.Logf("shutdown returned %v (expected ENOTCONN on unconnected UDP)", serr)
+	}
+
+	select {
+	case err := <-done:
+		if !errors.Is(err, net.ErrClosed) {
+			t.Errorf("recvmmsg after post-rx shutdown returned %v, want net.ErrClosed", err)
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatalf("HANG: recvmmsg did not return after shutdown following a received packet")
+	}
+	_ = unix.Close(sc.sysFd)
+}
+
+// TestListenOutTeardown_TrafficPatterns reproduces the field report: a blocking
+// reader must tear down cleanly on Close() regardless of what the socket has
+// carried. The three cases the report called out:
+//
+//	no traffic ever      -> works (shutdown wakes recvmmsg with n==0)
+//	ping once, then idle -> historically HUNG: once the socket has received a
+//	                        packet, shutdown(2) wakes recvmmsg with n>=1/Len==0,
+//	                        which an n==0-only teardown check misses
+//	continuous traffic   -> works (a real packet is always arriving)
+//
+// All three must return within the deadline; a hang dumps goroutines so the
+// stuck reader is visible.
+func TestListenOutTeardown_TrafficPatterns(t *testing.T) {
+	cases := []struct {
+		name    string
+		traffic func(send net.Conn, stop <-chan struct{})
+	}{
+		{"no_traffic_ever", func(net.Conn, <-chan struct{}) {}},
+		{"ping_once_then_idle", func(send net.Conn, _ <-chan struct{}) {
+			_, _ = send.Write([]byte("hello"))
+		}},
+		{"continuous", func(send net.Conn, stop <-chan struct{}) {
+			for {
+				select {
+				case <-stop:
+					return
+				default:
+					_, _ = send.Write([]byte("hello"))
+					time.Sleep(2 * time.Millisecond)
+				}
+			}
+		}},
+	}
+
+	// batch 1 exercises the recvmsg path, batch 64 the recvmmsg path; both must
+	// tear down cleanly.
+	for _, batch := range []int{1, 64} {
+		for _, tc := range cases {
+			t.Run(fmt.Sprintf("batch%d/%s", batch, tc.name), func(t *testing.T) {
+				runTeardownCase(t, batch, tc.name, tc.traffic)
+			})
+		}
+	}
+}
+
+func runTeardownCase(t *testing.T, batch int, name string, traffic func(send net.Conn, stop <-chan struct{})) {
+	c, err := NewListener(testLogger(), netip.MustParseAddr("127.0.0.1"), 0, true, batch)
+	if err != nil {
+		t.Fatalf("NewListener: %v", err)
+	}
+	sc := c.(*StdConn)
+	addr, err := sc.LocalAddr()
+	if err != nil {
+		t.Fatalf("LocalAddr: %v", err)
+	}
+
+	var received atomic.Int64
+	loopDone := make(chan error, 1)
+	go func() {
+		loopDone <- sc.ListenOut(func(netip.AddrPort, []byte) {
+			received.Add(1)
+		})
+	}()
+
+	send, err := net.Dial("udp", addr.String())
+	if err != nil {
+		t.Fatalf("dial: %v", err)
+	}
+	defer send.Close()
+
+	stop := make(chan struct{})
+	trafficDone := make(chan struct{})
+	go func() {
+		traffic(send, stop)
+		close(trafficDone)
+	}()
+
+	// Let the pattern run and, for the idle case, the reader park again on an
+	// empty queue with the socket already having received a packet.
+	time.Sleep(500 * time.Millisecond)
+
+	start := time.Now()
+	if err := sc.Close(); err != nil {
+		t.Fatalf("Close: %v", err)
+	}
+	close(stop)
+
+	select {
+	case err := <-loopDone:
+		// Clean teardown surfaces as net.ErrClosed (propagated like the other
+		// platforms); the caller absorbs it via its closed flag.
+		if err != nil && !errors.Is(err, net.ErrClosed) {
+			t.Fatalf("%s: ListenOut returned unexpected error on teardown: %v", name, err)
+		}
+		t.Logf("%s: closed in %v (received %d packets)", name, time.Since(start), received.Load())
+	case <-time.After(3 * time.Second):
+		buf := make([]byte, 1<<20)
+		n := runtime.Stack(buf, true)
+		t.Fatalf("%s: HANG, ListenOut did not return within 3s of Close\n%s", name, buf[:n])
+	}
+	<-trafficDone
+}