wgengine/magicsock: remove an alloc from ReceiveIPvN

We modified the standard net package to not allocate a *net.UDPAddr
during a call to (*net.UDPConn).ReadFromUDP if the caller's use
of the *net.UDPAddr does not cause it to escape.
That is https://golang.org/cl/291390.

This is the companion change to magicsock.
There are two changes required.
First, call ReadFromUDP instead of ReadFrom, if possible.
ReadFrom returns a net.Addr, which is an interface, which always allocates.
Second, reduce the lifetime of the returned *net.UDPAddr.
We do this by immediately converting it into a netaddr.IPPort.

We left the existing RebindingUDPConn.ReadFrom method in place,
as it is required to satisfy the net.PacketConn interface.

With the upstream change and both of these fixes in place,
we have removed one large allocation per packet received.

name           old time/op    new time/op    delta
ReceiveFrom-8    16.7µs ± 5%    16.4µs ± 8%     ~     (p=0.310 n=5+5)

name           old alloc/op   new alloc/op   delta
ReceiveFrom-8      112B ± 0%       64B ± 0%  -42.86%  (p=0.008 n=5+5)

name           old allocs/op  new allocs/op  delta
ReceiveFrom-8      3.00 ± 0%      2.00 ± 0%  -33.33%  (p=0.008 n=5+5)

Co-authored-by: Sonia Appasamy <sonia@tailscale.com>
Signed-off-by: Josh Bleecher Snyder <josh@tailscale.com>

wgengine/magicsock/magicsock_test.go

@@ -1512,16 +1512,16 @@ func addTestEndpoint(conn *Conn, sendConn net.PacketConn) (tailcfg.NodeKey, tail
 	return nodeKey, discoKey
 }
 
-func BenchmarkReceiveFrom(b *testing.B) {
-	conn := newNonLegacyTestConn(b)
-	defer conn.Close()
+func setUpReceiveFrom(tb testing.TB) (roundTrip func()) {
+	conn := newNonLegacyTestConn(tb)
+	tb.Cleanup(func() { conn.Close() })
 	conn.logf = logger.Discard
 
 	sendConn, err := net.ListenPacket("udp4", "127.0.0.1:0")
 	if err != nil {
-		b.Fatal(err)
+		tb.Fatal(err)
 	}
-	defer sendConn.Close()
+	tb.Cleanup(func() { sendConn.Close() })
 
 	addTestEndpoint(conn, sendConn)
 
@@ -1530,21 +1530,89 @@ func BenchmarkReceiveFrom(b *testing.B) {
 	for i := range sendBuf {
 		sendBuf[i] = 'x'
 	}
-
 	buf := make([]byte, 2<<10)
-	for i := 0; i < b.N; i++ {
+	return func() {
 		if _, err := sendConn.WriteTo(sendBuf, dstAddr); err != nil {
-			b.Fatalf("WriteTo: %v", err)
+			tb.Fatalf("WriteTo: %v", err)
 		}
 		n, ep, err := conn.ReceiveIPv4(buf)
 		if err != nil {
-			b.Fatal(err)
+			tb.Fatal(err)
 		}
 		_ = n
 		_ = ep
 	}
 }
 
+// goMajorVersion reports the major Go version and whether it is a Tailscale fork.
+// If parsing fails, goMajorVersion returns 0, false.
+func goMajorVersion(s string) (version int, isTS bool) {
+	if !strings.HasPrefix(s, "go1.") {
+		return 0, false
+	}
+	mm := s[len("go1."):]
+	var major, rest string
+	for _, sep := range []string{".", "rc", "beta"} {
+		i := strings.Index(mm, sep)
+		if i > 0 {
+			major, rest = mm[:i], mm[i:]
+			break
+		}
+	}
+	if major == "" {
+		major = mm
+	}
+	n, err := strconv.Atoi(major)
+	if err != nil {
+		return 0, false
+	}
+	return n, strings.Contains(rest, "ts")
+}
+
+func TestGoMajorVersion(t *testing.T) {
+	tests := []struct {
+		version string
+		wantN   int
+		wantTS  bool
+	}{
+		{"go1.15.8", 15, false},
+		{"go1.16rc1", 16, false},
+		{"go1.16rc1", 16, false},
+		{"go1.15.5-ts3bd89195a3", 15, true},
+		{"go1.15", 15, false},
+	}
+
+	for _, tt := range tests {
+		n, ts := goMajorVersion(tt.version)
+		if tt.wantN != n || tt.wantTS != ts {
+			t.Errorf("goMajorVersion(%s) = %v, %v, want %v, %v", tt.version, n, ts, tt.wantN, tt.wantTS)
+		}
+	}
+}
+
+func TestReceiveFromAllocs(t *testing.T) {
+	// Go 1.16 and before: allow 3 allocs.
+	// Go Tailscale fork, Go 1.17+: only allow 2 allocs.
+	major, ts := goMajorVersion(runtime.Version())
+	maxAllocs := 3
+	if major >= 17 || ts {
+		maxAllocs = 2
+	}
+	t.Logf("allowing %d allocs for Go version %q", maxAllocs, runtime.Version())
+	roundTrip := setUpReceiveFrom(t)
+	avg := int(testing.AllocsPerRun(100, roundTrip))
+	if avg > maxAllocs {
+		t.Fatalf("expected %d allocs in ReceiveFrom, got %v", maxAllocs, avg)
+	}
+}
+
+func BenchmarkReceiveFrom(b *testing.B) {
+	roundTrip := setUpReceiveFrom(b)
+	for i := 0; i < b.N; i++ {
+		roundTrip()
+	}
+}
+
 func BenchmarkReceiveFrom_Native(b *testing.B) {
 	recvConn, err := net.ListenPacket("udp4", "127.0.0.1:0")
 	if err != nil {