status-go/services/wallet/balance/balance_cache.go

package balance

import (
	"context"
	"math/big"
	"reflect"
	"sync"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/core/types"
)

// Reader interface for reading balance at a specified address.
type Reader interface {
	BalanceAt(ctx context.Context, account common.Address, blockNumber *big.Int) (*big.Int, error)
	NonceAt(ctx context.Context, account common.Address, blockNumber *big.Int) (uint64, error)
	HeaderByNumber(ctx context.Context, number *big.Int) (*types.Header, error)
	NetworkID() uint64
}

// Cacher interface for caching balance to BalanceCache. Requires BalanceReader to fetch balance.
type Cacher interface {
	BalanceAt(ctx context.Context, client Reader, account common.Address, blockNumber *big.Int) (*big.Int, error)
	NonceAt(ctx context.Context, client Reader, account common.Address, blockNumber *big.Int) (*int64, error)
	Clear()
	Cache() CacheIface
}

// Interface for cache of balances.
type CacheIface interface {
	GetBalance(account common.Address, chainID uint64, blockNumber *big.Int) *big.Int
	GetNonce(account common.Address, chainID uint64, blockNumber *big.Int) *int64
	AddBalance(account common.Address, chainID uint64, blockNumber *big.Int, balance *big.Int)
	AddNonce(account common.Address, chainID uint64, blockNumber *big.Int, nonce *int64)
	BalanceSize(account common.Address, chainID uint64) int
	NonceSize(account common.Address, chainID uint64) int
	Clear()
}

type addressChainMap[T any] map[common.Address]map[uint64]T // address->chainID

type cacheIface[K comparable, V any] interface {
	get(K) V
	set(K, V)
	len() int
	keys() []K
	clear()
	init()
}

// genericCache is a generic implementation of CacheIface
type genericCache[B cacheIface[uint64, *big.Int], N cacheIface[uint64, *int64], NR cacheIface[int64, nonceRange]] struct {
	nonceRangeCache[NR]

	// balances maps an address and chain to a cache of a block number and the balance of this particular address on the chain
	balances addressChainMap[B]
	nonces   addressChainMap[N]
	rw       sync.RWMutex
}

func (b *genericCache[_, _, _]) GetBalance(account common.Address, chainID uint64, blockNumber *big.Int) *big.Int {
	b.rw.RLock()
	defer b.rw.RUnlock()

	_, exists := b.balances[account]
	if !exists {
		return nil
	}

	_, exists = b.balances[account][chainID]
	if !exists {
		return nil
	}

	return b.balances[account][chainID].get(blockNumber.Uint64())
}

func (b *genericCache[B, _, _]) AddBalance(account common.Address, chainID uint64, blockNumber *big.Int, balance *big.Int) {
	b.rw.Lock()
	defer b.rw.Unlock()

	_, exists := b.balances[account]
	if !exists {
		b.balances[account] = make(map[uint64]B)
	}

	_, exists = b.balances[account][chainID]
	if !exists {
		b.balances[account][chainID] = reflect.New(reflect.TypeOf(b.balances[account][chainID]).Elem()).Interface().(B)
		b.balances[account][chainID].init()
	}

	b.balances[account][chainID].set(blockNumber.Uint64(), balance)
}

func (b *genericCache[_, _, _]) GetNonce(account common.Address, chainID uint64, blockNumber *big.Int) *int64 {
	b.rw.RLock()
	defer b.rw.RUnlock()

	_, exists := b.nonces[account]
	if !exists {
		return nil
	}

	_, exists = b.nonces[account][chainID]
	if !exists {
		return nil
	}

	nonce := b.nonces[account][chainID].get(blockNumber.Uint64())
	if nonce != nil {
		return nonce
	}

	return b.findNonceInRange(account, chainID, blockNumber)
}

func (b *genericCache[_, N, _]) AddNonce(account common.Address, chainID uint64, blockNumber *big.Int, nonce *int64) {
	b.rw.Lock()
	defer b.rw.Unlock()

	_, exists := b.nonces[account]
	if !exists {
		b.nonces[account] = make(map[uint64]N)
	}

	_, exists = b.nonces[account][chainID]
	if !exists {
		b.nonces[account][chainID] = reflect.New(reflect.TypeOf(b.nonces[account][chainID]).Elem()).Interface().(N)
		b.nonces[account][chainID].init()
	}

	b.nonces[account][chainID].set(blockNumber.Uint64(), nonce)
	b.updateNonceRange(account, chainID, blockNumber, nonce)
}

func (b *genericCache[_, _, _]) BalanceSize(account common.Address, chainID uint64) int {
	b.rw.RLock()
	defer b.rw.RUnlock()

	_, exists := b.balances[account]
	if !exists {
		return 0
	}

	_, exists = b.balances[account][chainID]
	if !exists {
		return 0
	}

	return b.balances[account][chainID].len()
}

func (b *genericCache[_, N, _]) NonceSize(account common.Address, chainID uint64) int {
	b.rw.RLock()
	defer b.rw.RUnlock()

	_, exists := b.nonces[account]
	if !exists {
		return 0
	}

	_, exists = b.nonces[account][chainID]
	if !exists {
		return 0
	}

	return b.nonces[account][chainID].len()
}

// implements Cacher interface that caches balance and nonce in memory.
type cacherImpl struct {
	cache CacheIface
}

func newCacherImpl(cache CacheIface) *cacherImpl {
	return &cacherImpl{
		cache: cache,
	}
}

func (b *cacherImpl) BalanceAt(ctx context.Context, client Reader, account common.Address, blockNumber *big.Int) (*big.Int, error) {
	cachedBalance := b.cache.GetBalance(account, client.NetworkID(), blockNumber)
	if cachedBalance != nil {
		return cachedBalance, nil
	}

	balance, err := client.BalanceAt(ctx, account, blockNumber)
	if err != nil {
		return nil, err
	}
	b.cache.AddBalance(account, client.NetworkID(), blockNumber, balance)
	return balance, nil
}

func (b *cacherImpl) NonceAt(ctx context.Context, client Reader, account common.Address, blockNumber *big.Int) (*int64, error) {
	cachedNonce := b.cache.GetNonce(account, client.NetworkID(), blockNumber)
	if cachedNonce != nil {
		return cachedNonce, nil
	}

	nonce, err := client.NonceAt(ctx, account, blockNumber)
	if err != nil {
		return nil, err
	}
	int64Nonce := int64(nonce)
	b.cache.AddNonce(account, client.NetworkID(), blockNumber, &int64Nonce)

	return &int64Nonce, nil
}

func (b *cacherImpl) Clear() {
	b.cache.Clear()
}

func (b *cacherImpl) Cache() CacheIface {
	return b.cache
}