feature/conn25: extend assignment expiry on use

When we use assigned addresses in response to a DNS request, extend the
expiry on the assignment.

Updates tailscale/corp#39975

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

feature/conn25/addrAssignments.go

@@ -34,16 +34,23 @@ type addrAssignments struct {
 
 const defaultExpiry = 48 * time.Hour
 
+func clampExpiryFromTTL(ttl time.Duration) time.Duration {
+	const minTTL = time.Minute * 1
+	const maxTTL = time.Hour * 72
+	expiry := max(minTTL, ttl)
+	return min(maxTTL, expiry)
+}
+
 func (a *addrAssignments) insert(as *addrs) error {
 	return a.insertWithExpiry(as, defaultExpiry)
 }
 
 func (a *addrAssignments) insertFromTTL(as *addrs, ttl time.Duration) error {
-	const minTTL = time.Minute * 1
+	return a.insertWithExpiry(as, clampExpiryFromTTL(ttl))
-	const maxTTL = time.Hour * 72
+}
-	expiry := max(minTTL, ttl)
+
-	expiry = min(maxTTL, expiry)
+func (a *addrAssignments) updateFromTTL(as *addrs, ttl time.Duration) {
-	return a.insertWithExpiry(as, expiry)
+	a.updateExpiry(as, clampExpiryFromTTL(ttl))
 }
 
 func (a *addrAssignments) insertWithExpiry(as *addrs, d time.Duration) error {
@@ -104,10 +111,10 @@ func (a *addrAssignments) lookupByTransitIP(tip netip.Addr) (*addrs, bool) {
 // 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{}
+		return nil
 	}
 	if !a.byExpiresAt.peek().expiresAt.Before(now) {
-		return &addrs{}
+		return nil
 	}
 	v := heap.Pop(&a.byExpiresAt).(*addrs)
 	delete(a.byMagicIP, v.magic)
@@ -117,6 +124,18 @@ func (a *addrAssignments) popExpired(now time.Time) *addrs {
 	return v
 }
 
+func (a *addrAssignments) updateExpiry(as *addrs, expiresIn time.Duration) {
+	now := a.clock.Now()
+	as.expiresAt = now.Add(expiresIn)
+	// TODO(fran) 2026-05-26 We can make this perform better.
+	//  * With a bit of extra effort, we can track the index so that heap.Fix can
+	//    be used.
+	//  * Alternatively, marking the heap dirty and waiting until the next
+	//    operation that requires it to be in the correct order would mean a
+	//    whole slew of updates can accumulate before paying for a heap.Init.
+	heap.Init(&a.byExpiresAt)
+}
+
 type addrsHeap []*addrs
 
 func (h addrsHeap) Len() int           { return len(h) }

feature/conn25/addrAssignments_test.go

@@ -107,20 +107,16 @@ func TestPopExpired(t *testing.T) {
 	}
 
 	nn := assignments.popExpired(clock.Now())
-	want := &addrs{} // invalid addr
+	if diff := doDiff(nil, nn); diff != "" {
-	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]
+	want := testAddrs[0]
 	nn = assignments.popExpired(clock.Now())
 	if diff := doDiff(want, nn); diff != "" {
 		t.Fatal(diff)
@@ -138,12 +134,62 @@ func TestPopExpired(t *testing.T) {
 		t.Fatalf("an assignment should have been removed")
 	}
 
-	want = &addrs{}
 	nn = assignments.popExpired(clock.Now())
-	if diff := doDiff(want, nn); diff != "" {
+	if diff := doDiff(nil, nn); diff != "" {
 		t.Fatal(diff)
 	}
 	if len(assignments.byMagicIP) != 0 {
 		t.Fatalf("there should have been no change")
 	}
 }
+
+func TestPopExpiredHandlesExpiresAtChanges(t *testing.T) {
+	expiryInterval := time.Second * 5
+	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.insertWithExpiry(as, expiryInterval)
+		if err != nil {
+			t.Fatal(err)
+		}
+		return as
+	}
+	addresses := []*addrs{}
+	// t = 0
+	for i := range 10 {
+		addresses = append(addresses, makeAndAddAddrs(i)) // expires at t=i+5
+		// We track the next addr to expire with a heap. updateExpiry changes the heap invariant.
+		// Twiddling the addrs in this particular way (updating item 1 after inserting 7) shows that
+		// we are fixing the heap after updating the invariant (if we weren't the test would fail).
+		if i == 6 {
+			assignments.updateExpiry(addresses[1], 20*time.Second) // addresses[1] expires at t=26
+		}
+		clock.Advance(time.Second)
+	}
+	// t = 10
+
+	expectedOrder := []int{0, 2, 3, 4, 5, 6, 7, 8, 9, 1}
+	i := 0
+	for tick := range 18 {
+		a := assignments.popExpired(clock.Now())
+		if a != nil {
+			expectedIdx := expectedOrder[i]
+			if a != addresses[expectedIdx] {
+				t.Fatalf("want %v, got %v at tick=%v", addresses[expectedIdx].magic, a.magic, tick)
+			}
+			i++
+		}
+		clock.Advance(time.Second)
+	}
+	if len(assignments.byMagicIP) != 0 {
+		t.Fatalf("expected assignments to be exhausted")
+	}
+}

feature/conn25/conn25.go

@@ -755,6 +755,7 @@ func (c *client) reserveAddresses(appName string, domain dnsname.FQDN, dst netip
 	c.mu.Lock()
 	defer c.mu.Unlock()
 	if existing, ok := c.assignments.lookupByDomainDst(domain, dst); ok {
+		c.assignments.updateFromTTL(existing, ttl)
 		return existing, nil
 	}
 
@@ -766,7 +767,7 @@ func (c *client) reserveAddresses(appName string, domain dnsname.FQDN, dst netip
 	now := c.assignments.clock.Now()
 	for range 10 {
 		a := c.assignments.popExpired(now)
-		if !a.isValid() {
+		if a == nil {
 			break
 		}
 		if a.is4() {

feature/conn25/conn25_test.go

@@ -1208,6 +1208,16 @@ func TestMapDNSResponseSetsExpiryBasedOnTTL(t *testing.T) {
 
 	assertExpiresAt(ipThree, clock.Now().Add(301*time.Second))
 	assertExpiresAt(ipFour, clock.Now().Add(61*time.Second))
+
+	elapsed := time.Second * 30
+	clock.Advance(elapsed)
+	assertExpiresAt(ipThree, clock.Now().Add(301*time.Second).Add(-1*elapsed))
+	assertExpiresAt(ipFour, clock.Now().Add(61*time.Second).Add(-1*elapsed))
+	c.mapDNSResponse(dnsRespV6)
+	// after seeing the addresses again, the expiry time is pushed out.
+	assertExpiresAt(ipThree, clock.Now().Add(301*time.Second))
+	assertExpiresAt(ipFour, clock.Now().Add(61*time.Second))
+
 }
 
 func TestNormalizedDNSNames(t *testing.T) {