[tailscale1.16] net: use mid-stack inlining with ReadFromUDP to avoid an allocation

This commit rewrites ReadFromUDP to be mid-stack inlined
and pass a UDPAddr for lower layers to fill in.

This lets performance-sensitive clients avoid an allocation.
It requires some care on their part to prevent the UDPAddr
from escaping, but it is now possible.
The UDPAddr trivially does not escape in the benchmark,
as it is immediately discarded.

name                  old time/op    new time/op    delta
WriteToReadFromUDP-8    17.2µs ± 6%    17.1µs ± 5%     ~     (p=0.387 n=9+9)

name                  old alloc/op   new alloc/op   delta
WriteToReadFromUDP-8      112B ± 0%       64B ± 0%  -42.86%  (p=0.000 n=10+10)

name                  old allocs/op  new allocs/op  delta
WriteToReadFromUDP-8      3.00 ± 0%      2.00 ± 0%  -33.33%  (p=0.000 n=10+10)

Updates golang#43451

Co-authored-by: Filippo Valsorda <[email protected]>
Change-Id: I1f9d2ab66bd7e4eff07fe39000cfa0b45717bd13

Author: 2 people

src/cmd/compile/internal/gc/inl_test.go

@@ -16,7 +16,7 @@ import (
 	"testing"
 )
 
-// TestIntendedInlining tests that specific runtime functions are inlined.
+// TestIntendedInlining tests that specific functions are inlined.
 // This allows refactoring for code clarity and re-use without fear that
 // changes to the compiler will cause silent performance regressions.
 func TestIntendedInlining(t *testing.T) {
@@ -155,6 +155,9 @@ func TestIntendedInlining(t *testing.T) {
 			"(*rngSource).Int63",
 			"(*rngSource).Uint64",
 		},
+		"net": {
+			"(*UDPConn).ReadFromUDP",
+		},
 	}
 
 	if runtime.GOARCH != "386" && runtime.GOARCH != "mips64" && runtime.GOARCH != "mips64le" && runtime.GOARCH != "riscv64" {

src/net/udpsock.go

@@ -99,11 +99,20 @@ func (c *UDPConn) SyscallConn() (syscall.RawConn, error) {
 }
 
 // ReadFromUDP acts like ReadFrom but returns a UDPAddr.
-func (c *UDPConn) ReadFromUDP(b []byte) (int, *UDPAddr, error) {
+func (c *UDPConn) ReadFromUDP(b []byte) (n int, addr *UDPAddr, err error) {
+	// This function is designed to allow the caller to control the lifetime
+	// of the returned *UDPAddr and thereby prevent an allocation.
+	// See https://blog.filippo.io/efficient-go-apis-with-the-inliner/.
+	// The real work is done by readFromUDP, below.
+	return c.readFromUDP(b, &UDPAddr{})
+}
+
+// readFromUDP implements ReadFromUDP.
+func (c *UDPConn) readFromUDP(b []byte, addr *UDPAddr) (int, *UDPAddr, error) {
 	if !c.ok() {
 		return 0, nil, syscall.EINVAL
 	}
-	n, addr, err := c.readFrom(b)
+	n, addr, err := c.readFrom(b, addr)
 	if err != nil {
 		err = &OpError{Op: "read", Net: c.fd.net, Source: c.fd.laddr, Addr: c.fd.raddr, Err: err}
 	}
@@ -115,7 +124,7 @@ func (c *UDPConn) ReadFrom(b []byte) (int, Addr, error) {
 	if !c.ok() {
 		return 0, nil, syscall.EINVAL
 	}
-	n, addr, err := c.readFrom(b)
+	n, addr, err := c.readFrom(b, &UDPAddr{})
 	if err != nil {
 		err = &OpError{Op: "read", Net: c.fd.net, Source: c.fd.laddr, Addr: c.fd.raddr, Err: err}
 	}

src/net/udpsock_plan9.go

@@ -11,7 +11,7 @@ import (
 	"syscall"
 )
 
-func (c *UDPConn) readFrom(b []byte) (n int, addr *UDPAddr, err error) {
+func (c *UDPConn) readFrom(b []byte, addr *UDPAddr) (int, *UDPAddr, error) {
 	buf := make([]byte, udpHeaderSize+len(b))
 	m, err := c.fd.Read(buf)
 	if err != nil {
@@ -23,8 +23,9 @@ func (c *UDPConn) readFrom(b []byte) (n int, addr *UDPAddr, err error) {
 	buf = buf[:m]
 
 	h, buf := unmarshalUDPHeader(buf)
-	n = copy(b, buf)
-	return n, &UDPAddr{IP: h.raddr, Port: int(h.rport)}, nil
+	n := copy(b, buf)
+	*addr = UDPAddr{IP: h.raddr, Port: int(h.rport)}
+	return n, addr, nil
 }
 
 func (c *UDPConn) readMsg(b, oob []byte) (n, oobn, flags int, addr *UDPAddr, err error) {

src/net/udpsock_posix.go

@@ -42,14 +42,13 @@ func (a *UDPAddr) toLocal(net string) sockaddr {
 	return &UDPAddr{loopbackIP(net), a.Port, a.Zone}
 }
 
-func (c *UDPConn) readFrom(b []byte) (int, *UDPAddr, error) {
-	var addr *UDPAddr
+func (c *UDPConn) readFrom(b []byte, addr *UDPAddr) (int, *UDPAddr, error) {
 	n, sa, err := c.fd.readFrom(b)
 	switch sa := sa.(type) {
 	case *syscall.SockaddrInet4:
-		addr = &UDPAddr{IP: sa.Addr[0:], Port: sa.Port}
+		*addr = UDPAddr{IP: sa.Addr[0:], Port: sa.Port}
 	case *syscall.SockaddrInet6:
-		addr = &UDPAddr{IP: sa.Addr[0:], Port: sa.Port, Zone: zoneCache.name(int(sa.ZoneId))}
+		*addr = UDPAddr{IP: sa.Addr[0:], Port: sa.Port, Zone: zoneCache.name(int(sa.ZoneId))}
 	}
 	return n, addr, err
 }

src/net/udpsock_test.go

@@ -444,3 +444,24 @@ func TestUDPReadSizeError(t *testing.T) {
 		}
 	}
 }
+
+func BenchmarkWriteToReadFromUDP(b *testing.B) {
+	conn, err := ListenUDP("udp4", new(UDPAddr))
+	if err != nil {
+		b.Fatal(err)
+	}
+	addr := conn.LocalAddr()
+	buf := make([]byte, 8)
+	b.ResetTimer()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		_, err := conn.WriteTo(buf, addr)
+		if err != nil {
+			b.Fatal(err)
+		}
+		_, _, err = conn.ReadFromUDP(buf)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}