github.com/decred/dcrlnd@v0.7.6/kvdb/prefetch_test.go (about)

     1  package kvdb
     2  
     3  import (
     4  	"fmt"
     5  	"testing"
     6  
     7  	"github.com/btcsuite/btcwallet/walletdb"
     8  	"github.com/davecgh/go-spew/spew"
     9  	"github.com/stretchr/testify/require"
    10  )
    11  
    12  func fetchBucket(t *testing.T, bucket walletdb.ReadBucket) map[string]string {
    13  	items := make(map[string]string)
    14  	err := bucket.ForEach(func(k, v []byte) error {
    15  		if v != nil {
    16  			items[string(k)] = string(v)
    17  		}
    18  
    19  		return nil
    20  	})
    21  	require.NoError(t, err)
    22  
    23  	return items
    24  }
    25  
    26  func alterBucket(t *testing.T, bucket walletdb.ReadWriteBucket,
    27  	put map[string]string, remove []string) {
    28  
    29  	for k, v := range put {
    30  		require.NoError(t, bucket.Put([]byte(k), []byte(v)))
    31  	}
    32  
    33  	for _, k := range remove {
    34  		require.NoError(t, bucket.Delete([]byte(k)))
    35  	}
    36  }
    37  
    38  func prefetchTest(t *testing.T, db walletdb.DB,
    39  	prefetchAt []bool, put map[string]string, remove []string) {
    40  
    41  	prefetch := func(i int, tx walletdb.ReadTx) {
    42  		require.Less(t, i, len(prefetchAt))
    43  		if prefetchAt[i] {
    44  			Prefetch(
    45  				RootBucket(tx),
    46  				[]string{"top"}, []string{"top", "bucket"},
    47  			)
    48  		}
    49  	}
    50  
    51  	items := map[string]string{
    52  		"a": "1",
    53  		"b": "2",
    54  		"c": "3",
    55  		"d": "4",
    56  		"e": "5",
    57  	}
    58  
    59  	err := Update(db, func(tx walletdb.ReadWriteTx) error {
    60  		top, err := tx.CreateTopLevelBucket([]byte("top"))
    61  		require.NoError(t, err)
    62  		require.NotNil(t, top)
    63  
    64  		for k, v := range items {
    65  			require.NoError(t, top.Put([]byte(k), []byte(v)))
    66  		}
    67  
    68  		bucket, err := top.CreateBucket([]byte("bucket"))
    69  		require.NoError(t, err)
    70  		require.NotNil(t, bucket)
    71  
    72  		for k, v := range items {
    73  			require.NoError(t, bucket.Put([]byte(k), []byte(v)))
    74  		}
    75  
    76  		return nil
    77  	}, func() {})
    78  	require.NoError(t, err)
    79  
    80  	for k, v := range put {
    81  		items[k] = v
    82  	}
    83  
    84  	for _, k := range remove {
    85  		delete(items, k)
    86  	}
    87  
    88  	err = Update(db, func(tx walletdb.ReadWriteTx) error {
    89  		prefetch(0, tx)
    90  		top := tx.ReadWriteBucket([]byte("top"))
    91  		require.NotNil(t, top)
    92  		alterBucket(t, top, put, remove)
    93  
    94  		prefetch(1, tx)
    95  		require.Equal(t, items, fetchBucket(t, top))
    96  
    97  		prefetch(2, tx)
    98  		bucket := top.NestedReadWriteBucket([]byte("bucket"))
    99  		require.NotNil(t, bucket)
   100  		alterBucket(t, bucket, put, remove)
   101  
   102  		prefetch(3, tx)
   103  		require.Equal(t, items, fetchBucket(t, bucket))
   104  
   105  		return nil
   106  	}, func() {})
   107  	require.NoError(t, err)
   108  
   109  	err = Update(db, func(tx walletdb.ReadWriteTx) error {
   110  		return tx.DeleteTopLevelBucket([]byte("top"))
   111  	}, func() {})
   112  	require.NoError(t, err)
   113  }
   114  
   115  // testPrefetch tests that prefetching buckets works as expected even when the
   116  // prefetch happens multiple times and the bucket contents change. Our expectation
   117  // is that with or without prefetches, the kvdb layer works accourding to the
   118  // interface specification.
   119  func testPrefetch(t *testing.T, db walletdb.DB) {
   120  	tests := []struct {
   121  		put    map[string]string
   122  		remove []string
   123  	}{
   124  		{
   125  			put:    nil,
   126  			remove: nil,
   127  		},
   128  		{
   129  			put: map[string]string{
   130  				"a":   "a",
   131  				"aa":  "aa",
   132  				"aaa": "aaa",
   133  				"x":   "x",
   134  				"y":   "y",
   135  			},
   136  			remove: nil,
   137  		},
   138  		{
   139  			put: map[string]string{
   140  				"a":   "a",
   141  				"aa":  "aa",
   142  				"aaa": "aaa",
   143  				"x":   "x",
   144  				"y":   "y",
   145  			},
   146  			remove: []string{"a", "c", "d"},
   147  		},
   148  		{
   149  			put:    nil,
   150  			remove: []string{"b", "d"},
   151  		},
   152  	}
   153  
   154  	prefetchAt := [][]bool{
   155  		{false, false, false, false},
   156  		{true, false, false, false},
   157  		{false, true, false, false},
   158  		{false, false, true, false},
   159  		{false, false, false, true},
   160  		{true, true, false, false},
   161  		{true, true, true, false},
   162  		{true, true, true, true},
   163  		{true, false, true, true},
   164  		{true, false, false, true},
   165  		{true, false, true, false},
   166  	}
   167  
   168  	for i, test := range tests {
   169  		test := test
   170  
   171  		for j := 0; j < len(prefetchAt); j++ {
   172  			if !t.Run(
   173  				fmt.Sprintf("prefetch %d %d", i, j),
   174  				func(t *testing.T) {
   175  					prefetchTest(
   176  						t, db, prefetchAt[j], test.put,
   177  						test.remove,
   178  					)
   179  				}) {
   180  
   181  				fmt.Printf("Prefetch test (%d, %d) failed:\n"+
   182  					"testcase=%v\n prefetch=%v\n",
   183  					i, j, spew.Sdump(test),
   184  					spew.Sdump(prefetchAt[j]))
   185  			}
   186  		}
   187  	}
   188  }