accounts_benchmark_test.go

// Copyright (c) 2025 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package wallet

import (
	"fmt"
	"testing"
	"time"

	"github.com/btcsuite/btcd/btcutil"
	"github.com/btcsuite/btcd/chaincfg/chainhash"
	"github.com/btcsuite/btcd/txscript"
	"github.com/btcsuite/btcd/wire"
	"github.com/btcsuite/btcwallet/waddrmgr"
	"github.com/btcsuite/btcwallet/walletdb"
	"github.com/btcsuite/btcwallet/wtxmgr"
	"github.com/stretchr/testify/require"
)

// benchmarkDataSize represents different test data sizes for stress testing.
type benchmarkDataSize struct {
	numAccounts int
	numUTXOs    int
}

// name returns a dynamically generated benchmark name based on numAccounts and
// numUTXOs.
func (b benchmarkDataSize) name() string {
	return fmt.Sprintf("%dAccounts_%dUTXOs", b.numAccounts, b.numUTXOs)
}

// generateBenchmarkSizes creates benchmark data sizes programmatically.
func generateBenchmarkSizes() []benchmarkDataSize {
	var sizes []benchmarkDataSize

	// Generate UTXO sizes from 2^0 to 2^13 and account sizes from 5 to 70.
	for i := 0; i <= 13; i++ {
		numUTXOs := 1 << i
		numAccounts := 5 + (i * 5)
		sizes = append(sizes, benchmarkDataSize{
			numAccounts: numAccounts,
			numUTXOs:    numUTXOs,
		})
	}

	return sizes
}

var benchmarkSizes = generateBenchmarkSizes()

// BenchmarkAccounts benchmarks the old deprecated Accounts method.
func BenchmarkAccounts(b *testing.B) {
	for _, size := range benchmarkSizes {
		b.Run(size.name(), func(b *testing.B) {
			w := setupBenchmarkWallet(
				b, size.numAccounts, size.numUTXOs,
			)

			scope := waddrmgr.KeyScopeBIP0044

			b.ReportAllocs()

			for b.Loop() {
				result, err := w.Accounts(scope)
				require.NoError(b, err)

				// Intentionally prevent compiler optimization
				// by using the result.
				_ = result.Accounts
			}
		})
	}
}

// BenchmarkListAccounts benchmarks the new optimized ListAccountsByScope
// method.
func BenchmarkListAccountsByScope(b *testing.B) {
	for _, size := range benchmarkSizes {
		b.Run(size.name(), func(b *testing.B) {
			w := setupBenchmarkWallet(
				b, size.numAccounts, size.numUTXOs,
			)

			scope := waddrmgr.KeyScopeBIP0044

			b.ReportAllocs()

			for b.Loop() {
				result, err := w.ListAccountsByScope(
					b.Context(), scope,
				)
				require.NoError(b, err)

				// Intentionally prevent compiler optimization
				// by using the result.
				_ = result.Accounts
			}
		})
	}
}

