tka: mark young AUMs as active even if the chain is long

Existing compaction logic seems to have had an assumption that
markActiveChain would cover a longer part of the chain than
markYoungAUMs. This prevented long, but fresh, chains, from being
compacted correctly.

Updates tailscale/corp#33537

Signed-off-by: Anton Tolchanov <anton@tailscale.com>

tka/tailchonk.go

@@ -668,7 +668,7 @@ const (
 )
 
 // markActiveChain marks AUMs in the active chain.
-// All AUMs that are within minChain ancestors of head are
+// All AUMs that are within minChain ancestors of head, or are marked as young, are
 // marked retainStateActive, and all remaining ancestors are
 // marked retainStateCandidate.
 //
@@ -700,19 +700,22 @@ func markActiveChain(storage Chonk, verdict map[AUMHash]retainState, minChain in
 
 	// If we got this far, we have at least minChain AUMs stored, and minChain number
 	// of ancestors have been marked for retention. We now continue to iterate backwards
-	// till we find an AUM which we can compact to (a Checkpoint AUM).
+	// till we find an AUM which we can compact to: either a Checkpoint AUM which is old
+	// enough, or the genesis AUM.
 	for {
 		h := next.Hash()
 		verdict[h] |= retainStateActive
+
+		parent, hasParent := next.Parent()
+		isYoung := verdict[h]&retainStateYoung != 0
+
 		if next.MessageKind == AUMCheckpoint {
 			lastActiveAncestor = h
-			break
+			if !isYoung || !hasParent {
+				break
+			}
 		}
 
-		parent, hasParent := next.Parent()
-		if !hasParent {
-			return AUMHash{}, errors.New("reached genesis AUM without finding an appropriate lastActiveAncestor")
-		}
 		if next, err = storage.AUM(parent); err != nil {
 			return AUMHash{}, fmt.Errorf("searching for compaction target (%v): %w", parent, err)
 		}
@@ -917,12 +920,12 @@ func Compact(storage CompactableChonk, head AUMHash, opts CompactionOptions) (la
 		verdict[h] = 0
 	}
 
-	if lastActiveAncestor, err = markActiveChain(storage, verdict, opts.MinChain, head); err != nil {
-		return AUMHash{}, fmt.Errorf("marking active chain: %w", err)
-	}
 	if err := markYoungAUMs(storage, verdict, opts.MinAge); err != nil {
 		return AUMHash{}, fmt.Errorf("marking young AUMs: %w", err)
 	}
+	if lastActiveAncestor, err = markActiveChain(storage, verdict, opts.MinChain, head); err != nil {
+		return AUMHash{}, fmt.Errorf("marking active chain: %w", err)
+	}
 	if err := markDescendantAUMs(storage, verdict); err != nil {
 		return AUMHash{}, fmt.Errorf("marking descendant AUMs: %w", err)
 	}

tka/tailchonk_test.go

@@ -15,6 +15,7 @@ import (
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
 	"golang.org/x/crypto/blake2s"
+	"tailscale.com/types/key"
 	"tailscale.com/util/must"
 )
 
@@ -601,3 +602,33 @@ func TestCompact(t *testing.T) {
 		}
 	}
 }
+
+func TestCompactLongButYoung(t *testing.T) {
+	ourPriv := key.NewNLPrivate()
+	ourKey := Key{Kind: Key25519, Public: ourPriv.Public().Verifier(), Votes: 1}
+	someOtherKey := Key{Kind: Key25519, Public: key.NewNLPrivate().Public().Verifier(), Votes: 1}
+
+	storage := &Mem{}
+	auth, _, err := Create(storage, State{
+		Keys:               []Key{ourKey, someOtherKey},
+		DisablementSecrets: [][]byte{DisablementKDF(bytes.Repeat([]byte{0xa5}, 32))},
+	}, ourPriv)
+	if err != nil {
+		t.Fatalf("tka.Create() failed: %v", err)
+	}
+
+	genesis := auth.Head()
+
+	for range 100 {
+		upd := auth.NewUpdater(ourPriv)
+		must.Do(upd.RemoveKey(someOtherKey.MustID()))
+		must.Do(upd.AddKey(someOtherKey))
+		aums := must.Get(upd.Finalize(storage))
+		must.Do(auth.Inform(storage, aums))
+	}
+
+	lastActiveAncestor := must.Get(Compact(storage, auth.Head(), CompactionOptions{MinChain: 5, MinAge: time.Hour}))
+	if lastActiveAncestor != genesis {
+		t.Errorf("last active ancestor = %v, want %v", lastActiveAncestor, genesis)
+	}
+}