[Access] Upgrade lru cache

engine/access/rpc/backend/backend.go

@@ -7,7 +7,8 @@ import (
 	"strconv"
 	"time"
 
-	lru "github.com/hashicorp/golang-lru"
+	lru "github.com/hashicorp/golang-lru/v2"
+	accessproto "github.com/onflow/flow/protobuf/go/flow/access"
 	"github.com/rs/zerolog"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -22,8 +23,6 @@ import (
 	"github.com/onflow/flow-go/module"
 	"github.com/onflow/flow-go/state/protocol"
 	"github.com/onflow/flow-go/storage"
-
-	accessproto "github.com/onflow/flow/protobuf/go/flow/access"
 )
 
 // minExecutionNodesCnt is the minimum number of execution nodes expected to have sent the execution receipt for a block
@@ -119,7 +118,7 @@ func New(
 		retry.Activate()
 	}
 
-	loggedScripts, err := lru.New(DefaultLoggedScriptsCacheSize)
+	loggedScripts, err := lru.New[[16]byte, time.Time](DefaultLoggedScriptsCacheSize)
 	if err != nil {
 		log.Fatal().Err(err).Msg("failed to initialize script logging cache")
 	}
@@ -228,14 +227,13 @@ func NewCache(
 	log zerolog.Logger,
 	accessMetrics module.AccessMetrics,
 	connectionPoolSize uint,
-) (*lru.Cache, uint, error) {
+) (*lru.Cache[string, *connection.CachedClient], uint, error) {
 
-	var cache *lru.Cache
+	var cache *lru.Cache[string, *connection.CachedClient]
 	cacheSize := connectionPoolSize
 	if cacheSize > 0 {
 		var err error
-		cache, err = lru.NewWithEvict(int(cacheSize), func(_, evictedValue interface{}) {
-			store := evictedValue.(*connection.CachedClient)
+		cache, err = lru.NewWithEvict[string, *connection.CachedClient](int(cacheSize), func(_ string, store *connection.CachedClient) {
 			store.Close()
 			log.Debug().Str("grpc_conn_evicted", store.Address).Msg("closing grpc connection evicted from pool")
 			if accessMetrics != nil {

engine/access/rpc/backend/backend_scripts.go

@@ -8,7 +8,7 @@ import (
 	"time"
 
 	"github.com/hashicorp/go-multierror"
-	lru "github.com/hashicorp/golang-lru"
+	lru "github.com/hashicorp/golang-lru/v2"
 	"github.com/onflow/flow/protobuf/go/flow/access"
 	execproto "github.com/onflow/flow/protobuf/go/flow/execution"
 	"github.com/rs/zerolog"
@@ -27,13 +27,14 @@ import (
 const uniqueScriptLoggingTimeWindow = 10 * time.Minute
 
 type backendScripts struct {
-	headers              storage.Headers
-	executionReceipts    storage.ExecutionReceipts
-	state                protocol.State
-	connFactory          connection.ConnectionFactory
-	log                  zerolog.Logger
-	metrics              module.BackendScriptsMetrics
-	loggedScripts        *lru.Cache
+	headers           storage.Headers
+	executionReceipts storage.ExecutionReceipts
+	state             protocol.State
+	connFactory       connection.ConnectionFactory
+	log               zerolog.Logger
+	metrics           module.BackendScriptsMetrics
+	loggedScripts     *lru.Cache[[16]byte, time.Time]
+
 	archiveAddressList   []string
 	archivePorts         []uint
 	scriptExecValidation bool
@@ -310,14 +311,11 @@ func (b *backendScripts) executeScriptOnAvailableExecutionNodes(
 
 // shouldLogScript checks if the script hash is unique in the time window
 func (b *backendScripts) shouldLogScript(execTime time.Time, scriptHash [16]byte) bool {
-	rawTimestamp, seen := b.loggedScripts.Get(scriptHash)
-	if !seen || rawTimestamp == nil {
+	timestamp, seen := b.loggedScripts.Get(scriptHash)
+	if !seen {
 		return true
-	} else {
-		// safe cast
-		timestamp := rawTimestamp.(time.Time)
-		return execTime.Sub(timestamp) >= uniqueScriptLoggingTimeWindow
 	}
+	return execTime.Sub(timestamp) >= uniqueScriptLoggingTimeWindow
 }
 
 func (b *backendScripts) tryExecuteScriptOnExecutionNode(

engine/access/rpc/connection/cache.go

@@ -4,7 +4,7 @@ import (
 	"sync"
 	"time"
 
-	lru "github.com/hashicorp/golang-lru"
+	lru "github.com/hashicorp/golang-lru/v2"
 	"go.uber.org/atomic"
 	"google.golang.org/grpc"
 )
@@ -38,12 +38,12 @@ func (s *CachedClient) Close() {
 
 // Cache represents a cache of CachedClient instances with a given maximum size.
 type Cache struct {
-	cache *lru.Cache
+	cache *lru.Cache[string, *CachedClient]
 	size  int
 }
 
 // NewCache creates a new Cache with the specified maximum size and the underlying LRU cache.
-func NewCache(cache *lru.Cache, size int) *Cache {
+func NewCache(cache *lru.Cache[string, *CachedClient], size int) *Cache {
 	return &Cache{
 		cache: cache,
 		size:  size,
@@ -57,7 +57,7 @@ func (c *Cache) Get(address string) (*CachedClient, bool) {
 	if !ok {
 		return nil, false
 	}
-	return val.(*CachedClient), true
+	return val, true
 }
 
 // GetOrAdd atomically gets the CachedClient for the given address from the cache, or adds a new one
@@ -71,7 +71,7 @@ func (c *Cache) GetOrAdd(address string, timeout time.Duration) (*CachedClient,
 
 	val, existed, _ := c.cache.PeekOrAdd(address, client)
 	if existed {
-		return val.(*CachedClient), true
+		return val, true
 	}
 
 	client.Address = address

engine/access/rpc/connection/connection_test.go

@@ -4,29 +4,24 @@ import (
 	"context"
 	"fmt"
 	"net"
-	"strconv"
-	"strings"
 	"sync"
 	"testing"
 	"time"
 
-	"github.com/sony/gobreaker"
-
-	"go.uber.org/atomic"
-	"pgregory.net/rapid"
-
-	lru "github.com/hashicorp/golang-lru"
+	lru "github.com/hashicorp/golang-lru/v2"
 	"github.com/onflow/flow/protobuf/go/flow/access"
 	"github.com/onflow/flow/protobuf/go/flow/execution"
+	"github.com/sony/gobreaker"
 	"github.com/stretchr/testify/assert"
 	testifymock "github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/atomic"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/status"
+	"pgregory.net/rapid"
 
-	"github.com/onflow/flow-go/engine/access/mock"
 	"github.com/onflow/flow-go/module/metrics"
 	"github.com/onflow/flow-go/utils/unittest"
 )
@@ -129,8 +124,8 @@ func TestProxyAccessAPIConnectionReuse(t *testing.T) {
 	connectionFactory.CollectionGRPCPort = cn.port
 	// set the connection pool cache size
 	cacheSize := 1
-	cache, _ := lru.NewWithEvict(cacheSize, func(_, evictedValue interface{}) {
-		evictedValue.(*CachedClient).Close()
+	cache, _ := lru.NewWithEvict[string, *CachedClient](cacheSize, func(_ string, client *CachedClient) {
+		client.Close()
 	})
 	connectionCache := NewCache(cache, cacheSize)
 
@@ -184,8 +179,8 @@ func TestProxyExecutionAPIConnectionReuse(t *testing.T) {
 	connectionFactory.ExecutionGRPCPort = en.port
 	// set the connection pool cache size
 	cacheSize := 5
-	cache, _ := lru.NewWithEvict(cacheSize, func(_, evictedValue interface{}) {
-		evictedValue.(*CachedClient).Close()
+	cache, _ := lru.NewWithEvict[string, *CachedClient](cacheSize, func(_ string, client *CachedClient) {
+		client.Close()
 	})
 	connectionCache := NewCache(cache, cacheSize)
 	// set metrics reporting
@@ -245,8 +240,8 @@ func TestExecutionNodeClientTimeout(t *testing.T) {
 	connectionFactory.ExecutionNodeGRPCTimeout = timeout
 	// set the connection pool cache size
 	cacheSize := 5
-	cache, _ := lru.NewWithEvict(cacheSize, func(_, evictedValue interface{}) {
-		evictedValue.(*CachedClient).Close()
+	cache, _ := lru.NewWithEvict[string, *CachedClient](cacheSize, func(_ string, client *CachedClient) {
+		client.Close()
 	})
 	connectionCache := NewCache(cache, cacheSize)
 	// set metrics reporting
@@ -294,8 +289,8 @@ func TestCollectionNodeClientTimeout(t *testing.T) {
 	connectionFactory.CollectionNodeGRPCTimeout = timeout
 	// set the connection pool cache size
 	cacheSize := 5
-	cache, _ := lru.NewWithEvict(cacheSize, func(_, evictedValue interface{}) {
-		evictedValue.(*CachedClient).Close()
+	cache, _ := lru.NewWithEvict[string, *CachedClient](cacheSize, func(_ string, client *CachedClient) {
+		client.Close()
 	})
 	connectionCache := NewCache(cache, cacheSize)
 	// set metrics reporting
@@ -343,9 +338,10 @@ func TestConnectionPoolFull(t *testing.T) {
 	connectionFactory.CollectionGRPCPort = cn1.port
 	// set the connection pool cache size
 	cacheSize := 2
-	cache, _ := lru.NewWithEvict(cacheSize, func(_, evictedValue interface{}) {
-		evictedValue.(*CachedClient).Close()
+	cache, _ := lru.NewWithEvict[string, *CachedClient](cacheSize, func(_ string, client *CachedClient) {
+		client.Close()
 	})
+
 	connectionCache := NewCache(cache, cacheSize)
 	// set metrics reporting
 	connectionFactory.AccessMetrics = metrics.NewNoopCollector()
@@ -419,9 +415,10 @@ func TestConnectionPoolStale(t *testing.T) {
 	connectionFactory.CollectionGRPCPort = cn.port
 	// set the connection pool cache size
 	cacheSize := 5
-	cache, _ := lru.NewWithEvict(cacheSize, func(_, evictedValue interface{}) {
-		evictedValue.(*CachedClient).Close()
+	cache, _ := lru.NewWithEvict[string, *CachedClient](cacheSize, func(_ string, client *CachedClient) {
+		client.Close()
 	})
+
 	connectionCache := NewCache(cache, cacheSize)
 	// set metrics reporting
 	connectionFactory.AccessMetrics = metrics.NewNoopCollector()
@@ -508,9 +505,10 @@ func TestExecutionNodeClientClosedGracefully(t *testing.T) {
 		connectionFactory.ExecutionNodeGRPCTimeout = time.Second
 		// set the connection pool cache size
 		cacheSize := 1
-		cache, _ := lru.NewWithEvict(cacheSize, func(_, evictedValue interface{}) {
-			evictedValue.(*CachedClient).Close()
+		cache, _ := lru.NewWithEvict[string, *CachedClient](cacheSize, func(_ string, client *CachedClient) {
+			client.Close()
 		})
+
 		connectionCache := NewCache(cache, cacheSize)
 		// set metrics reporting
 		connectionFactory.AccessMetrics = metrics.NewNoopCollector()
@@ -592,7 +590,7 @@ func TestExecutionEvictingCacheClients(t *testing.T) {
 	connectionFactory.CollectionNodeGRPCTimeout = 5 * time.Second
 	// Set the connection pool cache size
 	cacheSize := 1
-	cache, err := lru.New(cacheSize)
+	cache, err := lru.New[string, *CachedClient](cacheSize)
 	require.NoError(t, err)
 
 	connectionCache := NewCache(cache, cacheSize)
@@ -614,8 +612,7 @@ func TestExecutionEvictingCacheClients(t *testing.T) {
 	ctx := context.Background()
 
 	// Retrieve the cached client from the cache
-	result, _ := cache.Get(clientAddress)
-	cachedClient := result.(*CachedClient)
+	cachedClient, _ := cache.Get(clientAddress)
 
 	// Schedule the invalidation of the access API client after a delay
 	time.AfterFunc(250*time.Millisecond, func() {
@@ -671,7 +668,7 @@ func TestCircuitBreakerExecutionNode(t *testing.T) {
 
 	// Set the connection pool cache size.
 	cacheSize := 1
-	connectionCache, _ := lru.New(cacheSize)
+	connectionCache, _ := lru.New[string, *CachedClient](cacheSize)
 
 	connectionFactory.Manager = NewManager(
 		NewCache(connectionCache, cacheSize),
@@ -755,8 +752,8 @@ func TestCircuitBreakerCollectionNode(t *testing.T) {
 
 	// Set the connection pool cache size.
 	cacheSize := 1
-	connectionCache, _ := lru.New(cacheSize)
 
+	connectionCache, _ := lru.New[string, *CachedClient](cacheSize)
 	connectionFactory.Manager = NewManager(
 		NewCache(connectionCache, cacheSize),
 		unittest.Logger(),
@@ -812,79 +809,3 @@ func TestCircuitBreakerCollectionNode(t *testing.T) {
 	assert.Greater(t, requestTimeout, duration)
 	assert.Equal(t, nil, err)
 }
-
-// node mocks a flow node that runs a GRPC server
-type node struct {
-	server   *grpc.Server
-	listener net.Listener
-	port     uint
-}
-
-func (n *node) setupNode(t *testing.T) {
-	n.server = grpc.NewServer()
-	listener, err := net.Listen("tcp4", unittest.DefaultAddress)
-	assert.NoError(t, err)
-	n.listener = listener
-	assert.Eventually(t, func() bool {
-		return !strings.HasSuffix(listener.Addr().String(), ":0")
-	}, time.Second*4, 10*time.Millisecond)
-
-	_, port, err := net.SplitHostPort(listener.Addr().String())
-	assert.NoError(t, err)
-	portAsUint, err := strconv.ParseUint(port, 10, 32)
-	assert.NoError(t, err)
-	n.port = uint(portAsUint)
-}
-
-func (n *node) start(t *testing.T) {
-	// using a wait group here to ensure the goroutine has started before returning. Otherwise,
-	// there's a race condition where the server is sometimes stopped before it has started
-	wg := sync.WaitGroup{}
-	wg.Add(1)
-	go func() {
-		wg.Done()
-		err := n.server.Serve(n.listener)
-		assert.NoError(t, err)
-	}()
-	unittest.RequireReturnsBefore(t, wg.Wait, 10*time.Millisecond, "could not start goroutine on time")
-}
-
-func (n *node) stop(t *testing.T) {
-	if n.server != nil {
-		n.server.Stop()
-	}
-}
-
-type executionNode struct {
-	node
-	handler *mock.ExecutionAPIServer
-}
-
-func (en *executionNode) start(t *testing.T) {
-	en.setupNode(t)
-	handler := new(mock.ExecutionAPIServer)
-	execution.RegisterExecutionAPIServer(en.server, handler)
-	en.handler = handler
-	en.node.start(t)
-}
-
-func (en *executionNode) stop(t *testing.T) {
-	en.node.stop(t)
-}
-
-type collectionNode struct {
-	node
-	handler *mock.AccessAPIServer
-}
-
-func (cn *collectionNode) start(t *testing.T) {
-	cn.setupNode(t)
-	handler := new(mock.AccessAPIServer)
-	access.RegisterAccessAPIServer(cn.server, handler)
-	cn.handler = handler
-	cn.node.start(t)
-}
-
-func (cn *collectionNode) stop(t *testing.T) {
-	cn.node.stop(t)
-}