Evict blocks based on size

snow/engine/snowman/transitive.go

@@ -22,19 +22,25 @@ import (
 	"github.com/ava-labs/avalanchego/snow/events"
 	"github.com/ava-labs/avalanchego/snow/validators"
 	"github.com/ava-labs/avalanchego/utils/bag"
+	"github.com/ava-labs/avalanchego/utils/constants"
 	"github.com/ava-labs/avalanchego/utils/logging"
 	"github.com/ava-labs/avalanchego/utils/set"
+	"github.com/ava-labs/avalanchego/utils/units"
 	"github.com/ava-labs/avalanchego/utils/wrappers"
 )
 
-const nonVerifiedCacheSize = 128
+const nonVerifiedCacheSize = 128 * units.MiB
 
 var _ Engine = (*Transitive)(nil)
 
 func New(config Config) (Engine, error) {
 	return newTransitive(config)
 }
 
+func cachedBlockSize(_ ids.ID, blk snowman.Block) int {
+	return ids.IDLen + len(blk.Bytes()) + constants.PointerOverhead
+}
+
 // Transitive implements the Engine interface by attempting to fetch all
 // Transitive dependencies.
 type Transitive struct {
@@ -92,7 +98,10 @@ func newTransitive(config Config) (*Transitive, error) {
 	nonVerifiedCache, err := metercacher.New[ids.ID, snowman.Block](
 		"non_verified_cache",
 		config.Ctx.Registerer,
-		&cache.LRU[ids.ID, snowman.Block]{Size: nonVerifiedCacheSize},
+		cache.NewSizedLRU[ids.ID, snowman.Block](
+			nonVerifiedCacheSize,
+			cachedBlockSize,
+		),
 	)
 	if err != nil {
 		return nil, err

utils/constants/memory.go

@@ -0,0 +1,8 @@
+// Copyright (C) 2019-2023, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package constants
+
+// PointerOverhead is used to approximate the memory footprint from allocating a
+// pointer.
+const PointerOverhead = 8

vms/components/chain/state.go

@@ -17,8 +17,17 @@ import (
 	"github.com/ava-labs/avalanchego/snow/choices"
 	"github.com/ava-labs/avalanchego/snow/consensus/snowman"
 	"github.com/ava-labs/avalanchego/snow/engine/snowman/block"
+	"github.com/ava-labs/avalanchego/utils/constants"
 )
 
+func cachedBlockSize(_ ids.ID, bw *BlockWrapper) int {
+	return ids.IDLen + len(bw.Bytes()) + 2*constants.PointerOverhead
+}
+
+func cachedBlockBytesSize(blockBytes string, _ ids.ID) int {
+	return len(blockBytes) + ids.IDLen
+}
+
 // State implements an efficient caching layer used to wrap a VM
 // implementation.
 type State struct {
@@ -138,11 +147,20 @@ func (s *State) initialize(config *Config) {
 
 func NewState(config *Config) *State {
 	c := &State{
-		verifiedBlocks:   make(map[ids.ID]*BlockWrapper),
-		decidedBlocks:    &cache.LRU[ids.ID, *BlockWrapper]{Size: config.DecidedCacheSize},
-		missingBlocks:    &cache.LRU[ids.ID, struct{}]{Size: config.MissingCacheSize},
-		unverifiedBlocks: &cache.LRU[ids.ID, *BlockWrapper]{Size: config.UnverifiedCacheSize},
-		bytesToIDCache:   &cache.LRU[string, ids.ID]{Size: config.BytesToIDCacheSize},
+		verifiedBlocks: make(map[ids.ID]*BlockWrapper),
+		decidedBlocks: cache.NewSizedLRU[ids.ID, *BlockWrapper](
+			config.DecidedCacheSize,
+			cachedBlockSize,
+		),
+		missingBlocks: &cache.LRU[ids.ID, struct{}]{Size: config.MissingCacheSize},
+		unverifiedBlocks: cache.NewSizedLRU[ids.ID, *BlockWrapper](
+			config.UnverifiedCacheSize,
+			cachedBlockSize,
+		),
+		bytesToIDCache: cache.NewSizedLRU[string, ids.ID](
+			config.BytesToIDCacheSize,
+			cachedBlockBytesSize,
+		),
 	}
 	c.initialize(config)
 	return c
@@ -155,7 +173,10 @@ func NewMeteredState(
 	decidedCache, err := metercacher.New[ids.ID, *BlockWrapper](
 		"decided_cache",
 		registerer,
-		&cache.LRU[ids.ID, *BlockWrapper]{Size: config.DecidedCacheSize},
+		cache.NewSizedLRU[ids.ID, *BlockWrapper](
+			config.DecidedCacheSize,
+			cachedBlockSize,
+		),
 	)
 	if err != nil {
 		return nil, err
@@ -171,15 +192,21 @@ func NewMeteredState(
 	unverifiedCache, err := metercacher.New[ids.ID, *BlockWrapper](
 		"unverified_cache",
 		registerer,
-		&cache.LRU[ids.ID, *BlockWrapper]{Size: config.UnverifiedCacheSize},
+		cache.NewSizedLRU[ids.ID, *BlockWrapper](
+			config.UnverifiedCacheSize,
+			cachedBlockSize,
+		),
 	)
 	if err != nil {
 		return nil, err
 	}
 	bytesToIDCache, err := metercacher.New[string, ids.ID](
 		"bytes_to_id_cache",
 		registerer,
-		&cache.LRU[string, ids.ID]{Size: config.BytesToIDCacheSize},
+		cache.NewSizedLRU[string, ids.ID](
+			config.BytesToIDCacheSize,
+			cachedBlockBytesSize,
+		),
 	)
 	if err != nil {
 		return nil, err

vms/components/chain/state_test.go

@@ -585,7 +585,7 @@ func TestStateBytesToIDCache(t *testing.T) {
 		DecidedCacheSize:    0,
 		MissingCacheSize:    0,
 		UnverifiedCacheSize: 0,
-		BytesToIDCacheSize:  1,
+		BytesToIDCacheSize:  1 + ids.IDLen, // Size of one block
 		LastAcceptedBlock:   genesisBlock,
 		GetBlock:            getBlock,
 		UnmarshalBlock:      parseBlock,

vms/platformvm/state/state.go

@@ -40,8 +40,6 @@ import (
 	"github.com/ava-labs/avalanchego/vms/platformvm/txs"
 )
 
-const pointerOverhead = wrappers.LongLen
-
 var (
 	_ State = (*state)(nil)
 
@@ -383,23 +381,23 @@ type txAndStatus struct {
 
 func txSize(_ ids.ID, tx *txs.Tx) int {
 	if tx == nil {
-		return ids.IDLen + pointerOverhead
+		return ids.IDLen + constants.PointerOverhead
 	}
-	return ids.IDLen + len(tx.Bytes()) + pointerOverhead
+	return ids.IDLen + len(tx.Bytes()) + constants.PointerOverhead
 }
 
 func txAndStatusSize(_ ids.ID, t *txAndStatus) int {
 	if t == nil {
-		return ids.IDLen + pointerOverhead
+		return ids.IDLen + constants.PointerOverhead
 	}
-	return ids.IDLen + len(t.tx.Bytes()) + wrappers.IntLen + pointerOverhead
+	return ids.IDLen + len(t.tx.Bytes()) + wrappers.IntLen + 2*constants.PointerOverhead
 }
 
 func blockSize(_ ids.ID, blk blocks.Block) int {
 	if blk == nil {
-		return ids.IDLen + pointerOverhead
+		return ids.IDLen + constants.PointerOverhead
 	}
-	return ids.IDLen + len(blk.Bytes()) + pointerOverhead
+	return ids.IDLen + len(blk.Bytes()) + constants.PointerOverhead
 }
 
 func New(

vms/proposervm/state/block_state.go

@@ -14,10 +14,13 @@ import (
 	"github.com/ava-labs/avalanchego/database"
 	"github.com/ava-labs/avalanchego/ids"
 	"github.com/ava-labs/avalanchego/snow/choices"
+	"github.com/ava-labs/avalanchego/utils/constants"
+	"github.com/ava-labs/avalanchego/utils/units"
+	"github.com/ava-labs/avalanchego/utils/wrappers"
 	"github.com/ava-labs/avalanchego/vms/proposervm/block"
 )
 
-const blockCacheSize = 8192
+const blockCacheSize = 256 * units.MiB
 
 var (
 	errBlockWrongVersion = errors.New("wrong version")
@@ -45,18 +48,31 @@ type blockWrapper struct {
 	block block.Block
 }
 
+func cachedBlockSize(_ ids.ID, bw *blockWrapper) int {
+	if bw == nil {
+		return ids.IDLen + constants.PointerOverhead
+	}
+	return ids.IDLen + len(bw.Block) + wrappers.IntLen + 2*constants.PointerOverhead
+}
+
 func NewBlockState(db database.Database) BlockState {
 	return &blockState{
-		blkCache: &cache.LRU[ids.ID, *blockWrapper]{Size: blockCacheSize},
-		db:       db,
+		blkCache: cache.NewSizedLRU[ids.ID, *blockWrapper](
+			blockCacheSize,
+			cachedBlockSize,
+		),
+		db: db,
 	}
 }
 
 func NewMeteredBlockState(db database.Database, namespace string, metrics prometheus.Registerer) (BlockState, error) {
 	blkCache, err := metercacher.New[ids.ID, *blockWrapper](
 		fmt.Sprintf("%s_block_cache", namespace),
 		metrics,
-		&cache.LRU[ids.ID, *blockWrapper]{Size: blockCacheSize},
+		cache.NewSizedLRU[ids.ID, *blockWrapper](
+			blockCacheSize,
+			cachedBlockSize,
+		),
 	)
 
 	return &blockState{

vms/proposervm/vm.go

@@ -31,6 +31,7 @@ import (
 	"github.com/ava-labs/avalanchego/utils/constants"
 	"github.com/ava-labs/avalanchego/utils/math"
 	"github.com/ava-labs/avalanchego/utils/timer/mockable"
+	"github.com/ava-labs/avalanchego/utils/units"
 	"github.com/ava-labs/avalanchego/vms/proposervm/indexer"
 	"github.com/ava-labs/avalanchego/vms/proposervm/proposer"
 	"github.com/ava-labs/avalanchego/vms/proposervm/scheduler"
@@ -46,7 +47,7 @@ const (
 	DefaultMinBlockDelay = time.Second
 
 	checkIndexedFrequency = 10 * time.Second
-	innerBlkCacheSize     = 512
+	innerBlkCacheSize     = 128 * units.MiB
 )
 
 var (
@@ -75,6 +76,10 @@ func init() {
 	}
 }
 
+func cachedBlockSize(_ ids.ID, blk snowman.Block) int {
+	return ids.IDLen + len(blk.Bytes()) + constants.PointerOverhead
+}
+
 type VM struct {
 	block.ChainVM
 	blockBuilderVM block.BuildBlockWithContextChainVM
@@ -192,7 +197,10 @@ func (vm *VM) Initialize(
 	innerBlkCache, err := metercacher.New[ids.ID, snowman.Block](
 		"inner_block_cache",
 		registerer,
-		&cache.LRU[ids.ID, snowman.Block]{Size: innerBlkCacheSize},
+		cache.NewSizedLRU[ids.ID, snowman.Block](
+			innerBlkCacheSize,
+			cachedBlockSize,
+		),
 	)
 	if err != nil {
 		return err

vms/proposervm/vm_test.go

@@ -27,6 +27,7 @@ import (
 	"github.com/ava-labs/avalanchego/snow/engine/snowman/block/mocks"
 	"github.com/ava-labs/avalanchego/snow/validators"
 	"github.com/ava-labs/avalanchego/staking"
+	"github.com/ava-labs/avalanchego/utils"
 	"github.com/ava-labs/avalanchego/utils/timer/mockable"
 	"github.com/ava-labs/avalanchego/version"
 	"github.com/ava-labs/avalanchego/vms/proposervm/proposer"
@@ -2229,6 +2230,8 @@ func TestVMInnerBlkCache(t *testing.T) {
 	// We will ask the inner VM to parse.
 	mockInnerBlkNearTip := snowman.NewMockBlock(ctrl)
 	mockInnerBlkNearTip.EXPECT().Height().Return(uint64(1)).Times(2)
+	mockInnerBlkNearTip.EXPECT().Bytes().Return(blkNearTipInnerBytes).Times(1)
+
 	innerVM.EXPECT().ParseBlock(gomock.Any(), blkNearTipInnerBytes).Return(mockInnerBlkNearTip, nil).Times(2)
 	_, err = vm.ParseBlock(context.Background(), blkNearTip.Bytes())
 	require.NoError(err)
@@ -2444,6 +2447,7 @@ func TestVM_VerifyBlockWithContext(t *testing.T) {
 		).Return(nil)
 		innerBlk.MockBlock.EXPECT().Parent().Return(ids.GenerateTestID()).AnyTimes()
 		innerBlk.MockBlock.EXPECT().ID().Return(ids.GenerateTestID()).AnyTimes()
+		innerBlk.MockBlock.EXPECT().Bytes().Return(utils.RandomBytes(1024)).AnyTimes()
 
 		blk := NewMockPostForkBlock(ctrl)
 		blk.EXPECT().getInnerBlk().Return(innerBlk).AnyTimes()

vms/rpcchainvm/vm_client.go

@@ -41,6 +41,7 @@ import (
 	"github.com/ava-labs/avalanchego/snow/validators/gvalidators"
 	"github.com/ava-labs/avalanchego/utils/crypto/bls"
 	"github.com/ava-labs/avalanchego/utils/resource"
+	"github.com/ava-labs/avalanchego/utils/units"
 	"github.com/ava-labs/avalanchego/utils/wrappers"
 	"github.com/ava-labs/avalanchego/version"
 	"github.com/ava-labs/avalanchego/vms/components/chain"
@@ -62,11 +63,12 @@ import (
 	warppb "github.com/ava-labs/avalanchego/proto/pb/warp"
 )
 
+// TODO: Enable these to be configured by the user
 const (
-	decidedCacheSize    = 2048
+	decidedCacheSize    = 256 * units.MiB
 	missingCacheSize    = 2048
-	unverifiedCacheSize = 2048
-	bytesToIDCacheSize  = 2048
+	unverifiedCacheSize = 64 * units.MiB
+	bytesToIDCacheSize  = 64 * units.MiB
 )
 
 var (