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_test.go

@@ -107,20 +107,16 @@ func TestPopExpired(t *testing.T) {
 	}
 
 	nn := assignments.popExpired(clock.Now())
-	want := &addrs{} // invalid addr
-	if diff := doDiff(want, nn); diff != "" {
+	if diff := doDiff(nil, 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")
+	}
+}