kvdb: implement ApplyScanBatch

Add recovery scan batch types and route kvdb scan writes through
the legacy address and transaction managers so horizons, tx
notifications, and synced blocks remain atomic.

Author: yyforyongyu

wallet/internal/db/kvdb/txstore.go

@@ -373,6 +373,85 @@ func (s *Store) applyLegacyTxBatch(addrmgrNs walletdb.ReadWriteBucket,
 	return nil
 }
 
+// ApplyScanBatch atomically applies recovery scan writes through the legacy
+// walletdb managers.
+func (s *Store) ApplyScanBatch(_ context.Context,
+	params db.ScanBatchParams) error {
+
+	if s.addrStore == nil {
+		return fmt.Errorf("kvdb.Store.ApplyScanBatch: %w",
+			errMissingAddrStore)
+	}
+
+	err := walletdb.Update(s.db, func(tx walletdb.ReadWriteTx) error {
+		addrmgrNs := tx.ReadWriteBucket(waddrmgr.NamespaceKey)
+		if addrmgrNs == nil {
+			return errMissingAddrmgrNamespace
+		}
+
+		txmgrNs := tx.ReadWriteBucket(wtxmgrNamespaceKey)
+		if txmgrNs == nil {
+			return errMissingTxmgrNamespace
+		}
+
+		// Horizons are extended first because the transaction batch
+		// below resolves every scan-discovered credit against the
+		// address manager, so a freshly derived address must already be
+		// written to the addrmgr bucket before its crediting transaction
+		// is recorded.
+		err := s.applyLegacyScanHorizons(addrmgrNs, params.Horizons)
+		if err != nil {
+			return err
+		}
+
+		err = s.applyLegacyTxBatch(
+			addrmgrNs, txmgrNs, params.Transactions,
+		)
+		if err != nil {
+			return err
+		}
+
+		return s.applyLegacySyncedBlocks(addrmgrNs, params.SyncedBlocks)
+	})
+	if err != nil {
+		// ExtendAddresses advances the scoped manager's in-memory next
+		// indices and cached account info as a side effect of writing
+		// the addrmgr bucket. walletdb has now rolled that bucket back,
+		// but those in-memory mutations survive, so a failed batch would
+		// otherwise leave the live manager reporting an unpersisted
+		// horizon advance. Invalidate the touched account caches so the
+		// next read reloads the persisted (rolled-back) indices.
+		s.rollbackScanHorizonCaches(params.Horizons)
+
+		return fmt.Errorf("kvdb.Store.ApplyScanBatch: %w", err)
+	}
+
+	return nil
+}
+
+// rollbackScanHorizonCaches drops the scoped-manager account caches advanced by
+// a failed ApplyScanBatch horizon extension. ExtendAddresses bumps the
+// manager's in-memory next indices and cached account info before the
+// surrounding walletdb transaction commits; once that transaction rolls back,
+// invalidating these caches forces the persisted (rolled-back) indices to be
+// reloaded on the next access, so no unpersisted horizon advance stays
+// observable from the live manager.
+func (s *Store) rollbackScanHorizonCaches(horizons []db.ScanHorizon) {
+	for _, horizon := range horizons {
+		scopedMgr, err := s.addrStore.FetchScopedKeyManager(
+			waddrmgr.KeyScope(horizon.Scope),
+		)
+		if err != nil {
+			// The scope resolved while applying the batch, so a
+			// lookup miss here is unexpected; skip it as there is no
+			// cache to invalidate for an unknown scope.
+			continue
+		}
+
+		scopedMgr.InvalidateAccountCache(horizon.Account)
+	}
+}
+
 // GetTx retrieves one wallet-scoped transaction snapshot through the legacy
 // wtxmgr query path.
 func (s *Store) GetTx(_ context.Context, query db.GetTxQuery) (
@@ -606,6 +685,48 @@ func (s *Store) RollbackToBlock(_ context.Context, height uint32) error {
 	return nil
 }
 
+// applyLegacyScanHorizons extends legacy address branches for scan hits.
+func (s *Store) applyLegacyScanHorizons(ns walletdb.ReadWriteBucket,
+	horizons []db.ScanHorizon) error {
+
+	for _, horizon := range horizons {
+		scopedMgr, err := s.addrStore.FetchScopedKeyManager(
+			waddrmgr.KeyScope(horizon.Scope),
+		)
+		if err != nil {
+			return fmt.Errorf("fetch scoped manager: %w", err)
+		}
+
+		err = scopedMgr.ExtendAddresses(
+			ns, horizon.Account, horizon.Index, horizon.Branch,
+		)
+		if err != nil {
+			return fmt.Errorf("extend addresses: %w", err)
+		}
+	}
+
+	return nil
+}
+
+// applyLegacySyncedBlocks connects a sequence of legacy synced blocks.
+func (s *Store) applyLegacySyncedBlocks(ns walletdb.ReadWriteBucket,
+	blocks []db.Block) error {
+
+	for i := range blocks {
+		block, err := db.BlockStampFromBlock(&blocks[i])
+		if err != nil {
+			return err
+		}
+
+		err = s.addrStore.SetSyncedTo(ns, &block)
+		if err != nil {
+			return fmt.Errorf("set synced block %d: %w", i, err)
+		}
+	}
+
+	return nil
+}
+
 // applyLegacyTxNotification records one relevant transaction notification using
 // legacy wtxmgr semantics.
 func (s *Store) applyLegacyTxNotification(addrmgrNs walletdb.ReadWriteBucket,

wallet/internal/db/kvdb/txstore_test.go

@@ -1,6 +1,7 @@
 package kvdb
 
 import (
+	"math"
 	"testing"
 	"time"
 
@@ -323,6 +324,86 @@ func TestApplyTxBatchCreditNilFallback(t *testing.T) {
 	addrStore.AssertNotCalled(t, "MarkUsed", mock.Anything, mock.Anything)
 }
 
+// TestApplyScanBatchRecordsTxAndSyncedBlocks verifies that kvdb.Store applies
+// scan-discovered transactions and connected sync blocks in one write batch.
+func TestApplyScanBatchRecordsTxAndSyncedBlocks(t *testing.T) {
+	t.Parallel()
+
+	dbConn, cleanup := newTestDB(t)
+	t.Cleanup(cleanup)
+
+	newAddrmgrNamespace(t, dbConn)
+	txStore := newTxStore(t, dbConn)
+
+	addr, script := newTestAddressScript(t)
+	managedAddr := &bwmock.ManagedAddress{}
+	managedAddr.On("Internal").Return(false).Maybe()
+	managedAddr.On("Address").Return(addr).Maybe()
+	managedAddr.On("AddrType").Return(waddrmgr.WitnessPubKey).Maybe()
+	managedAddr.On("Imported").Return(false).Maybe()
+	managedAddr.On("InternalAccount").Return(uint32(0)).Maybe()
+	managedAddr.On("Compressed").Return(true).Maybe()
+	managedAddr.On("AddrHash").Return([]byte(nil)).Maybe()
+	managedAddr.On("Used", mock.Anything).Return(false).Maybe()
+
+	addrStore := &bwmock.AddrStore{}
+	addrStore.On("ChainParams").Return(&chaincfg.RegressionNetParams)
+	addrStore.On("Address", mock.Anything, mock.Anything).
+		Return(managedAddr, nil)
+	addrStore.On("MarkUsed", mock.Anything, mock.Anything).Return(nil)
+	addrStore.On("SetSyncedTo", mock.Anything, mock.Anything).Return(nil)
+	store := NewStore(dbConn, txStore, addrStore)
+
+	txMsg := &wire.MsgTx{Version: 1}
+	txMsg.AddTxIn(&wire.TxIn{PreviousOutPoint: wire.OutPoint{
+		Hash: chainhash.Hash{65},
+	}})
+	txMsg.AddTxOut(&wire.TxOut{Value: 11_000, PkScript: script})
+
+	syncedBlocks := []db.Block{{
+		Hash:      chainhash.Hash{66},
+		Height:    145,
+		Timestamp: time.Unix(1710003300, 0),
+	}, {
+		Hash:      chainhash.Hash{67},
+		Height:    146,
+		Timestamp: time.Unix(1710003400, 0),
+	}}
+	err := store.ApplyScanBatch(t.Context(), db.ScanBatchParams{
+		WalletID: 0,
+		Transactions: []db.CreateTxParams{{
+			WalletID: 0,
+			Tx:       txMsg,
+			Received: time.Unix(1710003350, 0),
+			Status:   db.TxStatusPublished,
+			Credits:  map[uint32]btcutil.Address{0: addr},
+		}},
+		SyncedBlocks: syncedBlocks,
+	})
+	require.NoError(t, err)
+
+	txid := txMsg.TxHash()
+	err = walletdb.View(dbConn, func(tx walletdb.ReadTx) error {
+		ns := tx.ReadBucket(wtxmgrNamespaceKey)
+		require.NotNil(t, ns)
+
+		details, err := txStore.TxDetails(ns, &txid)
+		require.NoError(t, err)
+		require.NotNil(t, details)
+		require.Len(t, details.Credits, 1)
+
+		return nil
+	})
+	require.NoError(t, err)
+	addrStore.AssertCalled(t, "MarkUsed", mock.Anything, mock.Anything)
+	addrStore.AssertCalled(
+		t, "SetSyncedTo", mock.Anything,
+		mock.MatchedBy(func(bs *waddrmgr.BlockStamp) bool {
+			return bs.Height == int32(146)
+		}),
+	)
+}
+
 // TestCreateTxCreditAddrMismatch verifies that crediting an output with an
 // address that the output script does not pay to is rejected, so a caller
 // cannot corrupt UTXO ownership by mislabeling a credit.
@@ -1169,3 +1250,111 @@ func newMultisigScript(t *testing.T) (btcutil.Address, []byte) {
 
 	return memberAddr, script
 }
+
+// TestApplyScanBatchHorizonRollbackSafety verifies that when a recovery scan
+// batch fails after its horizon extension has already advanced the scoped
+// manager's in-memory address state, the live manager does not keep reporting
+// the unpersisted horizon advance once walletdb rolls the batch back.
+func TestApplyScanBatchHorizonRollbackSafety(t *testing.T) {
+	t.Parallel()
+
+	dbConn, cleanup := newTestDB(t)
+	t.Cleanup(cleanup)
+
+	// newSpendableAddrMgr creates the addrmgr namespace and a real manager
+	// with the default account already present on every default scope.
+	addrStore := newSpendableAddrMgr(t, dbConn)
+	t.Cleanup(func() {
+		_ = addrStore.Lock()
+		addrStore.Close()
+	})
+	txStore := newTxStore(t, dbConn)
+	store := NewStore(dbConn, txStore, addrStore)
+
+	err := walletdb.View(dbConn, func(tx walletdb.ReadTx) error {
+		ns := tx.ReadBucket(waddrmgr.NamespaceKey)
+
+		return addrStore.Unlock(ns, testPrivPass)
+	})
+	require.NoError(t, err)
+
+	const account = waddrmgr.DefaultAccountNum
+
+	scope := waddrmgr.KeyScopeBIP0084
+	scopedMgr, err := addrStore.FetchScopedKeyManager(scope)
+	require.NoError(t, err)
+
+	// The fresh default account has not derived any external addresses yet.
+	require.Zero(t, externalKeyCount(t, dbConn, scopedMgr, account))
+
+	const horizonIndex = 9
+
+	horizon := db.ScanHorizon{
+		Scope:   db.KeyScope(scope),
+		Account: account,
+		Branch:  waddrmgr.ExternalBranch,
+		Index:   horizonIndex,
+	}
+
+	// A synced block whose height overflows the legacy int32 height domain
+	// fails the synced-block step, which runs after the horizon extension
+	// has already mutated the manager's in-memory address state.
+	overflowBlock := db.Block{
+		Hash:      chainhash.Hash{0xAB},
+		Height:    uint32(math.MaxInt32) + 1,
+		Timestamp: time.Unix(1_700_000_000, 0),
+	}
+
+	err = store.ApplyScanBatch(t.Context(), db.ScanBatchParams{
+		WalletID:     0,
+		Horizons:     []db.ScanHorizon{horizon},
+		SyncedBlocks: []db.Block{overflowBlock},
+	})
+	require.Error(t, err)
+
+	// The live manager must not observe the rolled-back horizon advance, and
+	// the persisted addrmgr bucket must not record any derived external
+	// addresses either.
+	require.Zero(t, externalKeyCount(t, dbConn, scopedMgr, account))
+
+	// As a positive control, the same horizon extension applied without the
+	// failing block must succeed and advance the external next index past the
+	// recovered child, proving the rollback path above did not simply skip the
+	// extension and that the invalidated cache reloads cleanly.
+	err = store.ApplyScanBatch(t.Context(), db.ScanBatchParams{
+		WalletID: 0,
+		Horizons: []db.ScanHorizon{horizon},
+	})
+	require.NoError(t, err)
+	require.Equal(
+		t, uint32(horizonIndex+1),
+		externalKeyCount(t, dbConn, scopedMgr, account),
+	)
+}
+
+// externalKeyCount reports the number of external keys the scoped manager
+// considers derived for the account, read through the live manager so any
+// unpersisted in-memory advance would be observable.
+func externalKeyCount(t *testing.T, dbConn walletdb.DB,
+	scopedMgr waddrmgr.AccountStore, account uint32) uint32 {
+
+	t.Helper()
+
+	var count uint32
+
+	err := walletdb.View(dbConn, func(tx walletdb.ReadTx) error {
+		ns := tx.ReadBucket(waddrmgr.NamespaceKey)
+
+		props, err := scopedMgr.AccountProperties(ns, account)
+		if err != nil {
+			return err
+		}
+
+		count = props.ExternalKeyCount
+
+		return nil
+	})
+	require.NoError(t, err)
+
+	return count
+}