headerfs: allow deleting multi indices from store

Author: mohamedawnallah

headerfs/index.go

@@ -3,6 +3,7 @@ package headerfs
 import (
 	"bytes"
 	"encoding/binary"
+	"errors"
 	"fmt"
 	"sort"
 
@@ -320,11 +321,25 @@ func (h *headerIndex) chainTipWithTx(tx walletdb.ReadTx) (*chainhash.Hash,
 	return tipHash, tipHeight, nil
 }
 
-// truncateIndex truncates the index for a particular header type by a single
-// header entry. The passed newTip pointer should point to the hash of the new
-// chain tip. Optionally, if the entry is to be deleted as well, then the
-// remove flag should be set to true.
-func (h *headerIndex) truncateIndex(newTip *chainhash.Hash, remove bool) error {
+// truncateIndices truncates the index for a particular header type by removing
+// a set of header entries. The passed newTip pointer should point to the hash
+// of the new chain tip. The blockHeadersToTruncate contains the hashes of all
+// block headers that should be removed from the index, which are the block
+// headers after the new tip. Optionally, if the entries are to be deleted as
+// well, then the remove flag should be set to true.
+func (h *headerIndex) truncateIndices(newTip *chainhash.Hash,
+	blockHeadersToTruncate []*chainhash.Hash, remove bool) error {
+
+	if remove && len(blockHeadersToTruncate) == 0 {
+		return errors.New("remove flag set but headers to truncate " +
+			"beyond new tip not provided")
+	}
+
+	if !remove && len(blockHeadersToTruncate) != 0 {
+		return errors.New("headers to truncate beyond new tip " +
+			"provided but remove flag not set")
+	}
+
 	return walletdb.Update(h.db, func(tx walletdb.ReadWriteTx) error {
 		rootBucket := tx.ReadWriteBucket(indexBucket)
 
@@ -338,13 +353,13 @@ func (h *headerIndex) truncateIndex(newTip *chainhash.Hash, remove bool) error {
 			return err
 		}
 
-		// If the remove flag is set, then we'll also delete this entry
-		// from the database as the primary index (block headers) is
-		// being rolled back.
+		// If the remove flag is set, then we'll also delete those
+		// entries from the database as the primary index
+		// (block headers) is being rolled back.
 		if remove {
-			prevTipHash := rootBucket.Get(tipKey)
-			err := delHeaderEntry(rootBucket, prevTipHash)
-			if err != nil {
+			if err := deleteHeaderEntries(
+				rootBucket, blockHeadersToTruncate,
+			); err != nil {
 				return err
 			}
 		}
@@ -417,26 +432,66 @@ func getHeaderEntryFallback(rootBucket walletdb.ReadBucket,
 	return binary.BigEndian.Uint32(heightBytes), nil
 }
 
-// delHeaderEntry tries to remove a header entry from the bbolt database. It
-// first looks if a key for it exists in the old place, the root bucket. If it
-// does, it's deleted from there. If not, it's attempted to be deleted from the
-// sub bucket instead.
-func delHeaderEntry(rootBucket walletdb.ReadWriteBucket, hashBytes []byte) error {
-	// In case this header was stored in the old place (the root bucket
-	// directly), let's remove it from there.
-	if len(rootBucket.Get(hashBytes)) == 4 {
-		return rootBucket.Delete(hashBytes)
+// deleteHeaderEntries tries to remove multiple header entries from the bbolt
+// database. For each header hash, it first looks if a key exists in the old
+// place, the root bucket. If it does, it's deleted from there. If not, it's
+// attempted to be deleted from the appropriate sub-bucket instead.
+func deleteHeaderEntries(rootBucket walletdb.ReadWriteBucket,
+	headerHashes []*chainhash.Hash) error {
+
+	if len(headerHashes) == 0 {
+		return nil
 	}
 
-	// The hash wasn't stored in the root bucket. So we try the sub bucket
-	// now. If that doesn't exist, something is wrong and we want to return
-	// an error here.
-	subBucket := rootBucket.NestedReadWriteBucket(
-		hashBytes[0:numSubBucketBytes],
-	)
-	if subBucket == nil {
-		return ErrHashNotFound
+	// Group hashes by their sub-bucket for more efficient deletion.
+	bySubBucket := make(map[string][]*chainhash.Hash)
+	rootBucketHashes := make([]*chainhash.Hash, 0, len(headerHashes))
+
+	// Check which hashes are in the root bucket and group the rest by their
+	// sub-bucket prefix.
+	for _, hash := range headerHashes {
+		// Convert hash to bytes for DB operations.
+		hashBytes := hash.CloneBytes()
+
+		// In case this header was stored in the old place
+		// (the root bucket directly), let's check and mark it for
+		// removal from there.
+		if len(rootBucket.Get(hashBytes)) == 4 {
+			rootBucketHashes = append(rootBucketHashes, hash)
+			continue
+		}
+		// The hash wasn't stored in the root bucket. So we need
+		// to use the sub-bucket. We extract the prefix to
+		// determine which sub-bucket to use.
+		prefix := string(hashBytes[0:numSubBucketBytes])
+		bySubBucket[prefix] = append(bySubBucket[prefix], hash)
+	}
+
+	// Delete entries from root bucket.
+	for _, hash := range rootBucketHashes {
+		hashBytes := hash.CloneBytes()
+		if err := rootBucket.Delete(hashBytes); err != nil {
+			return err
+		}
+	}
+
+	// Delete enties from sub-buckets.
+	for prefix, hashes := range bySubBucket {
+		// Try to get the sub-bucket for this prefix. If it doesn't
+		// exist, something is wrong and we want to return an error.
+		subBucket := rootBucket.NestedReadWriteBucket([]byte(prefix))
+		if subBucket == nil {
+			return fmt.Errorf("%w: sub-bucket for prefix %x not "+
+				"found", ErrHashNotFound, prefix)
+		}
+
+		for _, hash := range hashes {
+			hashBytes := hash.CloneBytes()
+			if err := subBucket.Delete(hashBytes); err != nil {
+				return err
+			}
+		}
 	}
 
-	return subBucket.Delete(hashBytes)
+	return nil
 }

headerfs/index_test.go

@@ -11,8 +11,10 @@ import (
 	"testing"
 	"time"
 
+	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcwallet/walletdb"
 	_ "github.com/btcsuite/btcwallet/walletdb/bdb"
+	"github.com/stretchr/testify/require"
 )
 
 func createTestIndex(t testing.TB) (func(), *headerIndex, error) {
@@ -43,6 +45,16 @@ func createTestIndex(t testing.TB) (func(), *headerIndex, error) {
 	return cleanUp, filterDB, nil
 }
 
+// TestAddHeadersIndexRetrieve tests the header index functionality by verifying
+// the writing of random headers, ensuring the database tip matches the last
+// inserted header, checking each header can be retrieved by hash, and testing
+// index truncation.
+// It specifically exercises the truncateIndices method with an explicit list of
+// headers to remove, ensuring proper index maintenance when headers are removed
+// from the chain. The test first writes a batch of headers, verifies the chain
+// tip, confirms retrieval by hash for all entries, truncates the last header,
+// and finally verifies the tip has been properly updated to the second-to-last
+// entry.
 func TestAddHeadersIndexRetrieve(t *testing.T) {
 	cleanUp, hIndex, err := createTestIndex(t)
 	defer cleanUp()
@@ -90,8 +102,12 @@ func TestAddHeadersIndexRetrieve(t *testing.T) {
 	// Next if we truncate the index by one, then we should end up at the
 	// second to last entry for the tip.
 	newTip := headerIndex[numHeaders-2]
-	if err := hIndex.truncateIndex(&newTip.hash, true); err != nil {
-		t.Fatalf("unable to truncate index: %v", err)
+
+	// Truncate just the last header.
+	headersToTruncate := []*chainhash.Hash{&lastEntry.hash}
+	err = hIndex.truncateIndices(&newTip.hash, headersToTruncate, true)
+	if err != nil {
+		t.Fatalf("unable to truncate indices: %v", err)
 	}
 
 	// This time the database tip should be the _second_ to last entry
@@ -113,7 +129,11 @@ func TestAddHeadersIndexRetrieve(t *testing.T) {
 
 // TestHeaderStorageFallback makes sure that the changes to the header storage
 // location in the bbolt database for reduced memory consumption don't impact
-// existing users that already have entries in their database.
+// existing users that already have entries in their database. The test verifies
+// compatibility with both old format headers stored directly in the root bucket
+// and new format headers (stored in sub-buckets). It tests reading from both
+// formats and ensures that the truncation functionality correctly handles
+// removing headers from either storage format.
 func TestHeaderStorageFallback(t *testing.T) {
 	cleanUp, hIndex, err := createTestIndex(t)
 	if err != nil {
@@ -184,49 +204,95 @@ func TestHeaderStorageFallback(t *testing.T) {
 		}
 	}
 
-	// And finally, we trim the chain all the way down to the first header.
-	// To do so, we first need to make sure the tip points to the last entry
-	// we added.
-	lastEntry := newHeaderEntries[len(newHeaderEntries)-1]
-	if err := hIndex.truncateIndex(&lastEntry.hash, false); err != nil {
-		t.Fatalf("error setting new tip: %v", err)
-	}
-	for _, header := range newHeaderEntries {
-		if err := hIndex.truncateIndex(&header.hash, true); err != nil {
-			t.Fatalf("error truncating tip: %v", err)
-		}
-	}
-	for _, header := range oldHeaderEntries {
-		if err := hIndex.truncateIndex(&header.hash, true); err != nil {
-			t.Fatalf("error truncating tip: %v", err)
-		}
+	// Now we'll test the truncation functionality by truncating all the way
+	// back to the first old header. We'll do this in steps to verify the
+	// truncation works properly on both new and old format headers.
+
+	// First, set the chain tip to the last new header without removing
+	// anything.
+	lastNewHeader := newHeaderEntries[len(newHeaderEntries)-1]
+	err = hIndex.truncateIndices(&lastNewHeader.hash, nil, false)
+	require.NoError(t, err)
+
+	// Next, truncate all new headers except the first one.
+	truncationPoint := newHeaderEntries[0]
+	headersToTruncate := make([]*chainhash.Hash, 0, len(newHeaderEntries)-1)
+	for i := 1; i < len(newHeaderEntries); i++ {
+		headersToTruncate = append(
+			headersToTruncate, &newHeaderEntries[i].hash,
+		)
 	}
+	err = hIndex.truncateIndices(
+		&truncationPoint.hash, headersToTruncate, true,
+	)
+	require.NoError(t, err)
 
-	// All the headers except the very last should now be deleted.
-	for i := 0; i < len(oldHeaderEntries)-1; i++ {
-		header := oldHeaderEntries[i]
-		if _, err := hIndex.heightFromHash(&header.hash); err == nil {
-			t.Fatalf("expected error reading old entry %x",
-				header.hash[:])
+	// Verify that only the first new header remains and all others
+	// are gone.
+	for i, header := range newHeaderEntries {
+		height, err := hIndex.heightFromHash(&header.hash)
+		if i == 0 {
+			// First header should still be there.
+			msg := "first new header should still exist"
+			require.NoError(t, err, msg)
+			require.Equal(t, header.height, height)
+			continue
 		}
-	}
-	for _, header := range newHeaderEntries {
-		if _, err := hIndex.heightFromHash(&header.hash); err == nil {
-			t.Fatalf("expected error reading old entry %x",
-				header.hash[:])
+
+		if err == nil {
+			// All other headers should be gone.
+			msg := fmt.Sprintf("header at index %d should be "+
+				"deleted, but still exists", i)
+			require.Fail(t, msg)
 		}
 	}
 
-	// The last entry should still be there.
-	lastEntry = oldHeaderEntries[len(oldHeaderEntries)-1]
-	height, err := hIndex.heightFromHash(&lastEntry.hash)
-	if err != nil {
-		t.Fatalf("error reading old entry: %v", err)
-	}
+	// Now truncate back to the last old header.
+	truncationPoint = oldHeaderEntries[len(oldHeaderEntries)-1]
+	headersToTruncate = []*chainhash.Hash{&newHeaderEntries[0].hash}
+	err = hIndex.truncateIndices(
+		&truncationPoint.hash, headersToTruncate, true,
+	)
+	require.NoError(t, err, "error truncating to old headers")
+
+	// Verify all new headers are gone.
+	for i, header := range newHeaderEntries {
+		_, err := hIndex.heightFromHash(&header.hash)
+		msg := fmt.Sprintf("new header at index %d should be deleted, "+
+			"but still exists", i)
+		require.Error(t, err, msg)
+	}
+
+	// Finally, truncate to the first old header.
+	truncationPoint = oldHeaderEntries[0]
+	headersToTruncate = make([]*chainhash.Hash, 0, len(oldHeaderEntries)-1)
+	for i := 1; i < len(oldHeaderEntries); i++ {
+		headersToTruncate = append(
+			headersToTruncate, &oldHeaderEntries[i].hash,
+		)
+	}
+	err = hIndex.truncateIndices(
+		&truncationPoint.hash, headersToTruncate, true,
+	)
+	require.NoError(t, err, "error truncating to old headers")
 
-	if height != lastEntry.height {
-		t.Fatalf("unexpected height, got %d wanted %d", height,
-			lastEntry.height)
+	// Verify only the first old header remains.
+	for i, header := range oldHeaderEntries {
+		height, err := hIndex.heightFromHash(&header.hash)
+		if i == 0 {
+			// First header should still be there.
+			msg := "first old header should still exist"
+			require.NoError(t, err, msg)
+			require.Equal(t, header.height, height)
+			continue
+		}
+
+		if err == nil {
+			// All other headers should be gone.
+			msg := fmt.Sprintf("old header at index %d should be "+
+				"deleted, but still exists", i)
+			require.Fail(t, msg)
+		}
 	}
 }
 

headerfs/store.go

@@ -405,7 +405,7 @@ func (h *blockHeaderStore) RollbackLastBlock() (*BlockStamp, error) {
 	defer h.mtx.Unlock()
 
 	// First, we'll obtain the latest height that the index knows of.
-	_, chainTipHeight, err := h.chainTip()
+	lastBlock, chainTipHeight, err := h.chainTip()
 	if err != nil {
 		return nil, err
 	}
@@ -419,13 +419,19 @@ func (h *blockHeaderStore) RollbackLastBlock() (*BlockStamp, error) {
 	}
 	prevHeaderHash := prevHeader.BlockHash()
 
+	// Compute the block headers to truncate.
+	headersToTruncate := []*chainhash.Hash{lastBlock}
+
 	// Now that we have the information we need to return from this
 	// function, we can now truncate the header file, and then use the hash
 	// of the prevHeader to set the proper index chain tip.
 	if err := h.truncateHeaders(1, h.indexType); err != nil {
 		return nil, err
 	}
-	if err := h.truncateIndex(&prevHeaderHash, true); err != nil {
+
+	if err := h.truncateIndices(
+		&prevHeaderHash, headersToTruncate, true,
+	); err != nil {
 		return nil, err
 	}
 
@@ -989,7 +995,7 @@ func (f *filterHeaderStore) WriteHeaders(hdrs ...FilterHeader) error {
 	// As the block headers should already be written, we only need to
 	// update the tip pointer for this particular header type.
 	newTip := hdrs[len(hdrs)-1].toIndexEntry().hash
-	return f.truncateIndex(&newTip, false)
+	return f.truncateIndices(&newTip, []*chainhash.Hash{}, false)
 }
 
 // ChainTip returns the latest filter header and height known to the
@@ -1047,7 +1053,9 @@ func (f *filterHeaderStore) RollbackLastBlock(
 	if err := f.truncateHeaders(1, f.indexType); err != nil {
 		return nil, err
 	}
-	if err := f.truncateIndex(newTip, false); err != nil {
+
+	err = f.truncateIndices(newTip, []*chainhash.Hash{}, false)
+	if err != nil {
 		return nil, err
 	}