Add cache for flan spans (#6199)

* flat span cache

* Plug cache into kv

* Add comments

* Fix

* Add doc.go

* Gaz

Co-authored-by: Ivan Martinez <ivanthegreatdev@gmail.com>
Co-authored-by: prylabs-bulldozer[bot] <58059840+prylabs-bulldozer[bot]@users.noreply.github.com>

slasher/cache/BUILD.bazel

@@ -3,6 +3,8 @@ load("@prysm//tools/go:def.bzl", "go_library")
 go_library(
     name = "go_default_library",
     srcs = [
+        "doc.go",
+        "flat_span_cache.go",
         "span_cache.go",
         "validators_cache.go",
     ],

slasher/cache/doc.go

@@ -0,0 +1,3 @@
+// Package cache contains critical caches necessary for the runtime
+// of the slasher service, such as a cache for maintaining validator history on a per epoch basis
+package cache

slasher/cache/flat_span_cache.go

@@ -0,0 +1,67 @@
+package cache
+
+import (
+	lru "github.com/hashicorp/golang-lru"
+	"github.com/prysmaticlabs/prysm/slasher/detection/attestations/types"
+	log "github.com/sirupsen/logrus"
+)
+
+// EpochFlatSpansCache is used to store the spans needed on a per-epoch basis for slashing detection.
+type EpochFlatSpansCache struct {
+	cache *lru.Cache
+}
+
+// NewEpochFlatSpansCache initializes the underlying cache with the given size and on evict function.
+func NewEpochFlatSpansCache(size int, onEvicted func(key interface{}, value interface{})) (*EpochFlatSpansCache, error) {
+	if size != 0 {
+		epochSpansCacheSize = size
+	}
+	cache, err := lru.NewWithEvict(epochSpansCacheSize, onEvicted)
+	if err != nil {
+		return nil, err
+	}
+	return &EpochFlatSpansCache{cache: cache}, nil
+}
+
+// Get returns an ok bool and the cached value for the requested epoch key, if any.
+func (c *EpochFlatSpansCache) Get(epoch uint64) (*types.EpochStore, bool) {
+	item, exists := c.cache.Get(epoch)
+	if exists && item != nil {
+		epochSpansCacheHit.Inc()
+		return item.(*types.EpochStore), true
+	}
+
+	epochSpansCacheMiss.Inc()
+	return &types.EpochStore{}, false
+}
+
+// Set the response in the cache.
+func (c *EpochFlatSpansCache) Set(epoch uint64, epochSpans *types.EpochStore) {
+	_ = c.cache.Add(epoch, epochSpans)
+}
+
+// Delete removes an epoch from the cache and returns if it existed or not.
+// Performs the onEviction function before removal.
+func (c *EpochFlatSpansCache) Delete(epoch uint64) bool {
+	return c.cache.Remove(epoch)
+}
+
+// PruneOldest removes the oldest key from the span cache, calling its OnEvict function.
+func (c *EpochFlatSpansCache) PruneOldest() uint64 {
+	if c.cache.Len() == epochSpansCacheSize {
+		epoch, _, _ := c.cache.RemoveOldest()
+		return epoch.(uint64)
+	}
+	return 0
+}
+
+// Has returns true if the key exists in the cache.
+func (c *EpochFlatSpansCache) Has(epoch uint64) bool {
+	return c.cache.Contains(epoch)
+}
+
+// Purge removes all keys from the SpanCache and evicts all current data.
+func (c *EpochFlatSpansCache) Purge() {
+	log.Info("Saving all cached data to DB, please wait for completion.")
+	c.cache.Purge()
+}

slasher/cache/span_cache.go

@@ -1,6 +1,3 @@
-// Package cache contains critical caches necessary for the runtime
-// of the slasher service, being able to keep attestation spans by epoch
-// for the validators active in the beacon chain.
 package cache
 
 import (

slasher/db/kv/kv.go

@@ -22,6 +22,7 @@ type Store struct {
 	db               *bolt.DB
 	databasePath     string
 	spanCache        *cache.EpochSpansCache
+	flatSpanCache    *cache.EpochFlatSpansCache
 	spanCacheEnabled bool
 }
 
@@ -104,6 +105,11 @@ func NewKVStore(dirPath string, cfg *Config) (*Store, error) {
 		return nil, errors.Wrap(err, "could not create new cache")
 	}
 	kv.spanCache = spanCache
+	flatSpanCache, err := cache.NewEpochFlatSpansCache(cfg.SpanCacheSize, persistFlatSpanMapsOnEviction(kv))
+	if err != nil {
+		return nil, errors.Wrap(err, "could not create new flat cache")
+	}
+	kv.flatSpanCache = flatSpanCache
 
 	if err := kv.db.Update(func(tx *bolt.Tx) error {
 		return createBuckets(

slasher/db/kv/spanner_new.go

@@ -2,8 +2,8 @@ package kv
 
 import (
 	"context"
-	"errors"
 
+	"github.com/pkg/errors"
 	"github.com/prysmaticlabs/prysm/shared/bytesutil"
 	"github.com/prysmaticlabs/prysm/slasher/detection/attestations/types"
 	log "github.com/sirupsen/logrus"