// setupBenchmarkWallet creates a wallet with test data for benchmarking
func setupBenchmarkWallet(b *testing.B, numAccounts, numUTXOs int) *Wallet {
	b.Helper()

	t := &testing.T{}

	// TODO: making sure the wallet db is deleted at the end of the test.
	// Perhaps a better way is modifying testWallet function to return a
	// cleanup callback function that deletes the wallet db.
	w, _ := testWallet(t)

	scopes := []waddrmgr.KeyScope{
		waddrmgr.KeyScopeBIP0044,
	}

	var (
		allAccountNumbers []uint32
		allAddresses      []waddrmgr.ManagedAddress
	)

	err := walletdb.Update(w.db, func(tx walletdb.ReadWriteTx) error {
		for _, scope := range scopes {
			manager, err := w.addrStore.FetchScopedKeyManager(scope)
			if err != nil {
				return err
			}

			addrmgrNs := tx.ReadWriteBucket(waddrmgrNamespaceKey)

			for i := range numAccounts {
				accountName := fmt.Sprintf("bench-account-%d", i)
				account, err := manager.NewAccount(
					addrmgrNs, accountName,
				)
				if err != nil {
					return err
				}
				allAccountNumbers = append(
					allAccountNumbers, account,
				)

				addrs, err := manager.NextExternalAddresses(
					addrmgrNs, account, 5,
				)
				if err != nil {
					return err
				}
				allAddresses = append(allAddresses, addrs...)
			}
		}
		return nil
	})
	require.NoError(t, err, "failed to create test accounts: %v", err)

	// Create valid UTXOs using existing test transaction data.
	err = walletdb.Update(w.db, func(tx walletdb.ReadWriteTx) error {
		txmgrNs := tx.ReadWriteBucket(wtxmgrNamespaceKey)

		msgTx := TstTx.MsgTx()

		for i := 0; i < numUTXOs && i < len(allAddresses); i++ {
			newMsgTx := wire.NewMsgTx(msgTx.Version)

			addr := allAddresses[i%len(allAddresses)]
			pkScript, err := txscript.PayToAddrScript(
				addr.Address(),
			)
			if err != nil {
				return err
			}

			amount := btcutil.Amount(100000 + i*1000)
			newMsgTx.AddTxOut(wire.NewTxOut(
				int64(amount), pkScript),
			)

			// Add a dummy input to make it valid.
			prevHash := chainhash.Hash{}
			prevHash[0] = byte(i)
			newMsgTx.AddTxIn(wire.NewTxIn(
				wire.NewOutPoint(&prevHash, 0), nil, nil,
			))

			rec, err := wtxmgr.NewTxRecordFromMsgTx(
				newMsgTx, time.Now(),
			)
			if err != nil {
				return err
			}

			blockMeta := &wtxmgr.BlockMeta{
				Block: wtxmgr.Block{
					Hash:   chainhash.Hash{},
					Height: 1,
				},
				Time: time.Now(),
			}

			if err = w.txStore.InsertTx(
				txmgrNs, rec, blockMeta,
			); err != nil {
				return err
			}

			// Mark the output as unspent.
			if err = w.txStore.AddCredit(
				txmgrNs, rec, blockMeta, 0, false,
			); err != nil {
				return err
			}
		}

		return nil
	})
	require.NoError(t, err, "failed to create test UTXOs: %v", err)

	return w
}

package wallet

import (
	"context"
	"fmt"
	"testing"
	"time"

	"github.com/btcsuite/btcd/btcutil"
	"github.com/btcsuite/btcd/chaincfg/chainhash"
	"github.com/btcsuite/btcd/txscript"
	"github.com/btcsuite/btcd/wire"
	"github.com/btcsuite/btcwallet/waddrmgr"
	"github.com/btcsuite/btcwallet/walletdb"
	"github.com/btcsuite/btcwallet/wtxmgr"
	"github.com/stretchr/testify/require"
)

// benchmarkDataSize represents different test data sizes for stress testing.
type benchmarkDataSize struct {
	numAccounts int
	numUTXOs    int
}

// name returns a dynamically generated benchmark name based on numAccounts and
// numUTXOs.
func (b benchmarkDataSize) name() string {
	return fmt.Sprintf("%d-Accounts_%d-UTXOs", b.numAccounts, b.numUTXOs)
}

// generateBenchmarkSizes creates benchmark data sizes programmatically.
func generateBenchmarkSizes() []benchmarkDataSize {
	var sizes []benchmarkDataSize

	// Generate UTXO sizes from 2^0 to 2^13 and account sizes from 5 to 70.
	for i := 0; i <= 13; i++ {
		numUTXOs := 1 << i
		numAccounts := 5 + (i * 5)
		sizes = append(sizes, benchmarkDataSize{
			numAccounts: numAccounts,
			numUTXOs:    numUTXOs,
		})
	}

	return sizes
}

// BenchmarkListAccountsByScope benchmarks the performance of the
// ListAccountsByScope method with proper setup/teardown isolation.
func BenchmarkListAccountsByScope(b *testing.B) {
	benchmarkSizes := generateBenchmarkSizes()

	for _, size := range benchmarkSizes {
		// Perform all expensive setup before starting the timer. This
		// includes wallet creation, account creation, and UTXO
		// generation.
		w, cleanup := setupBenchmarkWallet(
			b, size.numAccounts, size.numUTXOs,
		)

		scope := waddrmgr.KeyScopeBIP0044

		// Reset the timer to ensure the setup time is not included in
		// the benchmark results.
		b.ReportAllocs()
		b.ResetTimer()

		b.Run(size.name(), func(b *testing.B) {
			// The loop runs b.N times.
			for b.Loop() {
				_, err := w.ListAccountsByScope(
					context.Background(), scope,
				)
				require.NoError(b, err)
			}

			// Stop the timer before performing cleanup to ensure
			// that the teardown time is not included.
			b.StopTimer()
			cleanup()
		})
	}
}

