Improve node cache (#921)

Improves the implementation of the node cache under client.

With these changes, nodeCache is now thread-safe, and supports concurrent fetching of arbitrary nodes, potentially speeding up the process of building trees and proofs if the underlying storage has high latency.

Author: AlCutter

client/client.go

@@ -23,6 +23,7 @@ import (
 	"fmt"
 	"sync"
 
+	lru "github.com/hashicorp/golang-lru/v2"
 	"github.com/transparency-dev/formats/log"
 	"github.com/transparency-dev/merkle/compact"
 	"github.com/transparency-dev/merkle/proof"
@@ -32,6 +33,7 @@ import (
 	"github.com/transparency-dev/tessera/internal/otel"
 	"go.opentelemetry.io/otel/trace"
 	"golang.org/x/mod/sumdb/note"
+	"golang.org/x/sync/errgroup"
 )
 
 var (
@@ -125,15 +127,7 @@ func FetchRangeNodes(ctx context.Context, s uint64, f TileFetcherFunc) ([][]byte
 		nc := newNodeCache(f, s)
 		nIDs := make([]compact.NodeID, 0, compact.RangeSize(0, s))
 		nIDs = compact.RangeNodes(0, s, nIDs)
-		hashes := make([][]byte, 0, len(nIDs))
-		for _, n := range nIDs {
-			h, err := nc.GetNode(ctx, n)
-			if err != nil {
-				return nil, err
-			}
-			hashes = append(hashes, h)
-		}
-		return hashes, nil
+		return nc.GetNodes(ctx, nIDs)
 	})
 }
 
@@ -180,12 +174,12 @@ func GetEntryBundle(ctx context.Context, f EntryBundleFetcherFunc, i, logSize ui
 // at a given tree size.
 type ProofBuilder struct {
 	treeSize  uint64
-	nodeCache nodeCache
+	nodeCache *nodeCache
 }
 
 // NewProofBuilder creates a new ProofBuilder object for a given tree size.
-// The returned ProofBuilder can be re-used for proofs related to a given tree size, but
-// it is not thread-safe and should not be accessed concurrently.
+// The returned ProofBuilder can be re-used for proofs related to a given tree size, and is
+// thread-safe.
 func NewProofBuilder(ctx context.Context, treeSize uint64, f TileFetcherFunc) (*ProofBuilder, error) {
 	pb := &ProofBuilder{
 		treeSize:  treeSize,
@@ -204,7 +198,7 @@ func (pb *ProofBuilder) InclusionProof(ctx context.Context, index uint64) ([][]b
 		if err != nil {
 			return nil, fmt.Errorf("failed to calculate inclusion proof node list: %v", err)
 		}
-		return pb.fetchNodes(ctx, nodes)
+		return pb.materialiseProof(ctx, nodes)
 	})
 }
 
@@ -221,22 +215,16 @@ func (pb *ProofBuilder) ConsistencyProof(ctx context.Context, smaller, larger ui
 		if err != nil {
 			return nil, fmt.Errorf("failed to calculate consistency proof node list: %v", err)
 		}
-		return pb.fetchNodes(ctx, nodes)
+		return pb.materialiseProof(ctx, nodes)
 	})
 }
 
