kvdb: extra test coverage for kvdb.Prefetch

kvdb/bolt_test.go

@@ -63,6 +63,10 @@ func TestBolt(t *testing.T) {
 			name: "tx rollback",
 			test: testTxRollback,
 		},
+		{
+			name: "prefetch",
+			test: testPrefetch,
+		},
 	}
 
 	for _, test := range tests {

kvdb/etcd_test.go

@@ -135,6 +135,11 @@ func TestEtcd(t *testing.T) {
 			test:       testTxRollback,
 			expectedDb: map[string]string{},
 		},
+		{
+			name:       "prefetch",
+			test:       testPrefetch,
+			expectedDb: map[string]string{},
+		},
 	}
 
 	for _, test := range tests {

kvdb/go.mod

@@ -3,6 +3,7 @@ module github.com/lightningnetwork/lnd/kvdb
 require (
 	github.com/btcsuite/btclog v0.0.0-20170628155309-84c8d2346e9f
 	github.com/btcsuite/btcwallet/walletdb v1.3.6-0.20210803004036-eebed51155ec
+	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/google/btree v1.0.1
 	github.com/lightningnetwork/lnd/healthcheck v1.0.0
 	github.com/stretchr/testify v1.7.0

kvdb/prefetch_test.go

@@ -0,0 +1,188 @@
+package kvdb
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/btcsuite/btcwallet/walletdb"
+	"github.com/davecgh/go-spew/spew"
+	"github.com/stretchr/testify/require"
+)
+
+func fetchBucket(t *testing.T, bucket walletdb.ReadBucket) map[string]string {
+	items := make(map[string]string)
+	err := bucket.ForEach(func(k, v []byte) error {
+		if v != nil {
+			items[string(k)] = string(v)
+		}
+
+		return nil
+	})
+	require.NoError(t, err)
+
+	return items
+}
+
+func alterBucket(t *testing.T, bucket walletdb.ReadWriteBucket,
+	put map[string]string, remove []string) {
+
+	for k, v := range put {
+		require.NoError(t, bucket.Put([]byte(k), []byte(v)))
+	}
+
+	for _, k := range remove {
+		require.NoError(t, bucket.Delete([]byte(k)))
+	}
+}
+
+func prefetchTest(t *testing.T, db walletdb.DB,
+	prefetchAt []bool, put map[string]string, remove []string) {
+
+	prefetch := func(i int, tx walletdb.ReadTx) {
+		require.Less(t, i, len(prefetchAt))
+		if prefetchAt[i] {
+			Prefetch(
+				RootBucket(tx),
+				[]string{"top"}, []string{"top", "bucket"},
+			)
+		}
+	}
+
+	items := map[string]string{
+		"a": "1",
+		"b": "2",
+		"c": "3",
+		"d": "4",
+		"e": "5",
+	}
+
+	err := Update(db, func(tx walletdb.ReadWriteTx) error {
+		top, err := tx.CreateTopLevelBucket([]byte("top"))
+		require.NoError(t, err)
+		require.NotNil(t, top)
+
+		for k, v := range items {
+			require.NoError(t, top.Put([]byte(k), []byte(v)))
+		}
+
+		bucket, err := top.CreateBucket([]byte("bucket"))
+		require.NoError(t, err)
+		require.NotNil(t, bucket)
+
+		for k, v := range items {
+			require.NoError(t, bucket.Put([]byte(k), []byte(v)))
+		}
+
+		return nil
+	}, func() {})
+	require.NoError(t, err)
+
+	for k, v := range put {
+		items[k] = v
+	}
+
+	for _, k := range remove {
+		delete(items, k)
+	}
+
+	err = Update(db, func(tx walletdb.ReadWriteTx) error {
+		prefetch(0, tx)
+		top := tx.ReadWriteBucket([]byte("top"))
+		require.NotNil(t, top)
+		alterBucket(t, top, put, remove)
+
+		prefetch(1, tx)
+		require.Equal(t, items, fetchBucket(t, top))
+
+		prefetch(2, tx)
+		bucket := top.NestedReadWriteBucket([]byte("bucket"))
+		require.NotNil(t, bucket)
+		alterBucket(t, bucket, put, remove)
+
+		prefetch(3, tx)
+		require.Equal(t, items, fetchBucket(t, bucket))
+
+		return nil
+	}, func() {})
+	require.NoError(t, err)
+
+	err = Update(db, func(tx walletdb.ReadWriteTx) error {
+		return tx.DeleteTopLevelBucket([]byte("top"))
+	}, func() {})
+	require.NoError(t, err)
+}
+
+// testPrefetch tests that prefetching buckets works as expected even when the
+// prefetch happens multiple times and the bucket contents change. Our expectation
+// is that with or without prefetches, the kvdb layer works accourding to the
+// interface specification.
+func testPrefetch(t *testing.T, db walletdb.DB) {
+	tests := []struct {
+		put    map[string]string
+		remove []string
+	}{
+		{
+			put:    nil,
+			remove: nil,
+		},
+		{
+			put: map[string]string{
+				"a":   "a",
+				"aa":  "aa",
+				"aaa": "aaa",
+				"x":   "x",
+				"y":   "y",
+			},
+			remove: nil,
+		},
+		{
+			put: map[string]string{
+				"a":   "a",
+				"aa":  "aa",
+				"aaa": "aaa",
+				"x":   "x",
+				"y":   "y",
+			},
+			remove: []string{"a", "c", "d"},
+		},
+		{
+			put:    nil,
+			remove: []string{"b", "d"},
+		},
+	}
+
+	prefetchAt := [][]bool{
+		{false, false, false, false},
+		{true, false, false, false},
+		{false, true, false, false},
+		{false, false, true, false},
+		{false, false, false, true},
+		{true, true, false, false},
+		{true, true, true, false},
+		{true, true, true, true},
+		{true, false, true, true},
+		{true, false, false, true},
+		{true, false, true, false},
+	}
+
+	for i, test := range tests {
+		test := test
+
+		for j := 0; j < len(prefetchAt); j++ {
+			if !t.Run(
+				fmt.Sprintf("prefetch %d %d", i, j),
+				func(t *testing.T) {
+					prefetchTest(
+						t, db, prefetchAt[j], test.put,
+						test.remove,
+					)
+				}) {
+
+				fmt.Printf("Prefetch test (%d, %d) failed:\n"+
+					"testcase=%v\n prefetch=%v\n",
+					i, j, spew.Sdump(test),
+					spew.Sdump(prefetchAt[j]))
+			}
+		}
+	}
+}