// BenchmarkAccountsDeprecated benchmarks the performance of the deprecated
// Accounts method with proper setup/teardown isolation.
func BenchmarkAccountsDeprecated(b *testing.B) {
	benchmarkSizes := generateBenchmarkSizes()

	for _, size := range benchmarkSizes {
		// Perform all expensive setup before starting the timer.
		w, cleanup := setupBenchmarkWallet(
			b, size.numAccounts, size.numUTXOs,
		)

		scope := waddrmgr.KeyScopeBIP0044

		// Reset the timer to ensure the setup time is not included in
		// the benchmark results.
		b.ReportAllocs()
		b.ResetTimer()

		b.Run(size.name(), func(b *testing.B) {
			// The loop runs b.N times.
			for b.Loop() {
				_, err := w.Accounts(scope)
				require.NoError(b, err)
			}

			// Stop the timer before performing cleanup.
			b.StopTimer()
			cleanup()
		})
	}
}

// setupBenchmarkWallet creates a wallet with test data for benchmarking. It
// also returns a cleanup function to remove the wallet database.
func setupBenchmarkWallet(b *testing.B, numAccounts,
	numUTXOs int) (*Wallet, func()) {

	b.Helper()

	// Since testWallet requires a *testing.T, we can't pass the
	// benchmark's *testing.B. Instead, we create a dummy *testing.T and
	// manually fail the benchmark if the setup fails.
	t := &testing.T{}
	w, cleanup := testWallet(t)
	if t.Failed() {
		b.Fatalf("testWallet setup failed")
	}

	scopes := []waddrmgr.KeyScope{
		waddrmgr.KeyScopeBIP0044,
	}

	var (
		allAccountNumbers []uint32
		allAddresses      []waddrmgr.ManagedAddress
	)

	err := walletdb.Update(w.db, func(tx walletdb.ReadWriteTx) error {
		for _, scope := range scopes {
			manager, err := w.addrStore.FetchScopedKeyManager(scope)
			if err != nil {
				return err
			}

			addrmgrNs := tx.ReadWriteBucket(waddrmgrNamespaceKey)

			for i := range numAccounts {
				accountName := fmt.Sprintf("bench-account-%d", i)
				account, err := manager.NewAccount(
					addrmgrNs, accountName,
				)
				if err != nil {
					return err
				}
				allAccountNumbers = append(
					allAccountNumbers, account,
				)

				addrs, err := manager.NextExternalAddresses(
					addrmgrNs, account, 5,
				)
				if err != nil {
					return err
				}
				allAddresses = append(allAddresses, addrs...)
			}
		}
		return nil
	})
	require.NoError(b, err, "failed to create test accounts: %v", err)

	// Create valid UTXOs using existing test transaction data.
	err = walletdb.Update(w.db, func(tx walletdb.ReadWriteTx) error {
		txmgrNs := tx.ReadWriteBucket(wtxmgrNamespaceKey)

		// Use a transaction from the test suite as a template.
		msgTx := TstTx.MsgTx()

		for i := 0; i < numUTXOs && i < len(allAddresses); i++ {
			newMsgTx := wire.NewMsgTx(msgTx.Version)

			addr := allAddresses[i%len(allAddresses)]
			pkScript, err := txscript.PayToAddrScript(
				addr.Address(),
			)
			if err != nil {
				return err
			}

			amount := btcutil.Amount(100000 + i*1000)
			newMsgTx.AddTxOut(wire.NewTxOut(
				int64(amount), pkScript),
			)

			// Add a dummy input to make it valid.
			prevHash := chainhash.Hash{}
			prevHash[0] = byte(i)
			newMsgTx.AddTxIn(wire.NewTxIn(
				wire.NewOutPoint(&prevHash, 0), nil, nil,
			))

			rec, err := wtxmgr.NewTxRecordFromMsgTx(
				newMsgTx, time.Now(),
			)
			if err != nil {
				return err
			}

			blockMeta := &wtxmgr.BlockMeta{
				Block: wtxmgr.Block{
					Hash:   chainhash.Hash{},
					Height: 1,
				},
				Time: time.Now(),
			}

			if err = w.txStore.InsertTx(
				txmgrNs, rec, blockMeta,
			); err != nil {
				return err
			}

			// Mark the output as unspent.
			if err = w.txStore.AddCredit(
				txmgrNs, rec, blockMeta, 0, false,
			); err != nil {
				return err
			}
		}

		return nil
	})
	require.NoError(b, err, "failed to create test UTXOs: %v", err)

	return w, cleanup
}