net: add ReadFromUDPAddrPort

It is now possible to do completely allocation-free UDP.

This is implemented completely separately from ReadFromUDP
because ReadFromUDP exists in a delicate balance to allow
mid-stack inlining. After performance-sensitive callers have
migrated to ReadFromUDPAddrPort, we may be able to simplify
ReadFromUDP to call ReadFromUDPAddrPort.

name                          old time/op    new time/op    delta
WriteToReadFromUDPAddrPort-8    4.71µs ± 2%    4.81µs ± 5%    +2.18%  (p=0.000 n=14+14)

name                          old alloc/op   new alloc/op   delta
WriteToReadFromUDPAddrPort-8     4.00B ± 0%     0.00B       -100.00%  (p=0.000 n=15+15)

name                          old allocs/op  new allocs/op  delta
WriteToReadFromUDPAddrPort-8      1.00 ± 0%      0.00       -100.00%  (p=0.000 n=15+15)

Change-Id: I37f5ad9416a1d4333ed48d83474b2cf933b2a1be
Reviewed-on: https://go-review.googlesource.com/c/go/+/360600
Trust: Josh Bleecher Snyder <josharian@gmail.com>
Run-TryBot: Josh Bleecher Snyder <josharian@gmail.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>

src/net/udpsock.go

@@ -167,6 +167,18 @@ func (c *UDPConn) ReadFrom(b []byte) (int, Addr, error) {
 	return n, addr, err
 }
 
+// ReadFromUDPAddrPort acts like ReadFrom but returns a netip.AddrPort.
+func (c *UDPConn) ReadFromUDPAddrPort(b []byte) (n int, addr netip.AddrPort, err error) {
+	if !c.ok() {
+		return 0, netip.AddrPort{}, syscall.EINVAL
+	}
+	n, addr, err = c.readFromAddrPort(b)
+	if err != nil {
+		err = &OpError{Op: "read", Net: c.fd.net, Source: c.fd.laddr, Addr: c.fd.raddr, Err: err}
+	}
+	return n, addr, err
+}
+
 // ReadMsgUDP reads a message from c, copying the payload into b and
 // the associated out-of-band data into oob. It returns the number of
 // bytes copied into b, the number of bytes copied into oob, the flags

src/net/udpsock_plan9.go

@@ -29,6 +29,25 @@ func (c *UDPConn) readFrom(b []byte, addr *UDPAddr) (int, *UDPAddr, error) {
 	return n, addr, nil
 }
 
+func (c *UDPConn) readFromAddrPort(b []byte) (int, netip.AddrPort, error) {
+	// TODO: optimize. The equivalent code on posix is alloc-free.
+	buf := make([]byte, udpHeaderSize+len(b))
+	m, err := c.fd.Read(buf)
+	if err != nil {
+		return 0, netip.AddrPort{}, err
+	}
+	if m < udpHeaderSize {
+		return 0, netip.AddrPort{}, errors.New("short read reading UDP header")
+	}
+	buf = buf[:m]
+
+	h, buf := unmarshalUDPHeader(buf)
+	n := copy(b, buf)
+	ip, _ := netip.AddrFromSlice(h.raddr)
+	addr := netip.AddrPortFrom(ip, h.rport)
+	return n, addr, nil
+}
+
 func (c *UDPConn) readMsg(b, oob []byte) (n, oobn, flags int, addr netip.AddrPort, err error) {
 	return 0, 0, 0, netip.AddrPort{}, syscall.EPLAN9
 }

src/net/udpsock_posix.go

@@ -69,6 +69,31 @@ func (c *UDPConn) readFrom(b []byte, addr *UDPAddr) (int, *UDPAddr, error) {
 	return n, addr, err
 }
 
+func (c *UDPConn) readFromAddrPort(b []byte) (n int, addr netip.AddrPort, err error) {
+	var ip netip.Addr
+	var port int
+	switch c.fd.family {
+	case syscall.AF_INET:
+		var from syscall.SockaddrInet4
+		n, err = c.fd.readFromInet4(b, &from)
+		if err == nil {
+			ip = netip.AddrFrom4(from.Addr)
+			port = from.Port
+		}
+	case syscall.AF_INET6:
+		var from syscall.SockaddrInet6
+		n, err = c.fd.readFromInet6(b, &from)
+		if err == nil {
+			ip = netip.AddrFrom16(from.Addr).WithZone(zoneCache.name(int(from.ZoneId)))
+			port = from.Port
+		}
+	}
+	if err == nil {
+		addr = netip.AddrPortFrom(ip, uint16(port))
+	}
+	return n, addr, err
+}
+
 func (c *UDPConn) readMsg(b, oob []byte) (n, oobn, flags int, addr netip.AddrPort, err error) {
 	var sa syscall.Sockaddr
 	n, oobn, flags, sa, err = c.fd.readMsg(b, oob, 0)

src/net/udpsock_test.go

@@ -533,7 +533,7 @@ func BenchmarkWriteToReadFromUDPAddrPort(b *testing.B) {
 		if err != nil {
 			b.Fatal(err)
 		}
-		_, _, err = conn.ReadFromUDP(buf) // TODO: create and use ReadFromUDPAddrPort
+		_, _, err = conn.ReadFromUDPAddrPort(buf)
 		if err != nil {
 			b.Fatal(err)
 		}