From 4c66edf2c229ff1b4f171041fa1ee7e037a667b7 Mon Sep 17 00:00:00 2001
From: Shay Zluf <thezluf@gmail.com>
Date: Sun, 14 Jun 2020 05:51:54 +0300
Subject: [PATCH] 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        |  2 +
 slasher/cache/doc.go             |  3 ++
 slasher/cache/flat_span_cache.go | 67 ++++++++++++++++++++++++++++++++
 slasher/cache/span_cache.go      |  3 --
 slasher/db/kv/kv.go              |  6 +++
 slasher/db/kv/spanner_new.go     |  2 +-
 6 files changed, 79 insertions(+), 4 deletions(-)
 create mode 100644 slasher/cache/doc.go
 create mode 100644 slasher/cache/flat_span_cache.go

diff --git a/slasher/cache/BUILD.bazel b/slasher/cache/BUILD.bazel
index d6f67040c..c7453ca76 100644
--- a/slasher/cache/BUILD.bazel
+++ b/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",
     ],
diff --git a/slasher/cache/doc.go b/slasher/cache/doc.go
new file mode 100644
index 000000000..7fcc59c1d
--- /dev/null
+++ b/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
diff --git a/slasher/cache/flat_span_cache.go b/slasher/cache/flat_span_cache.go
new file mode 100644
index 000000000..0e2fe44bf
--- /dev/null
+++ b/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()
+}
diff --git a/slasher/cache/span_cache.go b/slasher/cache/span_cache.go
index f3a39690e..381156261 100644
--- a/slasher/cache/span_cache.go
+++ b/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 (
diff --git a/slasher/db/kv/kv.go b/slasher/db/kv/kv.go
index 03b388be8..b18711a7b 100644
--- a/slasher/db/kv/kv.go
+++ b/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(
diff --git a/slasher/db/kv/spanner_new.go b/slasher/db/kv/spanner_new.go
index 4ddfcdcae..34d5df4ab 100644
--- a/slasher/db/kv/spanner_new.go
+++ b/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"