github.com/treeverse/lakefs@v1.24.1-0.20240520134607-95648127bfb0/pkg/graveler/sstable/iterator_test.go (about)

     1  package sstable_test
     2  
     3  import (
     4  	"os"
     5  	"sort"
     6  	"testing"
     7  
     8  	"github.com/cockroachdb/pebble"
     9  	pebblesst "github.com/cockroachdb/pebble/sstable"
    10  	"github.com/golang/mock/gomock"
    11  	"github.com/stretchr/testify/require"
    12  	"github.com/treeverse/lakefs/pkg/graveler/committed"
    13  	"github.com/treeverse/lakefs/pkg/graveler/sstable"
    14  )
    15  
    16  func TestIteratorSuccess(t *testing.T) {
    17  	ctrl := gomock.NewController(t)
    18  	defer ctrl.Finish()
    19  
    20  	count := 1000
    21  	keys := randomStrings(count)
    22  	sort.Strings(keys)
    23  	vals := randomStrings(count)
    24  	iter := createSStableIterator(t, keys, vals)
    25  
    26  	called := 0
    27  	sut := sstable.NewIterator(iter, func() error {
    28  		called++
    29  		return nil
    30  	})
    31  	require.NotNil(t, sut)
    32  
    33  	// read first -> nothing to read
    34  	require.Nil(t, sut.Value())
    35  	require.NoError(t, sut.Err())
    36  
    37  	// advance by one and read
    38  	require.True(t, sut.Next())
    39  	val := sut.Value()
    40  	require.NoError(t, sut.Err())
    41  	require.NotNil(t, val)
    42  	require.Equal(t, committed.Key(keys[0]), val.Key)
    43  	require.NotNil(t, val.Value)
    44  
    45  	// advance by one and read
    46  	require.True(t, sut.Next())
    47  	val = sut.Value()
    48  	require.NoError(t, sut.Err())
    49  	require.NotNil(t, val)
    50  	require.Equal(t, committed.Key(keys[1]), val.Key)
    51  	require.NotNil(t, val.Value)
    52  
    53  	// seek to a random offset
    54  	seekedKeyIndex := count / 3
    55  	seekedKey := committed.Key(keys[seekedKeyIndex])
    56  	sut.SeekGE(seekedKey)
    57  	require.NoError(t, sut.Err())
    58  	// value should be nil until next is called
    59  	require.Nil(t, sut.Value())
    60  	require.True(t, sut.Next())
    61  	val = sut.Value()
    62  	require.NoError(t, sut.Err())
    63  	require.NotNil(t, val)
    64  	require.Equal(t, seekedKey, val.Key)
    65  	require.NotNil(t, val.Value)
    66  
    67  	// read till the end
    68  	for i := seekedKeyIndex + 1; i < count; i++ {
    69  		require.True(t, sut.Next())
    70  		val = sut.Value()
    71  		require.NoError(t, sut.Err())
    72  		require.NotNil(t, val)
    73  		require.Equal(t, committed.Key(keys[i]), val.Key)
    74  		require.NotNil(t, val.Value)
    75  	}
    76  
    77  	// reached the end
    78  	require.False(t, sut.Next())
    79  	require.NoError(t, sut.Err())
    80  
    81  	sut.Close()
    82  	require.NoError(t, sut.Err())
    83  	require.Equal(t, 1, called)
    84  }
    85  
    86  // createSStableIterator creates the iterator from keys, vals passed to it
    87  func createSStableIterator(t *testing.T, keys, vals []string) pebblesst.Iterator {
    88  	ssReader := createSStableReader(t, keys, vals)
    89  
    90  	iter, err := ssReader.NewIter(nil, nil)
    91  	require.NoError(t, err)
    92  
    93  	t.Cleanup(func() {
    94  		_ = iter.Close()
    95  	})
    96  	return iter
    97  }
    98  
    99  // wrapReadableFile wraps os.File to count how many times it has been closed.
   100  type wrapReadableFile struct {
   101  	*os.File
   102  	NumClosed int
   103  }
   104  
   105  func (f *wrapReadableFile) ReadAt(p []byte, off int64) (n int, err error) {
   106  	return f.File.ReadAt(p, off)
   107  }
   108  
   109  func (f *wrapReadableFile) Close() error {
   110  	f.NumClosed++
   111  	return f.File.Close()
   112  }
   113  
   114  func (f *wrapReadableFile) Stat() (os.FileInfo, error) {
   115  	return f.File.Stat()
   116  }
   117  
   118  type fakeReader struct {
   119  	*pebblesst.Reader
   120  	GetNumClosed func() int
   121  }
   122  
   123  // createSStableReader creates the table from keys, vals passed to it
   124  func createSStableReader(t *testing.T, keys []string, vals []string) fakeReader {
   125  	f, err := os.CreateTemp(os.TempDir(), "test file")
   126  	require.NoError(t, err)
   127  	w := pebblesst.NewWriter(f, pebblesst.WriterOptions{
   128  		Compression: pebblesst.SnappyCompression,
   129  	})
   130  	for i, key := range keys {
   131  		require.NoError(t, w.Set([]byte(key), []byte(vals[i])))
   132  	}
   133  	require.NoError(t, w.Close())
   134  
   135  	cache := pebble.NewCache(0)
   136  	t.Cleanup(func() {
   137  		cache.Unref()
   138  	})
   139  
   140  	readF, err := os.Open(f.Name())
   141  	require.NoError(t, err)
   142  	wf := &wrapReadableFile{readF, 0}
   143  	ssReader, err := pebblesst.NewReader(wf, pebblesst.ReaderOptions{Cache: cache})
   144  	require.NoError(t, err)
   145  	t.Cleanup(func() {
   146  		if wf.NumClosed == 0 {
   147  			_ = ssReader.Close()
   148  		}
   149  	})
   150  
   151  	return fakeReader{ssReader, func() int { return wf.NumClosed }}
   152  }