Cache Block Ancestor for Fork Choice (#2269)

beacon-chain/cache/BUILD.bazel

@@ -2,10 +2,14 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "go_default_library",
-    srcs = ["committee.go"],
+    srcs = [
+        "block.go",
+        "committee.go",
+    ],
     importpath = "github.com/prysmaticlabs/prysm/beacon-chain/cache",
     visibility = ["//beacon-chain:__subpackages__"],
     deps = [
+        "//proto/beacon/p2p/v1:go_default_library",
         "//shared/params:go_default_library",
         "@com_github_prometheus_client_golang//prometheus:go_default_library",
         "@com_github_prometheus_client_golang//prometheus/promauto:go_default_library",
@@ -15,6 +19,10 @@ go_library(
 
 go_test(
     name = "go_default_test",
-    srcs = ["committee_test.go"],
+    srcs = [
+        "block_test.go",
+        "committee_test.go",
+    ],
     embed = [":go_default_library"],
+    deps = ["//proto/beacon/p2p/v1:go_default_library"],
 )

beacon-chain/cache/block.go

@@ -0,0 +1,104 @@
+package cache
+
+import (
+	"errors"
+	"strconv"
+	"sync"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
+	pb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
+	"k8s.io/client-go/tools/cache"
+)
+
+var (
+	// ErrNotAncestorCacheObj will be returned when a cache object is not a pointer to
+	// block ancestor cache obj.
+	ErrNotAncestorCacheObj = errors.New("object is not an ancestor object for cache")
+	// Metrics
+	ancestorBlockCacheMiss = promauto.NewCounter(prometheus.CounterOpts{
+		Name: "ancestor_block_cache_miss",
+		Help: "The number of ancestor block requests that aren't present in the cache.",
+	})
+	ancestorBlockCacheHit = promauto.NewCounter(prometheus.CounterOpts{
+		Name: "ancestor_block_cache_hit",
+		Help: "The number of ancestor block requests that are present in the cache.",
+	})
+	ancestorBlockCacheSize = promauto.NewGauge(prometheus.GaugeOpts{
+		Name: "ancestor_block_cache_size",
+		Help: "The number of ancestor blocks in the ancestorBlock cache",
+	})
+)
+
+// AncestorInfo defines the cached ancestor block object for height.
+type AncestorInfo struct {
+	Hash   []byte
+	Height uint64
+	Block  *pb.BeaconBlock
+}
+
+// AncestorBlockCache structs with 1 queue for looking up block ancestor by height.
+type AncestorBlockCache struct {
+	ancestorBlockCache *cache.FIFO
+	lock               sync.RWMutex
+}
+
+// heightKeyFn takes the string representation of the block hash + height as the key
+// for the ancestor of a given block (AncestorInfo).
+func heightKeyFn(obj interface{}) (string, error) {
+	aInfo, ok := obj.(*AncestorInfo)
+	if !ok {
+		return "", ErrNotAncestorCacheObj
+	}
+
+	return string(aInfo.Hash) + strconv.Itoa(int(aInfo.Height)), nil
+}
+
+// NewBlockAncestorCache creates a new block ancestor cache for storing/accessing block ancestor
+// from memory.
+func NewBlockAncestorCache() *AncestorBlockCache {
+	return &AncestorBlockCache{
+		ancestorBlockCache: cache.NewFIFO(heightKeyFn),
+	}
+}
+
+// AncestorBySlot fetches block's ancestor by height. Returns true with a
+// reference to the ancestor block, if exists. Otherwise returns false, nil.
+func (a *AncestorBlockCache) AncestorBySlot(blockHash []byte, height uint64) (*AncestorInfo, error) {
+	a.lock.RLock()
+	defer a.lock.RUnlock()
+
+	obj, exists, err := a.ancestorBlockCache.GetByKey(string(blockHash) + strconv.Itoa(int(height)))
+	if err != nil {
+		return nil, err
+	}
+
+	if exists {
+		ancestorBlockCacheHit.Inc()
+	} else {
+		ancestorBlockCacheMiss.Inc()
+		return nil, nil
+	}
+
+	aInfo, ok := obj.(*AncestorInfo)
+	if !ok {
+		return nil, ErrNotACommitteeInfo
+	}
+
+	return aInfo, nil
+}
+
+// AddBlockAncestor adds block ancestor object to the cache. This method also trims the least
+// recently added ancestor if the cache size has ready the max cache size limit.
+func (a *AncestorBlockCache) AddBlockAncestor(ancestorInfo *AncestorInfo) error {
+	a.lock.Lock()
+	defer a.lock.Unlock()
+
+	if err := a.ancestorBlockCache.AddIfNotPresent(ancestorInfo); err != nil {
+		return err
+	}
+
+	trim(a.ancestorBlockCache, maxCacheSize)
+	ancestorBlockCacheSize.Set(float64(len(a.ancestorBlockCache.ListKeys())))
+	return nil
+}

beacon-chain/cache/block_test.go

@@ -0,0 +1,101 @@
+package cache
+
+import (
+	"reflect"
+	"strconv"
+	"testing"
+
+	pb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
+)
+
+func TestHeightHeightFn_OK(t *testing.T) {
+	height := uint64(999)
+	hash := []byte{'A'}
+	aInfo := &AncestorInfo{
+		Height: height,
+		Hash:   hash,
+	}
+
+	key, err := heightKeyFn(aInfo)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	strHeightKey := string(aInfo.Hash) + strconv.Itoa(int(aInfo.Height))
+	if key != strHeightKey {
+		t.Errorf("Incorrect hash key: %s, expected %s", key, strHeightKey)
+	}
+}
+
+func TestHeightKeyFn_InvalidObj(t *testing.T) {
+	_, err := heightKeyFn("bad")
+	if err != ErrNotAncestorCacheObj {
+		t.Errorf("Expected error %v, got %v", ErrNotAncestorCacheObj, err)
+	}
+}
+
+func TestAncestorCache_AncestorInfoByHeight(t *testing.T) {
+	cache := NewBlockAncestorCache()
+
+	height := uint64(123)
+	hash := []byte{'B'}
+	aInfo := &AncestorInfo{
+		Height: height,
+		Hash:   hash,
+		Block:  &pb.BeaconBlock{Slot: height},
+	}
+
+	fetchedInfo, err := cache.AncestorBySlot(hash, height)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if fetchedInfo != nil {
+		t.Error("Expected ancestor info not to exist in empty cache")
+	}
+
+	if err := cache.AddBlockAncestor(aInfo); err != nil {
+		t.Fatal(err)
+	}
+	fetchedInfo, err = cache.AncestorBySlot(hash, height)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if fetchedInfo == nil {
+		t.Error("Expected ancestor info to exist")
+	}
+	if fetchedInfo.Height != height {
+		t.Errorf(
+			"Expected fetched slot number to be %d, got %d",
+			aInfo.Height,
+			fetchedInfo.Height,
+		)
+	}
+	if !reflect.DeepEqual(fetchedInfo.Block, aInfo.Block) {
+		t.Errorf(
+			"Expected fetched info committee to be %v, got %v",
+			aInfo.Block,
+			fetchedInfo.Block,
+		)
+	}
+}
+
+func TestBlockAncestor_maxSize(t *testing.T) {
+	cache := NewBlockAncestorCache()
+
+	for i := 0; i < maxCacheSize+10; i++ {
+		aInfo := &AncestorInfo{
+			Height: uint64(i),
+		}
+		if err := cache.AddBlockAncestor(aInfo); err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	if len(cache.ancestorBlockCache.ListKeys()) != maxCacheSize {
+		t.Errorf(
+			"Expected hash cache key size to be %d, got %d",
+			maxCacheSize,
+			len(cache.ancestorBlockCache.ListKeys()),
+		)
+	}
+}