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]
+}

feature/conn25/addrAssignments_test.go

@@ -4,17 +4,20 @@
 package conn25
 
 import (
+	"fmt"
 	"net/netip"
 	"testing"
 	"time"
 
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
 	"tailscale.com/tstest"
 )
 
 func TestAssignmentsExpire(t *testing.T) {
 	clock := tstest.NewClock(tstest.ClockOpts{Start: time.Now()})
 	assignments := addrAssignments{clock: clock}
-	as := addrs{
+	as := &addrs{
 		dst:     netip.MustParseAddr("0.0.0.1"),
 		magic:   netip.MustParseAddr("0.0.0.2"),
 		transit: netip.MustParseAddr("0.0.0.3"),
@@ -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(clock.Now())
+	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(clock.Now())
+	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(clock.Now())
+	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(clock.Now())
+	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

@@ -748,9 +748,9 @@ func (cfg *config) getAppsForConnectorDomain(domain dnsname.FQDN, prefsAdvertise
 // the app name refers to a configured app.
 // It checks that this domain should be routed and that this client is not itself a connector for the domain
 // and generally if it is valid to make the assignment.
-func (c *client) reserveAddresses(appName string, domain dnsname.FQDN, dst netip.Addr) (addrs, error) {
+func (c *client) reserveAddresses(appName string, domain dnsname.FQDN, dst netip.Addr) (*addrs, error) {
 	if !dst.IsValid() {
-		return addrs{}, errors.New("dst is not valid")
+		return nil, errors.New("dst is not valid")
 	}
 	c.mu.Lock()
 	defer c.mu.Unlock()
@@ -758,30 +758,52 @@ func (c *client) reserveAddresses(appName string, domain dnsname.FQDN, dst netip
 		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.
+	now := c.assignments.clock.Now()
+	for range 10 {
+		a := c.assignments.popExpired(now)
+		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)
+		} else {
+			return nil, errors.New("unexpected neither 4 nor 6")
+		}
+	}
+
 	var mip, tip netip.Addr
 	var err error
 	if dst.Is4() {
 		mip, err = c.v4MagicIPPool.next()
 		if err != nil {
-			return addrs{}, err
+			return nil, err
 		}
 		tip, err = c.v4TransitIPPool.next()
 		if err != nil {
-			return addrs{}, err
+			return nil, err
 		}
 	} else if dst.Is6() {
 		mip, err = c.v6MagicIPPool.next()
 		if err != nil {
-			return addrs{}, err
+			return nil, err
 		}
 		tip, err = c.v6TransitIPPool.next()
 		if err != nil {
-			return addrs{}, err
+			return nil, err
 		}
 	} else {
-		return addrs{}, errors.New("unexpected neither 4 nor 6")
+		return nil, errors.New("unexpected neither 4 nor 6")
 	}
-	as := addrs{
+	as := &addrs{
 		dst:     dst,
 		magic:   mip,
 		transit: tip,
@@ -789,11 +811,11 @@ func (c *client) reserveAddresses(appName string, domain dnsname.FQDN, dst netip
 		domain:  domain,
 	}
 	if err := c.assignments.insert(as); err != nil {
-		return addrs{}, err
+		return nil, err
 	}
 	err = c.enqueueAddressAssignment(as)
 	if err != nil {
-		return addrs{}, err
+		return nil, err
 	}
 	return as, nil
 }
@@ -835,11 +857,11 @@ func (e *extension) handleAddressAssignment(ctx context.Context, as addrs) error
 	return nil
 }
 
-func (c *client) enqueueAddressAssignment(addrs addrs) error {
+func (c *client) enqueueAddressAssignment(addrs *addrs) error {
 	select {
 	// TODO(fran) investigate the value of waiting for multiple addresses and sending them
 	// in one ConnectorTransitIPRequest
-	case c.addrsCh <- addrs:
+	case c.addrsCh <- *addrs:
 		return nil
 	default:
 		c.logf("address assignment queue full, dropping transit assignment for %v", addrs.domain)
@@ -1243,8 +1265,16 @@ type addrs struct {
 	expiresAt time.Time
 }
 
-func (c addrs) isValid() bool {
-	return c.dst.IsValid()
+func (as addrs) isValid() bool {
+	return as.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

feature/conn25/conn25_test.go

@@ -914,7 +914,7 @@ func TestMapDNSResponseAssignsAddrs(t *testing.T) {
 		v6Addrs             []*dnsmessage.AAAAResource
 		selfTags            []string
 		isEligibleConnector bool
-		wantByMagicIP       map[netip.Addr]addrs
+		wantByMagicIP       map[netip.Addr]*addrs
 	}{
 		{
 			name:       "one-ip-matches",
@@ -922,7 +922,7 @@ func TestMapDNSResponseAssignsAddrs(t *testing.T) {
 			domain:     "example.com.",
 			v4Addrs:    []*dnsmessage.AResource{{A: [4]byte{1, 0, 0, 0}}},
 			// these are 'expected' because they are the beginning of the provided pools
-			wantByMagicIP: map[netip.Addr]addrs{
+			wantByMagicIP: map[netip.Addr]*addrs{
 				netip.MustParseAddr("100.64.0.0"): {
 					domain:  "example.com.",
 					dst:     netip.MustParseAddr("1.0.0.0"),
@@ -940,7 +940,7 @@ func TestMapDNSResponseAssignsAddrs(t *testing.T) {
 				{AAAA: [16]byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1}},
 				{AAAA: [16]byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2}},
 			},
-			wantByMagicIP: map[netip.Addr]addrs{
+			wantByMagicIP: map[netip.Addr]*addrs{
 				netip.MustParseAddr("fd7a:115c:a1e0:a99c::"): {
 					domain:  "example.com.",
 					dst:     netip.MustParseAddr("::1"),
@@ -965,7 +965,7 @@ func TestMapDNSResponseAssignsAddrs(t *testing.T) {
 				{A: [4]byte{1, 0, 0, 0}},
 				{A: [4]byte{2, 0, 0, 0}},
 			},
-			wantByMagicIP: map[netip.Addr]addrs{
+			wantByMagicIP: map[netip.Addr]*addrs{
 				netip.MustParseAddr("100.64.0.0"): {
 					domain:  "example.com.",
 					dst:     netip.MustParseAddr("1.0.0.0"),
@@ -1007,7 +1007,7 @@ func TestMapDNSResponseAssignsAddrs(t *testing.T) {
 			selfTags:   []string{"tag:woo"},
 			// isEligibleConnector is false: tag matches but prefs not set,
 			// so DNS response should be rewritten normally.
-			wantByMagicIP: map[netip.Addr]addrs{
+			wantByMagicIP: map[netip.Addr]*addrs{
 				netip.MustParseAddr("100.64.0.0"): {
 					domain:  "example.com.",
 					dst:     netip.MustParseAddr("1.0.0.0"),
@@ -1026,7 +1026,7 @@ func TestMapDNSResponseAssignsAddrs(t *testing.T) {
 			isEligibleConnector: true,
 			// isEligibleConnector is true but tag doesn't match the app,
 			// so DNS response should be rewritten normally.
-			wantByMagicIP: map[netip.Addr]addrs{
+			wantByMagicIP: map[netip.Addr]*addrs{
 				netip.MustParseAddr("100.64.0.0"): {
 					domain:  "example.com.",
 					dst:     netip.MustParseAddr("1.0.0.0"),
@@ -1042,7 +1042,7 @@ func TestMapDNSResponseAssignsAddrs(t *testing.T) {
 			domain:     "sub.example.com.",
 			v4Addrs:    []*dnsmessage.AResource{{A: [4]byte{1, 0, 0, 0}}},
 			// these are 'expected' because they are the beginning of the provided pools
-			wantByMagicIP: map[netip.Addr]addrs{
+			wantByMagicIP: map[netip.Addr]*addrs{
 				netip.MustParseAddr("100.64.0.0"): {
 					domain:  "sub.example.com.",
 					dst:     netip.MustParseAddr("1.0.0.0"),
@@ -1058,7 +1058,7 @@ func TestMapDNSResponseAssignsAddrs(t *testing.T) {
 			domain:     "sub.example.com.",
 			v4Addrs:    []*dnsmessage.AResource{{A: [4]byte{1, 0, 0, 0}}},
 			// these are 'expected' because they are the beginning of the provided pools
-			wantByMagicIP: map[netip.Addr]addrs{
+			wantByMagicIP: map[netip.Addr]*addrs{
 				netip.MustParseAddr("100.64.0.0"): {
 					domain:  "sub.example.com.",
 					dst:     netip.MustParseAddr("1.0.0.0"),
@@ -1074,7 +1074,7 @@ func TestMapDNSResponseAssignsAddrs(t *testing.T) {
 			domain:     "a.sub.example.com.",
 			v4Addrs:    []*dnsmessage.AResource{{A: [4]byte{1, 0, 0, 0}}},
 			// these are 'expected' because they are the beginning of the provided pools
-			wantByMagicIP: map[netip.Addr]addrs{
+			wantByMagicIP: map[netip.Addr]*addrs{
 				netip.MustParseAddr("100.64.0.0"): {
 					domain:  "a.sub.example.com.",
 					dst:     netip.MustParseAddr("1.0.0.0"),
@@ -1177,9 +1177,11 @@ func TestReserveAddressesDeduplicated(t *testing.T) {
 				t.Fatal(err)
 			}
 
-			if first.magic != second.magic {
-				t.Errorf("expected same magic addrs on repeated call, got first=%v second=%v", first.magic, second.magic)
+			if first != second {
+				// reserveAddresses should return the existing entry when called for a domain that already has assigned addrs
+				t.Fatalf("want first==second, got first: %v, second: %v", first, second)
 			}
+
 			if got := len(c.assignments.byMagicIP); got != 1 {
 				t.Errorf("want 1 entry in byMagicIP, got %d", got)
 			}
@@ -1314,7 +1316,7 @@ func TestAddressAssignmentIsHandled(t *testing.T) {
 	cfg := mustConfig(t, sn)
 	ext.conn25.reconfig(cfg)
 
-	as := addrs{
+	as := &addrs{
 		dst:     netip.MustParseAddr("1.2.3.4"),
 		magic:   netip.MustParseAddr("100.64.0.0"),
 		transit: netip.MustParseAddr("169.254.0.1"),
@@ -1930,12 +1932,12 @@ func TestHandleAddressAssignmentStoresTransitIPs(t *testing.T) {
 	// and then does the lookups.
 	steps := []struct {
 		name    string
-		as      addrs
+		as      *addrs
 		lookups []lookup
 	}{
 		{
 			name: "step-1-conn1-tip1",
-			as: addrs{
+			as: &addrs{
 				dst:     netip.MustParseAddr("1.2.3.1"),
 				magic:   netip.MustParseAddr("100.64.0.1"),
 				transit: transitIPs[0].Addr(),
@@ -1959,7 +1961,7 @@ func TestHandleAddressAssignmentStoresTransitIPs(t *testing.T) {
 		},
 		{
 			name: "step-2-conn1-tip2",
-			as: addrs{
+			as: &addrs{
 				dst:     netip.MustParseAddr("1.2.3.2"),
 				magic:   netip.MustParseAddr("100.64.0.2"),
 				transit: transitIPs[1].Addr(),
@@ -1979,7 +1981,7 @@ func TestHandleAddressAssignmentStoresTransitIPs(t *testing.T) {
 		},
 		{
 			name: "step-3-conn2-tip1",
-			as: addrs{
+			as: &addrs{
 				dst:     netip.MustParseAddr("1.2.3.3"),
 				magic:   netip.MustParseAddr("100.64.0.3"),
 				transit: transitIPs[2].Addr(),
@@ -2040,7 +2042,7 @@ func TestHandleAddressAssignmentStoresTransitIPs(t *testing.T) {
 func TestTransitIPConnMapping(t *testing.T) {
 	conn25 := newConn25(t.Logf)
 
-	as := addrs{
+	as := &addrs{
 		dst:     netip.MustParseAddr("1.2.3.1"),
 		magic:   netip.MustParseAddr("100.64.0.1"),
 		transit: netip.MustParseAddr("169.254.0.1"),
@@ -2151,14 +2153,14 @@ func TestClientTransitIPForMagicIP(t *testing.T) {
 			c := newConn25(t.Logf)
 			c.reconfig(cfg)
 
-			if err := c.client.assignments.insert(addrs{
+			if err := c.client.assignments.insert(&addrs{
 				magic:   mappedMip,
 				transit: mappedTip,
 				dst:     dst,
 			}); err != nil {
 				t.Fatal(err)
 			}
-			if err := c.client.assignments.insert(addrs{
+			if err := c.client.assignments.insert(&addrs{
 				magic:   v6MappedMip,
 				transit: v6MappedTip,
 				dst:     v6Dst,
@@ -2253,7 +2255,7 @@ func TestIsKnownTransitIP(t *testing.T) {
 	unknownTip := netip.MustParseAddr("100.64.0.42")
 
 	c := newConn25(t.Logf)
-	c.client.assignments.insert(addrs{
+	c.client.assignments.insert(&addrs{
 		transit: knownTip,
 	})
 
@@ -2269,7 +2271,7 @@ func TestLinkLocalAllow(t *testing.T) {
 	knownTip := netip.MustParseAddr("100.64.0.41")
 
 	c := newConn25(t.Logf)
-	c.client.assignments.insert(addrs{
+	c.client.assignments.insert(&addrs{
 		transit: knownTip,
 	})