feat: batch computation of ECChain Keys in chainexchange (#918)

* feat: batch computation of ECChain Keys in chainexchange

Adjust the for loop not to repeat the full chain
The EqualExportedValues is required as we never call explicit Key these
values, leaving cached key empty.

Signed-off-by: Jakub Sztandera <oss@kubuxu.com>

* fix: make fake MarshalPayloadForSigning use real signing

It conflicts with partial message validation

Signed-off-by: Jakub Sztandera <oss@kubuxu.com>

* Address review

Signed-off-by: Jakub Sztandera <oss@kubuxu.com>

---------

Signed-off-by: Jakub Sztandera <oss@kubuxu.com>

Author: Kubuxu

chainexchange/options.go

@@ -5,6 +5,7 @@ import (
 	"time"
 
 	"github.com/filecoin-project/go-f3/gpbft"
+	"github.com/filecoin-project/go-f3/internal/clock"
 	"github.com/filecoin-project/go-f3/internal/psutil"
 	"github.com/filecoin-project/go-f3/manifest"
 	pubsub "github.com/libp2p/go-libp2p-pubsub"
@@ -25,6 +26,7 @@ type options struct {
 	listener                       Listener
 	maxTimestampAge                time.Duration
 	compression                    bool
+	clk                            clock.Clock
 }
 
 func newOptions(o ...Option) (*options, error) {
@@ -35,6 +37,7 @@ func newOptions(o ...Option) (*options, error) {
 		maxInstanceLookahead:           manifest.DefaultCommitteeLookback,
 		maxDiscoveredChainsPerInstance: 1000,
 		maxWantedChainsPerInstance:     1000,
+		clk:                            clock.RealClock,
 	}
 	for _, apply := range o {
 		if err := apply(opts); err != nil {
@@ -163,3 +166,13 @@ func WithMaxTimestampAge(max time.Duration) Option {
 		return nil
 	}
 }
+
+func WithClock(clk clock.Clock) Option {
+	return func(o *options) error {
+		if clk == nil {
+			return errors.New("clock cannot be nil")
+		}
+		o.clk = clk
+		return nil
+	}
+}

chainexchange/pubsub.go

@@ -245,7 +245,7 @@ func (p *PubSubChainExchange) validatePubSubMessage(ctx context.Context, _ peer.
 			return pubsub.ValidationReject
 		}
 	}
-	now := time.Now().UnixMilli()
+	now := p.clk.Now().UnixMilli()
 	lowerBound := now - p.maxTimestampAge.Milliseconds()
 	if lowerBound > cmsg.Timestamp || cmsg.Timestamp > now {
 		// The timestamp is too old or too far ahead. Ignore the message to avoid
@@ -263,11 +263,11 @@ func (p *PubSubChainExchange) cacheAsDiscoveredChain(ctx context.Context, cmsg M
 	wanted := p.getChainsDiscoveredAt(ctx, cmsg.Instance)
 	discovered := p.getChainsDiscoveredAt(ctx, cmsg.Instance)
 
-	for offset := cmsg.Chain.Len(); offset >= 0 && ctx.Err() == nil; offset-- {
-		// TODO: Expose internals of merkle.go so that keys can be generated
-		//       cumulatively for a more efficient prefix chain key generation.
-		prefix := cmsg.Chain.Prefix(offset)
+	allPrefixes := cmsg.Chain.AllPrefixes()
+	for i := len(allPrefixes) - 1; i >= 0 && ctx.Err() == nil; i-- {
+		prefix := allPrefixes[i]
 		key := prefix.Key()
+
 		if portion, found := wanted.Peek(key); !found {
 			// Not a wanted key; add it to discovered chains if they are not there already,
 			// i.e. without modifying the recent-ness of any of the discovered values.
@@ -329,11 +329,12 @@ type discovery struct {
 func (p *PubSubChainExchange) cacheAsWantedChain(ctx context.Context, cmsg Message) {
 	var notifications []discovery
 	wanted := p.getChainsWantedAt(ctx, cmsg.Instance)
-	for offset := cmsg.Chain.Len(); offset >= 0 && ctx.Err() == nil; offset-- {
-		// TODO: Expose internals of merkle.go so that keys can be generated
-		//       cumulatively for a more efficient prefix chain key generation.
-		prefix := cmsg.Chain.Prefix(offset)
+
+	allPrefixes := cmsg.Chain.AllPrefixes()
+	for i := len(allPrefixes) - 1; i >= 0 && ctx.Err() == nil; i-- {
+		prefix := allPrefixes[i]
 		key := prefix.Key()
+
 		if portion, found := wanted.Peek(key); !found || portion.IsPlaceholder() {
 			wanted.Add(key, &chainPortion{
 				chain: prefix,

chainexchange/pubsub_test.go

@@ -88,25 +88,25 @@ func TestPubSubChainExchange_Broadcast(t *testing.T) {
 				chain, found = subject.GetChainByInstance(ctx, instance, key)
 				return found
 			}, time.Second, 100*time.Millisecond)
-			require.Equal(t, ecChain, chain)
+			require.EqualExportedValues(t, ecChain, chain)
 
 			baseChain := ecChain.BaseChain()
 			baseKey := baseChain.Key()
 			require.Eventually(t, func() bool {
 				chain, found = subject.GetChainByInstance(ctx, instance, baseKey)
 				return found
 			}, time.Second, 100*time.Millisecond)
-			require.Equal(t, baseChain, chain)
+			require.EqualExportedValues(t, baseChain, chain)
 
 			// Assert that we have received 2 notifications, because ecChain has 2 tipsets.
 			// First should be the ecChain, second should be the baseChain.
 
 			notifications := testListener.getNotifications()
 			require.Len(t, notifications, 2)
 			require.Equal(t, instance, notifications[1].instance)
-			require.Equal(t, baseChain, notifications[1].chain)
+			require.EqualExportedValues(t, baseChain, notifications[1].chain)
 			require.Equal(t, instance, notifications[0].instance)
-			require.Equal(t, ecChain, notifications[0].chain)
+			require.EqualExportedValues(t, ecChain, notifications[0].chain)
 
 			require.NoError(t, subject.Shutdown(ctx))
 		})

gpbft/chain.go

@@ -425,6 +425,51 @@ func (c *ECChain) Key() ECChainKey {
 	return c.key
 }
 
+// KeysForPrefixes return batch of keys for all prefixes
+// the indexes to this array correspond to Prefix(i) calls
+func (c *ECChain) KeysForPrefixes() []ECChainKey {
+	if c.IsZero() {
+		return nil
+	}
+	values := make([][]byte, c.Len())
+	for i, ts := range c.TipSets {
+		values[i] = ts.MarshalForSigning()
+	}
+
+	batch := merkle.BatchTree(values)
+	res := make([]ECChainKey, c.Len())
+	for i := range c.Len() {
+		res[i] = batch[i]
+	}
+	return res
+}
+
+// AllPrefixes returns an array of all prefix chain including the c itself.
+// It precomputes keys for them as well, populating the key cache.
+func (c *ECChain) AllPrefixes() []*ECChain {
+	if c.IsZero() {
+		return nil
+	}
+	values := make([][]byte, c.Len())
+	for i, ts := range c.TipSets {
+		values[i] = ts.MarshalForSigning()
+	}
+	batch := merkle.BatchTree(values)
+
+	res := make([]*ECChain, len(c.TipSets))
+	for i := range len(c.TipSets) {
+		var prefix ECChain
+		prefix.TipSets = c.TipSets[: i+1 : i+1]
+
+		copy(prefix.key[:], batch[i][:]) // populate the key cache
+		prefix.keyLazyLoader.Do(func() {})
+
+		res[i] = &prefix
+	}
+
+	return res
+}
+
 func (c *ECChain) String() string {
 	if c.IsZero() {
 		return "丄"

gpbft/chain_test.go

@@ -189,6 +189,55 @@ func TestECChain(t *testing.T) {
 	})
 }
 
+func TestECChainKeysForPrefixes(t *testing.T) {
+	t.Parallel()
+
+	pt1Cid := gpbft.MakeCid([]byte("pt1"))
+	pt2Cid := gpbft.MakeCid([]byte("pt2"))
+	pt3Cid := gpbft.MakeCid([]byte("pt3"))
+	tipSets := []*gpbft.TipSet{
+		{Epoch: 0, Key: []byte{0}, PowerTable: pt1Cid},
+		{Epoch: 1, Key: []byte{1}, PowerTable: pt2Cid},
+		{Epoch: 2, Key: []byte{2}, PowerTable: pt3Cid},
+	}
+
+	chain := &gpbft.ECChain{TipSets: tipSets}
+	keys := chain.KeysForPrefixes()
+
+	require.Equal(t, len(keys), len(tipSets), "KeysForPrefixes should return a key for each prefix")
+
+	for i, key := range keys {
+		prefixChain := chain.Prefix(i)
+		require.Equal(t, key, prefixChain.Key(), "Key for prefix %d should match", i)
+	}
+}
+
+func TestECChainAllPrefixes(t *testing.T) {
+	t.Parallel()
+
+	pt1Cid := gpbft.MakeCid([]byte("pt1"))
+	pt2Cid := gpbft.MakeCid([]byte("pt2"))
+	pt3Cid := gpbft.MakeCid([]byte("pt3"))
+	tipSets := []*gpbft.TipSet{
+		{Epoch: 0, Key: []byte{0}, PowerTable: pt1Cid},
+		{Epoch: 1, Key: []byte{1}, PowerTable: pt2Cid},
+		{Epoch: 2, Key: []byte{2}, PowerTable: pt3Cid},
+	}
+
+	chain := &gpbft.ECChain{TipSets: tipSets}
+	prefixes := chain.AllPrefixes()
+
+	require.Equal(t, len(prefixes), len(tipSets), "AllPrefixes should return a prefix for each tipset")
+
+	for i, prefix := range prefixes {
+		require.Equal(t, prefix.Len(), i+1, "Prefix %d should have length %d", i, i+1)
+		expected := chain.Prefix(i)
+		require.True(t, expected.Eq(prefix))
+		require.Equal(t, expected.Key(), prefix.Key())
+		require.NotSame(t, expected, prefix)
+	}
+}
+
 func TestECChain_Eq(t *testing.T) {
 	t.Parallel()
 	var (

host.go

@@ -147,7 +147,7 @@ func newRunner(
 		runner.msgEncoding = encoding.NewCBOR[*PartialGMessage]()
 	}
 
-	runner.pmm, err = newPartialMessageManager(runner.Progress, ps, m)
+	runner.pmm, err = newPartialMessageManager(runner.Progress, ps, m, runner.clock)
 	if err != nil {
 		return nil, fmt.Errorf("creating partial message manager: %w", err)
 	}

internal/clock/clock.go

@@ -26,13 +26,13 @@ func WithMockClock(ctx context.Context) (context.Context, *Mock) {
 	return context.WithValue(ctx, clockKey, (Clock)(clk)), clk
 }
 
-var realClock = clock.New()
+var RealClock = clock.New()
 
 // GetClock either retrieves a mock clock from the context or returns a realtime clock.
 func GetClock(ctx context.Context) Clock {
 	clk := ctx.Value(clockKey)
 	if clk == nil {
-		return realClock
+		return RealClock
 	}
 	return clk.(Clock)
 }

merkle/merkle.go

@@ -29,7 +29,7 @@ func TreeWithProofs(values [][]byte) (Digest, [][]Digest) {
 
 // Tree returns a the root of the merkle-tree of the given values.
 func Tree(values [][]byte) Digest {
-	return buildTree(bits.Len(uint(len(values))-1), values, nil)
+	return buildTree(depth(values), values, nil)
 }
 
 // VerifyProof verifies that the given value maps to the given index in the merkle-tree with the
@@ -123,3 +123,32 @@ func buildTree(depth int, values [][]byte, proofs [][]Digest) Digest {
 
 	return internalHash(leftHash, rightHash)
 }
+
+// BatchTree creates a batch of prefixes of values: [[0], [0, 1], [0, 1, 2], ...]
+// and provides digests for all of them.
+func BatchTree(values [][]byte) []Digest {
+	n := len(values)
+	roots := make([]Digest, n+1) // roots[0] is empty
+
+	for k := 1; k <= n; k++ {
+		if k == 1 {
+			roots[k] = leafHash(values[0])
+			continue
+		}
+
+		depth := bits.Len(uint(k - 1))
+		split := 1 << (depth - 1)
+
+		// Reuse left subtree root from previous computation
+		leftRoot := roots[split]
+
+		// compute right subtree root
+		// this could be made more optimal but leads to more messy code
+		rightValues := values[split:k]
+		rightRoot := buildTree(depth-1, rightValues, nil)
+
+		roots[k] = internalHash(leftRoot, rightRoot)
+	}
+
+	return roots[1:]
+}

merkle/merkle_test.go

@@ -1,7 +1,9 @@
 package merkle
 
 import (
+	"encoding/hex"
 	"fmt"
+	"runtime"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -65,3 +67,87 @@ func TestHashZero(t *testing.T) {
 	valid, _ := VerifyProof(root, 0, nil, nil)
 	assert.False(t, valid)
 }
+
+func generateInputs(dst string, N int) [][]byte {
+	var res [][]byte
+	for i := 0; i < N; i++ {
+		res = append(res, []byte(dst+fmt.Sprintf("%09d", i)))
+	}
+	return res
+}
+
+func TestHashBatch(t *testing.T) {
+	for N := 1; N < 300; {
+		inputs := generateInputs(fmt.Sprintf("batch-%d", N), N)
+		expected := make([]Digest, 0, N-1)
+
+		for i := 0; i < N; i++ {
+			expected = append(expected, Tree(inputs[0:i+1]))
+		}
+
+		batchRes := BatchTree(inputs)
+		require.Equal(t, expected, batchRes, "failed at size %d", N)
+		if N < 32 {
+			N++
+		} else {
+			N += 13
+		}
+	}
+}
+
+func TestHashTreeGolden(t *testing.T) {
+	expectedHex := []string{
+		"3d4395573ce4d2acbce4fe8a4be67ca5e7cdfb8ee2e85b2f6733c16b24c3b175",
+		"91b7c899421ca7f3228e10265c6970a03bc2ccba44367b1d44a9d8597b20a32e",
+		"69abe78dc2390b4666b60d0582e1799e73e48766f6e502c515e79d6cd2ae3c45",
+		"bc4ce8dbf993eb2e87c02bbf19cd4faeb3a0672188bc6be6c8d867cef9b08917",
+		"538cfd0c1f6b7ab4c3d20466d4e01b438972212fe5257eae213ae0a040da977f",
+		"e28aa108b0263820dfe2c7f051ddc8794ab48ebd3c1813db28bf9f06bedc52f3",
+		"875cb1d5027522b344b8adc62cd6bd110d97eaedd40a35bcb2fe142a9cb4612b",
+		"63804e8b6cb16993d5d43d9d7faf17ba967365dac141a4afbce1d794157a1b8e",
+		"07105bd8716bebc90036c8ebfe23a92bd775c09664b076ffa1d9a29d30647f91",
+		"960b7eb6440789f76f5d53965e8b208e34777bc4aab78edf6827d71c7eea4933",
+		"d55e07222c786722e1ad1b5bcc2ebaf04b2b4e92c07f3f7b61b0fbf0fd78fb9b",
+		"ee5a34dfae748e088a1b99386274158266f44ceeb2c5190f4e9bbc39cd8a4d26",
+		"15def4fc077ccfb0e48b32bc07ea3b91acecc5b73ed9caf13b10adf17052c371",
+		"07cfe4ec2efa9075763f921e9f29794ec6b945694e41cc19911101270d8b1087",
+		"84cdf541cbb3b9b3f26dbdeb9ca3a2721d15447a8b47074c3b08b560f79e5d85",
+		"af8e9fc2f15aaedadb96da1afb339b93e3174661327dcc6aad70ea67e183369d",
+	}
+	var results []string
+	N := 16
+	for i := 0; i < N; i++ {
+		inputs := generateInputs("golden", i+1)
+		res := Tree(inputs)
+		resHex := hex.EncodeToString(res[:])
+		assert.Equal(t, expectedHex[i], resHex)
+		results = append(results, resHex)
+	}
+	t.Logf("results: %#v", results)
+
+	batchRes := BatchTree(generateInputs("golden", N))
+	batchResHash := make([]string, N)
+	for i := 0; i < N; i++ {
+		batchResHash[i] = hex.EncodeToString(batchRes[i][:])
+	}
+	assert.Equal(t, expectedHex, batchResHash)
+}
+
+func BenchmarkPrefixes(b *testing.B) {
+	K := 128
+	inputs := generateInputs("golden", K)
+	b.Run("IndividualPrefix", func(b *testing.B) {
+		b.ReportAllocs()
+		for range b.N {
+			for i := range len(inputs) {
+				runtime.KeepAlive(Tree(inputs[:i+1]))
+			}
+		}
+	})
+	b.Run("BatchPrefix", func(b *testing.B) {
+		b.ReportAllocs()
+		for range b.N {
+			runtime.KeepAlive(BatchTree(inputs))
+		}
+	})
+}