feature/conn25: return expired assignments to address pools

Make it possible to remove the least recently used expired address
assignment from addrAssignments.
Before checking out a new address from the IP pools, return a handful of
expired addresses.

Updates tailscale/corp#39975

Signed-off-by: Fran Bull <fran@tailscale.com>

feature/conn25/addrAssignments.go

@@ -4,6 +4,7 @@
 package conn25
 
 import (
+	"container/heap"
 	"errors"
 	"net/netip"
 	"time"
@@ -23,27 +24,25 @@ type domainDst struct {
 // addrAssignments is the collection of addrs assigned by this client
 // supporting lookup by magic IP, transit IP or domain+dst, or to lookup all
 // transit IPs associated with a given connector (identified by its node key).
-// byConnKey stores netip.Prefix versions of the transit IPs for use in the
-// WireGuard hooks.
 type addrAssignments struct {
-	byMagicIP   map[netip.Addr]addrs
-	byTransitIP map[netip.Addr]addrs
-	byDomainDst map[domainDst]addrs
+	byMagicIP   map[netip.Addr]*addrs
+	byTransitIP map[netip.Addr]*addrs
+	byDomainDst map[domainDst]*addrs
+	byExpiresAt addrsHeap
 	clock       tstime.Clock
 }
 
 const defaultExpiry = 48 * time.Hour
 
-func (a *addrAssignments) insert(as addrs) error {
+func (a *addrAssignments) insert(as *addrs) error {
 	return a.insertWithExpiry(as, defaultExpiry)
 }
 
-func (a *addrAssignments) insertWithExpiry(as addrs, d time.Duration) error {
-	if !as.expiresAt.IsZero() {
+func (a *addrAssignments) insertWithExpiry(as *addrs, d time.Duration) error {
+	now := a.clock.Now()
+	if !as.expiresAt.IsZero() && !as.expiresAt.Before(now) {
 		return errors.New("expiresAt already set")
 	}
-	now := a.clock.Now()
-	as.expiresAt = now.Add(d)
 	// we don't expect for addresses to be reused before expiry
 	if existing, ok := a.byMagicIP[as.magic]; ok {
 		if !existing.expiresAt.Before(now) {
@@ -61,32 +60,74 @@ func (a *addrAssignments) insertWithExpiry(as addrs, d time.Duration) error {
 			return errors.New("byTransitIP key exists")
 		}
 	}
+	as.expiresAt = now.Add(d)
 	mak.Set(&a.byMagicIP, as.magic, as)
 	mak.Set(&a.byTransitIP, as.transit, as)
 	mak.Set(&a.byDomainDst, ddst, as)
+	heap.Push(&a.byExpiresAt, as)
 	return nil
 }
 
-func (a *addrAssignments) lookupByDomainDst(domain dnsname.FQDN, dst netip.Addr) (addrs, bool) {
+func (a *addrAssignments) lookupByDomainDst(domain dnsname.FQDN, dst netip.Addr) (*addrs, bool) {
 	v, ok := a.byDomainDst[domainDst{domain: domain, dst: dst}]
 	if !ok || v.expiresAt.Before(a.clock.Now()) {
-		return addrs{}, false
+		return &addrs{}, false
 	}
 	return v, true
 }
 
-func (a *addrAssignments) lookupByMagicIP(mip netip.Addr) (addrs, bool) {
+func (a *addrAssignments) lookupByMagicIP(mip netip.Addr) (*addrs, bool) {
 	v, ok := a.byMagicIP[mip]
 	if !ok || v.expiresAt.Before(a.clock.Now()) {
-		return addrs{}, false
+		return &addrs{}, false
 	}
 	return v, true
 }
 
-func (a *addrAssignments) lookupByTransitIP(tip netip.Addr) (addrs, bool) {
+func (a *addrAssignments) lookupByTransitIP(tip netip.Addr) (*addrs, bool) {
 	v, ok := a.byTransitIP[tip]
 	if !ok || v.expiresAt.Before(a.clock.Now()) {
-		return addrs{}, false
+		return &addrs{}, false
 	}
 	return v, true
 }
+
+// popExpired returns the member of addrAssignments that expired earliest,
+// or an invalid addrs if there are no expired members of addrAssignments.
+func (a *addrAssignments) popExpired(now time.Time) *addrs {
+	if a.byExpiresAt.Len() == 0 {
+		return &addrs{}
+	}
+	if !a.byExpiresAt.peek().expiresAt.Before(now) {
+		return &addrs{}
+	}
+	v := heap.Pop(&a.byExpiresAt).(*addrs)
+	delete(a.byMagicIP, v.magic)
+	delete(a.byTransitIP, v.transit)
+	dd := domainDst{domain: v.domain, dst: v.dst}
+	delete(a.byDomainDst, dd)
+	return v
+}
+
+type addrsHeap []*addrs
+
+func (h addrsHeap) Len() int           { return len(h) }
+func (h addrsHeap) Less(i, j int) bool { return h[i].expiresAt.Before(h[j].expiresAt) }
+func (h addrsHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
+func (h *addrsHeap) Push(x any) {
+	as, ok := x.(*addrs)
+	if !ok {
+		panic("unexpected not an addrs")
+	}
+	*h = append(*h, as)
+}
+func (h *addrsHeap) Pop() any {
+	old := *h
+	n := len(old)
+	x := old[n-1]
+	*h = old[0 : n-1]
+	return x
+}
+func (h addrsHeap) peek() *addrs {
+	return (h)[0]
+}