diff --git a/storage/modules/balance_storage.go b/storage/modules/balance_storage.go index 54e940a9..122bcadd 100644 --- a/storage/modules/balance_storage.go +++ b/storage/modules/balance_storage.go @@ -1097,7 +1097,16 @@ func (b *BalanceStorage) BootstrapBalances( dbTransaction := b.db.Transaction(ctx) defer dbTransaction.Discard(ctx) - for _, balance := range balances { + for i, balance := range balances { + // Commit transaction batch by batch rather than commit at one time. + // This helps reduce memory usage and improve running time when bootstrap_balances.json + // contains huge number of accounts. + if i != 0 && i%utils.MaxEntrySizePerTxn == 0 { + if err := dbTransaction.Commit(ctx); err != nil { + return err + } + dbTransaction = b.db.Transaction(ctx) + } // Ensure change.Difference is valid amountValue, ok := new(big.Int).SetString(balance.Value, 10) if !ok { diff --git a/storage/modules/balance_storage_test.go b/storage/modules/balance_storage_test.go index 8b3a16a3..23c4b3cc 100644 --- a/storage/modules/balance_storage_test.go +++ b/storage/modules/balance_storage_test.go @@ -1020,6 +1020,14 @@ func TestBootstrapBalances(t *testing.T) { account = &types.AccountIdentifier{ Address: "hello", } + + account2 = &types.AccountIdentifier{ + Address: "hello world", + } + + account3 = &types.AccountIdentifier{ + Address: "hello world new", + } ) ctx := context.Background() @@ -1064,6 +1072,16 @@ func TestBootstrapBalances(t *testing.T) { Value: amount.Value, Currency: amount.Currency, }, + { + Account: account2, + Value: amount.Value, + Currency: amount.Currency, + }, + { + Account: account3, + Value: amount.Value, + Currency: amount.Currency, + }, }, "", " ") assert.NoError(t, err) @@ -1083,7 +1101,7 @@ func TestBootstrapBalances(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).Once() + mockHandler.On("AccountsSeen", ctx, mock.Anything, 1).Return(nil).Times(3) err = storage.BootstrapBalances( ctx, bootstrapBalancesFile, @@ -1101,6 +1119,26 @@ func TestBootstrapBalances(t *testing.T) { assert.Equal(t, amount, retrievedAmount) assert.NoError(t, err) + retrievedAmount2, err := storage.GetOrSetBalance( + ctx, + account2, + amount.Currency, + genesisBlockIdentifier, + ) + + assert.Equal(t, amount, retrievedAmount2) + assert.NoError(t, err) + + retrievedAmount3, err := storage.GetOrSetBalance( + ctx, + account3, + amount.Currency, + genesisBlockIdentifier, + ) + + assert.Equal(t, amount, retrievedAmount3) + assert.NoError(t, err) + // Attempt to update balance txn := storage.db.Transaction(ctx) newAccount, err := storage.UpdateBalance( diff --git a/utils/utils.go b/utils/utils.go index 1726ded3..da286da9 100644 --- a/utils/utils.go +++ b/utils/utils.go @@ -63,6 +63,12 @@ const ( // to consider when estimating time to tip if the provided // estimate is 0. minBlocksPerSecond = 0.0001 + + // MaxEntrySizePerTxn is the maximum number of entries + // in one transaction object. This is used for bootstrap + // balances process to avoid TxnTooBig error when memory_limit_disabled=false + // as well as reduce the running time. + MaxEntrySizePerTxn = 600 ) var (