milvus/internal/rootcoord/suffix_snapshot_test.go

package rootcoord

import (
	"errors"
	"fmt"
	"math/rand"
	"testing"
	"time"

	etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
	"github.com/milvus-io/milvus/internal/util/typeutil"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

var (
	snapshotPrefix = "snapshots"
)

func Test_binarySearchRecords(t *testing.T) {
	type testcase struct {
		records     []tsv
		ts          typeutil.Timestamp
		expected    string
		shouldFound bool
	}

	cases := []testcase{
		{
			records:     []tsv{},
			ts:          0,
			expected:    "",
			shouldFound: false,
		},
		{
			records: []tsv{
				{
					ts:    101,
					value: "abc",
				},
			},
			ts:          100,
			expected:    "",
			shouldFound: false,
		},
		{
			records: []tsv{
				{
					ts:    100,
					value: "a",
				},
			},
			ts:          100,
			expected:    "a",
			shouldFound: true,
		},
		{
			records: []tsv{
				{
					ts:    100,
					value: "a",
				},
				{
					ts:    200,
					value: "b",
				},
			},
			ts:          100,
			expected:    "a",
			shouldFound: true,
		},
		{
			records: []tsv{
				{
					ts:    100,
					value: "a",
				},
				{
					ts:    200,
					value: "b",
				},
				{
					ts:    300,
					value: "c",
				},
			},
			ts:          150,
			expected:    "a",
			shouldFound: true,
		},
		{
			records: []tsv{
				{
					ts:    100,
					value: "a",
				},
				{
					ts:    200,
					value: "b",
				},
				{
					ts:    300,
					value: "c",
				},
			},
			ts:          300,
			expected:    "c",
			shouldFound: true,
		},
		{
			records: []tsv{
				{
					ts:    100,
					value: "a",
				},
				{
					ts:    200,
					value: "b",
				},
				{
					ts:    300,
					value: "c",
				},
			},
			ts:          201,
			expected:    "b",
			shouldFound: true,
		},
		{
			records: []tsv{
				{
					ts:    100,
					value: "a",
				},
				{
					ts:    200,
					value: "b",
				},
				{
					ts:    300,
					value: "c",
				},
			},
			ts:          301,
			expected:    "c",
			shouldFound: true,
		},
	}
	for _, c := range cases {
		result, found := binarySearchRecords(c.records, c.ts)
		assert.Equal(t, c.expected, result)
		assert.Equal(t, c.shouldFound, found)
	}
}

func Test_ComposeIsTsKey(t *testing.T) {
	sep := "_ts"
	ss, err := newSuffixSnapshot((*etcdkv.EtcdKV)(nil), sep, "", snapshotPrefix)
	require.Nil(t, err)
	type testcase struct {
		key         string
		expected    uint64
		shouldFound bool
	}
	testCases := []testcase{
		{
			key:         ss.composeTSKey("key", 100),
			expected:    100,
			shouldFound: true,
		},
		{
			key:         ss.composeTSKey("other-key", 65536),
			expected:    65536,
			shouldFound: true,
		},
		{
			key:         "snapshots/test/1000",
			expected:    0,
			shouldFound: false,
		},
		{
			key:         "snapshots",
			expected:    0,
			shouldFound: false,
		},
	}
	for _, c := range testCases {
		ts, found := ss.isTSKey(c.key)
		assert.EqualValues(t, c.expected, ts)
		assert.Equal(t, c.shouldFound, found)
	}
}

func Test_SuffixSnaphotIsTSOfKey(t *testing.T) {
	sep := "_ts"
	ss, err := newSuffixSnapshot((*etcdkv.EtcdKV)(nil), sep, "", snapshotPrefix)
	require.Nil(t, err)
	type testcase struct {
		key         string
		target      string
		expected    uint64
		shouldFound bool
	}
	testCases := []testcase{
		{
			key:         ss.composeTSKey("key", 100),
			target:      "key",
			expected:    100,
			shouldFound: true,
		},
		{
			key:         ss.composeTSKey("other-key", 65536),
			target:      "other-key",
			expected:    65536,
			shouldFound: true,
		},
		{
			key:         ss.composeTSKey("other-key", 65536),
			target:      "key",
			expected:    0,
			shouldFound: false,
		},
		{
			key:         "snapshots/test/1000",
			expected:    0,
			shouldFound: false,
		},
		{
			key:         "snapshots",
			expected:    0,
			shouldFound: false,
		},
	}
	for _, c := range testCases {
		ts, found := ss.isTSOfKey(c.key, c.target)
		assert.EqualValues(t, c.expected, ts)
		assert.Equal(t, c.shouldFound, found)
	}

}

func Test_SuffixSnapshotLoad(t *testing.T) {
	rand.Seed(time.Now().UnixNano())
	randVal := rand.Int()

	Params.Init()
	rootPath := fmt.Sprintf("/test/meta/%d", randVal)
	sep := "_ts"

	etcdkv, err := etcdkv.NewEtcdKV(Params.EtcdEndpoints, rootPath)
	require.Nil(t, err)
	defer etcdkv.Close()

	var vtso typeutil.Timestamp
	ftso := func() typeutil.Timestamp {
		return vtso
	}

	ss, err := newSuffixSnapshot(etcdkv, sep, rootPath, snapshotPrefix)
	assert.Nil(t, err)
	assert.NotNil(t, ss)

	for i := 0; i < 20; i++ {
		vtso = typeutil.Timestamp(100 + i*5)
		ts := ftso()
		err = ss.Save("key", fmt.Sprintf("value-%d", i), ts)
		assert.Nil(t, err)
		assert.Equal(t, vtso, ts)
	}
	for i := 0; i < 20; i++ {
		val, err := ss.Load("key", typeutil.Timestamp(100+i*5+2))
		t.Log("ts:", typeutil.Timestamp(100+i*5+2), i, val)
		assert.Nil(t, err)
		assert.Equal(t, fmt.Sprintf("value-%d", i), val)
	}
	val, err := ss.Load("key", 0)
	assert.Nil(t, err)
	assert.Equal(t, "value-19", val)

	ss, err = newSuffixSnapshot(etcdkv, sep, rootPath, snapshotPrefix)
	assert.Nil(t, err)
	assert.NotNil(t, ss)

	for i := 0; i < 20; i++ {
		val, err := ss.Load("key", typeutil.Timestamp(100+i*5+2))
		assert.Nil(t, err)
		assert.Equal(t, val, fmt.Sprintf("value-%d", i))
	}

	ss.RemoveWithPrefix("")
}

func Test_SuffixSnapshotMultiSave(t *testing.T) {
	rand.Seed(time.Now().UnixNano())
	randVal := rand.Int()

	Params.Init()
	rootPath := fmt.Sprintf("/test/meta/%d", randVal)
	sep := "_ts"
	etcdkv, err := etcdkv.NewEtcdKV(Params.EtcdEndpoints, rootPath)
	require.Nil(t, err)
	defer etcdkv.Close()

	var vtso typeutil.Timestamp
	ftso := func() typeutil.Timestamp {
		return vtso
	}

	ss, err := newSuffixSnapshot(etcdkv, sep, rootPath, snapshotPrefix)
	assert.Nil(t, err)
	assert.NotNil(t, ss)

	for i := 0; i < 20; i++ {
		saves := map[string]string{"k1": fmt.Sprintf("v1-%d", i), "k2": fmt.Sprintf("v2-%d", i)}
		vtso = typeutil.Timestamp(100 + i*5)
		ts := ftso()
		err = ss.MultiSave(saves, ts, func(ts typeutil.Timestamp) (string, string, error) {
			return "extra", "extra-value", nil
		}, func(ts typeutil.Timestamp) (string, string, error) {
			return "extra", "extra-value", errors.New("out of range")
		})
		assert.Nil(t, err)
		assert.Equal(t, vtso, ts)
	}
	for i := 0; i < 20; i++ {
		keys, vals, err := ss.LoadWithPrefix("k", typeutil.Timestamp(100+i*5+2))
		t.Log(i, keys, vals)
		assert.Nil(t, err)
		assert.Equal(t, len(keys), len(vals))
		assert.Equal(t, len(keys), 2)
		assert.Equal(t, keys[0], "k1")
		assert.Equal(t, keys[1], "k2")
		assert.Equal(t, vals[0], fmt.Sprintf("v1-%d", i))
		assert.Equal(t, vals[1], fmt.Sprintf("v2-%d", i))
	}
	keys, vals, err := ss.LoadWithPrefix("k", 0)
	assert.Nil(t, err)
	assert.Equal(t, len(keys), len(vals))
	assert.Equal(t, len(keys), 2)
	assert.Equal(t, keys[0], "k1")
	assert.Equal(t, keys[1], "k2")
	assert.Equal(t, vals[0], "v1-19")
	assert.Equal(t, vals[1], "v2-19")

	ss, err = newSuffixSnapshot(etcdkv, sep, rootPath, snapshotPrefix)
	assert.Nil(t, err)
	assert.NotNil(t, ss)
	for i := 0; i < 20; i++ {
		keys, vals, err := ss.LoadWithPrefix("k", typeutil.Timestamp(100+i*5+2))
		assert.Nil(t, err)
		assert.Equal(t, len(keys), len(vals))
		assert.Equal(t, len(keys), 2)
		assert.ElementsMatch(t, keys, []string{"k1", "k2"})
		assert.ElementsMatch(t, vals, []string{fmt.Sprintf("v1-%d", i), fmt.Sprintf("v2-%d", i)})
	}
	// mix non ts k-v
	err = ss.Save("kextra", "extra-value", 0)
	assert.Nil(t, err)
	keys, vals, err = ss.LoadWithPrefix("k", typeutil.Timestamp(300))
	assert.Nil(t, err)
	assert.Equal(t, len(keys), len(vals))
	assert.Equal(t, len(keys), 3)
	assert.ElementsMatch(t, keys, []string{"k1", "k2", "kextra"})
	assert.ElementsMatch(t, vals, []string{"v1-19", "v2-19", "extra-value"})

	// clean up
	ss.RemoveWithPrefix("")
}

func Test_SuffixSnapshotMultiSaveAndRemoveWithPrefix(t *testing.T) {
	rand.Seed(time.Now().UnixNano())
	randVal := rand.Int()

	Params.Init()
	rootPath := fmt.Sprintf("/test/meta/%d", randVal)
	sep := "_ts"

	etcdkv, err := etcdkv.NewEtcdKV(Params.EtcdEndpoints, rootPath)
	require.Nil(t, err)
	defer etcdkv.Close()

	var vtso typeutil.Timestamp
	ftso := func() typeutil.Timestamp {
		return vtso
	}

	ss, err := newSuffixSnapshot(etcdkv, sep, rootPath, snapshotPrefix)
	assert.Nil(t, err)
	assert.NotNil(t, ss)

	for i := 0; i < 20; i++ {
		vtso = typeutil.Timestamp(100 + i*5)
		ts := ftso()
		err = ss.Save(fmt.Sprintf("kd-%04d", i), fmt.Sprintf("value-%d", i), ts)
		assert.Nil(t, err)
		assert.Equal(t, vtso, ts)
	}
	for i := 20; i < 40; i++ {
		sm := map[string]string{"ks": fmt.Sprintf("value-%d", i)}
		dm := []string{fmt.Sprintf("kd-%04d", i-20)}
		vtso = typeutil.Timestamp(100 + i*5)
		ts := ftso()
		err = ss.MultiSaveAndRemoveWithPrefix(sm, dm, ts, func(ts typeutil.Timestamp) (string, string, error) {
			return "extra", "extra-value", nil
		}, func(ts typeutil.Timestamp) (string, string, error) {
			return "extra", "extra-value", errors.New("out of range")
		})
		assert.Nil(t, err)
		assert.Equal(t, vtso, ts)
	}
	for i := 0; i < 20; i++ {
		val, err := ss.Load(fmt.Sprintf("kd-%04d", i), typeutil.Timestamp(100+i*5+2))
		assert.Nil(t, err)
		assert.Equal(t, fmt.Sprintf("value-%d", i), val)
		_, vals, err := ss.LoadWithPrefix("kd-", typeutil.Timestamp(100+i*5+2))
		assert.Nil(t, err)
		assert.Equal(t, i+1, len(vals))
	}
	for i := 20; i < 40; i++ {
		val, err := ss.Load("ks", typeutil.Timestamp(100+i*5+2))
		assert.Nil(t, err)
		assert.Equal(t, fmt.Sprintf("value-%d", i), val)
		_, vals, err := ss.LoadWithPrefix("kd-", typeutil.Timestamp(100+i*5+2))
		assert.Nil(t, err)
		assert.Equal(t, 39-i, len(vals))
	}

	ss, err = newSuffixSnapshot(etcdkv, sep, rootPath, snapshotPrefix)
	assert.Nil(t, err)
	assert.NotNil(t, ss)

	for i := 0; i < 20; i++ {
		val, err := ss.Load(fmt.Sprintf("kd-%04d", i), typeutil.Timestamp(100+i*5+2))
		assert.Nil(t, err)
		assert.Equal(t, fmt.Sprintf("value-%d", i), val)
		_, vals, err := ss.LoadWithPrefix("kd-", typeutil.Timestamp(100+i*5+2))
		assert.Nil(t, err)
		assert.Equal(t, i+1, len(vals))
	}
	for i := 20; i < 40; i++ {
		val, err := ss.Load("ks", typeutil.Timestamp(100+i*5+2))
		assert.Nil(t, err)
		assert.Equal(t, fmt.Sprintf("value-%d", i), val)
		_, vals, err := ss.LoadWithPrefix("kd-", typeutil.Timestamp(100+i*5+2))
		assert.Nil(t, err)
		assert.Equal(t, 39-i, len(vals))
	}
	// try to load
	_, err = ss.Load("kd-0000", 500)
	assert.NotNil(t, err)
	_, err = ss.Load("kd-0000", 0)
	assert.NotNil(t, err)
	_, err = ss.Load("kd-0000", 1)
	assert.NotNil(t, err)

	// cleanup
	ss.MultiSaveAndRemoveWithPrefix(map[string]string{}, []string{""}, 0)
}