tka: optimize common case of processing updates built from head

Signed-off-by: Tom DNetto <tom@tailscale.com>

tka/tka.go

@@ -553,13 +553,31 @@ func Bootstrap(storage Chonk, bootstrap AUM) (*Authority, error) {
 // should be ordered oldest to newest. An error is returned if any
 // of the updates could not be processed.
 func (a *Authority) Inform(updates []AUM) error {
+	if len(updates) == 0 {
+		return errors.New("inform called with empty slice")
+	}
 	stateAt := make(map[AUMHash]State, len(updates)+1)
 	toCommit := make([]AUM, 0, len(updates))
+	prevHash := a.Head()
+
+	// The state at HEAD is the current state of the authority. Its likely
+	// to be needed, so we prefill it rather than computing it.
+	stateAt[prevHash] = a.state
+
+	// Optimization: If the set of updates is a chain building from
+	// the current head, EG:
+	//   <a.Head()> ==> updates[0] ==> updates[1] ...
+	// Then theres no need to recompute the resulting state from the
+	// stored ancestor, because the last state computed during iteration
+	// is the new state. This should be the common case.
+	// isHeadChain keeps track of this.
+	isHeadChain := true
 
 	for i, update := range updates {
 		hash := update.Hash()
+		// Check if we already have this AUM thus don't need to process it.
 		if _, err := a.storage.AUM(hash); err == nil {
-			// Already have this AUM.
+			isHeadChain = false // Disable the head-chain optimization.
 			continue
 		}
 
@@ -583,6 +601,11 @@ func (a *Authority) Inform(updates []AUM) error {
 		if stateAt[hash], err = state.applyVerifiedAUM(update); err != nil {
 			return fmt.Errorf("update %d cannot be applied: %v", i, err)
 		}
+
+		if isHeadChain && parent != prevHash {
+			isHeadChain = false
+		}
+		prevHash = hash
 		toCommit = append(toCommit, update)
 	}
 
@@ -590,18 +613,22 @@ func (a *Authority) Inform(updates []AUM) error {
 		return fmt.Errorf("commit: %v", err)
 	}
 
-	// TODO(tom): Theres no need to recompute the state from scratch
-	//            in every case. We should detect when updates were
-	//            a linear, non-forking series applied to head, and
-	//            just use the last State we computed.
-	oldestAncestor := a.oldestAncestor.Hash()
-	c, err := computeActiveChain(a.storage, &oldestAncestor, 2000)
-	if err != nil {
-		return fmt.Errorf("recomputing active chain: %v", err)
+	if isHeadChain {
+		// Head-chain fastpath: We can use the state we computed
+		// in the last iteration.
+		a.head = updates[len(updates)-1]
+		a.state = stateAt[prevHash]
+	} else {
+		oldestAncestor := a.oldestAncestor.Hash()
+		c, err := computeActiveChain(a.storage, &oldestAncestor, 2000)
+		if err != nil {
+			return fmt.Errorf("recomputing active chain: %v", err)
+		}
+		a.head = c.Head
+		a.oldestAncestor = c.Oldest
+		a.state = c.state
 	}
-	a.head = c.Head
-	a.oldestAncestor = c.Oldest
-	a.state = c.state
+
 	return nil
 }
 

tka/tka_test.go

@@ -325,7 +325,7 @@ func TestCreateBootstrapAuthority(t *testing.T) {
 	}
 }
 
-func TestAuthorityInform(t *testing.T) {
+func TestAuthorityInformNonLinear(t *testing.T) {
 	pub, priv := testingKey25519(t, 1)
 	key := Key{Kind: Key25519, Public: pub, Votes: 2}
 
@@ -351,6 +351,8 @@ func TestAuthorityInform(t *testing.T) {
 		t.Fatalf("Bootstrap() failed: %v", err)
 	}
 
+	// L2 does not chain from L1, disabling the isHeadChain optimization
+	// and forcing Inform() to take the slow path.
 	informAUMs := []AUM{c.AUMs["L1"], c.AUMs["L2"], c.AUMs["L3"], c.AUMs["L4"], c.AUMs["L5"]}
 
 	if err := a.Inform(informAUMs); err != nil {
@@ -371,3 +373,46 @@ func TestAuthorityInform(t *testing.T) {
 		t.Fatal("authority did not converge to correct AUM")
 	}
 }
+
+func TestAuthorityInformLinear(t *testing.T) {
+	pub, priv := testingKey25519(t, 1)
+	key := Key{Kind: Key25519, Public: pub, Votes: 2}
+
+	c := newTestchain(t, `
+        G1 -> L1 -> L2 -> L3
+
+        G1.template = genesis
+    `,
+		optTemplate("genesis", AUM{MessageKind: AUMCheckpoint, State: &State{
+			Keys:               []Key{key},
+			DisablementSecrets: [][]byte{disablementKDF([]byte{1, 2, 3})},
+		}}),
+		optKey("key", key, priv),
+		optSignAllUsing("key"))
+
+	storage := &Mem{}
+	a, err := Bootstrap(storage, c.AUMs["G1"])
+	if err != nil {
+		t.Fatalf("Bootstrap() failed: %v", err)
+	}
+
+	informAUMs := []AUM{c.AUMs["L1"], c.AUMs["L2"], c.AUMs["L3"]}
+
+	if err := a.Inform(informAUMs); err != nil {
+		t.Fatalf("Inform() failed: %v", err)
+	}
+	for i, update := range informAUMs {
+		stored, err := storage.AUM(update.Hash())
+		if err != nil {
+			t.Errorf("reading stored update %d: %v", i, err)
+			continue
+		}
+		if diff := cmp.Diff(update, stored); diff != "" {
+			t.Errorf("update %d differs (-want, +got):\n%s", i, diff)
+		}
+	}
+
+	if a.Head() != c.AUMHashes["L3"] {
+		t.Fatal("authority did not converge to correct AUM")
+	}
+}