From ef5c587f782d944005971c7ba5da88405f9b000b Mon Sep 17 00:00:00 2001
From: Josh Bleecher Snyder <josh@tailscale.com>
Date: Tue, 22 Mar 2022 11:23:56 -0700
Subject: [PATCH] conn: remove the final alloc per packet receive
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This does bind_std only; other platforms remain.

The remaining alloc per iteration in the Throughput benchmark
comes from the tuntest package, and should not appear in regular use.

name           old time/op      new time/op      delta
Latency-10         25.2µs ± 1%      25.0µs ± 0%   -0.58%  (p=0.006 n=10+10)
Throughput-10      2.44µs ± 3%      2.41µs ± 2%     ~     (p=0.140 n=10+8)

name           old alloc/op     new alloc/op     delta
Latency-10           854B ± 5%        741B ± 3%  -13.22%  (p=0.000 n=10+10)
Throughput-10        265B ±34%        267B ±39%     ~     (p=0.670 n=10+10)

name           old allocs/op    new allocs/op    delta
Latency-10           16.0 ± 0%        14.0 ± 0%  -12.50%  (p=0.000 n=10+10)
Throughput-10        2.00 ± 0%        1.00 ± 0%  -50.00%  (p=0.000 n=10+10)

name           old packet-loss  new packet-loss  delta
Throughput-10        0.01 ±82%       0.01 ±282%     ~     (p=0.321 n=9+8)

Signed-off-by: Josh Bleecher Snyder <josh@tailscale.com>
Signed-off-by: Jason A. Donenfeld <Jason@zx2c4.com>
---
 conn/bind_std.go | 53 +++++++++++++++++++++++++++++++++---------------
 1 file changed, 37 insertions(+), 16 deletions(-)

diff --git a/conn/bind_std.go b/conn/bind_std.go
index e0f6cdd..b6a7ab3 100644
--- a/conn/bind_std.go
+++ b/conn/bind_std.go
@@ -31,34 +31,34 @@ type StdNetEndpoint netip.AddrPort
 
 var (
 	_ Bind     = (*StdNetBind)(nil)
-	_ Endpoint = (*StdNetEndpoint)(nil)
+	_ Endpoint = StdNetEndpoint{}
 )
 
 func (*StdNetBind) ParseEndpoint(s string) (Endpoint, error) {
 	e, err := netip.ParseAddrPort(s)
-	return (*StdNetEndpoint)(&e), err
+	return asEndpoint(e), err
 }
 
-func (*StdNetEndpoint) ClearSrc() {}
+func (StdNetEndpoint) ClearSrc() {}
 
-func (e *StdNetEndpoint) DstIP() netip.Addr {
-	return (*netip.AddrPort)(e).Addr()
+func (e StdNetEndpoint) DstIP() netip.Addr {
+	return (netip.AddrPort)(e).Addr()
 }
 
-func (e *StdNetEndpoint) SrcIP() netip.Addr {
+func (e StdNetEndpoint) SrcIP() netip.Addr {
 	return netip.Addr{} // not supported
 }
 
-func (e *StdNetEndpoint) DstToBytes() []byte {
-	b, _ := (*netip.AddrPort)(e).MarshalBinary()
+func (e StdNetEndpoint) DstToBytes() []byte {
+	b, _ := (netip.AddrPort)(e).MarshalBinary()
 	return b
 }
 
-func (e *StdNetEndpoint) DstToString() string {
-	return (*netip.AddrPort)(e).String()
+func (e StdNetEndpoint) DstToString() string {
+	return (netip.AddrPort)(e).String()
 }
 
-func (e *StdNetEndpoint) SrcToString() string {
+func (e StdNetEndpoint) SrcToString() string {
 	return ""
 }
 
@@ -152,24 +152,24 @@ func (bind *StdNetBind) Close() error {
 func (*StdNetBind) makeReceiveIPv4(conn *net.UDPConn) ReceiveFunc {
 	return func(buff []byte) (int, Endpoint, error) {
 		n, endpoint, err := conn.ReadFromUDPAddrPort(buff)
-		return n, (*StdNetEndpoint)(&endpoint), err
+		return n, asEndpoint(endpoint), err
 	}
 }
 
 func (*StdNetBind) makeReceiveIPv6(conn *net.UDPConn) ReceiveFunc {
 	return func(buff []byte) (int, Endpoint, error) {
 		n, endpoint, err := conn.ReadFromUDPAddrPort(buff)
-		return n, (*StdNetEndpoint)(&endpoint), err
+		return n, asEndpoint(endpoint), err
 	}
 }
 
 func (bind *StdNetBind) Send(buff []byte, endpoint Endpoint) error {
 	var err error
-	nend, ok := endpoint.(*StdNetEndpoint)
+	nend, ok := endpoint.(StdNetEndpoint)
 	if !ok {
 		return ErrWrongEndpointType
 	}
-	addrPort := (*netip.AddrPort)(nend)
+	addrPort := netip.AddrPort(nend)
 
 	bind.mu.Lock()
 	blackhole := bind.blackhole4
@@ -186,6 +186,27 @@ func (bind *StdNetBind) Send(buff []byte, endpoint Endpoint) error {
 	if conn == nil {
 		return syscall.EAFNOSUPPORT
 	}
-	_, err = conn.WriteToUDPAddrPort(buff, *addrPort)
+	_, err = conn.WriteToUDPAddrPort(buff, addrPort)
 	return err
 }
+
+// endpointPool contains a re-usable set of mapping from netip.AddrPort to Endpoint.
+// This exists to reduce allocations: Putting a netip.AddrPort in an Endpoint allocates,
+// but Endpoints are immutable, so we can re-use them.
+var endpointPool = sync.Pool{
+	New: func() any {
+		return make(map[netip.AddrPort]Endpoint)
+	},
+}
+
+// asEndpoint returns an Endpoint containing ap.
+func asEndpoint(ap netip.AddrPort) Endpoint {
+	m := endpointPool.Get().(map[netip.AddrPort]Endpoint)
+	defer endpointPool.Put(m)
+	e, ok := m[ap]
+	if !ok {
+		e = Endpoint(StdNetEndpoint(ap))
+		m[ap] = e
+	}
+	return e
+}