refactor(balance): refactor balance indexing + add test coverage

refactor(evm/batcher): create reusable batcher struct
remove code

Author: Tangui-Bitfly

backend/cmd/eth1indexer/main.go

@@ -19,6 +19,7 @@ import (
 	"github.com/gobitfly/beaconchain/pkg/commons/db"
 	"github.com/gobitfly/beaconchain/pkg/commons/db2/data"
 	"github.com/gobitfly/beaconchain/pkg/commons/db2/database"
+	"github.com/gobitfly/beaconchain/pkg/commons/db2/metadata"
 	"github.com/gobitfly/beaconchain/pkg/commons/db2/metadataupdates"
 	"github.com/gobitfly/beaconchain/pkg/commons/erc20"
 	"github.com/gobitfly/beaconchain/pkg/commons/log"
@@ -29,6 +30,7 @@ import (
 	"github.com/gobitfly/beaconchain/pkg/commons/utils"
 	"github.com/gobitfly/beaconchain/pkg/commons/version"
 	"github.com/gobitfly/beaconchain/pkg/executionlayer"
+	"github.com/gobitfly/beaconchain/pkg/executionlayer/evm"
 
 	"github.com/coocood/freecache"
 	"github.com/ethereum/go-ethereum/common"
@@ -145,8 +147,6 @@ func Run() {
 
 	chainId := strconv.FormatUint(utils.Config.Chain.ClConfig.DepositChainID, 10)
 
-	balanceUpdaterPrefix := chainId + ":B:"
-
 	nodeChainId, err := client.GetNativeClient().ChainID(context.Background())
 	if err != nil {
 		log.Fatal(err, "node chain id error", 0)
@@ -179,26 +179,36 @@ func Run() {
 		go ImportEnsUpdatesLoop(bt, client, *ensBatchSize)
 	}
 
-	if *enableFullBalanceUpdater {
-		ProcessMetadataUpdates(bt, client, balanceUpdaterPrefix, *balanceUpdaterBatchSize, -1)
-		return
-	}
-
-	cache := freecache.NewCache(100 * 1024 * 1024) // 100 MB limit
-
 	bigtable, err := database.NewBigTable(utils.Config.Bigtable.Project, utils.Config.Bigtable.Instance, nil)
 	if err != nil {
 		log.Fatal(err, "error connecting to bigtable", 0)
 	}
 
+	cache := freecache.NewCache(100 * 1024 * 1024) // 100 MB limit
+	metadataUpdatesStore := metadataupdates.NewStore(database.Wrap(bigtable, metadataupdates.Table), cache)
+	dataStore := data.NewStore(database.Wrap(bigtable, data.Table))
+	metadataStore := metadata.NewStore(database.Wrap(bigtable, metadata.Table))
+
 	indexer := executionlayer.NewIndexer(
 		executionlayer.NewAdaptorV1(
-			data.NewStore(database.Wrap(bigtable, data.Table)),
-			metadataupdates.NewStore(database.Wrap(bigtable, metadataupdates.Table), cache),
+			dataStore,
+			metadataUpdatesStore,
 		),
 		executionlayer.AllTransformers...,
 	)
 
+	balanceUpdater := executionlayer.NewBalanceUpdater(
+		chainId,
+		metadataUpdatesStore,
+		metadataStore,
+		evm.NewBatcher(chainId, client.GetNativeClient()),
+	)
+
+	if *enableFullBalanceUpdater {
+		ProcessBalanceUpdates(balanceUpdater, *balanceUpdaterBatchSize, -1)
+		return
+	}
+
 	if *block != 0 {
 		err = IndexFromNode(bt, client, *block, *block, *concurrencyBlocks, *traceMode)
 		if err != nil {
@@ -361,7 +371,7 @@ func Run() {
 		}
 
 		if *enableBalanceUpdater {
-			ProcessMetadataUpdates(bt, client, balanceUpdaterPrefix, *balanceUpdaterBatchSize, 10)
+			ProcessBalanceUpdates(balanceUpdater, *balanceUpdaterBatchSize, 10)
 		}
 
 		log.Infof("index run completed")
@@ -545,55 +555,15 @@ func HandleChainReorgs(bt *db.Bigtable, client *rpc.ErigonClient, depth int) err
 	return nil
 }
 
-func ProcessMetadataUpdates(bt *db.Bigtable, client *rpc.ErigonClient, prefix string, batchSize int, iterations int) {
-	lastKey := prefix
-
-	its := 0
-	for {
+func ProcessBalanceUpdates(balanceUpdater executionlayer.BalanceUpdater, batchSize int, iterations int) {
+	for its := 0; iterations == -1 || its < iterations; its++ {
 		start := time.Now()
-		keys, pairs, err := bt.GetMetadataUpdates(prefix, lastKey, batchSize)
+		balances, err := balanceUpdater.UpdateBalances(int64(batchSize))
 		if err != nil {
-			log.Error(err, "error retrieving metadata updates from bigtable", 0)
-			return
-		}
-
-		if len(keys) == 0 {
-			return
-		}
-
-		balances := make([]*types.Eth1AddressBalance, 0, len(pairs))
-		for b := 0; b < len(pairs); b += batchSize {
-			start := b
-			end := b + batchSize
-			if len(pairs) < end {
-				end = len(pairs)
-			}
-
-			log.Infof("processing batch %v with start %v and end %v", b, start, end)
-
-			b, err := client.GetBalances(pairs[start:end], 2, 4)
-
-			if err != nil {
-				log.Error(err, "error retrieving balances from node", 0)
-				return
-			}
-			balances = append(balances, b...)
-		}
-
-		err = bt.SaveBalances(balances, keys)
-		if err != nil {
-			log.Error(err, "error saving balances to bigtable", 0)
-			return
-		}
-
-		lastKey = keys[len(keys)-1]
-		log.Infof("retrieved %v balances in %v, currently at %v", len(balances), time.Since(start), lastKey)
-
-		its++
-
-		if iterations != -1 && its > iterations {
+			log.Error(err, "error updating balances", 0)
 			return
 		}
+		log.Infof("retrieved %v balances in %v, currently at %s", len(balances), time.Since(start), balances[len(balances)-1].Address)
 	}
 }
 

backend/pkg/commons/db/bigtable_eth1.go

@@ -2339,53 +2339,6 @@ func (bigtable *Bigtable) SaveContractMetadata(address []byte, metadata *types.C
 	return bigtable.tableMetadata.Apply(ctx, fmt.Sprintf("%s:%x", bigtable.chainId, address), mut)
 }
 
-func (bigtable *Bigtable) SaveBalances(balances []*types.Eth1AddressBalance, deleteKeys []string) error {
-	if len(balances) == 0 {
-		return nil
-	}
-
-	mutsWrite := &types.BulkMutations{
-		Keys: make([]string, 0, len(balances)),
-		Muts: make([]*gcp_bigtable.Mutation, 0, len(balances)),
-	}
-
-	for _, balance := range balances {
-		mutWrite := gcp_bigtable.NewMutation()
-
-		mutWrite.Set(ACCOUNT_METADATA_FAMILY, fmt.Sprintf("B:%x", balance.Token), gcp_bigtable.Timestamp(0), balance.Balance)
-		mutsWrite.Keys = append(mutsWrite.Keys, fmt.Sprintf("%s:%x", bigtable.chainId, balance.Address))
-		mutsWrite.Muts = append(mutsWrite.Muts, mutWrite)
-	}
-
-	err := bigtable.WriteBulk(mutsWrite, bigtable.tableMetadata, DEFAULT_BATCH_INSERTS)
-
-	if err != nil {
-		return err
-	}
-
-	if len(deleteKeys) == 0 {
-		return nil
-	}
-	mutsDelete := &types.BulkMutations{
-		Keys: make([]string, 0, len(balances)),
-		Muts: make([]*gcp_bigtable.Mutation, 0, len(balances)),
-	}
-	for _, key := range deleteKeys {
-		mutDelete := gcp_bigtable.NewMutation()
-		mutDelete.DeleteRow()
-		mutsDelete.Keys = append(mutsDelete.Keys, key)
-		mutsDelete.Muts = append(mutsDelete.Muts, mutDelete)
-	}
-
-	err = bigtable.WriteBulk(mutsDelete, bigtable.tableMetadataUpdates, DEFAULT_BATCH_INSERTS)
-
-	if err != nil {
-		return err
-	}
-
-	return nil
-}
-
 func (bigtable *Bigtable) SaveERC20TokenPrices(prices []*types.ERC20TokenPrice) error {
 	if len(prices) == 0 {
 		return nil

backend/pkg/commons/db2/database/bigtable.go

@@ -82,6 +82,10 @@ func (w TableWrapper) GetRowsWithKeys(keys []string) ([]Row, error) {
 	return res, nil
 }
 
+func (w TableWrapper) DeleteRowsWithKeys(keys []string, opts ...Option) error {
+	return w.BigTable.DeleteRowsWithKeys(w.table, keys, opts...)
+}
+
 // BigTable is a wrapper around Google Cloud Bigtable for storing and retrieving data
 type BigTable struct {
 	client *bigtable.Client
@@ -162,10 +166,6 @@ func createTableAndFamilies(ctx context.Context, admin *bigtable.AdminClient, ta
 func (b BigTable) BulkAdd(table string, itemsByKey map[string][]Item, opts ...Option) error {
 	options := apply(opts)
 
-	tbl := b.client.Open(table)
-	ctx, cancel := context.WithTimeout(context.Background(), timeout)
-	defer cancel()
-
 	var keys []string
 	var muts []*bigtable.Mutation
 	for key, items := range itemsByKey {
@@ -180,9 +180,17 @@ func (b BigTable) BulkAdd(table string, itemsByKey map[string][]Item, opts ...Op
 		keys = append(keys, key)
 		muts = append(muts, mut)
 	}
+	return b.bulk(table, keys, muts, options)
+}
+
+func (b BigTable) bulk(table string, keys []string, mutations []*bigtable.Mutation, options options) error {
+	tbl := b.client.Open(table)
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+
 	bulk := &bulkMutations{
 		Keys: keys,
-		Muts: muts,
+		Muts: mutations,
 	}
 	sort.Sort(bulk)
 	for i := int64(0); i < int64(bulk.Len()); i = i + options.BatchSize {
@@ -337,6 +345,23 @@ func (b BigTable) GetRowsWithKeys(table string, keys []string) ([]Row, error) {
 	return data, nil
 }
 
+func (b BigTable) DeleteRowsWithKeys(table string, source []string, opts ...Option) error {
+	options := apply(opts)
+
+	var muts []*bigtable.Mutation
+	var keys []string
+	for _, key := range source {
+		if key == "" {
+			continue
+		}
+		mut := bigtable.NewMutation()
+		mut.DeleteRow()
+		muts = append(muts, mut)
+		keys = append(keys, key)
+	}
+	return b.bulk(table, keys, muts, options)
+}
+
 func (b BigTable) Clear() error {
 	ctx, cancel := context.WithTimeout(context.Background(), timeout)
 	defer cancel()

backend/pkg/commons/db2/database/bigtable_test.go

@@ -158,6 +158,19 @@ func TestBigTable(t *testing.T) {
 					}
 				}
 			})
+
+			t.Run("DeleteRowsWithKeys", func(t *testing.T) {
+				if err := db.DeleteRowsWithKeys(maps.Keys(tt.items)); err != nil {
+					t.Error(err)
+				}
+				res, err := db.Read("")
+				if err != nil {
+					t.Error(err)
+				}
+				if len(res) != 0 {
+					t.Errorf("not empty after delete %v", res)
+				}
+			})
 		})
 	}
 

backend/pkg/commons/db2/database/database.go

@@ -6,6 +6,7 @@ type Database interface {
 	GetRow(key string) (*Row, error)
 	GetRowsWithKeys(keys []string) ([]Row, error)
 	GetRowsRange(high, low string, opts ...Option) ([]Row, error)
+	DeleteRowsWithKeys(keys []string, opts ...Option) error
 
 	Close() error
 	Clear() error

backend/pkg/commons/db2/database/option.go

@@ -1,7 +1,7 @@
 package database
 
 const (
-	defaultBatchSize = 10000
+	defaultBatchSize = 10_000
 	defaultLimit     = 100
 )
 

backend/pkg/commons/db2/database/remote.go

@@ -300,6 +300,10 @@ func (r RemoteClient) GetRowsWithKeys(keys []string) ([]Row, error) {
 	return rows, nil
 }
 
+func (r RemoteClient) DeleteRowsWithKeys(keys []string, opts ...Option) error {
+	panic("implement me")
+}
+
 func (r RemoteClient) Close() error {
 	//TODO implement me
 	panic("implement me")

backend/pkg/commons/db2/metadata/metadata.go

@@ -0,0 +1,36 @@
+package metadata
+
+import (
+	"fmt"
+	"math/big"
+
+	"github.com/gobitfly/beaconchain/pkg/commons/db2/database"
+	"github.com/gobitfly/beaconchain/pkg/commons/db2/metadataupdates"
+)
+
+type Balance struct {
+	metadataupdates.Pair
+	Value *big.Int
+}
+type Store struct {
+	db database.Database
+}
+
+func NewStore(db database.Database) Store {
+	return Store{
+		db: db,
+	}
+}
+
+func (store Store) UpdateBalance(chainID string, balances []Balance) error {
+	updates := make(map[string][]database.Item)
+	for _, balance := range balances {
+		key := fmt.Sprintf("%s:%x", chainID, balance.Address)
+		updates[key] = append(updates[key], database.Item{
+			Family: accountFamily,
+			Column: fmt.Sprintf("B:%x", balance.Token),
+			Data:   balance.Value.Bytes(),
+		})
+	}
+	return store.db.BulkAdd(updates)
+}

backend/pkg/commons/db2/metadata/tables.go

@@ -0,0 +1,13 @@
+package metadata
+
+const Table = "metadata"
+
+var Schema = map[string][]string{
+	Table: {
+		accountFamily,
+	},
+}
+
+const (
+	accountFamily = "a"
+)

backend/pkg/commons/db2/metadataupdates/keys.go

@@ -67,8 +67,8 @@ func ContractUpdate(blockNumber uint64, chainID string, updates []ContractUpdate
 func MarkBalanceUpdate(chainID string, address []byte, token []byte, cache Cache) map[string][]database.Item {
 	items := make(map[string][]database.Item)
 
-	key := fmt.Sprintf("%s:B:%x", chainID, address) // format is B: for balance update as chainid:prefix:address (token id will be encoded as column name)
-	keyCache := []byte(fmt.Sprintf("%s:B:%x:%x", chainID, address, token))
+	key := fmt.Sprintf("%s:%s:%x", chainID, balanceKey, address) // format is B: for balance update as chainid:prefix:address (token id will be encoded as column name)
+	keyCache := []byte(fmt.Sprintf("%s:%s:%x:%x", chainID, balanceKey, address, token))
 	if _, err := cache.Get(keyCache); err != nil {
 		items[key] = []database.Item{
 			{