lnd/kvdb/prefetch_test.go

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]))
			}
		}
	}
}