fix(aggregation): recover from over-budget cost estimates

cmd/gean/main.go

@@ -92,6 +92,7 @@ func run(cfg config) error {
 		VerifyAggregatedSignatures: cfg.ShadowVerifyAggregatedSignaturesRate,
 	}
 	n := node.New(s, fc, p2pHost, inputs.keyManager, aggCtl, cfg.CommitteeCount, shadowRates)
+	n.AggregateSubnetIDs = cfg.AggregateSubnetIDs
 	startNodeNetworking(ctx, n, s, p2pHost, inputs.bootnodes)
 
 	apiAddr, metricsAddr := startHTTPServers(cfg, s, fc, aggCtl)

internal/aggregation/aggregate_test.go

@@ -4,6 +4,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/geanlabs/gean/internal/metrics"
 	"github.com/geanlabs/gean/internal/shadow"
 	"github.com/geanlabs/gean/internal/store"
 	"github.com/geanlabs/gean/internal/types"
@@ -78,7 +79,7 @@ func TestAggregateFromSnapshotExpiredDeadlineReportsTruncation(t *testing.T) {
 	snap := aggregateTestSnapshot(5)
 	cache := xmss.NewPubKeyCache()
 
-	aggs, payloads, deletes, truncated, _ := aggregateFromSnapshot(snap, cache, time.Now().Add(-time.Second), shadow.Rates{}, newUnitCostEstimator())
+	aggs, payloads, deletes, truncated, _ := aggregateFromSnapshot(snap, cache, time.Now().Add(-time.Second), MaxGroupsPerSession, shadow.Rates{}, newUnitCostEstimator())
 
 	if !truncated {
 		t.Fatal("expected truncation with expired deadline")
@@ -91,47 +92,101 @@ func TestAggregateFromSnapshotExpiredDeadlineReportsTruncation(t *testing.T) {
 func TestAggregateFromSnapshotZeroDeadlineProcessesAll(t *testing.T) {
 	snap := aggregateTestSnapshot(5)
 
-	_, _, _, truncated, _ := aggregateFromSnapshot(snap, xmss.NewPubKeyCache(), time.Time{}, shadow.Rates{}, newUnitCostEstimator())
+	_, _, _, truncated, _ := aggregateFromSnapshot(snap, xmss.NewPubKeyCache(), time.Time{}, MaxGroupsPerSession, shadow.Rates{}, newUnitCostEstimator())
 
 	if truncated {
 		t.Fatal("zero deadline must never truncate")
 	}
 }
 
-func TestUnitCostEstimatorMaxUnitsWithin(t *testing.T) {
-	e := newUnitCostEstimator() // seed 0.1s/unit
+// The cost of a proof is the proof, not what it covers: two-signature groups
+// measured 2.0-5.2s on a 16-core host. Dividing that by the signature count is
+// what previously concluded a signature costs seconds and pinned every later
+// group at the two-signature floor.
+func TestUnitCostEstimatorLearnsFixedCostPerProof(t *testing.T) {
+	e := newUnitCostEstimator()
 
-	if got := e.maxUnitsWithin(time.Second); got != 10 {
-		t.Fatalf("maxUnitsWithin(1s)=%d, want 10", got)
+	for range 10 {
+		e.observeGroup(4*time.Second, 0)
 	}
-	// Never below the spec minimum of two, even for a tiny or expired budget.
-	if got := e.maxUnitsWithin(time.Millisecond); got != 2 {
-		t.Fatalf("maxUnitsWithin(1ms)=%d, want 2 (floor)", got)
+
+	if got := e.nextGroupDuration(); got < 3800*time.Millisecond || got > 4200*time.Millisecond {
+		t.Fatalf("nextGroupDuration=%v, want about 4s (the whole proof, not a share of it)", got)
+	}
+}
+
+// Children are the part that does scale, so they are charged whatever the fixed
+// cost does not explain.
+func TestUnitCostEstimatorChargesChildrenTheResidual(t *testing.T) {
+	e := newUnitCostEstimator()
+
+	for range 10 {
+		e.observeGroup(2*time.Second, 0)
 	}
-	if got := e.maxUnitsWithin(-time.Second); got != 2 {
-		t.Fatalf("maxUnitsWithin(-1s)=%d, want 2 (floor)", got)
+	for range 10 {
+		e.observeGroup(5*time.Second, 1)
+	}
+
+	if got := e.childDuration(); got < 2500*time.Millisecond || got > 3500*time.Millisecond {
+		t.Fatalf("childDuration=%v, want about 3s (5s group less the 2s baseline)", got)
+	}
+	if got := e.nextGroupDuration(); got > 2500*time.Millisecond {
+		t.Fatalf("nextGroupDuration=%v, want the recursive group kept out of the fixed cost", got)
+	}
+
+	// A group cheaper than a raw-only one says nothing about its children.
+	steady := e.childDuration()
+	e.observeGroup(time.Millisecond, 1)
+	if e.childDuration() != steady {
+		t.Fatalf("under-cost group moved the child estimate: %v -> %v", steady, e.childDuration())
+	}
+
+	// With no baseline yet there is nothing to subtract, so a recursive group is
+	// ignored rather than charged the whole duration.
+	fresh := newUnitCostEstimator()
+	before := fresh.childDuration()
+	fresh.observeGroup(9*time.Second, 1)
+	if fresh.childDuration() != before {
+		t.Fatalf("child estimate moved without a fixed-cost baseline: %v -> %v", before, fresh.childDuration())
 	}
 }
 
-func TestUnitCostEstimatorObserveConverges(t *testing.T) {
-	e := newUnitCostEstimator() // seed 0.1s/unit
+// A budget stop defers every group still queued, not only the one it examined.
+// Counting a single skip understated the backlog and made a session that dropped
+// a long queue look like one that dropped a single group.
+func TestAggregateFromSnapshotBudgetStopCountsEveryDeferredGroup(t *testing.T) {
+	snap := aggregateTestSnapshot(5, 6, 7)
+	cache := xmss.NewPubKeyCache()
+
+	_, _, _, truncated, skips := aggregateFromSnapshot(snap, cache, time.Now().Add(-time.Second), MaxGroupsPerSession, shadow.Rates{}, newUnitCostEstimator())
 
-	// A cheaper-than-seed observation must pull the estimate down, letting more
-	// units fit the budget on the next pass.
-	before := e.maxUnitsWithin(time.Second)
-	for range 20 {
-		e.observe(200*time.Millisecond, 10) // 0.02s/unit
+	if !truncated {
+		t.Fatal("expected truncation with expired deadline")
 	}
-	after := e.maxUnitsWithin(time.Second)
-	if after <= before {
-		t.Fatalf("estimate did not converge down: before=%d after=%d units/sec", before, after)
+	if got := skips[metrics.AggGroupSkipBudget]; got != 3 {
+		t.Fatalf("budget skips = %d, want 3 (one per deferred group)", got)
 	}
+}
 
-	// Degenerate inputs are ignored, not divided by.
-	steady := e.perUnitSeconds
-	e.observe(0, 10)
-	e.observe(time.Second, 0)
-	if e.perUnitSeconds != steady {
-		t.Fatalf("degenerate observe mutated estimate: %v -> %v", steady, e.perUnitSeconds)
+// This slot's votes are the only ones with a deadline: they must be aggregated
+// in time to reach the next block, while backlog entries lose nothing by waiting
+// a slot. With a session capped at two groups, ordering purely by target slot
+// would spend both on the oldest backlog and leave the current slot unaggregated.
+func TestOrderedGroupsPutsCurrentSlotFirst(t *testing.T) {
+	snap := aggregateTestSnapshot(10, 11, 12)
+	snap.slot = 12
+
+	groups := orderedGroups(snap, groupSkips{})
+	if len(groups) != 3 {
+		t.Fatalf("groups=%d, want 3", len(groups))
+	}
+	if !groups[0].currentSlot || groups[0].targetSlot != 12 {
+		t.Fatalf("first group targets slot %d (current=%v), want the current slot",
+			groups[0].targetSlot, groups[0].currentSlot)
+	}
+	// Behind it, the frontier rule still holds: oldest unjustified target first.
+	if groups[1].targetSlot != 10 || groups[2].targetSlot != 11 {
+		t.Fatalf("backlog order = %d,%d, want ascending target 10,11",
+			groups[1].targetSlot, groups[2].targetSlot)
 	}
 }

internal/aggregation/budget_recovery_test.go

@@ -0,0 +1,121 @@
+package aggregation
+
+import (
+	"errors"
+	"testing"
+	"testing/synctest"
+	"time"
+
+	"github.com/geanlabs/gean/internal/metrics"
+	"github.com/geanlabs/gean/internal/shadow"
+	"github.com/geanlabs/gean/internal/store"
+	"github.com/geanlabs/gean/internal/types"
+	"github.com/geanlabs/gean/xmss"
+)
+
+// The signatures are only decoded: the injected prover tests scheduling and
+// result retention, not cryptographic validity.
+func budgetTestSnapshot() *Snapshot {
+	snap := aggregateTestSnapshot(5, 6, 7)
+	snap.headState.Validators = []*types.Validator{{Index: 0}, {Index: 1}, {Index: 2}}
+	for _, dr := range [][32]byte{rootByte(2), rootByte(3)} {
+		snap.attSigs[dr].Signatures = []store.AttestationSignatureEntry{{ValidatorID: 0}, {ValidatorID: 1}, {ValidatorID: 2}}
+	}
+	return snap
+}
+
+func TestAggregationBudgetRecovery(t *testing.T) {
+	for _, tc := range []struct {
+		name         string
+		observations []time.Duration
+		fail         bool
+	}{
+		{"first_slow_group", []time.Duration{2 * time.Second}, false},
+		{"steady_then_spike", []time.Duration{1400 * time.Millisecond, 2500 * time.Millisecond}, false},
+		{"failed_attempt", []time.Duration{2 * time.Second}, true},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			synctest.Test(t, func(t *testing.T) {
+				e := newUnitCostEstimator()
+				for _, d := range tc.observations {
+					e.observeGroup(d, 0)
+				}
+				cache := xmss.NewPubKeyCache()
+				defer cache.Close()
+				for session := 0; session < 5; session++ {
+					before := e.nextGroupDuration()
+					calls := 0
+					prove := func(pks []xmss.CPubKey, sigs []xmss.CSig, children []xmss.ChildProof, root [32]byte, slot uint32) ([]byte, error) {
+						calls++
+						// All three of the group's signatures, not the two the
+						// old per-signature budget floor allowed: the proof costs
+						// the same either way.
+						if slot != 6 || len(pks) != 3 || len(sigs) != 3 || len(children) != 0 {
+							t.Fatalf("unexpected inputs: slot=%d raw=%d children=%d", slot, len(sigs), len(children))
+						}
+						time.Sleep(time.Second)
+						if tc.fail {
+							return nil, errors.New("test prover failed")
+						}
+						return []byte{1}, nil
+					}
+					aggs, payloads, deletes, truncated, skips := aggregateFromSnapshotWithProver(budgetTestSnapshot(), cache, time.Now().Add(SessionBudget), MaxGroupsPerSession, shadow.Rates{}, e, prove)
+					if calls != 1 || !truncated || skips[metrics.AggGroupSkipBudget] != 1 || skips[metrics.AggGroupSkipTooFewSigners] != 1 {
+						t.Fatalf("session=%d calls=%d truncated=%v skips=%v", session, calls, truncated, skips)
+					}
+					if tc.fail {
+						if len(aggs) != 0 || len(payloads) != 0 || len(deletes) != 0 || skips[metrics.AggGroupSkipError] != 1 || e.nextGroupDuration() != before {
+							t.Fatalf("failure changed results/estimate or lost skip: %v", skips)
+						}
+					} else {
+						if len(aggs) != 1 || len(payloads) != 1 || len(deletes) != 3 {
+							t.Fatalf("lost partial results: aggs=%d payloads=%d deletes=%d", len(aggs), len(payloads), len(deletes))
+						}
+						if e.nextGroupDuration() >= before {
+							t.Fatal("successful attempt did not recalibrate")
+						}
+						// Every signature the group held is covered and retired:
+						// the proof costs the same whether it carries two or
+						// three, so none is left behind for a later session.
+						if types.BitlistCount(aggs[0].Proof.Participants) != 3 ||
+							deletes[0].ValidatorID != 0 || deletes[1].ValidatorID != 1 || deletes[2].ValidatorID != 2 {
+							t.Fatal("group did not cover and retire every signature it held")
+						}
+					}
+				}
+			})
+		})
+	}
+}
+
+func TestAggregationBudgetDeadline(t *testing.T) {
+	for _, tc := range []struct {
+		name          string
+		deadline      time.Duration
+		proofTime     time.Duration
+		wantCalls     int
+		wantTruncated bool
+	}{
+		{"expired", -time.Second, 0, 0, true},
+		{"at_deadline", 0, 0, 0, true},
+		{"first_proof_overruns", SessionBudget, 2 * time.Second, 1, true},
+		{"enough_for_later_group", SessionBudget, 100 * time.Millisecond, 2, false},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			synctest.Test(t, func(t *testing.T) {
+				cache := xmss.NewPubKeyCache()
+				defer cache.Close()
+				calls := 0
+				prove := func([]xmss.CPubKey, []xmss.CSig, []xmss.ChildProof, [32]byte, uint32) ([]byte, error) {
+					calls++
+					time.Sleep(tc.proofTime)
+					return []byte{1}, nil
+				}
+				aggs, _, _, truncated, skips := aggregateFromSnapshotWithProver(budgetTestSnapshot(), cache, time.Now().Add(tc.deadline), MaxGroupsPerSession, shadow.Rates{}, newUnitCostEstimator(), prove)
+				if calls != tc.wantCalls || len(aggs) != calls || truncated != tc.wantTruncated {
+					t.Fatalf("calls=%d aggs=%d truncated=%v skips=%v", calls, len(aggs), truncated, skips)
+				}
+			})
+		})
+	}
+}