-// fetchNodes retrieves the specified proof nodes via pb's nodeCache.
-func (pb *ProofBuilder) fetchNodes(ctx context.Context, nodes proof.Nodes) ([][]byte, error) {
-	hashes := make([][]byte, 0, len(nodes.IDs))
-	// TODO(al) parallelise this.
-	for _, id := range nodes.IDs {
-		h, err := pb.nodeCache.GetNode(ctx, id)
-		if err != nil {
-			return nil, fmt.Errorf("failed to get node (%v): %v", id, err)
-		}
-		hashes = append(hashes, h)
+// materialiseProof retrieves the specified proof nodes via pb's nodeCache, recreating ephemeral nodes if necessary.
+func (pb *ProofBuilder) materialiseProof(ctx context.Context, nodes proof.Nodes) ([][]byte, error) {
+	hashes, err := pb.nodeCache.GetNodes(ctx, nodes.IDs)
+	if err != nil {
+		return nil, err
 	}
-	var err error
 	if hashes, err = nodes.Rehash(hashes, hasher.HashChildren); err != nil {
 		return nil, fmt.Errorf("failed to rehash proof: %v", err)
 	}
@@ -344,82 +332,179 @@ func (lst *LogStateTracker) Latest() log.Checkpoint {
 	return lst.latestConsistent
 }
 
-// tileKey is used as a key in nodeCache's tile map.
-type tileKey struct {
-	tileLevel uint64
-	tileIndex uint64
-}
-
 // nodeCache hides the tiles abstraction away, and improves
 // performance by caching tiles it's seen.
-// Not threadsafe, and intended to be only used throughout the course
-// of a single request.
+// Threadsafe.
 type nodeCache struct {
 	logSize   uint64
-	ephemeral map[compact.NodeID][]byte
-	tiles     map[tileKey]api.HashTile
+	nodes     *lru.Cache[compact.NodeID, []byte]
 	getTile   TileFetcherFunc
+	tileLocks *shardedMutex[compact.NodeID]
 }
 
 // newNodeCache creates a new nodeCache instance for a given log size.
-func newNodeCache(f TileFetcherFunc, logSize uint64) nodeCache {
-	return nodeCache{
+func newNodeCache(f TileFetcherFunc, logSize uint64) *nodeCache {
+	c, err := lru.New[compact.NodeID, []byte](64 << 10)
+	if err != nil {
+		panic(fmt.Errorf("lru.New: %v", err))
+	}
+	return &nodeCache{
 		logSize:   logSize,
-		ephemeral: make(map[compact.NodeID][]byte),
-		tiles:     make(map[tileKey]api.HashTile),
+		nodes:     c,
 		getTile:   f,
+		tileLocks: newShardedMutex[compact.NodeID](),
 	}
 }
 
-// SetEphemeralNode stored a derived "ephemeral" tree node.
-func (n *nodeCache) SetEphemeralNode(id compact.NodeID, h []byte) {
-	n.ephemeral[id] = h
-}
-
 // GetNode returns the internal log tree node hash for the specified node ID.
-// A previously set ephemeral node will be returned if id matches, otherwise
-// the tile containing the requested node will be fetched and cached, and the
-// node hash returned.
+// The tile containing the node will be fetched if necessary.
 func (n *nodeCache) GetNode(ctx context.Context, id compact.NodeID) ([]byte, error) {
 	return otel.Trace(ctx, "tessera.client.nodecache.GetNode", tracer, func(ctx context.Context, span trace.Span) ([]byte, error) {
 		span.SetAttributes(indexKey.Int64(otel.Clamp64(id.Index)), levelKey.Int64(int64(id.Level)))
+		// Fast-path: check to see we have this node in the cache and return it directly if so, otherwise we'll need to fetch it.
+		if e, ok := n.nodes.Get(id); ok {
+			return e, nil
+		}
 
-		// First check for ephemeral nodes:
-		if e := n.ephemeral[id]; len(e) != 0 {
+		// No dice, so we need to fetch the tile and use the contents to populate the cache.
+		// We only want to do this once per tile, so lock keyed by the _tile_ ID here.
+		tileLevel, tileIndex, _, _ := layout.NodeCoordsToTileAddress(uint64(id.Level), uint64(id.Index))
+		k := compact.NodeID{Level: uint(tileLevel), Index: tileIndex}
+		n.tileLocks.Lock(k)
+		defer n.tileLocks.Unlock(k)
+		// Re-check if we have the node cached - since we're under lock here it's possible that another goroutine
+		// managed to get into this section before us and populate the cache.
+		if e, ok := n.nodes.Get(id); ok {
 			return e, nil
 		}
-		// Otherwise look in fetched tiles:
-		tileLevel, tileIndex, nodeLevel, nodeIndex := layout.NodeCoordsToTileAddress(uint64(id.Level), uint64(id.Index))
-		tKey := tileKey{tileLevel, tileIndex}
-		t, ok := n.tiles[tKey]
-		if !ok {
-			span.AddEvent("cache miss")
-			p := layout.PartialTileSize(tileLevel, tileIndex, n.logSize)
-			tileRaw, err := n.getTile(ctx, tileLevel, tileIndex, p)
+		span.AddEvent("cache miss")
+
+		p := layout.PartialTileSize(tileLevel, tileIndex, n.logSize)
+		nodes, err := n.fetchTileNodes(ctx, tileLevel, tileIndex, p)
+		if err != nil {
+			return nil, fmt.Errorf("failed to fetch and populate node cache: %v", err)
+		}
+		for k, v := range nodes {
+			n.nodes.Add(k, v)
+		}
+		if e, ok := nodes[id]; ok {
+			return e, nil
+		}
+		return nil, fmt.Errorf("internal error: missing node %+v", id)
+	})
+}
+
+// GetNodes returns the tree hashes at the provided locations.
+func (n *nodeCache) GetNodes(ctx context.Context, nIDs []compact.NodeID) ([][]byte, error) {
+	hashes := make([][]byte, len(nIDs))
+	g, ctx := errgroup.WithContext(ctx)
+	for i, id := range nIDs {
+		g.Go(func() error {
+			h, err := n.GetNode(ctx, id)
 			if err != nil {
-				return nil, fmt.Errorf("failed to fetch tile: %v", err)
-			}
-			var tile api.HashTile
-			if err := tile.UnmarshalText(tileRaw); err != nil {
-				return nil, fmt.Errorf("failed to parse tile: %v", err)
+				return err
 			}
-			t = tile
-			n.tiles[tKey] = tile
+			hashes[i] = h
+			return nil
+		})
+	}
+	if err := g.Wait(); err != nil {
+		return nil, err
+	}
+	return hashes, nil
+}
+
+// fetchTileNodes retrieves the specified tile, parses it, and returns a map of tree-space-coordinate to node hash.
+func (n *nodeCache) fetchTileNodes(ctx context.Context, tileLevel, tileIndex uint64, p uint8) (map[compact.NodeID][]byte, error) {
+	return otel.Trace(ctx, "tessera.client.nodecache.fetchTileNodes", tracer, func(ctx context.Context, span trace.Span) (map[compact.NodeID][]byte, error) {
+		tileRaw, err := n.getTile(ctx, tileLevel, tileIndex, p)
+		if err != nil {
+			return nil, fmt.Errorf("failed to fetch tile: %v", err)
 		}
-		// We've got the tile, now we need to look up (or calculate) the node inside of it
-		numLeaves := 1 << nodeLevel
-		firstLeaf := int(nodeIndex) * numLeaves
-		lastLeaf := firstLeaf + numLeaves
-		if lastLeaf > len(t.Nodes) {
-			return nil, fmt.Errorf("require leaf nodes [%d, %d) but only got %d leaves", firstLeaf, lastLeaf, len(t.Nodes))
+
+		var tile api.HashTile
+		if err := tile.UnmarshalText(tileRaw); err != nil {
+			return nil, fmt.Errorf("failed to parse tile: %v", err)
+		}
+
+		ret := make(map[compact.NodeID][]byte, 256*2-1)
+		// visitFn is a visitor callback which populates the nodes cache.
+		// Used by the calls to compact range below.
+		visitFn := func(intID compact.NodeID, h []byte) {
+			// Figure out the "global" nodeID for the node intID in the requested tile.
+			i := compact.NodeID{
+				Level: uint(tileLevel*layout.TileHeight) + intID.Level,
+				Index: (tileIndex*layout.TileWidth)>>intID.Level + intID.Index,
+			}
+			ret[i] = h
 		}
 		rf := compact.RangeFactory{Hash: hasher.HashChildren}
 		r := rf.NewEmptyRange(0)
-		for _, l := range t.Nodes[firstLeaf:lastLeaf] {
-			if err := r.Append(l, nil); err != nil {
+		for _, l := range tile.Nodes {
+			if err := r.Append(l, visitFn); err != nil {
 				return nil, fmt.Errorf("failed to Append: %v", err)
 			}
 		}
-		return r.GetRootHash(nil)
+		if _, err := r.GetRootHash(visitFn); err != nil {
+			return nil, fmt.Errorf("failed to visit all nodes: %v", err)
+		}
+		return ret, nil
 	})
 }
+
+// cLock is a mutex which keeps track of the number of goroutines attempting to acquire a lock.
+type cLock struct {
+	sync.Mutex
+	n int64
+}
+
+// shardedMutex is a set of mutexes sharded by key.
+//
+// For a given key, it acts as a regular mutex with the exception that it also tracks the number
+// of blocked goroutines waiting to acquire the lock.
+//
+// If a mutex doesn't exist for a given key at the point that Lock is called, one will be created.
+// To help guard against unbounded growth, mutexes with zero pending waiters at the point they're unlocked are deleted.
+type shardedMutex[K comparable] struct {
+	// m protects the locks map and the waiter counts it contains.
+	m     *sync.Mutex
+	locks map[K]*cLock
+}
+
+// newShardedMutex creates a new shardedLock instance.
+func newShardedMutex[K comparable]() *shardedMutex[K] {
+	return &shardedMutex[K]{
+		m:     &sync.Mutex{},
+		locks: make(map[K]*cLock),
+	}
+}
+
+// Lock locks the given key.
+func (sl *shardedMutex[K]) Lock(k K) {
+	sl.m.Lock()
+	l, ok := sl.locks[k]
+	if !ok {
+		l = &cLock{}
+		sl.locks[k] = l
+	}
+	l.n++
+	sl.m.Unlock()
+
+	l.Lock()
+}
+
+// Unlock unlocks the given key.
+func (sl *shardedMutex[K]) Unlock(k K) {
+	sl.m.Lock()
+	l, ok := sl.locks[k]
+	if !ok {
+		panic("unlock on non-existent key")
+	}
+	l.n--
+	if l.n == 0 {
+		delete(sl.locks, k)
+	}
+	sl.m.Unlock()
+
+	l.Unlock()
+}