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

220 lines
5.8 KiB
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)
CallBlockHashByTransaction(ctx context.Context, blockNumber *big.Int, index uint) (common.Hash, 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
}