Repository: coinbase/rosetta-sdk-go
Pull request #252 (2020-11-25 19:32:00)

[storage] Balance Storage Overhaul
Closes: #247 

This PR introduces significant optimizations to `BalanceStorage` that dramatically speed up sync speed in the `rosetta-cli` (mostly through the reduction of lock contention). Special shoutout to @yflinn, @sidhujag, and @HelloKashif for debugging the issues that led to this PR.

### Changes
- [x] Use `big.Int` bytes format to store amount
- [x] Don't use `txn.Scan` when updating balance
- [x] cache most recent balance
  - [x] update during orphan
- [x] Don't store pruning index, last reconciled index, and current balance on same record
- [x] Count accounts seen and accounts reconciled
  - [x] Use in `EstimatedReconciliationCoverage`
  - [x] Update accounts reconciled in block syncing loop so that account txs don't conflict
- [x] Rebase on https://github.com/coinbase/rosetta-sdk-go/pull/254
  - [x] squash commits
- [x] Add tests (blocked by https://github.com/coinbase/rosetta-sdk-go/pull/254)
  - [x] Pass existing tests
  - [x] Add test for estimated reconciliation coverage
- [x] Rollback counts when deleting account records (conditionally remove reconciliation if exists)
- [x] Don't error if reconciled orphaned account (could be a race where we delete account while reconciling)

### Future Work (before release)
* limit concurrent goroutines in syncing to `NumCPU` (pretty good default) instead of unbounded (can cause CPU to freeze or OOM)
  * Use: https://github.com/neilotoole/errgroup
* pre-allocate `allChanges` in `parser` (#247 )

coveralls: ## Pull Request Test Coverage Report for [Build 12234](https://coveralls.io/builds/35248833)

* **0** of **0**   changed or added relevant lines in **0** files are covered.
* No unchanged relevant lines lost coverage.
* Overall coverage increased (+**1.5%**) to **80.141%**

---



|  Totals | [![Coverage Status](https://coveralls.io/builds/35248833/badge)](https://coveralls.io/builds/35248833) |
| :-- | --: |
| Change from base [Build 12213](https://coveralls.io/builds/35246694): |  1.5% |
| Covered Lines: | 5331 |
| Relevant Lines: | 6652 |

---
##### 💛  - [Coveralls](https://coveralls.io)

HelloKashif: Hi @patrick-ogrady 
Do you have timeline on this PR completion?
patrick-ogrady: > Hi @patrick-ogrady
Do you have timeline on this PR completion?

Looking to have this all merged and released by Tuesday! Thanks for your patience!

Patch:
diff --git a/mocks/storage/modules/balance_storage_handler.go b/mocks/storage/modules/balance_storage_handler.go
index ae76fe5..ec9d411 100644
--- a/mocks/storage/modules/balance_storage_handler.go
+++ b/mocks/storage/modules/balance_storage_handler.go
@@ -8,6 +8,7 @@ import (
 	mock "github.com/stretchr/testify/mock"
 
 	parser "github.com/coinbase/rosetta-sdk-go/parser"
+	database "github.com/coinbase/rosetta-sdk-go/storage/database"
 	types "github.com/coinbase/rosetta-sdk-go/types"
 )
 
@@ -16,6 +17,34 @@ type BalanceStorageHandler struct {
 	mock.Mock
 }
 
+// AccountsReconciled provides a mock function with given fields: ctx, dbTx, count
+func (_m *BalanceStorageHandler) AccountsReconciled(ctx context.Context, dbTx database.Transaction, count int) error {
+	ret := _m.Called(ctx, dbTx, count)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, database.Transaction, int) error); ok {
+		r0 = rf(ctx, dbTx, count)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// AccountsSeen provides a mock function with given fields: ctx, dbTx, count
+func (_m *BalanceStorageHandler) AccountsSeen(ctx context.Context, dbTx database.Transaction, count int) error {
+	ret := _m.Called(ctx, dbTx, count)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, database.Transaction, int) error); ok {
+		r0 = rf(ctx, dbTx, count)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
 // BlockAdded provides a mock function with given fields: ctx, block, changes
 func (_m *BalanceStorageHandler) BlockAdded(ctx context.Context, block *types.Block, changes []*parser.BalanceChange) error {
 	ret := _m.Called(ctx, block, changes)
diff --git a/mocks/storage/modules/balance_storage_helper.go b/mocks/storage/modules/balance_storage_helper.go
index cf59a69..b65e50d 100644
--- a/mocks/storage/modules/balance_storage_helper.go
+++ b/mocks/storage/modules/balance_storage_helper.go
@@ -4,11 +4,13 @@ package modules
 
 import (
 	context "context"
+	big "math/big"
 
 	mock "github.com/stretchr/testify/mock"
 
 	asserter "github.com/coinbase/rosetta-sdk-go/asserter"
 	parser "github.com/coinbase/rosetta-sdk-go/parser"
+	database "github.com/coinbase/rosetta-sdk-go/storage/database"
 	types "github.com/coinbase/rosetta-sdk-go/types"
 )
 
@@ -40,6 +42,52 @@ func (_m *BalanceStorageHelper) AccountBalance(ctx context.Context, account *typ
 	return r0, r1
 }
 
+// AccountsReconciled provides a mock function with given fields: ctx, dbTx
+func (_m *BalanceStorageHelper) AccountsReconciled(ctx context.Context, dbTx database.Transaction) (*big.Int, error) {
+	ret := _m.Called(ctx, dbTx)
+
+	var r0 *big.Int
+	if rf, ok := ret.Get(0).(func(context.Context, database.Transaction) *big.Int); ok {
+		r0 = rf(ctx, dbTx)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*big.Int)
+		}
+	}
+
+	var r1 error
+	if rf, ok := ret.Get(1).(func(context.Context, database.Transaction) error); ok {
+		r1 = rf(ctx, dbTx)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// AccountsSeen provides a mock function with given fields: ctx, dbTx
+func (_m *BalanceStorageHelper) AccountsSeen(ctx context.Context, dbTx database.Transaction) (*big.Int, error) {
+	ret := _m.Called(ctx, dbTx)
+
+	var r0 *big.Int
+	if rf, ok := ret.Get(0).(func(context.Context, database.Transaction) *big.Int); ok {
+		r0 = rf(ctx, dbTx)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*big.Int)
+		}
+	}
+
+	var r1 error
+	if rf, ok := ret.Get(1).(func(context.Context, database.Transaction) error); ok {
+		r1 = rf(ctx, dbTx)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
 // Asserter provides a mock function with given fields:
 func (_m *BalanceStorageHelper) Asserter() *asserter.Asserter {
 	ret := _m.Called()
diff --git a/parser/balance_changes.go b/parser/balance_changes.go
index a876424..9b6c01b 100644
--- a/parser/balance_changes.go
+++ b/parser/balance_changes.go
@@ -99,7 +99,6 @@ func (p *Parser) BalanceChanges(
 					return nil, err
 				}
 				amountValue = negatedValue
-				blockIdentifier = block.ParentBlockIdentifier
 			}
 
 			// Merge values by account and currency
diff --git a/storage/errors/errors.go b/storage/errors/errors.go
index 60f9b38..e49190b 100644
--- a/storage/errors/errors.go
+++ b/storage/errors/errors.go
@@ -324,12 +324,22 @@ var (
 	// account does not exist.
 	ErrAccountMissing = errors.New("block nil")
 
+	// ErrInvalidChangeValue is returned when the change value
+	// cannot be parsed.
+	ErrInvalidChangeValue = errors.New("invalid change value")
+
+	// ErrInvalidValue is returned when the value we are trying
+	// to save cannot be parsed.
+	ErrInvalidValue = errors.New("invalid value")
+
 	BalanceStorageErrs = []error{
 		ErrNegativeBalance,
 		ErrInvalidLiveBalance,
 		ErrBalancePruned,
 		ErrBlockNil,
 		ErrAccountMissing,
+		ErrInvalidChangeValue,
+		ErrInvalidValue,
 	}
 )
 
diff --git a/storage/modules/balance_storage.go b/storage/modules/balance_storage.go
index 2b27c36..c470545 100644
--- a/storage/modules/balance_storage.go
+++ b/storage/modules/balance_storage.go
@@ -20,6 +20,7 @@ import (
 	"fmt"
 	"log"
 	"math/big"
+	"sync"
 
 	"golang.org/x/sync/errgroup"
 
@@ -35,15 +36,31 @@ var _ BlockWorker = (*BalanceStorage)(nil)
 
 const (
 	// accountNamespace is prepended to any stored account.
-	accountNamespace = "account"
+	accountNamespace = "acc"
+
+	// balanceNamespace is prepended to any stored balance.
+	balanceNamespace = "bal"
 
 	// historicalBalanceNamespace is prepended to any stored
 	// historical balance.
-	historicalBalanceNamespace = "balance"
+	historicalBalanceNamespace = "hbal"
+
+	// reconciliationNamespace is prepended to any stored
+	// reconciliation.
+	reconciliationNamepace = "recacc"
+
+	// pruneNamespace is prepended to any stored pruning
+	// record.
+	pruneNamespace = "pruneacc"
+
+	// maxBalancePruneSize is the maximum number of balances
+	// we should consider pruning at one time.
+	maxBalancePruneSize = 5000
 )
 
 var (
 	errAccountFound = errors.New("account found")
+	errTooManyKeys  = errors.New("too many keys")
 )
 
 /*
@@ -51,9 +68,13 @@ var (
 */
 
 // GetAccountKey returns a deterministic hash of a types.Account + types.Currency.
-func GetAccountKey(account *types.AccountIdentifier, currency *types.Currency) []byte {
+func GetAccountKey(
+	namespace string,
+	account *types.AccountIdentifier,
+	currency *types.Currency,
+) []byte {
 	return []byte(
-		fmt.Sprintf("%s/%s/%s", accountNamespace, types.Hash(account), types.Hash(currency)),
+		fmt.Sprintf("%s/%s/%s", namespace, types.Hash(account), types.Hash(currency)),
 	)
 }
 
@@ -92,6 +113,9 @@ func GetHistoricalBalancePrefix(account *types.AccountIdentifier, currency *type
 type BalanceStorageHandler interface {
 	BlockAdded(ctx context.Context, block *types.Block, changes []*parser.BalanceChange) error
 	BlockRemoved(ctx context.Context, block *types.Block, changes []*parser.BalanceChange) error
+
+	AccountsReconciled(ctx context.Context, dbTx database.Transaction, count int) error
+	AccountsSeen(ctx context.Context, dbTx database.Transaction, count int) error
 }
 
 // BalanceStorageHelper functions are used by BalanceStorage to process balances. Defining an
@@ -108,6 +132,9 @@ type BalanceStorageHelper interface {
 	ExemptFunc() parser.ExemptOperation
 	BalanceExemptions() []*types.BalanceExemption
 	Asserter() *asserter.Asserter
+
+	AccountsReconciled(ctx context.Context, dbTx database.Transaction) (*big.Int, error)
+	AccountsSeen(ctx context.Context, dbTx database.Transaction) (*big.Int, error)
 }
 
 // BalanceStorage implements block specific storage methods
@@ -117,6 +144,12 @@ type BalanceStorage struct {
 	helper  BalanceStorageHelper
 	handler BalanceStorageHandler
 
+	// To scale up write concurrency on reconciliation,
+	// we don't update the global reconciliation counter
+	// in the account-scoped reconciliatiion transaction.
+	pendingReconciliations     int
+	pendingReconciliationMutex *utils.PriorityMutex
+
 	parser *parser.Parser
 }
 
@@ -125,7 +158,8 @@ func NewBalanceStorage(
 	db database.Database,
 ) *BalanceStorage {
 	return &BalanceStorage{
-		db: db,
+		db:                         db,
+		pendingReconciliationMutex: new(utils.PriorityMutex),
 	}
 }
 
@@ -155,6 +189,11 @@ func (b *BalanceStorage) AddingBlock(
 		return nil, fmt.Errorf("%w: unable to calculate balance changes", err)
 	}
 
+	// Keep track of how many new accounts have been seen so that the counter
+	// can be updated in a single op.
+	var newAccounts int
+	var newAccountsLock sync.Mutex
+
 	g, gctx := errgroup.WithContext(ctx)
 	for i := range changes {
 		// We need to set variable before calling goroutine
@@ -162,7 +201,25 @@ func (b *BalanceStorage) AddingBlock(
 		// continues.
 		change := changes[i]
 		g.Go(func() error {
-			return b.UpdateBalance(gctx, transaction, change, block.ParentBlockIdentifier)
+			newAccount, err := b.UpdateBalance(
+				gctx,
+				transaction,
+				change,
+				block.ParentBlockIdentifier,
+			)
+			if err != nil {
+				return err
+			}
+
+			if !newAccount {
+				return nil
+			}
+
+			newAccountsLock.Lock()
+			newAccounts++
+			newAccountsLock.Unlock()
+
+			return nil
 		})
 	}
 
@@ -170,6 +227,26 @@ func (b *BalanceStorage) AddingBlock(
 		return nil, err
 	}
 
+	// Update accounts seen
+	if newAccounts > 0 {
+		if err := b.handler.AccountsSeen(ctx, transaction, newAccounts); err != nil {
+			return nil, err
+		}
+	}
+
+	// Update accounts reconciled
+	var pending int
+	b.pendingReconciliationMutex.Lock(true)
+	pending = b.pendingReconciliations
+	b.pendingReconciliations = 0
+	b.pendingReconciliationMutex.Unlock()
+
+	if pending > 0 {
+		if err := b.handler.AccountsReconciled(ctx, transaction, pending); err != nil {
+			return nil, err
+		}
+	}
+
 	return func(ctx context.Context) error {
 		return b.handler.BlockAdded(ctx, block, changes)
 	}, nil
@@ -186,6 +263,11 @@ func (b *BalanceStorage) RemovingBlock(
 		return nil, fmt.Errorf("%w: unable to calculate balance changes", err)
 	}
 
+	// staleAccounts should be removed because the orphaned
+	// balance was the last stored balance.
+	staleAccounts := []*types.AccountCurrency{}
+	var staleAccountsMutex sync.Mutex
+
 	g, gctx := errgroup.WithContext(ctx)
 	for i := range changes {
 		// We need to set variable before calling goroutine
@@ -193,13 +275,27 @@ func (b *BalanceStorage) RemovingBlock(
 		// continues.
 		change := changes[i]
 		g.Go(func() error {
-			return b.OrphanBalance(
+			shouldRemove, err := b.OrphanBalance(
 				gctx,
 				transaction,
-				change.Account,
-				change.Currency,
-				block.BlockIdentifier,
+				change,
 			)
+			if err != nil {
+				return err
+			}
+
+			if !shouldRemove {
+				return nil
+			}
+
+			staleAccountsMutex.Lock()
+			staleAccounts = append(staleAccounts, &types.AccountCurrency{
+				Account:  change.Account,
+				Currency: change.Currency,
+			})
+			staleAccountsMutex.Unlock()
+
+			return nil
 		})
 	}
 
@@ -208,15 +304,29 @@ func (b *BalanceStorage) RemovingBlock(
 	}
 
 	return func(ctx context.Context) error {
-		return b.handler.BlockRemoved(ctx, block, changes)
+		if err := b.handler.BlockRemoved(ctx, block, changes); err != nil {
+			return err
+		}
+
+		if len(staleAccounts) == 0 {
+			return nil
+		}
+
+		dbTx := b.db.Transaction(ctx)
+		defer dbTx.Discard(ctx)
+		for _, account := range staleAccounts {
+			if err := b.deleteAccountRecords(ctx, dbTx, account.Account, account.Currency); err != nil {
+				return err
+			}
+		}
+
+		return dbTx.Commit(ctx)
 	}, nil
 }
 
 type accountEntry struct {
-	Account        *types.AccountIdentifier `json:"account"`
-	Currency       *types.Currency          `json:"currency"`
-	LastReconciled *types.BlockIdentifier   `json:"last_reconciled"`
-	LastPruned     *int64                   `json:"last_pruned"`
+	Account  *types.AccountIdentifier `json:"account"`
+	Currency *types.Currency          `json:"currency"`
 }
 
 // SetBalance allows a client to set the balance of an account in a database
@@ -241,7 +351,23 @@ func (b *BalanceStorage) SetBalance(
 		return err
 	}
 
-	serialAcc, err := b.db.Encoder().Encode(historicalBalanceNamespace, accountEntry{
+	// Remove all account keys
+	if err := b.deleteAccountRecords(
+		ctx,
+		dbTransaction,
+		account,
+		amount.Currency,
+	); err != nil {
+		return err
+	}
+
+	// Mark as new account seen
+	if err := b.handler.AccountsSeen(ctx, dbTransaction, 1); err != nil {
+		return err
+	}
+
+	// Serialize account entry
+	serialAcc, err := b.db.Encoder().Encode(accountNamespace, accountEntry{
 		Account:  account,
 		Currency: amount.Currency,
 	})
@@ -249,15 +375,27 @@ func (b *BalanceStorage) SetBalance(
 		return err
 	}
 
-	// Set current record
-	key := GetAccountKey(account, amount.Currency)
+	// Set account
+	key := GetAccountKey(accountNamespace, account, amount.Currency)
 	if err := dbTransaction.Set(ctx, key, serialAcc, true); err != nil {
 		return err
 	}
 
-	// Set historical record
+	// Set current balance
+	key = GetAccountKey(balanceNamespace, account, amount.Currency)
+	value, ok := new(big.Int).SetString(amount.Value, 10)
+	if !ok {
+		return storageErrs.ErrInvalidValue
+	}
+
+	valueBytes := value.Bytes()
+	if err := dbTransaction.Set(ctx, key, valueBytes, false); err != nil {
+		return err
+	}
+
+	// Set historical balance
 	key = GetHistoricalBalanceKey(account, amount.Currency, block.Index)
-	if err := dbTransaction.Set(ctx, key, []byte(amount.Value), true); err != nil {
+	if err := dbTransaction.Set(ctx, key, valueBytes, true); err != nil {
 		return err
 	}
 
@@ -273,31 +411,42 @@ func (b *BalanceStorage) Reconciled(
 	currency *types.Currency,
 	block *types.BlockIdentifier,
 ) error {
-	dbTx := b.db.Transaction(ctx)
+	key := GetAccountKey(reconciliationNamepace, account, currency)
+	dbTx := b.db.WriteTransaction(ctx, string(key), false)
 	defer dbTx.Discard(ctx)
 
-	err := b.updateAccountEntry(
-		ctx,
-		dbTx,
-		account,
-		currency,
-		func(accEntry accountEntry) *accountEntry {
-			// Don't update last reconciled if the most recent reconciliation was
-			// lower than the last reconciliation. This can occur when inactive
-			// reconciliation gets ahead of the active reconciliation backlog.
-			if accEntry.LastReconciled != nil && accEntry.LastReconciled.Index > block.Index {
-				return nil
-			}
-
-			accEntry.LastReconciled = block
+	// Return nil if account record does not exist (could have
+	// occurred during a reorg at tip
+	acctKey := GetAccountKey(accountNamespace, account, currency)
+	acctExists, _, err := dbTx.Get(ctx, acctKey)
+	if err != nil {
+		return err
+	}
+	if !acctExists {
+		return nil
+	}
 
-			return &accEntry
-		},
-	)
+	exists, lastReconciled, err := BigIntGet(ctx, key, dbTx)
 	if err != nil {
 		return err
 	}
 
+	if exists {
+		// Skip update if we already reconciled the account
+		// after the index.
+		if block.Index <= lastReconciled.Int64() {
+			return nil
+		}
+	} else {
+		b.pendingReconciliationMutex.Lock(false)
+		b.pendingReconciliations++
+		b.pendingReconciliationMutex.Unlock()
+	}
+
+	if err := dbTx.Set(ctx, key, new(big.Int).SetInt64(block.Index).Bytes(), true); err != nil {
+		return err
+	}
+
 	if err := dbTx.Commit(ctx); err != nil {
 		return fmt.Errorf("%w: unable to commit last reconciliation update", err)
 	}
@@ -305,6 +454,31 @@ func (b *BalanceStorage) Reconciled(
 	return nil
 }
 
+// EstimatedReconciliationCoverage returns an estimated
+// reconciliation coverage metric. This can be used to
+// get an idea of the reconciliation coverage without
+// doing an expensive DB scan across all accounts.
+func (b *BalanceStorage) EstimatedReconciliationCoverage(ctx context.Context) (float64, error) {
+	dbTx := b.db.ReadTransaction(ctx)
+	defer dbTx.Discard(ctx)
+
+	reconciled, err := b.helper.AccountsReconciled(ctx, dbTx)
+	if err != nil {
+		return -1, err
+	}
+
+	accounts, err := b.helper.AccountsSeen(ctx, dbTx)
+	if err != nil {
+		return -1, err
+	}
+
+	if accounts.Sign() == 0 {
+		return 0, nil
+	}
+
+	return float64(reconciled.Int64()) / float64(accounts.Int64()), nil
+}
+
 // ReconciliationCoverage returns the proportion of accounts [0.0, 1.0] that
 // have been reconciled at an index >= to a minimumIndex.
 func (b *BalanceStorage) ReconciliationCoverage(
@@ -313,15 +487,25 @@ func (b *BalanceStorage) ReconciliationCoverage(
 ) (float64, error) {
 	seen := 0
 	validCoverage := 0
-	err := b.getAllAccountEntries(ctx, func(entry accountEntry) {
+	err := b.getAllAccountEntries(ctx, func(txn database.Transaction, entry accountEntry) error {
 		seen++
-		if entry.LastReconciled == nil {
-			return
+
+		// Fetch last reconciliation index in same database.Transaction
+		key := GetAccountKey(reconciliationNamepace, entry.Account, entry.Currency)
+		exists, lastReconciled, err := BigIntGet(ctx, key, txn)
+		if err != nil {
+			return err
+		}
+
+		if !exists {
+			return nil
 		}
 
-		if entry.LastReconciled.Index >= minimumIndex {
+		if lastReconciled.Int64() >= minimumIndex {
 			validCoverage++
 		}
+
+		return nil
 	})
 	if err != nil {
 		return -1, fmt.Errorf("%w: unable to get all account entries", err)
@@ -341,12 +525,11 @@ func (b *BalanceStorage) ReconciliationCoverage(
 // we update the passed in existing value.
 func (b *BalanceStorage) existingValue(
 	ctx context.Context,
+	exists bool,
 	change *parser.BalanceChange,
 	parentBlock *types.BlockIdentifier,
 	existingValue string,
 ) (string, error) {
-	exists := len(existingValue) > 0
-
 	if exists {
 		return existingValue, nil
 	}
@@ -446,26 +629,83 @@ func (b *BalanceStorage) applyExemptions(
 	return liveAmount.Value, nil
 }
 
+// deleteAccountRecords is a convenience method that deletes
+// all traces of an account. This method should be called
+// within a global write transaction as it could contend with
+// other database.Transactions.
+//
+// WARNING: DO NOT RUN THIS CONCURRENTLY! IT CAN LEAD TO COUNTER
+// CORRUPTION!
+func (b *BalanceStorage) deleteAccountRecords(
+	ctx context.Context,
+	dbTx database.Transaction,
+	account *types.AccountIdentifier,
+	currency *types.Currency,
+) error {
+	for _, namespace := range []string{
+		accountNamespace,
+		reconciliationNamepace,
+		pruneNamespace,
+		balanceNamespace,
+	} {
+		key := GetAccountKey(namespace, account, currency)
+
+		// Determine if we should decrement the accounts
+		// seen counter
+		if namespace == accountNamespace {
+			exists, _, err := dbTx.Get(ctx, key)
+			if err != nil {
+				return err
+			}
+
+			if exists {
+				if err := b.handler.AccountsSeen(ctx, dbTx, -1); err != nil {
+					return err
+				}
+			}
+		}
+
+		// Determine if we should decrement the reconciled
+		// accounts counter
+		if namespace == reconciliationNamepace {
+			exists, _, err := dbTx.Get(ctx, key)
+			if err != nil {
+				return err
+			}
+
+			if exists {
+				if err := b.handler.AccountsReconciled(ctx, dbTx, -1); err != nil {
+					return err
+				}
+			}
+		}
+
+		if err := dbTx.Delete(ctx, key); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 // OrphanBalance removes all saved
 // states for a *types.Account and *types.Currency
 // at blocks >= the provided block.
 func (b *BalanceStorage) OrphanBalance(
 	ctx context.Context,
 	dbTransaction database.Transaction,
-	account *types.AccountIdentifier,
-	currency *types.Currency,
-	block *types.BlockIdentifier,
-) error {
+	change *parser.BalanceChange,
+) (bool, error) {
 	err := b.removeHistoricalBalances(
 		ctx,
 		dbTransaction,
-		account,
-		currency,
-		block.Index,
+		change.Account,
+		change.Currency,
+		change.Block.Index,
 		true,
 	)
 	if err != nil {
-		return err
+		return false, err
 	}
 
 	// Check if we should remove the account record
@@ -473,19 +713,39 @@ func (b *BalanceStorage) OrphanBalance(
 	_, err = b.getHistoricalBalance(
 		ctx,
 		dbTransaction,
-		account,
-		currency,
-		block.Index,
+		change.Account,
+		change.Currency,
+		change.Block.Index,
 	)
 	switch {
 	case errors.Is(err, storageErrs.ErrAccountMissing):
-		key := GetAccountKey(account, currency)
-		return dbTransaction.Delete(ctx, key)
+		return true, nil
 	case err != nil:
-		return err
-	default:
-		return nil
+		return false, err
+	}
+
+	// Update current balance
+	key := GetAccountKey(balanceNamespace, change.Account, change.Currency)
+	exists, lastBalance, err := BigIntGet(ctx, key, dbTransaction)
+	if err != nil {
+		return false, err
+	}
+
+	if !exists {
+		return false, storageErrs.ErrAccountMissing
 	}
+
+	difference, ok := new(big.Int).SetString(change.Difference, 10)
+	if !ok {
+		return false, storageErrs.ErrInvalidChangeValue
+	}
+
+	newBalance := new(big.Int).Add(lastBalance, difference)
+	if err := dbTransaction.Set(ctx, key, newBalance.Bytes(), true); err != nil {
+		return false, err
+	}
+
+	return false, nil
 }
 
 // PruneBalances removes all historical balance states
@@ -498,7 +758,8 @@ func (b *BalanceStorage) PruneBalances(
 	currency *types.Currency,
 	index int64,
 ) error {
-	dbTx := b.db.Transaction(ctx)
+	key := GetAccountKey(pruneNamespace, account, currency)
+	dbTx := b.db.WriteTransaction(ctx, string(key), false)
 	defer dbTx.Discard(ctx)
 
 	err := b.removeHistoricalBalances(
@@ -528,55 +789,39 @@ func (b *BalanceStorage) UpdateBalance(
 	dbTransaction database.Transaction,
 	change *parser.BalanceChange,
 	parentBlock *types.BlockIdentifier,
-) error {
+) (bool, error) {
 	if change.Currency == nil {
-		return errors.New("invalid currency")
+		return false, errors.New("invalid currency")
 	}
 
-	// Get existing account key to determine if
-	// balance should be fetched.
-	key := GetAccountKey(change.Account, change.Currency)
-	exists, _, err := dbTransaction.Get(ctx, key)
+	// If the balance key does not exist, the account
+	// does not exist.
+	key := GetAccountKey(balanceNamespace, change.Account, change.Currency)
+	exists, currentBalance, err := BigIntGet(ctx, key, dbTransaction)
 	if err != nil {
-		return err
-	}
-
-	var storedValue string
-	if exists {
-		// Get most recent historical balance
-		balance, err := b.getHistoricalBalance(
-			ctx,
-			dbTransaction,
-			change.Account,
-			change.Currency,
-			change.Block.Index,
-		)
-		switch {
-		case errors.Is(err, storageErrs.ErrAccountMissing):
-			storedValue = "0"
-		case err != nil:
-			return err
-		default:
-			storedValue = balance.Value
-		}
+		return false, err
 	}
 
 	// Find account existing value whether the account is new, has an
 	// existing balance, or is subject to additional accounting from
 	// a balance exemption.
+	//
+	// currentBalance is equal to 0 when it doesn't exist
+	// so we don't need to apply any conditional logic here.
 	existingValue, err := b.existingValue(
 		ctx,
+		exists,
 		change,
 		parentBlock,
-		storedValue,
+		currentBalance.String(),
 	)
 	if err != nil {
-		return err
+		return false, err
 	}
 
 	newVal, err := types.AddValues(change.Difference, existingValue)
 	if err != nil {
-		return err
+		return false, err
 	}
 
 	// If any exemptions apply, the returned new value will
@@ -584,16 +829,16 @@ func (b *BalanceStorage) UpdateBalance(
 	// and *types.Currency.
 	newVal, err = b.applyExemptions(ctx, change, newVal)
 	if err != nil {
-		return err
+		return false, err
 	}
 
 	bigNewVal, ok := new(big.Int).SetString(newVal, 10)
 	if !ok {
-		return fmt.Errorf("%s is not an integer", newVal)
+		return false, fmt.Errorf("%s is not an integer", newVal)
 	}
 
 	if bigNewVal.Sign() == -1 {
-		return fmt.Errorf(
+		return false, fmt.Errorf(
 			"%w %s:%+v for %+v at %+v",
 			storageErrs.ErrNegativeBalance,
 			newVal,
@@ -604,30 +849,38 @@ func (b *BalanceStorage) UpdateBalance(
 	}
 
 	// Add account entry if doesn't exist
+	var newAccount bool
 	if !exists {
+		newAccount = true
+		key := GetAccountKey(accountNamespace, change.Account, change.Currency)
 		serialAcc, err := b.db.Encoder().Encode(accountNamespace, accountEntry{
 			Account:  change.Account,
 			Currency: change.Currency,
 		})
 		if err != nil {
-			return err
+			return false, err
 		}
 		if err := dbTransaction.Set(ctx, key, serialAcc, true); err != nil {
-			return err
+			return false, err
 		}
 	}
 
+	// Update current balance
+	if err := dbTransaction.Set(ctx, key, bigNewVal.Bytes(), true); err != nil {
+		return false, err
+	}
+
 	// Add a new historical record for the balance.
 	historicalKey := GetHistoricalBalanceKey(
 		change.Account,
 		change.Currency,
 		change.Block.Index,
 	)
-	if err := dbTransaction.Set(ctx, historicalKey, []byte(newVal), true); err != nil {
-		return err
+	if err := dbTransaction.Set(ctx, historicalKey, bigNewVal.Bytes(), true); err != nil {
+		return false, err
 	}
 
-	return nil
+	return newAccount, nil
 }
 
 // GetBalance returns the balance of a types.AccountIdentifier
@@ -664,8 +917,8 @@ func (b *BalanceStorage) GetBalanceTransactional(
 	currency *types.Currency,
 	index int64,
 ) (*types.Amount, error) {
-	key := GetAccountKey(account, currency)
-	exists, acct, err := dbTx.Get(ctx, key)
+	key := GetAccountKey(balanceNamespace, account, currency)
+	exists, _, err := dbTx.Get(ctx, key)
 	if err != nil {
 		return nil, err
 	}
@@ -674,20 +927,17 @@ func (b *BalanceStorage) GetBalanceTransactional(
 		return nil, storageErrs.ErrAccountMissing
 	}
 
-	var accEntry accountEntry
-	if err := b.db.Encoder().Decode(accountNamespace, acct, &accEntry, true); err != nil {
-		return nil, fmt.Errorf(
-			"%w: unable to parse balance entry for %s",
-			err,
-			string(acct),
-		)
+	key = GetAccountKey(pruneNamespace, account, currency)
+	exists, lastPruned, err := BigIntGet(ctx, key, dbTx)
+	if err != nil {
+		return nil, err
 	}
 
-	if accEntry.LastPruned != nil && *accEntry.LastPruned >= index {
+	if exists && lastPruned.Int64() >= index {
 		return nil, fmt.Errorf(
 			"%w: last pruned %d",
 			storageErrs.ErrBalancePruned,
-			*accEntry.LastPruned,
+			lastPruned.Int64(),
 		)
 	}
 
@@ -870,8 +1120,7 @@ func (b *BalanceStorage) BootstrapBalances(
 		}
 	}
 
-	err := dbTransaction.Commit(ctx)
-	if err != nil {
+	if err := dbTransaction.Commit(ctx); err != nil {
 		return err
 	}
 
@@ -881,7 +1130,7 @@ func (b *BalanceStorage) BootstrapBalances(
 
 func (b *BalanceStorage) getAllAccountEntries(
 	ctx context.Context,
-	handler func(accountEntry),
+	handler func(database.Transaction, accountEntry) error,
 ) error {
 	txn := b.db.ReadTransaction(ctx)
 	defer txn.Discard(ctx)
@@ -901,9 +1150,7 @@ func (b *BalanceStorage) getAllAccountEntries(
 				)
 			}
 
-			handler(accEntry)
-
-			return nil
+			return handler(txn, accEntry)
 		},
 		false,
 		false,
@@ -924,8 +1171,9 @@ func (b *BalanceStorage) GetAllAccountCurrency(
 	log.Println("Loading previously seen accounts (this could take a while)...")
 
 	accountEntries := []*accountEntry{}
-	if err := b.getAllAccountEntries(ctx, func(entry accountEntry) {
+	if err := b.getAllAccountEntries(ctx, func(_ database.Transaction, entry accountEntry) error {
 		accountEntries = append(accountEntries, &entry)
+		return nil
 	}); err != nil {
 		return nil, fmt.Errorf("%w: unable to get all balance entries", err)
 	}
@@ -996,7 +1244,7 @@ func (b *BalanceStorage) getHistoricalBalance(
 		GetHistoricalBalancePrefix(account, currency),
 		GetHistoricalBalanceKey(account, currency, index),
 		func(k []byte, v []byte) error {
-			foundValue = string(v)
+			foundValue = new(big.Int).SetBytes(v).String()
 			return errAccountFound
 		},
 		false,
@@ -1015,54 +1263,6 @@ func (b *BalanceStorage) getHistoricalBalance(
 	return nil, storageErrs.ErrAccountMissing
 }
 
-func (b *BalanceStorage) updateAccountEntry(
-	ctx context.Context,
-	dbTx database.Transaction,
-	account *types.AccountIdentifier,
-	currency *types.Currency,
-	handler func(accountEntry) *accountEntry,
-) error {
-	key := GetAccountKey(account, currency)
-	exists, acc, err := dbTx.Get(ctx, key)
-	if err != nil {
-		return fmt.Errorf(
-			"%w: unable to get account entry for %s:%s",
-			err,
-			types.PrettyPrintStruct(account),
-			types.PrettyPrintStruct(currency),
-		)
-	}
-
-	if !exists {
-		return fmt.Errorf(
-			"account entry is missing for %s:%s",
-			types.PrettyPrintStruct(account),
-			types.PrettyPrintStruct(currency),
-		)
-	}
-
-	var accEntry accountEntry
-	if err := b.db.Encoder().Decode(accountNamespace, acc, &accEntry, true); err != nil {
-		return fmt.Errorf("%w: unable to decode account entry", err)
-	}
-
-	newAcctEntry := handler(accEntry)
-	if newAcctEntry == nil {
-		return nil
-	}
-
-	serialAcc, err := b.db.Encoder().Encode(accountNamespace, newAcctEntry)
-	if err != nil {
-		return fmt.Errorf("%w: unable to encode account entry", err)
-	}
-
-	if err := dbTx.Set(ctx, key, serialAcc, true); err != nil {
-		return fmt.Errorf("%w: unable to set account entry", err)
-	}
-
-	return nil
-}
-
 // removeHistoricalBalances deletes all historical balances
 // >= (used during reorg) or <= (used during pruning) a particular
 // index.
@@ -1083,6 +1283,13 @@ func (b *BalanceStorage) removeHistoricalBalances(
 			thisK := make([]byte, len(k))
 			copy(thisK, k)
 
+			// If we don't limit the max pruning size,
+			// we may accidentally create a commit too large for
+			// the database.
+			if !orphan && len(foundKeys) >= maxBalancePruneSize {
+				return errTooManyKeys
+			}
+
 			foundKeys = append(foundKeys, thisK)
 			return nil
 		},
@@ -1093,7 +1300,7 @@ func (b *BalanceStorage) removeHistoricalBalances(
 		// are pruning, we want to sort least to greatest.
 		!orphan,
 	)
-	if err != nil {
+	if err != nil && !errors.Is(err, errTooManyKeys) {
 		return fmt.Errorf("%w: database scan failed", err)
 	}
 
@@ -1103,29 +1310,20 @@ func (b *BalanceStorage) removeHistoricalBalances(
 		}
 	}
 
-	// Update the last pruned index
-	if !orphan {
-		err := b.updateAccountEntry(
-			ctx,
-			dbTx,
-			account,
-			currency,
-			func(accEntry accountEntry) *accountEntry {
-				// Don't update last pruned if the most recent pruning was
-				// lower than the last prune.
-				if accEntry.LastPruned != nil && *accEntry.LastPruned > index {
-					return nil
-				}
+	if orphan {
+		return nil
+	}
 
-				accEntry.LastPruned = &index
+	// Update the last pruned index
+	key := GetAccountKey(pruneNamespace, account, currency)
+	exists, lastPruned, err := BigIntGet(ctx, key, dbTx)
+	if err != nil {
+		return err
+	}
 
-				return &accEntry
-			},
-		)
-		if err != nil {
-			return err
-		}
+	if exists && lastPruned.Int64() > index {
+		return nil
 	}
 
-	return nil
+	return dbTx.Set(ctx, key, big.NewInt(index).Bytes(), true)
 }
diff --git a/storage/modules/counter_storage.go b/storage/modules/counter_storage.go
index af9d65f..7cedd5a 100644
--- a/storage/modules/counter_storage.go
+++ b/storage/modules/counter_storage.go
@@ -21,6 +21,7 @@ import (
 
 	"github.com/coinbase/rosetta-sdk-go/storage/database"
 	"github.com/coinbase/rosetta-sdk-go/types"
+	"github.com/coinbase/rosetta-sdk-go/utils"
 )
 
 const (
@@ -75,6 +76,12 @@ const (
 	// or the block where an account was updated has been orphaned.
 	SkippedReconciliationsCounter = "skipped_reconciliations"
 
+	// SeenAccounts is the total number of accounts seen.
+	SeenAccounts = "seen_accounts"
+
+	// ReconciledAccounts is the total number of accounts seen.
+	ReconciledAccounts = "reconciled_accounts"
+
 	// counterNamespace is preprended to any counter.
 	counterNamespace = "counter"
 )
@@ -85,6 +92,8 @@ var _ BlockWorker = (*CounterStorage)(nil)
 // on top of a database.Database and database.Transaction interface.
 type CounterStorage struct {
 	db database.Database
+
+	m *utils.MutexMap
 }
 
 // NewCounterStorage returns a new CounterStorage.
@@ -93,6 +102,7 @@ func NewCounterStorage(
 ) *CounterStorage {
 	return &CounterStorage{
 		db: db,
+		m:  utils.NewMutexMap(utils.DefaultShards),
 	}
 }
 
@@ -100,21 +110,21 @@ func getCounterKey(counter string) []byte {
 	return []byte(fmt.Sprintf("%s/%s", counterNamespace, counter))
 }
 
-func transactionalGet(
+func BigIntGet(
 	ctx context.Context,
-	counter string,
+	key []byte,
 	txn database.Transaction,
-) (*big.Int, error) {
-	exists, val, err := txn.Get(ctx, getCounterKey(counter))
+) (bool, *big.Int, error) {
+	exists, val, err := txn.Get(ctx, key)
 	if err != nil {
-		return nil, err
+		return false, nil, err
 	}
 
 	if !exists {
-		return big.NewInt(0), nil
+		return false, big.NewInt(0), nil
 	}
 
-	return new(big.Int).SetBytes(val), nil
+	return true, new(big.Int).SetBytes(val), nil
 }
 
 // UpdateTransactional updates the value of a counter by amount and returns the new
@@ -125,7 +135,13 @@ func (c *CounterStorage) UpdateTransactional(
 	counter string,
 	amount *big.Int,
 ) (*big.Int, error) {
-	val, err := transactionalGet(ctx, counter, dbTx)
+	// Ensure that same counter is not incremented
+	// concurrently. This is necessary because we concurrently
+	// store account balances.
+	c.m.Lock(counter, false)
+	defer c.m.Unlock(counter)
+
+	_, val, err := BigIntGet(ctx, getCounterKey(counter), dbTx)
 	if err != nil {
 		return nil, err
 	}
@@ -165,7 +181,8 @@ func (c *CounterStorage) Get(ctx context.Context, counter string) (*big.Int, err
 	transaction := c.db.ReadTransaction(ctx)
 	defer transaction.Discard(ctx)
 
-	return transactionalGet(ctx, counter, transaction)
+	_, value, err := BigIntGet(ctx, getCounterKey(counter), transaction)
+	return value, err
 }
 
 // AddingBlock is called by BlockStorage when adding a block.


Test patch:
diff --git a/parser/balance_changes_test.go b/parser/balance_changes_test.go
index ae789b9..8f50c00 100644
--- a/parser/balance_changes_test.go
+++ b/parser/balance_changes_test.go
@@ -233,8 +233,8 @@ func TestBalanceChanges(t *testing.T) {
 					},
 					Currency: currency,
 					Block: &types.BlockIdentifier{
-						Hash:  "0",
-						Index: 0,
+						Hash:  "1",
+						Index: 1,
 					},
 					Difference: "-250",
 				},
@@ -244,8 +244,8 @@ func TestBalanceChanges(t *testing.T) {
 					},
 					Currency: currency,
 					Block: &types.BlockIdentifier{
-						Hash:  "0",
-						Index: 0,
+						Hash:  "1",
+						Index: 1,
 					},
 					Difference: "-150",
 				},
diff --git a/storage/modules/balance_storage_test.go b/storage/modules/balance_storage_test.go
index 11757e3..c5111e1 100644
--- a/storage/modules/balance_storage_test.go
+++ b/storage/modules/balance_storage_test.go
@@ -224,6 +224,7 @@ func TestBalance(t *testing.T) {
 			&types.Amount{Value: "10", Currency: currency},
 			nil,
 		).Once()
+		mockHandler.On("AccountsSeen", ctx, mock.Anything, 1).Return(nil).Once()
 		amount, err := storage.GetOrSetBalance(ctx, account, currency, newBlock)
 		assert.NoError(t, err)
 		assert.Equal(t, &types.Amount{
@@ -241,6 +242,7 @@ func TestBalance(t *testing.T) {
 
 	t.Run("Set and get genesis balance", func(t *testing.T) {
 		txn := storage.db.Transaction(ctx)
+		mockHandler.On("AccountsSeen", ctx, txn, 1).Return(nil).Once()
 		err := storage.SetBalance(
 			ctx,
 			txn,
@@ -264,6 +266,13 @@ func TestBalance(t *testing.T) {
 
 	t.Run("Set and get balance", func(t *testing.T) {
 		txn := storage.db.Transaction(ctx)
+
+		// When deleting records, we decrement the account seen count.
+		mockHandler.On("AccountsSeen", ctx, txn, -1).Return(nil).Once()
+
+		// When adding the account, we increment the account seen count.
+		mockHandler.On("AccountsSeen", ctx, txn, 1).Return(nil).Once()
+
 		err := storage.SetBalance(
 			ctx,
 			txn,
@@ -295,7 +304,7 @@ func TestBalance(t *testing.T) {
 			nil,
 		).Once()
 		txn := storage.db.Transaction(ctx)
-		err := storage.UpdateBalance(
+		newAccount, err := storage.UpdateBalance(
 			ctx,
 			txn,
 			&parser.BalanceChange{
@@ -306,6 +315,7 @@ func TestBalance(t *testing.T) {
 			},
 			nil,
 		)
+		assert.True(t, newAccount)
 		assert.NoError(t, err)
 		assert.NoError(t, txn.Commit(ctx))
 
@@ -316,7 +326,7 @@ func TestBalance(t *testing.T) {
 
 	t.Run("Set balance with nil currency", func(t *testing.T) {
 		txn := storage.db.Transaction(ctx)
-		err := storage.UpdateBalance(
+		newAccount, err := storage.UpdateBalance(
 			ctx,
 			txn,
 			&parser.BalanceChange{
@@ -327,6 +337,7 @@ func TestBalance(t *testing.T) {
 			},
 			nil,
 		)
+		assert.False(t, newAccount)
 		assert.EqualError(t, err, "invalid currency")
 		txn.Discard(ctx)
 
@@ -337,7 +348,7 @@ func TestBalance(t *testing.T) {
 
 	t.Run("Modify existing balance", func(t *testing.T) {
 		txn := storage.db.Transaction(ctx)
-		err = storage.UpdateBalance(
+		newAccount, err := storage.UpdateBalance(
 			ctx,
 			txn,
 			&parser.BalanceChange{
@@ -348,6 +359,7 @@ func TestBalance(t *testing.T) {
 			},
 			nil,
 		)
+		assert.False(t, newAccount)
 		assert.NoError(t, err)
 		assert.NoError(t, txn.Commit(ctx))
 
@@ -358,7 +370,7 @@ func TestBalance(t *testing.T) {
 
 	t.Run("Discard transaction", func(t *testing.T) {
 		txn := storage.db.Transaction(ctx)
-		err := storage.UpdateBalance(
+		newAccount, err := storage.UpdateBalance(
 			ctx,
 			txn,
 			&parser.BalanceChange{
@@ -369,6 +381,7 @@ func TestBalance(t *testing.T) {
 			},
 			nil,
 		)
+		assert.False(t, newAccount)
 		assert.NoError(t, err)
 
 		// Get balance during transaction
@@ -389,7 +402,7 @@ func TestBalance(t *testing.T) {
 
 	t.Run("Attempt modification to push balance negative on existing account", func(t *testing.T) {
 		txn := storage.db.Transaction(ctx)
-		err := storage.UpdateBalance(
+		newAccount, err := storage.UpdateBalance(
 			ctx,
 			txn,
 			&parser.BalanceChange{
@@ -401,6 +414,7 @@ func TestBalance(t *testing.T) {
 			nil,
 		)
 		assert.True(t, errors.Is(err, storageErrs.ErrNegativeBalance))
+		assert.False(t, newAccount)
 		txn.Discard(ctx)
 	})
 
@@ -416,7 +430,7 @@ func TestBalance(t *testing.T) {
 			&types.Amount{Value: "0", Currency: largeDeduction.Currency},
 			nil,
 		).Once()
-		err := storage.UpdateBalance(
+		newAccount, err := storage.UpdateBalance(
 			ctx,
 			txn,
 			&parser.BalanceChange{
@@ -427,6 +441,7 @@ func TestBalance(t *testing.T) {
 			},
 			nil,
 		)
+		assert.False(t, newAccount)
 		assert.Error(t, err)
 		assert.True(t, errors.Is(err, storageErrs.ErrNegativeBalance))
 		txn.Discard(ctx)
@@ -444,7 +459,7 @@ func TestBalance(t *testing.T) {
 			&types.Amount{Value: "0", Currency: amount.Currency},
 			nil,
 		).Once()
-		err := storage.UpdateBalance(
+		newAccount, err := storage.UpdateBalance(
 			ctx,
 			txn,
 			&parser.BalanceChange{
@@ -455,6 +470,7 @@ func TestBalance(t *testing.T) {
 			},
 			nil,
 		)
+		assert.True(t, newAccount)
 		assert.NoError(t, err)
 		assert.NoError(t, txn.Commit(ctx))
 
@@ -480,7 +496,7 @@ func TestBalance(t *testing.T) {
 			&types.Amount{Value: "0", Currency: amount.Currency},
 			nil,
 		).Once()
-		err := storage.UpdateBalance(
+		newAccount, err := storage.UpdateBalance(
 			ctx,
 			txn,
 			&parser.BalanceChange{
@@ -491,6 +507,7 @@ func TestBalance(t *testing.T) {
 			},
 			nil,
 		)
+		assert.True(t, newAccount)
 		assert.NoError(t, err)
 		assert.NoError(t, txn.Commit(ctx))
 
@@ -516,7 +533,7 @@ func TestBalance(t *testing.T) {
 			&types.Amount{Value: "0", Currency: amount.Currency},
 			nil,
 		).Once()
-		err := storage.UpdateBalance(
+		newAccount, err := storage.UpdateBalance(
 			ctx,
 			txn,
 			&parser.BalanceChange{
@@ -527,6 +544,7 @@ func TestBalance(t *testing.T) {
 			},
 			nil,
 		)
+		assert.True(t, newAccount)
 		assert.NoError(t, err)
 		assert.NoError(t, txn.Commit(ctx))
 
@@ -542,6 +560,7 @@ func TestBalance(t *testing.T) {
 
 	t.Run("balance exemption update", func(t *testing.T) {
 		txn := storage.db.Transaction(ctx)
+		mockHandler.On("AccountsSeen", ctx, txn, 1).Return(nil).Once()
 		err := storage.SetBalance(
 			ctx,
 			txn,
@@ -567,7 +586,7 @@ func TestBalance(t *testing.T) {
 			nil,
 		).Once()
 		txn = storage.db.Transaction(ctx)
-		err = storage.UpdateBalance(
+		newAccount, err := storage.UpdateBalance(
 			ctx,
 			txn,
 			&parser.BalanceChange{
@@ -578,6 +597,7 @@ func TestBalance(t *testing.T) {
 			},
 			nil,
 		)
+		assert.False(t, newAccount)
 		assert.NoError(t, err)
 		assert.NoError(t, txn.Commit(ctx))
 
@@ -602,7 +622,7 @@ func TestBalance(t *testing.T) {
 			nil,
 		).Once()
 		txn = storage.db.Transaction(ctx)
-		err = storage.UpdateBalance(
+		newAccount, err = storage.UpdateBalance(
 			ctx,
 			txn,
 			&parser.BalanceChange{
@@ -613,6 +633,7 @@ func TestBalance(t *testing.T) {
 			},
 			nil,
 		)
+		assert.False(t, newAccount)
 		assert.NoError(t, err)
 		assert.NoError(t, txn.Commit(ctx))
 
@@ -637,7 +658,7 @@ func TestBalance(t *testing.T) {
 			nil,
 		).Once()
 		txn = storage.db.Transaction(ctx)
-		err = storage.UpdateBalance(
+		newAccount, err = storage.UpdateBalance(
 			ctx,
 			txn,
 			&parser.BalanceChange{
@@ -648,6 +669,7 @@ func TestBalance(t *testing.T) {
 			},
 			nil,
 		)
+		assert.False(t, newAccount)
 		assert.Error(t, err)
 		txn.Discard(ctx)
 
@@ -743,7 +765,7 @@ func TestBalance(t *testing.T) {
 	t.Run("update existing balance", func(t *testing.T) {
 		txn := storage.db.Transaction(ctx)
 		orphanValue, _ := new(big.Int).SetString(largeDeduction.Value, 10)
-		err := storage.UpdateBalance(
+		newAccount, err := storage.UpdateBalance(
 			ctx,
 			txn,
 			&parser.BalanceChange{
@@ -754,6 +776,7 @@ func TestBalance(t *testing.T) {
 			},
 			nil,
 		)
+		assert.False(t, newAccount)
 		assert.NoError(t, err)
 		retrievedAmount, err := storage.GetBalanceTransactional(
 			ctx,
@@ -796,13 +819,16 @@ func TestBalance(t *testing.T) {
 		}, retrievedAmount)
 
 		txn := storage.db.Transaction(ctx)
-		err = storage.OrphanBalance(
+		shouldRemove, err := storage.OrphanBalance(
 			ctx,
 			txn,
-			account,
-			largeDeduction.Currency,
-			newBlock2,
+			&parser.BalanceChange{
+				Account:  account,
+				Currency: largeDeduction.Currency,
+				Block:    newBlock,
+			},
 		)
+		assert.True(t, shouldRemove)
 		assert.NoError(t, err)
 		assert.NoError(t, txn.Commit(ctx))
 
@@ -810,11 +836,11 @@ func TestBalance(t *testing.T) {
 			ctx,
 			account,
 			largeDeduction.Currency,
-			newBlock3,
+			newBlock,
 		)
 		assert.NoError(t, err)
 		assert.Equal(t, &types.Amount{
-			Value:    "100",
+			Value:    "0",
 			Currency: largeDeduction.Currency,
 		}, retrievedAmount)
 	})
@@ -944,6 +970,7 @@ func TestSetBalanceImported(t *testing.T) {
 	storage.Initialize(mockHelper, mockHandler)
 
 	t.Run("Set balance successfully", func(t *testing.T) {
+		mockHandler.On("AccountsSeen", ctx, mock.Anything, 1).Return(nil).Twice()
 		err = storage.SetBalanceImported(
 			ctx,
 			nil,
@@ -1042,6 +1069,7 @@ func TestBootstrapBalances(t *testing.T) {
 			ioutil.WriteFile(bootstrapBalancesFile, file, utils.DefaultFilePermissions),
 		)
 
+		mockHandler.On("AccountsSeen", ctx, mock.Anything, 1).Return(nil).Once()
 		err = storage.BootstrapBalances(
 			ctx,
 			bootstrapBalancesFile,
@@ -1061,7 +1089,7 @@ func TestBootstrapBalances(t *testing.T) {
 
 		// Attempt to update balance
 		txn := storage.db.Transaction(ctx)
-		err = storage.UpdateBalance(
+		newAccount, err := storage.UpdateBalance(
 			ctx,
 			txn,
 			&parser.BalanceChange{
@@ -1072,6 +1100,7 @@ func TestBootstrapBalances(t *testing.T) {
 			},
 			newBlock,
 		)
+		assert.False(t, newAccount)
 		assert.NoError(t, err)
 		assert.NoError(t, txn.Commit(ctx))
 
@@ -1222,7 +1251,7 @@ func TestBalanceReconciliation(t *testing.T) {
 
 	t.Run("attempt to store reconciliation for non-existent account", func(t *testing.T) {
 		err := storage.Reconciled(ctx, account, currency, genesisBlock)
-		assert.Error(t, err)
+		assert.NoError(t, err)
 
 		coverage, err := storage.ReconciliationCoverage(ctx, 0)
 		assert.NoError(t, err)
@@ -1231,7 +1260,7 @@ func TestBalanceReconciliation(t *testing.T) {
 
 	t.Run("set balance", func(t *testing.T) {
 		txn := storage.db.Transaction(ctx)
-		err := storage.UpdateBalance(
+		newAccount, err := storage.UpdateBalance(
 			ctx,
 			txn,
 			&parser.BalanceChange{
@@ -1242,6 +1271,7 @@ func TestBalanceReconciliation(t *testing.T) {
 			},
 			genesisBlock,
 		)
+		assert.True(t, newAccount)
 		assert.NoError(t, err)
 		assert.NoError(t, txn.Commit(ctx))
 
@@ -1255,7 +1285,7 @@ func TestBalanceReconciliation(t *testing.T) {
 		assert.NoError(t, err)
 
 		txn := storage.db.Transaction(ctx)
-		err = storage.UpdateBalance(
+		newAccount, err := storage.UpdateBalance(
 			ctx,
 			txn,
 			&parser.BalanceChange{
@@ -1266,6 +1296,7 @@ func TestBalanceReconciliation(t *testing.T) {
 			},
 			genesisBlock,
 		)
+		assert.True(t, newAccount)
 		assert.NoError(t, err)
 		assert.NoError(t, txn.Commit(ctx))
 
@@ -1307,7 +1338,7 @@ func TestBalanceReconciliation(t *testing.T) {
 
 	t.Run("add unreconciled", func(t *testing.T) {
 		txn := storage.db.Transaction(ctx)
-		err = storage.UpdateBalance(
+		newAccount, err := storage.UpdateBalance(
 			ctx,
 			txn,
 			&parser.BalanceChange{
@@ -1318,6 +1349,7 @@ func TestBalanceReconciliation(t *testing.T) {
 			},
 			newBlock,
 		)
+		assert.True(t, newAccount)
 		assert.NoError(t, err)
 		assert.NoError(t, txn.Commit(ctx))
 
@@ -1326,6 +1358,22 @@ func TestBalanceReconciliation(t *testing.T) {
 		assert.Equal(t, float64(1)/float64(3), coverage)
 	})
 
+	t.Run("test estimated no reconciliations", func(t *testing.T) {
+		mockHelper.On("AccountsReconciled", ctx, mock.Anything).Return(big.NewInt(0), nil).Once()
+		mockHelper.On("AccountsSeen", ctx, mock.Anything).Return(big.NewInt(0), nil).Once()
+		coverage, err := storage.EstimatedReconciliationCoverage(ctx)
+		assert.Equal(t, float64(0), coverage)
+		assert.NoError(t, err)
+	})
+
+	t.Run("test estimated some reconciliations", func(t *testing.T) {
+		mockHelper.On("AccountsReconciled", ctx, mock.Anything).Return(big.NewInt(1), nil).Once()
+		mockHelper.On("AccountsSeen", ctx, mock.Anything).Return(big.NewInt(2), nil).Once()
+		coverage, err := storage.EstimatedReconciliationCoverage(ctx)
+		assert.Equal(t, float64(0.5), coverage)
+		assert.NoError(t, err)
+	})
+
 	mockHelper.AssertExpectations(t)
 	mockHandler.AssertExpectations(t)
 }
@@ -1533,6 +1581,7 @@ func TestBlockSyncing(t *testing.T) {
 			&types.Amount{Value: "1", Currency: curr},
 			nil,
 		).Once()
+		mockHandler.On("AccountsSeen", ctx, dbTx, 1).Return(nil).Once()
 		_, err = storage.AddingBlock(ctx, b1, dbTx)
 		assert.NoError(t, err)
 		assert.NoError(t, dbTx.Commit(ctx))
@@ -1555,6 +1604,10 @@ func TestBlockSyncing(t *testing.T) {
 	})
 
 	t.Run("add block 2", func(t *testing.T) {
+		// Reconcile a previously seen account
+		err := storage.Reconciled(ctx, addr1, curr, b1.BlockIdentifier)
+		assert.NoError(t, err)
+
 		dbTx := database.Transaction(ctx)
 		mockHelper.On(
 			"AccountBalance",
@@ -1566,6 +1619,8 @@ func TestBlockSyncing(t *testing.T) {
 			&types.Amount{Value: "0", Currency: curr},
 			nil,
 		).Once()
+		mockHandler.On("AccountsSeen", ctx, dbTx, 1).Return(nil).Once()
+		mockHandler.On("AccountsReconciled", ctx, dbTx, 1).Return(nil).Once()
 		_, err = storage.AddingBlock(ctx, b2, dbTx)
 		assert.NoError(t, err)
 		assert.NoError(t, dbTx.Commit(ctx))
@@ -1604,9 +1659,12 @@ func TestBlockSyncing(t *testing.T) {
 
 	t.Run("orphan block 2", func(t *testing.T) {
 		dbTx := database.Transaction(ctx)
-		_, err = storage.RemovingBlock(ctx, b2, dbTx)
+		commitWorker, err := storage.RemovingBlock(ctx, b2, dbTx)
 		assert.NoError(t, err)
 		assert.NoError(t, dbTx.Commit(ctx))
+		mockHandler.On("BlockRemoved", ctx, b2, mock.Anything).Return(nil).Once()
+		mockHandler.On("AccountsSeen", ctx, mock.Anything, -1).Return(nil).Once()
+		assert.NoError(t, commitWorker(ctx))
 
 		amount, err := storage.GetBalance(ctx, addr1, curr, b0.BlockIdentifier.Index)
 		assert.NoError(t, err)
@@ -1636,9 +1694,13 @@ func TestBlockSyncing(t *testing.T) {
 
 	t.Run("orphan block 1", func(t *testing.T) {
 		dbTx := database.Transaction(ctx)
-		_, err = storage.RemovingBlock(ctx, b1, dbTx)
+		commitWorker, err := storage.RemovingBlock(ctx, b1, dbTx)
 		assert.NoError(t, err)
 		assert.NoError(t, dbTx.Commit(ctx))
+		mockHandler.On("BlockRemoved", ctx, b1, mock.Anything).Return(nil).Once()
+		mockHandler.On("AccountsSeen", ctx, mock.Anything, -1).Return(nil).Once()
+		mockHandler.On("AccountsReconciled", ctx, mock.Anything, -1).Return(nil).Once()
+		assert.NoError(t, commitWorker(ctx))
 
 		amount, err := storage.GetBalance(ctx, addr1, curr, b0.BlockIdentifier.Index)
 		assert.True(t, errors.Is(err, storageErrs.ErrAccountMissing))
@@ -1672,6 +1734,7 @@ func TestBlockSyncing(t *testing.T) {
 			&types.Amount{Value: "1", Currency: curr},
 			nil,
 		).Once()
+		mockHandler.On("AccountsSeen", ctx, dbTx, 1).Return(nil).Once()
 		_, err = storage.AddingBlock(ctx, b1, dbTx)
 		assert.NoError(t, err)
 		assert.NoError(t, dbTx.Commit(ctx))
@@ -1714,6 +1777,7 @@ func TestBlockSyncing(t *testing.T) {
 			&types.Amount{Value: "0", Currency: curr},
 			nil,
 		).Once()
+		mockHandler.On("AccountsSeen", ctx, dbTx, 1).Return(nil).Once()
 		_, err = storage.AddingBlock(ctx, b2a, dbTx)
 		assert.NoError(t, err)
 		assert.NoError(t, dbTx.Commit(ctx))