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"
@@ -29,6 +30,7 @@ type addrAssignments struct {
 	byMagicIP   map[netip.Addr]addrs
 	byTransitIP map[netip.Addr]addrs
 	byDomainDst map[domainDst]addrs
+	byExpiresAt addrsHeap
 	clock       tstime.Clock
 }
 
@@ -64,6 +66,7 @@ func (a *addrAssignments) insertWithExpiry(as addrs, d time.Duration) error {
 	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
 }
 
@@ -90,3 +93,37 @@ func (a *addrAssignments) lookupByTransitIP(tip netip.Addr) (addrs, bool) {
 	}
 	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() addrs {
+	if a.byExpiresAt.Len() == 0 {
+		return addrs{}
+	}
+	if !a.byExpiresAt.peek().expiresAt.Before(a.clock.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)        { *h = append(*h, x.(addrs)) }
+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]
+}

feature/conn25/addrAssignments_test.go

@@ -4,10 +4,13 @@
 package conn25
 
 import (
+	"fmt"
 	"net/netip"
 	"testing"
 	"time"
 
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
 	"tailscale.com/tstest"
 )
 
@@ -64,3 +67,83 @@ func TestAssignmentsExpire(t *testing.T) {
 		t.Fatalf("expected foundAs to expire after now")
 	}
 }
+
+func TestPopExpired(t *testing.T) {
+	clock := tstest.NewClock(tstest.ClockOpts{Start: time.Now()})
+	assignments := addrAssignments{clock: clock}
+	makeAndAddAddrs := func(n int) addrs {
+		t.Helper()
+		as := addrs{
+			dst:     netip.MustParseAddr(fmt.Sprintf("0.0.1.%d", n)),
+			magic:   netip.MustParseAddr(fmt.Sprintf("0.0.2.%d", n)),
+			transit: netip.MustParseAddr(fmt.Sprintf("0.0.3.%d", n)),
+			app:     "a",
+			domain:  "example.com.",
+		}
+		err := assignments.insert(as)
+		if err != nil {
+			t.Fatal(err)
+		}
+		return as
+	}
+	// cmp.Diff addrs ignoring expiresAt
+	doDiff := func(want, got addrs) string {
+		t.Helper()
+		return cmp.Diff(
+			want,
+			got,
+			cmp.AllowUnexported(addrs{}),
+			cmpopts.EquateComparable(netip.Addr{}),
+			cmpopts.IgnoreFields(addrs{}, "expiresAt"),
+		)
+	}
+	testAddrs := []addrs{}
+	for i := range 2 {
+		testAddrs = append(testAddrs, makeAndAddAddrs(i+1))
+		clock.Advance(1 * time.Second)
+	}
+	if len(assignments.byMagicIP) != 2 {
+		t.Fatalf("test setup wrong")
+	}
+
+	nn := assignments.popExpired()
+	want := addrs{} // invalid addr
+	if diff := doDiff(want, nn); diff != "" {
+		t.Fatalf("only expired addresses are removed: %s", diff)
+	}
+	if len(assignments.byMagicIP) != 2 {
+		t.Fatalf("nothing should have been removed")
+	}
+	if nn.isValid() {
+		t.Fatal("empty addrs should be invalid")
+	}
+
+	clock.Advance(2 * defaultExpiry) // all addrs are now expired
+
+	want = testAddrs[0]
+	nn = assignments.popExpired()
+	if diff := doDiff(want, nn); diff != "" {
+		t.Fatal(diff)
+	}
+	if len(assignments.byMagicIP) != 1 {
+		t.Fatalf("an assignment should have been removed")
+	}
+
+	want = testAddrs[1]
+	nn = assignments.popExpired()
+	if diff := doDiff(want, nn); diff != "" {
+		t.Fatal(diff)
+	}
+	if len(assignments.byMagicIP) != 0 {
+		t.Fatalf("an assignment should have been removed")
+	}
+
+	want = addrs{}
+	nn = assignments.popExpired()
+	if diff := doDiff(want, nn); diff != "" {
+		t.Fatal(diff)
+	}
+	if len(assignments.byMagicIP) != 0 {
+		t.Fatalf("there should have been no change")
+	}
+}

feature/conn25/conn25.go

@@ -729,6 +729,25 @@ func (c *client) reserveAddresses(app string, domain dnsname.FQDN, dst netip.Add
 		return existing, nil
 	}
 
+	// Before we check out more addresses from the pools try to return some.
+	// Trying to return any number greater than 1 will cause the number of
+	// addresses used to trend down in general. But as we have 2 different
+	// pools for the different IP versions, use a number a bit higher than
+	// 2 to try and process bursty behavior faster.
+	for range 10 {
+		a := c.assignments.popExpired()
+		if !a.isValid() {
+			break
+		}
+		if a.is4() {
+			c.v4MagicIPPool.returnAddr(a.magic)
+			c.v4TransitIPPool.returnAddr(a.transit)
+		} else if a.is6() {
+			c.v6MagicIPPool.returnAddr(a.magic)
+			c.v6TransitIPPool.returnAddr(a.transit)
+		}
+	}
+
 	var mip, tip netip.Addr
 	var err error
 	if dst.Is4() {
@@ -1182,6 +1201,14 @@ func (c addrs) isValid() bool {
 	return c.dst.IsValid()
 }
 
+func (as addrs) is4() bool {
+	return as.dst.Is4()
+}
+
+func (as addrs) is6() bool {
+	return as.dst.Is6()
+}
+
 // insertTransitConnMapping adds an entry to the byConnKey map
 // for the provided transitIP (as a prefix).
 // The provided transitIP must already be present in the byTransitIP map.