go-kv/memtable/skiplist_test.go

package memtable

import (
	"fmt"
	"math/rand"
	"sync"
	"testing"

	"github.com/stretchr/testify/require"
)

func TestSkipListEmpty(t *testing.T) {
	sl := NewSkipList(NewArena(1024))

	found, value, sequence := sl.Get([]byte("missing"))
	require.False(t, found)
	require.Nil(t, value)
	require.Zero(t, sequence)
}

func TestSkipListOrdered(t *testing.T) {
	sl := NewSkipList(NewArena(64 << 20))

	keys := []string{"c", "a", "e", "b", "d"}
	for i, key := range keys {
		require.NoError(t, sl.Put([]byte(key), []byte("value-"+key), uint64(i+1), false))
	}

	for i, key := range keys {
		found, value, sequence := sl.Get([]byte(key))
		require.True(t, found, "key %q", key)
		require.Equal(t, []byte("value-"+key), value)
		require.Equal(t, uint64(i+1), sequence)
	}

	it := sl.NewIterator()
	var got []string
	for it.Valid() {
		got = append(got, string(it.Key()))
		it.Next()
	}
	require.Equal(t, []string{"a", "b", "c", "d", "e"}, got)
}

func TestSkipListOverwrite(t *testing.T) {
	sl := NewSkipList(NewArena(1024))

	require.NoError(t, sl.Put([]byte("key"), []byte("old"), 1, false))
	require.NoError(t, sl.Put([]byte("key"), []byte("new"), 2, false))

	found, value, sequence := sl.Get([]byte("key"))
	require.True(t, found)
	require.Equal(t, []byte("new"), value)
	require.Equal(t, uint64(2), sequence)

	it := sl.NewIterator()
	require.True(t, it.Valid())
	require.Equal(t, []byte("key"), it.Key())
	require.Equal(t, []byte("new"), it.Value())
	it.Next()
	require.False(t, it.Valid())
}

func TestSkipListPendingEntriesAreHidden(t *testing.T) {
	sl := NewSkipList(NewArena(1024))

	require.NoError(t, sl.Put([]byte("a"), []byte("pending"), 1, true))
	require.NoError(t, sl.Put([]byte("b"), []byte("published"), 2, false))

	found, _, _ := sl.Get([]byte("a"))
	require.False(t, found)

	it := sl.NewIterator()
	require.True(t, it.Valid())
	require.Equal(t, []byte("b"), it.Key())
	it.Next()
	require.False(t, it.Valid())
}

func TestSkipListConcurrent(t *testing.T) {
	sl := NewSkipList(NewArena(64 << 20))

	const (
		writers        = 4
		readers        = 4
		keysPerWriter  = 100
		readsPerReader = 1000
	)

	var wg sync.WaitGroup
	start := make(chan struct{})
	errCh := make(chan error, writers+readers)

	for writer := range writers {
		wg.Go(func() {
			<-start
			for i := range keysPerWriter {
				key := fmt.Appendf(nil, "writer-%d-key-%03d", writer, i)
				value := fmt.Appendf(nil, "value-%d-%03d", writer, i)
				sequence := uint64(writer*keysPerWriter + i + 1)
				if err := sl.Put(key, value, sequence, false); err != nil {
					errCh <- err
					return
				}
			}
		})
	}

	for reader := range readers {
		wg.Go(func() {
			<-start
			rng := rand.New(rand.NewSource(int64(reader)))
			for range readsPerReader {
				writer := rng.Intn(writers)
				keyID := rng.Intn(keysPerWriter)
				key := fmt.Appendf(nil, "writer-%d-key-%03d", writer, keyID)
				found, value, sequence := sl.Get(key)
				if found {
					expectedValue := fmt.Appendf(nil, "value-%d-%03d", writer, keyID)
					expectedSequence := uint64(writer*keysPerWriter + keyID + 1)
					if err := compareEntry(value, sequence, expectedValue, expectedSequence); err != nil {
						errCh <- err
						return
					}
				}
			}
		})
	}

	close(start)
	wg.Wait()
	close(errCh)
	for err := range errCh {
		require.NoError(t, err)
	}

	for writer := range writers {
		for i := range keysPerWriter {
			key := fmt.Appendf(nil, "writer-%d-key-%03d", writer, i)
			found, value, sequence := sl.Get(key)
			require.True(t, found, "key %q", key)
			require.Equal(t, fmt.Appendf(nil, "value-%d-%03d", writer, i), value)
			require.Equal(t, uint64(writer*keysPerWriter+i+1), sequence)
		}
	}
}

func compareEntry(value []byte, sequence uint64, expectedValue []byte, expectedSequence uint64) error {
	if string(value) != string(expectedValue) {
		return fmt.Errorf("value = %q, want %q", value, expectedValue)
	}
	if sequence != expectedSequence {
		return fmt.Errorf("sequence = %d, want %d", sequence, expectedSequence)
	}
	return nil
}