github.com/weaviate/weaviate@v1.24.6/adapters/repos/db/lsmkv/strategies_set_integration_test.go (about)

     1  //                           _       _
     2  // __      _____  __ ___   ___  __ _| |_ ___
     3  // \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
     4  //  \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
     5  //   \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
     6  //
     7  //  Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
     8  //
     9  //  CONTACT: hello@weaviate.io
    10  //
    11  
    12  //go:build integrationTest
    13  // +build integrationTest
    14  
    15  package lsmkv
    16  
    17  import (
    18  	"context"
    19  	"fmt"
    20  	"testing"
    21  
    22  	"github.com/stretchr/testify/assert"
    23  	"github.com/stretchr/testify/require"
    24  	"github.com/weaviate/weaviate/entities/cyclemanager"
    25  )
    26  
    27  func TestSetCollectionStrategy(t *testing.T) {
    28  	ctx := testCtx()
    29  	tests := bucketIntegrationTests{
    30  		{
    31  			name: "collectionInsertAndSetAdd",
    32  			f:    collectionInsertAndSetAdd,
    33  			opts: []BucketOption{
    34  				WithStrategy(StrategySetCollection),
    35  			},
    36  		},
    37  		{
    38  			name: "collectionInsertAndSetAddInsertAndDelete",
    39  			f:    collectionInsertAndSetAddInsertAndDelete,
    40  			opts: []BucketOption{
    41  				WithStrategy(StrategySetCollection),
    42  			},
    43  		},
    44  		{
    45  			name: "collectionCursors",
    46  			f:    collectionCursors,
    47  			opts: []BucketOption{
    48  				WithStrategy(StrategySetCollection),
    49  			},
    50  		},
    51  	}
    52  	tests.run(ctx, t)
    53  }
    54  
    55  func collectionInsertAndSetAdd(ctx context.Context, t *testing.T, opts []BucketOption) {
    56  	dirName := t.TempDir()
    57  
    58  	t.Run("memtable-only", func(t *testing.T) {
    59  		b, err := NewBucket(ctx, dirName, "", nullLogger(), nil,
    60  			cyclemanager.NewCallbackGroupNoop(), cyclemanager.NewCallbackGroupNoop(), opts...)
    61  		require.Nil(t, err)
    62  
    63  		// so big it effectively never triggers as part of this test
    64  		b.SetMemtableThreshold(1e9)
    65  
    66  		key1 := []byte("test1-key-1")
    67  		key2 := []byte("test1-key-2")
    68  		key3 := []byte("test1-key-3")
    69  
    70  		t.Run("set original values and verify", func(t *testing.T) {
    71  			orig1 := [][]byte{[]byte("value 1.1"), []byte("value 1.2")}
    72  			orig2 := [][]byte{[]byte("value 2.1"), []byte("value 2.2")}
    73  			orig3 := [][]byte{[]byte("value 3.1"), []byte("value 3.2")}
    74  
    75  			err = b.SetAdd(key1, orig1)
    76  			require.Nil(t, err)
    77  			err = b.SetAdd(key2, orig2)
    78  			require.Nil(t, err)
    79  			err = b.SetAdd(key3, orig3)
    80  			require.Nil(t, err)
    81  
    82  			res, err := b.SetList(key1)
    83  			require.Nil(t, err)
    84  			assert.Equal(t, res, orig1)
    85  			res, err = b.SetList(key2)
    86  			require.Nil(t, err)
    87  			assert.Equal(t, res, orig2)
    88  			res, err = b.SetList(key3)
    89  			require.Nil(t, err)
    90  			assert.Equal(t, res, orig3)
    91  		})
    92  
    93  		t.Run("replace some, keep one", func(t *testing.T) {
    94  			orig1 := [][]byte{[]byte("value 1.1"), []byte("value 1.2")}
    95  			orig2 := [][]byte{[]byte("value 2.1"), []byte("value 2.2")}
    96  			orig3 := [][]byte{[]byte("value 3.1"), []byte("value 3.2")}
    97  			append2 := [][]byte{[]byte("value 2.3")}
    98  			append3 := [][]byte{[]byte("value 3.3")}
    99  
   100  			err = b.SetAdd(key2, append2)
   101  			require.Nil(t, err)
   102  			err = b.SetAdd(key3, append3)
   103  			require.Nil(t, err)
   104  
   105  			res, err := b.SetList(key1)
   106  			require.Nil(t, err)
   107  			assert.Equal(t, orig1, res)
   108  			res, err = b.SetList(key2)
   109  			require.Nil(t, err)
   110  			assert.Equal(t, append(orig2, append2...), res)
   111  			res, err = b.SetList(key3)
   112  			require.Nil(t, err)
   113  			assert.Equal(t, append(orig3, append3...), res)
   114  		})
   115  	})
   116  
   117  	t.Run("with a single flush between updates", func(t *testing.T) {
   118  		b, err := NewBucket(ctx, dirName, "", nullLogger(), nil,
   119  			cyclemanager.NewCallbackGroupNoop(), cyclemanager.NewCallbackGroupNoop(), opts...)
   120  		require.Nil(t, err)
   121  
   122  		// so big it effectively never triggers as part of this test
   123  		b.SetMemtableThreshold(1e9)
   124  
   125  		key1 := []byte("test2-key-1")
   126  		key2 := []byte("test2-key-2")
   127  		key3 := []byte("test2-key-3")
   128  
   129  		t.Run("set original values and verify", func(t *testing.T) {
   130  			orig1 := [][]byte{[]byte("value 1.1"), []byte("value 1.2")}
   131  			orig2 := [][]byte{[]byte("value 2.1"), []byte("value 2.2")}
   132  			orig3 := [][]byte{[]byte("value 3.1"), []byte("value 3.2")}
   133  
   134  			err = b.SetAdd(key1, orig1)
   135  			require.Nil(t, err)
   136  			err = b.SetAdd(key2, orig2)
   137  			require.Nil(t, err)
   138  			err = b.SetAdd(key3, orig3)
   139  			require.Nil(t, err)
   140  
   141  			res, err := b.SetList(key1)
   142  			require.Nil(t, err)
   143  			assert.Equal(t, res, orig1)
   144  			res, err = b.SetList(key2)
   145  			require.Nil(t, err)
   146  			assert.Equal(t, res, orig2)
   147  			res, err = b.SetList(key3)
   148  			require.Nil(t, err)
   149  			assert.Equal(t, res, orig3)
   150  		})
   151  
   152  		t.Run("flush to disk", func(t *testing.T) {
   153  			require.Nil(t, b.FlushAndSwitch())
   154  		})
   155  
   156  		t.Run("replace some, keep one", func(t *testing.T) {
   157  			orig1 := [][]byte{[]byte("value 1.1"), []byte("value 1.2")}
   158  			orig2 := [][]byte{[]byte("value 2.1"), []byte("value 2.2")}
   159  			orig3 := [][]byte{[]byte("value 3.1"), []byte("value 3.2")}
   160  			append2 := [][]byte{[]byte("value 2.3")}
   161  			append3 := [][]byte{[]byte("value 3.3")}
   162  
   163  			err = b.SetAdd(key2, append2)
   164  			require.Nil(t, err)
   165  			err = b.SetAdd(key3, append3)
   166  			require.Nil(t, err)
   167  
   168  			res, err := b.SetList(key1)
   169  			require.Nil(t, err)
   170  			assert.Equal(t, orig1, res)
   171  			res, err = b.SetList(key2)
   172  			require.Nil(t, err)
   173  			assert.Equal(t, append(orig2, append2...), res)
   174  			res, err = b.SetList(key3)
   175  			require.Nil(t, err)
   176  			assert.Equal(t, append(orig3, append3...), res)
   177  		})
   178  	})
   179  
   180  	t.Run("with flushes after initial and update", func(t *testing.T) {
   181  		b, err := NewBucket(ctx, dirName, "", nullLogger(), nil,
   182  			cyclemanager.NewCallbackGroupNoop(), cyclemanager.NewCallbackGroupNoop(), opts...)
   183  		require.Nil(t, err)
   184  
   185  		// so big it effectively never triggers as part of this test
   186  		b.SetMemtableThreshold(1e9)
   187  		key1 := []byte("test-3-key-1")
   188  		key2 := []byte("test-3-key-2")
   189  		key3 := []byte("test-3-key-3")
   190  
   191  		t.Run("set original values and verify", func(t *testing.T) {
   192  			orig1 := [][]byte{[]byte("value 1.1"), []byte("value 1.2")}
   193  			orig2 := [][]byte{[]byte("value 2.1"), []byte("value 2.2")}
   194  			orig3 := [][]byte{[]byte("value 3.1"), []byte("value 3.2")}
   195  
   196  			err = b.SetAdd(key1, orig1)
   197  			require.Nil(t, err)
   198  			err = b.SetAdd(key2, orig2)
   199  			require.Nil(t, err)
   200  			err = b.SetAdd(key3, orig3)
   201  			require.Nil(t, err)
   202  
   203  			res, err := b.SetList(key1)
   204  			require.Nil(t, err)
   205  			assert.Equal(t, res, orig1)
   206  			res, err = b.SetList(key2)
   207  			require.Nil(t, err)
   208  			assert.Equal(t, res, orig2)
   209  			res, err = b.SetList(key3)
   210  			require.Nil(t, err)
   211  			assert.Equal(t, res, orig3)
   212  		})
   213  
   214  		t.Run("flush to disk", func(t *testing.T) {
   215  			require.Nil(t, b.FlushAndSwitch())
   216  		})
   217  
   218  		t.Run("replace some, keep one", func(t *testing.T) {
   219  			orig1 := [][]byte{[]byte("value 1.1"), []byte("value 1.2")}
   220  			orig2 := [][]byte{[]byte("value 2.1"), []byte("value 2.2")}
   221  			orig3 := [][]byte{[]byte("value 3.1"), []byte("value 3.2")}
   222  			append2 := [][]byte{[]byte("value 2.3")}
   223  			append3 := [][]byte{[]byte("value 3.3")}
   224  
   225  			err = b.SetAdd(key2, append2)
   226  			require.Nil(t, err)
   227  			err = b.SetAdd(key3, append3)
   228  			require.Nil(t, err)
   229  
   230  			// Flush again!
   231  			require.Nil(t, b.FlushAndSwitch())
   232  
   233  			res, err := b.SetList(key1)
   234  			require.Nil(t, err)
   235  			assert.Equal(t, orig1, res)
   236  			res, err = b.SetList(key2)
   237  			require.Nil(t, err)
   238  			assert.Equal(t, append(orig2, append2...), res)
   239  			res, err = b.SetList(key3)
   240  			require.Nil(t, err)
   241  			assert.Equal(t, append(orig3, append3...), res)
   242  		})
   243  	})
   244  
   245  	t.Run("update in memtable, then do an orderly shutdown, and re-init", func(t *testing.T) {
   246  		b, err := NewBucket(ctx, dirName, "", nullLogger(), nil,
   247  			cyclemanager.NewCallbackGroupNoop(), cyclemanager.NewCallbackGroupNoop(), opts...)
   248  		require.Nil(t, err)
   249  
   250  		// so big it effectively never triggers as part of this test
   251  		b.SetMemtableThreshold(1e9)
   252  
   253  		key1 := []byte("test4-key-1")
   254  		key2 := []byte("test4-key-2")
   255  		key3 := []byte("test4-key-3")
   256  
   257  		t.Run("set original values and verify", func(t *testing.T) {
   258  			orig1 := [][]byte{[]byte("value 1.1"), []byte("value 1.2")}
   259  			orig2 := [][]byte{[]byte("value 2.1"), []byte("value 2.2")}
   260  			orig3 := [][]byte{[]byte("value 3.1"), []byte("value 3.2")}
   261  
   262  			err = b.SetAdd(key1, orig1)
   263  			require.Nil(t, err)
   264  			err = b.SetAdd(key2, orig2)
   265  			require.Nil(t, err)
   266  			err = b.SetAdd(key3, orig3)
   267  			require.Nil(t, err)
   268  
   269  			res, err := b.SetList(key1)
   270  			require.Nil(t, err)
   271  			assert.Equal(t, res, orig1)
   272  			res, err = b.SetList(key2)
   273  			require.Nil(t, err)
   274  			assert.Equal(t, res, orig2)
   275  			res, err = b.SetList(key3)
   276  			require.Nil(t, err)
   277  			assert.Equal(t, res, orig3)
   278  		})
   279  
   280  		t.Run("replace some, keep one", func(t *testing.T) {
   281  			orig1 := [][]byte{[]byte("value 1.1"), []byte("value 1.2")}
   282  			orig2 := [][]byte{[]byte("value 2.1"), []byte("value 2.2")}
   283  			orig3 := [][]byte{[]byte("value 3.1"), []byte("value 3.2")}
   284  			append2 := [][]byte{[]byte("value 2.3")}
   285  			append3 := [][]byte{[]byte("value 3.3")}
   286  
   287  			err = b.SetAdd(key2, append2)
   288  			require.Nil(t, err)
   289  			err = b.SetAdd(key3, append3)
   290  			require.Nil(t, err)
   291  
   292  			res, err := b.SetList(key1)
   293  			require.Nil(t, err)
   294  			assert.Equal(t, orig1, res)
   295  			res, err = b.SetList(key2)
   296  			require.Nil(t, err)
   297  			assert.Equal(t, append(orig2, append2...), res)
   298  			res, err = b.SetList(key3)
   299  			require.Nil(t, err)
   300  			assert.Equal(t, append(orig3, append3...), res)
   301  		})
   302  
   303  		t.Run("orderly shutdown", func(t *testing.T) {
   304  			b.Shutdown(context.Background())
   305  		})
   306  
   307  		t.Run("init another bucket on the same files", func(t *testing.T) {
   308  			b2, err := NewBucket(ctx, dirName, "", nullLogger(), nil,
   309  				cyclemanager.NewCallbackGroupNoop(), cyclemanager.NewCallbackGroupNoop(), opts...)
   310  			require.Nil(t, err)
   311  
   312  			orig1 := [][]byte{[]byte("value 1.1"), []byte("value 1.2")}
   313  			orig2 := [][]byte{[]byte("value 2.1"), []byte("value 2.2")}
   314  			orig3 := [][]byte{[]byte("value 3.1"), []byte("value 3.2")}
   315  			append2 := [][]byte{[]byte("value 2.3")}
   316  			append3 := [][]byte{[]byte("value 3.3")}
   317  
   318  			res, err := b2.SetList(key1)
   319  			require.Nil(t, err)
   320  			assert.Equal(t, orig1, res)
   321  			res, err = b2.SetList(key2)
   322  			require.Nil(t, err)
   323  			assert.Equal(t, append(orig2, append2...), res)
   324  			res, err = b2.SetList(key3)
   325  			require.Nil(t, err)
   326  			assert.Equal(t, append(orig3, append3...), res)
   327  		})
   328  	})
   329  }
   330  
   331  func collectionInsertAndSetAddInsertAndDelete(ctx context.Context, t *testing.T, opts []BucketOption) {
   332  	dirName := t.TempDir()
   333  
   334  	t.Run("memtable-only", func(t *testing.T) {
   335  		b, err := NewBucket(ctx, dirName, "", nullLogger(), nil,
   336  			cyclemanager.NewCallbackGroupNoop(), cyclemanager.NewCallbackGroupNoop(), opts...)
   337  		require.Nil(t, err)
   338  
   339  		// so big it effectively never triggers as part of this test
   340  		b.SetMemtableThreshold(1e9)
   341  
   342  		key1 := []byte("test1-key-1")
   343  		key2 := []byte("test1-key-2")
   344  		key3 := []byte("test1-key-3")
   345  
   346  		t.Run("set original values and verify", func(t *testing.T) {
   347  			orig1 := [][]byte{[]byte("value 1.1"), []byte("value 1.2")}
   348  			orig2 := [][]byte{[]byte("value 2.1"), []byte("value 2.2")}
   349  			orig3 := [][]byte{[]byte("value 3.1"), []byte("value 3.2")}
   350  
   351  			err = b.SetAdd(key1, orig1)
   352  			require.Nil(t, err)
   353  			err = b.SetAdd(key2, orig2)
   354  			require.Nil(t, err)
   355  			err = b.SetAdd(key3, orig3)
   356  			require.Nil(t, err)
   357  
   358  			res, err := b.SetList(key1)
   359  			require.Nil(t, err)
   360  			assert.Equal(t, orig1, res)
   361  			res, err = b.SetList(key2)
   362  			require.Nil(t, err)
   363  			assert.Equal(t, orig2, res)
   364  			res, err = b.SetList(key3)
   365  			require.Nil(t, err)
   366  			assert.Equal(t, orig3, res)
   367  		})
   368  
   369  		t.Run("delete individual keys", func(t *testing.T) {
   370  			delete2 := []byte("value 2.1")
   371  			delete3 := []byte("value 3.2")
   372  
   373  			err = b.SetDeleteSingle(key2, delete2)
   374  			require.Nil(t, err)
   375  			err = b.SetDeleteSingle(key3, delete3)
   376  			require.Nil(t, err)
   377  		})
   378  
   379  		t.Run("validate the results", func(t *testing.T) {
   380  			expected1 := [][]byte{[]byte("value 1.1"), []byte("value 1.2")} // unchanged
   381  			expected2 := [][]byte{[]byte("value 2.2")}                      // value1 deleted
   382  			expected3 := [][]byte{[]byte("value 3.1")}                      // value2 deleted
   383  
   384  			res, err := b.SetList(key1)
   385  			require.Nil(t, err)
   386  			assert.Equal(t, expected1, res)
   387  			res, err = b.SetList(key2)
   388  			require.Nil(t, err)
   389  			assert.Equal(t, expected2, res)
   390  			res, err = b.SetList(key3)
   391  			require.Nil(t, err)
   392  			assert.Equal(t, expected3, res)
   393  		})
   394  
   395  		t.Run("re-add keys which were previously deleted and new ones", func(t *testing.T) {
   396  			readd2 := [][]byte{[]byte("value 2.1"), []byte("value 2.3")}
   397  			readd3 := [][]byte{[]byte("value 3.2"), []byte("value 3.3")}
   398  
   399  			err = b.SetAdd(key2, readd2)
   400  			require.Nil(t, err)
   401  			err = b.SetAdd(key3, readd3)
   402  			require.Nil(t, err)
   403  		})
   404  
   405  		t.Run("validate the results again", func(t *testing.T) {
   406  			expected1 := [][]byte{[]byte("value 1.1"), []byte("value 1.2")} // unchanged
   407  			expected2 := [][]byte{
   408  				[]byte("value 2.2"), // from original import
   409  				[]byte("value 2.1"), // added again after initial deletion
   410  				[]byte("value 2.3"), // newly added
   411  			}
   412  			expected3 := [][]byte{
   413  				[]byte("value 3.1"), // form original import
   414  				[]byte("value 3.2"), // added again after initial deletion
   415  				[]byte("value 3.3"), // newly added
   416  			} // value2 deleted
   417  
   418  			res, err := b.SetList(key1)
   419  			require.Nil(t, err)
   420  			assert.Equal(t, expected1, res)
   421  			res, err = b.SetList(key2)
   422  			require.Nil(t, err)
   423  			assert.Equal(t, expected2, res)
   424  			res, err = b.SetList(key3)
   425  			require.Nil(t, err)
   426  			assert.Equal(t, expected3, res)
   427  		})
   428  	})
   429  
   430  	t.Run("with a single flush between updates", func(t *testing.T) {
   431  		b, err := NewBucket(ctx, dirName, "", nullLogger(), nil,
   432  			cyclemanager.NewCallbackGroupNoop(), cyclemanager.NewCallbackGroupNoop(), opts...)
   433  		require.Nil(t, err)
   434  
   435  		// so big it effectively never triggers as part of this test
   436  		b.SetMemtableThreshold(1e9)
   437  
   438  		key1 := []byte("test2-key-1")
   439  		key2 := []byte("test2-key-2")
   440  		key3 := []byte("test2-key-3")
   441  
   442  		t.Run("set original values and verify", func(t *testing.T) {
   443  			orig1 := [][]byte{[]byte("value 1.1"), []byte("value 1.2")}
   444  			orig2 := [][]byte{[]byte("value 2.1"), []byte("value 2.2")}
   445  			orig3 := [][]byte{[]byte("value 3.1"), []byte("value 3.2")}
   446  
   447  			err = b.SetAdd(key1, orig1)
   448  			require.Nil(t, err)
   449  			err = b.SetAdd(key2, orig2)
   450  			require.Nil(t, err)
   451  			err = b.SetAdd(key3, orig3)
   452  			require.Nil(t, err)
   453  
   454  			res, err := b.SetList(key1)
   455  			require.Nil(t, err)
   456  			assert.Equal(t, res, orig1)
   457  			res, err = b.SetList(key2)
   458  			require.Nil(t, err)
   459  			assert.Equal(t, res, orig2)
   460  			res, err = b.SetList(key3)
   461  			require.Nil(t, err)
   462  			assert.Equal(t, res, orig3)
   463  		})
   464  
   465  		t.Run("flush to disk", func(t *testing.T) {
   466  			require.Nil(t, b.FlushAndSwitch())
   467  		})
   468  
   469  		t.Run("delete individual keys", func(t *testing.T) {
   470  			delete2 := []byte("value 2.1")
   471  			delete3 := []byte("value 3.2")
   472  
   473  			err = b.SetDeleteSingle(key2, delete2)
   474  			require.Nil(t, err)
   475  			err = b.SetDeleteSingle(key3, delete3)
   476  			require.Nil(t, err)
   477  		})
   478  
   479  		t.Run("validate the results", func(t *testing.T) {
   480  			expected1 := [][]byte{[]byte("value 1.1"), []byte("value 1.2")} // unchanged
   481  			expected2 := [][]byte{[]byte("value 2.2")}                      // value1 deleted
   482  			expected3 := [][]byte{[]byte("value 3.1")}                      // value2 deleted
   483  
   484  			res, err := b.SetList(key1)
   485  			require.Nil(t, err)
   486  			assert.Equal(t, expected1, res)
   487  			res, err = b.SetList(key2)
   488  			require.Nil(t, err)
   489  			assert.Equal(t, expected2, res)
   490  			res, err = b.SetList(key3)
   491  			require.Nil(t, err)
   492  			assert.Equal(t, expected3, res)
   493  		})
   494  
   495  		t.Run("re-add keys which were previously deleted and new ones", func(t *testing.T) {
   496  			readd2 := [][]byte{[]byte("value 2.1"), []byte("value 2.3")}
   497  			readd3 := [][]byte{[]byte("value 3.2"), []byte("value 3.3")}
   498  
   499  			err = b.SetAdd(key2, readd2)
   500  			require.Nil(t, err)
   501  			err = b.SetAdd(key3, readd3)
   502  			require.Nil(t, err)
   503  		})
   504  
   505  		t.Run("validate the results again", func(t *testing.T) {
   506  			expected1 := [][]byte{[]byte("value 1.1"), []byte("value 1.2")} // unchanged
   507  			expected2 := [][]byte{
   508  				[]byte("value 2.2"), // from original import
   509  				[]byte("value 2.1"), // added again after initial deletion
   510  				[]byte("value 2.3"), // newly added
   511  			}
   512  			expected3 := [][]byte{
   513  				[]byte("value 3.1"), // form original import
   514  				[]byte("value 3.2"), // added again after initial deletion
   515  				[]byte("value 3.3"), // newly added
   516  			} // value2 deleted
   517  
   518  			res, err := b.SetList(key1)
   519  			require.Nil(t, err)
   520  			assert.Equal(t, expected1, res)
   521  			res, err = b.SetList(key2)
   522  			require.Nil(t, err)
   523  			assert.Equal(t, expected2, res)
   524  			res, err = b.SetList(key3)
   525  			require.Nil(t, err)
   526  			assert.Equal(t, expected3, res)
   527  		})
   528  	})
   529  
   530  	t.Run("with flushes in between and after the update", func(t *testing.T) {
   531  		b, err := NewBucket(ctx, dirName, "", nullLogger(), nil,
   532  			cyclemanager.NewCallbackGroupNoop(), cyclemanager.NewCallbackGroupNoop(), opts...)
   533  		require.Nil(t, err)
   534  
   535  		// so big it effectively never triggers as part of this test
   536  		b.SetMemtableThreshold(1e9)
   537  
   538  		key1 := []byte("test3-key-1")
   539  		key2 := []byte("test3-key-2")
   540  		key3 := []byte("test3-key-3")
   541  
   542  		t.Run("set original values and verify", func(t *testing.T) {
   543  			orig1 := [][]byte{[]byte("value 1.1"), []byte("value 1.2")}
   544  			orig2 := [][]byte{[]byte("value 2.1"), []byte("value 2.2")}
   545  			orig3 := [][]byte{[]byte("value 3.1"), []byte("value 3.2")}
   546  
   547  			err = b.SetAdd(key1, orig1)
   548  			require.Nil(t, err)
   549  			err = b.SetAdd(key2, orig2)
   550  			require.Nil(t, err)
   551  			err = b.SetAdd(key3, orig3)
   552  			require.Nil(t, err)
   553  
   554  			res, err := b.SetList(key1)
   555  			require.Nil(t, err)
   556  			assert.Equal(t, res, orig1)
   557  			res, err = b.SetList(key2)
   558  			require.Nil(t, err)
   559  			assert.Equal(t, res, orig2)
   560  			res, err = b.SetList(key3)
   561  			require.Nil(t, err)
   562  			assert.Equal(t, res, orig3)
   563  		})
   564  
   565  		t.Run("flush to disk", func(t *testing.T) {
   566  			require.Nil(t, b.FlushAndSwitch())
   567  		})
   568  
   569  		t.Run("delete individual keys", func(t *testing.T) {
   570  			delete2 := []byte("value 2.1")
   571  			delete3 := []byte("value 3.2")
   572  
   573  			err = b.SetDeleteSingle(key2, delete2)
   574  			require.Nil(t, err)
   575  			err = b.SetDeleteSingle(key3, delete3)
   576  			require.Nil(t, err)
   577  		})
   578  
   579  		t.Run("flush to disk - again!", func(t *testing.T) {
   580  			require.Nil(t, b.FlushAndSwitch())
   581  		})
   582  
   583  		t.Run("validate", func(t *testing.T) {
   584  			expected1 := [][]byte{[]byte("value 1.1"), []byte("value 1.2")} // unchanged
   585  			expected2 := [][]byte{[]byte("value 2.2")}                      // value1 deleted
   586  			expected3 := [][]byte{[]byte("value 3.1")}                      // value2 deleted
   587  
   588  			res, err := b.SetList(key1)
   589  			require.Nil(t, err)
   590  			assert.Equal(t, expected1, res)
   591  			res, err = b.SetList(key2)
   592  			require.Nil(t, err)
   593  			assert.Equal(t, expected2, res)
   594  			res, err = b.SetList(key3)
   595  			require.Nil(t, err)
   596  			assert.Equal(t, expected3, res)
   597  		})
   598  
   599  		t.Run("re-add keys which were previously deleted and new ones", func(t *testing.T) {
   600  			readd2 := [][]byte{[]byte("value 2.1"), []byte("value 2.3")}
   601  			readd3 := [][]byte{[]byte("value 3.2"), []byte("value 3.3")}
   602  
   603  			err = b.SetAdd(key2, readd2)
   604  			require.Nil(t, err)
   605  			err = b.SetAdd(key3, readd3)
   606  			require.Nil(t, err)
   607  		})
   608  
   609  		t.Run("flush to disk - yet again!", func(t *testing.T) {
   610  			require.Nil(t, b.FlushAndSwitch())
   611  		})
   612  
   613  		t.Run("validate the results again", func(t *testing.T) {
   614  			expected1 := [][]byte{[]byte("value 1.1"), []byte("value 1.2")} // unchanged
   615  			expected2 := [][]byte{
   616  				[]byte("value 2.2"), // from original import
   617  				[]byte("value 2.1"), // added again after initial deletion
   618  				[]byte("value 2.3"), // newly added
   619  			}
   620  			expected3 := [][]byte{
   621  				[]byte("value 3.1"), // form original import
   622  				[]byte("value 3.2"), // added again after initial deletion
   623  				[]byte("value 3.3"), // newly added
   624  			} // value2 deleted
   625  
   626  			res, err := b.SetList(key1)
   627  			require.Nil(t, err)
   628  			assert.Equal(t, expected1, res)
   629  			res, err = b.SetList(key2)
   630  			require.Nil(t, err)
   631  			assert.Equal(t, expected2, res)
   632  			res, err = b.SetList(key3)
   633  			require.Nil(t, err)
   634  			assert.Equal(t, expected3, res)
   635  		})
   636  	})
   637  
   638  	t.Run("update in memtable, make orderly shutdown, then create a new bucket from disk",
   639  		func(t *testing.T) {
   640  			b, err := NewBucket(ctx, dirName, "", nullLogger(), nil,
   641  				cyclemanager.NewCallbackGroupNoop(), cyclemanager.NewCallbackGroupNoop(), opts...)
   642  			require.Nil(t, err)
   643  
   644  			// so big it effectively never triggers as part of this test
   645  			b.SetMemtableThreshold(1e9)
   646  
   647  			key1 := []byte("test4-key-1")
   648  			key2 := []byte("test4-key-2")
   649  			key3 := []byte("test4-key-3")
   650  
   651  			t.Run("set original values and verify", func(t *testing.T) {
   652  				orig1 := [][]byte{[]byte("value 1.1"), []byte("value 1.2")}
   653  				orig2 := [][]byte{[]byte("value 2.1"), []byte("value 2.2")}
   654  				orig3 := [][]byte{[]byte("value 3.1"), []byte("value 3.2")}
   655  
   656  				err = b.SetAdd(key1, orig1)
   657  				require.Nil(t, err)
   658  				err = b.SetAdd(key2, orig2)
   659  				require.Nil(t, err)
   660  				err = b.SetAdd(key3, orig3)
   661  				require.Nil(t, err)
   662  
   663  				res, err := b.SetList(key1)
   664  				require.Nil(t, err)
   665  				assert.Equal(t, res, orig1)
   666  				res, err = b.SetList(key2)
   667  				require.Nil(t, err)
   668  				assert.Equal(t, res, orig2)
   669  				res, err = b.SetList(key3)
   670  				require.Nil(t, err)
   671  				assert.Equal(t, res, orig3)
   672  			})
   673  
   674  			t.Run("delete individual keys", func(t *testing.T) {
   675  				delete2 := []byte("value 2.1")
   676  				delete3 := []byte("value 3.2")
   677  
   678  				err = b.SetDeleteSingle(key2, delete2)
   679  				require.Nil(t, err)
   680  				err = b.SetDeleteSingle(key3, delete3)
   681  				require.Nil(t, err)
   682  			})
   683  
   684  			t.Run("orderly shutdown", func(t *testing.T) {
   685  				b.Shutdown(context.Background())
   686  			})
   687  
   688  			t.Run("init another bucket on the same files", func(t *testing.T) {
   689  				b2, err := NewBucket(ctx, dirName, "", nullLogger(), nil,
   690  					cyclemanager.NewCallbackGroupNoop(), cyclemanager.NewCallbackGroupNoop(), opts...)
   691  				require.Nil(t, err)
   692  
   693  				expected1 := [][]byte{[]byte("value 1.1"), []byte("value 1.2")} // unchanged
   694  				expected2 := [][]byte{[]byte("value 2.2")}                      // value1 deleted
   695  				expected3 := [][]byte{[]byte("value 3.1")}                      // value2 deleted
   696  
   697  				res, err := b2.SetList(key1)
   698  				require.Nil(t, err)
   699  				assert.Equal(t, expected1, res)
   700  				res, err = b2.SetList(key2)
   701  				require.Nil(t, err)
   702  				assert.Equal(t, expected2, res)
   703  				res, err = b2.SetList(key3)
   704  				require.Nil(t, err)
   705  				assert.Equal(t, expected3, res)
   706  			})
   707  		})
   708  }
   709  
   710  func collectionCursors(ctx context.Context, t *testing.T, opts []BucketOption) {
   711  	t.Run("memtable-only", func(t *testing.T) {
   712  		r := getRandomSeed()
   713  		dirName := t.TempDir()
   714  
   715  		b, err := NewBucket(ctx, dirName, "", nullLogger(), nil,
   716  			cyclemanager.NewCallbackGroupNoop(), cyclemanager.NewCallbackGroupNoop(), opts...)
   717  		require.Nil(t, err)
   718  
   719  		// so big it effectively never triggers as part of this test
   720  		b.SetMemtableThreshold(1e9)
   721  
   722  		t.Run("set original values", func(t *testing.T) {
   723  			pairs := 20
   724  			valuesPerPair := 3
   725  			keys := make([][]byte, pairs)
   726  			values := make([][][]byte, pairs)
   727  
   728  			for i := range keys {
   729  				keys[i] = []byte(fmt.Sprintf("key-%03d", i))
   730  				values[i] = make([][]byte, valuesPerPair)
   731  				for j := range values[i] {
   732  					values[i][j] = []byte(fmt.Sprintf("value-%03d.%d", i, j))
   733  				}
   734  			}
   735  
   736  			// shuffle to make sure the BST isn't accidentally in order
   737  			r.Shuffle(len(keys), func(i, j int) {
   738  				keys[i], keys[j] = keys[j], keys[i]
   739  				values[i], values[j] = values[j], values[i]
   740  			})
   741  
   742  			for i := range keys {
   743  				err = b.SetAdd(keys[i], values[i])
   744  				require.Nil(t, err)
   745  			}
   746  		})
   747  
   748  		t.Run("seek from somewhere in the middle", func(t *testing.T) {
   749  			expectedKeys := [][]byte{
   750  				[]byte("key-016"),
   751  				[]byte("key-017"),
   752  				[]byte("key-018"),
   753  				[]byte("key-019"),
   754  			}
   755  			expectedValues := [][][]byte{
   756  				{[]byte("value-016.0"), []byte("value-016.1"), []byte("value-016.2")},
   757  				{[]byte("value-017.0"), []byte("value-017.1"), []byte("value-017.2")},
   758  				{[]byte("value-018.0"), []byte("value-018.1"), []byte("value-018.2")},
   759  				{[]byte("value-019.0"), []byte("value-019.1"), []byte("value-019.2")},
   760  			}
   761  
   762  			var retrievedKeys [][]byte
   763  			var retrievedValues [][][]byte
   764  			c := b.SetCursor()
   765  			defer c.Close()
   766  			for k, v := c.Seek([]byte("key-016")); k != nil; k, v = c.Next() {
   767  				retrievedKeys = append(retrievedKeys, k)
   768  				retrievedValues = append(retrievedValues, v)
   769  			}
   770  
   771  			assert.Equal(t, expectedKeys, retrievedKeys)
   772  			assert.Equal(t, expectedValues, retrievedValues)
   773  		})
   774  
   775  		t.Run("start from the beginning", func(t *testing.T) {
   776  			expectedKeys := [][]byte{
   777  				[]byte("key-000"),
   778  				[]byte("key-001"),
   779  				[]byte("key-002"),
   780  			}
   781  			expectedValues := [][][]byte{
   782  				{[]byte("value-000.0"), []byte("value-000.1"), []byte("value-000.2")},
   783  				{[]byte("value-001.0"), []byte("value-001.1"), []byte("value-001.2")},
   784  				{[]byte("value-002.0"), []byte("value-002.1"), []byte("value-002.2")},
   785  			}
   786  
   787  			var retrievedKeys [][]byte
   788  			var retrievedValues [][][]byte
   789  			c := b.SetCursor()
   790  			defer c.Close()
   791  			retrieved := 0
   792  			for k, v := c.First(); k != nil && retrieved < 3; k, v = c.Next() {
   793  				retrieved++
   794  				retrievedKeys = append(retrievedKeys, k)
   795  				retrievedValues = append(retrievedValues, v)
   796  			}
   797  
   798  			assert.Equal(t, expectedKeys, retrievedKeys)
   799  			assert.Equal(t, expectedValues, retrievedValues)
   800  		})
   801  
   802  		t.Run("extend an existing key", func(t *testing.T) {
   803  			key := []byte("key-002")
   804  			extend := [][]byte{[]byte("value-002.3")}
   805  
   806  			require.Nil(t, b.SetAdd(key, extend))
   807  		})
   808  
   809  		t.Run("verify the extension is contained", func(t *testing.T) {
   810  			expectedKeys := [][]byte{
   811  				[]byte("key-001"),
   812  				[]byte("key-002"),
   813  			}
   814  			expectedValues := [][][]byte{
   815  				{[]byte("value-001.0"), []byte("value-001.1"), []byte("value-001.2")},
   816  				{
   817  					[]byte("value-002.0"), []byte("value-002.1"), []byte("value-002.2"),
   818  					[]byte("value-002.3"),
   819  				},
   820  			}
   821  
   822  			var retrievedKeys [][]byte
   823  			var retrievedValues [][][]byte
   824  			c := b.SetCursor()
   825  			defer c.Close()
   826  			retrieved := 0
   827  			for k, v := c.Seek([]byte("key-001")); k != nil && retrieved < 2; k, v = c.Next() {
   828  				retrieved++
   829  				retrievedKeys = append(retrievedKeys, k)
   830  				retrievedValues = append(retrievedValues, v)
   831  			}
   832  
   833  			assert.Equal(t, expectedKeys, retrievedKeys)
   834  			assert.Equal(t, expectedValues, retrievedValues)
   835  		})
   836  	})
   837  
   838  	t.Run("with flushes", func(t *testing.T) {
   839  		r := getRandomSeed()
   840  		dirName := t.TempDir()
   841  
   842  		b, err := NewBucket(ctx, dirName, "", nullLogger(), nil,
   843  			cyclemanager.NewCallbackGroupNoop(), cyclemanager.NewCallbackGroupNoop(), opts...)
   844  		require.Nil(t, err)
   845  
   846  		// so big it effectively never triggers as part of this test
   847  		b.SetMemtableThreshold(1e9)
   848  
   849  		t.Run("first third (%3==0)", func(t *testing.T) {
   850  			pairs := 20
   851  			valuesPerPair := 3
   852  			var keys [][]byte
   853  			var values [][][]byte
   854  
   855  			for i := 0; i < pairs; i++ {
   856  				if i%3 != 0 {
   857  					continue
   858  				}
   859  				keys = append(keys, []byte(fmt.Sprintf("key-%03d", i)))
   860  				curValues := make([][]byte, valuesPerPair)
   861  				for j := range curValues {
   862  					curValues[j] = []byte(fmt.Sprintf("value-%03d.%d", i, j))
   863  				}
   864  				values = append(values, curValues)
   865  			}
   866  
   867  			// shuffle to make sure the BST isn't accidentally in order
   868  			r.Shuffle(len(keys), func(i, j int) {
   869  				keys[i], keys[j] = keys[j], keys[i]
   870  				values[i], values[j] = values[j], values[i]
   871  			})
   872  
   873  			for i := range keys {
   874  				err = b.SetAdd(keys[i], values[i])
   875  				require.Nil(t, err)
   876  			}
   877  		})
   878  
   879  		t.Run("flush to disk", func(t *testing.T) {
   880  			require.Nil(t, b.FlushAndSwitch())
   881  		})
   882  
   883  		t.Run("second third (%3==1)", func(t *testing.T) {
   884  			pairs := 20
   885  			valuesPerPair := 3
   886  			var keys [][]byte
   887  			var values [][][]byte
   888  
   889  			for i := 0; i < pairs; i++ {
   890  				if i%3 != 1 {
   891  					continue
   892  				}
   893  				keys = append(keys, []byte(fmt.Sprintf("key-%03d", i)))
   894  				curValues := make([][]byte, valuesPerPair)
   895  				for j := range curValues {
   896  					curValues[j] = []byte(fmt.Sprintf("value-%03d.%d", i, j))
   897  				}
   898  				values = append(values, curValues)
   899  			}
   900  
   901  			// shuffle to make sure the BST isn't accidentally in order
   902  			r.Shuffle(len(keys), func(i, j int) {
   903  				keys[i], keys[j] = keys[j], keys[i]
   904  				values[i], values[j] = values[j], values[i]
   905  			})
   906  
   907  			for i := range keys {
   908  				err = b.SetAdd(keys[i], values[i])
   909  				require.Nil(t, err)
   910  			}
   911  		})
   912  
   913  		t.Run("flush to disk", func(t *testing.T) {
   914  			require.Nil(t, b.FlushAndSwitch())
   915  		})
   916  
   917  		t.Run("third (%3==2) memtable-only", func(t *testing.T) {
   918  			pairs := 20
   919  			valuesPerPair := 3
   920  			var keys [][]byte
   921  			var values [][][]byte
   922  
   923  			for i := 0; i < pairs; i++ {
   924  				if i%3 != 2 {
   925  					continue
   926  				}
   927  				keys = append(keys, []byte(fmt.Sprintf("key-%03d", i)))
   928  				curValues := make([][]byte, valuesPerPair)
   929  				for j := range curValues {
   930  					curValues[j] = []byte(fmt.Sprintf("value-%03d.%d", i, j))
   931  				}
   932  				values = append(values, curValues)
   933  			}
   934  
   935  			// shuffle to make sure the BST isn't accidentally in order
   936  			r.Shuffle(len(keys), func(i, j int) {
   937  				keys[i], keys[j] = keys[j], keys[i]
   938  				values[i], values[j] = values[j], values[i]
   939  			})
   940  
   941  			for i := range keys {
   942  				err = b.SetAdd(keys[i], values[i])
   943  				require.Nil(t, err)
   944  			}
   945  
   946  			// no flush for this one, so this segment stays in the memtable
   947  		})
   948  
   949  		t.Run("seek from somewhere in the middle", func(t *testing.T) {
   950  			expectedKeys := [][]byte{
   951  				[]byte("key-016"),
   952  				[]byte("key-017"),
   953  				[]byte("key-018"),
   954  				[]byte("key-019"),
   955  			}
   956  			expectedValues := [][][]byte{
   957  				{[]byte("value-016.0"), []byte("value-016.1"), []byte("value-016.2")},
   958  				{[]byte("value-017.0"), []byte("value-017.1"), []byte("value-017.2")},
   959  				{[]byte("value-018.0"), []byte("value-018.1"), []byte("value-018.2")},
   960  				{[]byte("value-019.0"), []byte("value-019.1"), []byte("value-019.2")},
   961  			}
   962  
   963  			var retrievedKeys [][]byte
   964  			var retrievedValues [][][]byte
   965  			c := b.SetCursor()
   966  			defer c.Close()
   967  			for k, v := c.Seek([]byte("key-016")); k != nil; k, v = c.Next() {
   968  				retrievedKeys = append(retrievedKeys, k)
   969  				retrievedValues = append(retrievedValues, v)
   970  			}
   971  
   972  			assert.Equal(t, expectedKeys, retrievedKeys)
   973  			assert.Equal(t, expectedValues, retrievedValues)
   974  		})
   975  
   976  		t.Run("start from the beginning", func(t *testing.T) {
   977  			expectedKeys := [][]byte{
   978  				[]byte("key-000"),
   979  				[]byte("key-001"),
   980  				[]byte("key-002"),
   981  			}
   982  			expectedValues := [][][]byte{
   983  				{[]byte("value-000.0"), []byte("value-000.1"), []byte("value-000.2")},
   984  				{[]byte("value-001.0"), []byte("value-001.1"), []byte("value-001.2")},
   985  				{[]byte("value-002.0"), []byte("value-002.1"), []byte("value-002.2")},
   986  			}
   987  
   988  			var retrievedKeys [][]byte
   989  			var retrievedValues [][][]byte
   990  			c := b.SetCursor()
   991  			defer c.Close()
   992  			retrieved := 0
   993  			for k, v := c.First(); k != nil && retrieved < 3; k, v = c.Next() {
   994  				retrieved++
   995  				retrievedKeys = append(retrievedKeys, k)
   996  				retrievedValues = append(retrievedValues, v)
   997  			}
   998  
   999  			assert.Equal(t, expectedKeys, retrievedKeys)
  1000  			assert.Equal(t, expectedValues, retrievedValues)
  1001  		})
  1002  
  1003  		t.Run("delete & extend an existing key", func(t *testing.T) {
  1004  			key := []byte("key-002")
  1005  			extend := [][]byte{[]byte("value-002.3")}
  1006  
  1007  			require.Nil(t, b.SetAdd(key, extend))
  1008  
  1009  			key = []byte("key-001")
  1010  			deleteValue := []byte("value-001.1")
  1011  			require.Nil(t, b.SetDeleteSingle(key, deleteValue))
  1012  		})
  1013  
  1014  		t.Run("verify the extension is contained", func(t *testing.T) {
  1015  			expectedKeys := [][]byte{
  1016  				[]byte("key-001"),
  1017  				[]byte("key-002"),
  1018  			}
  1019  			expectedValues := [][][]byte{
  1020  				{
  1021  					[]byte("value-001.0"),
  1022  					// "value-001.1" deleted
  1023  					[]byte("value-001.2"),
  1024  				},
  1025  				{
  1026  					[]byte("value-002.0"), []byte("value-002.1"), []byte("value-002.2"),
  1027  					[]byte("value-002.3"),
  1028  				},
  1029  			}
  1030  
  1031  			var retrievedKeys [][]byte
  1032  			var retrievedValues [][][]byte
  1033  			c := b.SetCursor()
  1034  			defer c.Close()
  1035  			retrieved := 0
  1036  			for k, v := c.Seek([]byte("key-001")); k != nil && retrieved < 2; k, v = c.Next() {
  1037  				retrieved++
  1038  				retrievedKeys = append(retrievedKeys, k)
  1039  				retrievedValues = append(retrievedValues, v)
  1040  			}
  1041  
  1042  			assert.Equal(t, expectedKeys, retrievedKeys)
  1043  			assert.Equal(t, expectedValues, retrievedValues)
  1044  		})
  1045  
  1046  		t.Run("flush to disk", func(t *testing.T) {
  1047  			require.Nil(t, b.FlushAndSwitch())
  1048  		})
  1049  
  1050  		t.Run("verify again after flush", func(t *testing.T) {
  1051  			expectedKeys := [][]byte{
  1052  				[]byte("key-001"),
  1053  				[]byte("key-002"),
  1054  			}
  1055  			expectedValues := [][][]byte{
  1056  				{
  1057  					[]byte("value-001.0"),
  1058  					// "value-001.1" deleted
  1059  					[]byte("value-001.2"),
  1060  				},
  1061  				{
  1062  					[]byte("value-002.0"), []byte("value-002.1"), []byte("value-002.2"),
  1063  					[]byte("value-002.3"),
  1064  				},
  1065  			}
  1066  
  1067  			var retrievedKeys [][]byte
  1068  			var retrievedValues [][][]byte
  1069  			c := b.SetCursor()
  1070  			defer c.Close()
  1071  			retrieved := 0
  1072  			for k, v := c.Seek([]byte("key-001")); k != nil && retrieved < 2; k, v = c.Next() {
  1073  				retrieved++
  1074  				retrievedKeys = append(retrievedKeys, k)
  1075  				retrievedValues = append(retrievedValues, v)
  1076  			}
  1077  
  1078  			assert.Equal(t, expectedKeys, retrievedKeys)
  1079  			assert.Equal(t, expectedValues, retrievedValues)
  1080  		})
  1081  	})
  1082  }