test/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2024-11-30 02:48:45 +08:00
Code Issues Actions 11 Packages Projects Releases Wiki Activity

Refactor NewEtcdKV API (#7075)

Browse Source 
This PR mainly did the following 3 things:

- Refactor NewEtcdKV API
- Add unittests for `kv/etcd/`, raise the coverage
up to 94%
- Correct some APIs from version to revision

The old NewEtcdKV() API took `clientv3.Client` as
one of the input, making everyone who using this API
had to create a new `clientv3.Client`. So there're a lot
of repeated codes everywhere.

Meanwhile, someplace didn't even check whether the
`clientv3.New()` returned error or not, and the old API
also didn't check whether the clientv3.Client was Nil.
This caused issue#6955.

Ectd's version and revision have nothing alike, they
represent different concepts, but our APIs kind of
mixed them all up. So I correct those APIs' names.

Resolves: #6955

Signed-off-by: yangxuan <xuan.yang@zilliz.com>
This commit is contained in:
XuanYang-cn 2021-08-13 11:04:09 +08:00 committed by GitHub
parent e1b3273d1c
commit 6c5be83ea1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
20 changed files with 554 additions and 432 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
internal/allocator/global_id_test.go
View File

@ -28,7 +28,10 @@ func TestGlobalTSOAllocator_All(t *testing.T) {
endpoints = "localhost:2379"
}
etcdEndpoints := strings.Split(endpoints, ",")
gTestIDAllocator = NewGlobalIDAllocator("idTimestamp", tsoutil.NewTSOKVBase(etcdEndpoints, "/test/root/kv", "gidTest"))
etcdKV, err := tsoutil.NewTSOKVBase(etcdEndpoints, "/test/root/kv", "gidTest")
assert.NoError(t, err)
gTestIDAllocator = NewGlobalIDAllocator("idTimestamp", etcdKV)
t.Run("Initialize", func(t *testing.T) {
err := gTestIDAllocator.Initialize()

6
internal/datacoord/server.go
View File

@ -20,7 +20,6 @@ import (
datanodeclient "github.com/milvus-io/milvus/internal/distributed/datanode/client"
rootcoordclient "github.com/milvus-io/milvus/internal/distributed/rootcoord/client"
"github.com/milvus-io/milvus/internal/logutil"
"go.etcd.io/etcd/clientv3"
"go.uber.org/zap"
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
@ -246,11 +245,12 @@ func (s *Server) startSegmentManager() {
func (s *Server) initMeta() error {
connectEtcdFn := func() error {
etcdClient, err := clientv3.New(clientv3.Config{Endpoints: Params.EtcdEndpoints})
etcdKV, err := etcdkv.NewEtcdKV(Params.EtcdEndpoints, Params.MetaRootPath)
if err != nil {
return err
}
s.kvClient = etcdkv.NewEtcdKV(etcdClient, Params.MetaRootPath)
s.kvClient = etcdKV
s.meta, err = newMeta(s.kvClient)
if err != nil {
return err

6
internal/datanode/mock_test.go
View File

@ -23,7 +23,6 @@ import (
"sync"
"time"
"go.etcd.io/etcd/clientv3"
"go.uber.org/zap"
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
@ -118,13 +117,10 @@ func refreshChannelNames() {
}
func clearEtcd(rootPath string) error {
etcdEndpoints := Params.EtcdEndpoints
log.Info("etcd tests address", zap.Any("endpoints", etcdEndpoints))
etcdClient, err := clientv3.New(clientv3.Config{Endpoints: etcdEndpoints})
etcdKV, err := etcdkv.NewEtcdKV(Params.EtcdEndpoints, rootPath)
if err != nil {
return err
}
etcdKV := etcdkv.NewEtcdKV(etcdClient, rootPath)
err = etcdKV.RemoveWithPrefix("writer/segment")
if err != nil {

11
internal/indexcoord/index_coord.go
View File

@ -25,7 +25,6 @@ import (
"github.com/coreos/etcd/mvcc/mvccpb"
"github.com/golang/protobuf/proto"
"go.etcd.io/etcd/clientv3"
"go.uber.org/zap"
"github.com/milvus-io/milvus/internal/allocator"
@ -104,11 +103,10 @@ func (i *IndexCoord) Init() error {
log.Debug("IndexCoord", zap.Any("etcd endpoints", Params.EtcdEndpoints))
connectEtcdFn := func() error {
etcdClient, err := clientv3.New(clientv3.Config{Endpoints: Params.EtcdEndpoints})
etcdKV, err := etcdkv.NewEtcdKV(Params.EtcdEndpoints, Params.MetaRootPath)
if err != nil {
return err
}
etcdKV := etcdkv.NewEtcdKV(etcdClient, Params.MetaRootPath)
metakv, err := NewMetaTable(etcdKV)
if err != nil {
return err
@ -145,7 +143,12 @@ func (i *IndexCoord) Init() error {
//init idAllocator
kvRootPath := Params.KvRootPath
i.idAllocator = allocator.NewGlobalIDAllocator("idTimestamp", tsoutil.NewTSOKVBase(Params.EtcdEndpoints, kvRootPath, "index_gid"))
etcdKV, err := tsoutil.NewTSOKVBase(Params.EtcdEndpoints, kvRootPath, "index_gid")
if err != nil {
log.Debug("IndexCoord TSOKVBase initialize failed", zap.Error(err))
}
i.idAllocator = allocator.NewGlobalIDAllocator("idTimestamp", etcdKV)
if err := i.idAllocator.Initialize(); err != nil {
log.Debug("IndexCoord idAllocator initialize failed", zap.Error(err))
return err

13
internal/indexcoord/index_coord_test.go
View File

@ -18,14 +18,14 @@ import (
"time"
"github.com/golang/protobuf/proto"
"github.com/stretchr/testify/assert"
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
"github.com/milvus-io/milvus/internal/proto/internalpb"
"go.etcd.io/etcd/clientv3"
"github.com/milvus-io/milvus/internal/types"
"github.com/milvus-io/milvus/internal/proto/commonpb"
"github.com/milvus-io/milvus/internal/proto/indexpb"
"github.com/milvus-io/milvus/internal/types"
"github.com/stretchr/testify/assert"
"github.com/milvus-io/milvus/internal/proto/internalpb"
)
type indexNodeMock struct {
@ -33,14 +33,13 @@ type indexNodeMock struct {
}
func (in *indexNodeMock) CreateIndex(ctx context.Context, req *indexpb.CreateIndexRequest) (*commonpb.Status, error) {
etcdClient, err := clientv3.New(clientv3.Config{Endpoints: Params.EtcdEndpoints})
indexMeta := indexpb.IndexMeta{}
etcdKV, err := etcdkv.NewEtcdKV(Params.EtcdEndpoints, Params.MetaRootPath)
if err != nil {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
}, err
}
indexMeta := indexpb.IndexMeta{}
etcdKV := etcdkv.NewEtcdKV(etcdClient, Params.MetaRootPath)
_, values, versions, err := etcdKV.LoadWithPrefix2(req.MetaPath)
if err != nil {
return &commonpb.Status{

5
internal/indexnode/indexnode.go
View File

@ -19,7 +19,6 @@ import (
"sync/atomic"
"time"
"go.etcd.io/etcd/clientv3"
"go.uber.org/zap"
"github.com/milvus-io/milvus/internal/kv"
@ -90,8 +89,8 @@ func (i *IndexNode) Register() error {
func (i *IndexNode) Init() error {
connectEtcdFn := func() error {
etcdClient, err := clientv3.New(clientv3.Config{Endpoints: Params.EtcdEndpoints})
i.etcdKV = etcdkv.NewEtcdKV(etcdClient, Params.MetaRootPath)
etcdKV, err := etcdkv.NewEtcdKV(Params.EtcdEndpoints, Params.MetaRootPath)
i.etcdKV = etcdKV
return err
}
err := retry.Do(i.loopCtx, connectEtcdFn, retry.Attempts(300))

20
internal/kv/etcd/etcd_kv.go
View File

@ -33,12 +33,22 @@ type EtcdKV struct {
}
// NewEtcdKV creates a new etcd kv.
func NewEtcdKV(client *clientv3.Client, rootPath string) *EtcdKV {
// func NewEtcdKV(etcdEndpints []string, rootPath string) (*EtcdKV, error)
func NewEtcdKV(etcdEndpoints []string, rootPath string) (*EtcdKV, error) {
client, err := clientv3.New(clientv3.Config{
Endpoints: etcdEndpoints,
DialTimeout: 5 * time.Second,
})
if err != nil {
return nil, err
}
kv := &EtcdKV{
client: client,
rootPath: rootPath,
}
return kv
return kv, nil
}
func (kv *EtcdKV) Close() {
@ -139,7 +149,7 @@ func (kv *EtcdKV) MultiLoad(keys []string) ([]string, error) {
return result, nil
}
func (kv *EtcdKV) LoadWithVersion(key string) ([]string, []string, int64, error) {
func (kv *EtcdKV) LoadWithRevision(key string) ([]string, []string, int64, error) {
key = path.Join(kv.rootPath, key)
log.Debug("LoadWithPrefix ", zap.String("prefix", key))
ctx, cancel := context.WithTimeout(context.TODO(), RequestTimeout)
@ -220,7 +230,7 @@ func (kv *EtcdKV) MultiRemove(keys []string) error {
}
func (kv *EtcdKV) MultiSaveAndRemove(saves map[string]string, removals []string) error {
ops := make([]clientv3.Op, 0, len(saves))
ops := make([]clientv3.Op, 0, len(saves)+len(removals))
for key, value := range saves {
ops = append(ops, clientv3.OpPut(path.Join(kv.rootPath, key), value))
}
@ -249,7 +259,7 @@ func (kv *EtcdKV) WatchWithPrefix(key string) clientv3.WatchChan {
return rch
}
func (kv *EtcdKV) WatchWithVersion(key string, revision int64) clientv3.WatchChan {
func (kv *EtcdKV) WatchWithRevision(key string, revision int64) clientv3.WatchChan {
key = path.Join(kv.rootPath, key)
rch := kv.client.Watch(context.Background(), key, clientv3.WithPrefix(), clientv3.WithPrevKV(), clientv3.WithRev(revision))
return rch

798
internal/kv/etcd/etcd_kv_test.go
View File

@ -16,10 +16,12 @@ import (
"strings"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.etcd.io/etcd/clientv3"
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
"github.com/milvus-io/milvus/internal/util/paramtable"
"github.com/stretchr/testify/assert"
"go.etcd.io/etcd/clientv3"
)
var Params paramtable.BaseTable
@ -30,348 +32,460 @@ func TestMain(m *testing.M) {
os.Exit(code)
}
func newEtcdClient() (*clientv3.Client, error) {
func TestEtcdKV_Load(te *testing.T) {
endpoints, err := Params.Load("_EtcdEndpoints")
if err != nil {
panic(err)
}
etcdEndpoints := strings.Split(endpoints, ",")
return clientv3.New(clientv3.Config{Endpoints: etcdEndpoints})
}
func TestEtcdKV_Load(t *testing.T) {
cli, err := newEtcdClient()
assert.Nil(t, err)
rootPath := "/etcd/test/root"
etcdKV := etcdkv.NewEtcdKV(cli, rootPath)
defer etcdKV.Close()
defer etcdKV.RemoveWithPrefix("")
err = etcdKV.Save("abc", "123")
assert.Nil(t, err)
err = etcdKV.Save("abcd", "1234")
assert.Nil(t, err)
val, err := etcdKV.Load("abc")
assert.Nil(t, err)
assert.Equal(t, val, "123")
keys, vals, err := etcdKV.LoadWithPrefix("abc")
assert.Nil(t, err)
assert.Equal(t, len(keys), len(vals))
assert.Equal(t, len(keys), 2)
assert.Equal(t, keys[0], etcdKV.GetPath("abc"))
assert.Equal(t, keys[1], etcdKV.GetPath("abcd"))
assert.Equal(t, vals[0], "123")
assert.Equal(t, vals[1], "1234")
err = etcdKV.Save("key_1", "123")
assert.Nil(t, err)
err = etcdKV.Save("key_2", "456")
assert.Nil(t, err)
err = etcdKV.Save("key_3", "789")
assert.Nil(t, err)
keys = []string{"key_1", "key_100"}
vals, err = etcdKV.MultiLoad(keys)
assert.NotNil(t, err)
assert.Equal(t, len(vals), len(keys))
assert.Equal(t, vals[0], "123")
assert.Empty(t, vals[1])
keys = []string{"key_1", "key_2"}
vals, err = etcdKV.MultiLoad(keys)
assert.Nil(t, err)
assert.Equal(t, len(vals), len(keys))
assert.Equal(t, vals[0], "123")
assert.Equal(t, vals[1], "456")
}
func TestEtcdKV_MultiSave(t *testing.T) {
cli, err := newEtcdClient()
assert.Nil(t, err)
rootPath := "/etcd/test/root"
etcdKV := etcdkv.NewEtcdKV(cli, rootPath)
defer etcdKV.Close()
defer etcdKV.RemoveWithPrefix("")
err = etcdKV.Save("key_1", "111")
assert.Nil(t, err)
kvs := map[string]string{
"key_1": "123",
"key_2": "456",
}
err = etcdKV.MultiSave(kvs)
assert.Nil(t, err)
val, err := etcdKV.Load("key_1")
assert.Nil(t, err)
assert.Equal(t, val, "123")
}
func TestEtcdKV_Remove(t *testing.T) {
cli, err := newEtcdClient()
assert.Nil(t, err)
rootPath := "/etcd/test/root"
etcdKV := etcdkv.NewEtcdKV(cli, rootPath)
defer etcdKV.Close()
defer etcdKV.RemoveWithPrefix("")
err = etcdKV.Save("key_1", "123")
assert.Nil(t, err)
err = etcdKV.Save("key_2", "456")
assert.Nil(t, err)
val, err := etcdKV.Load("key_1")
assert.Nil(t, err)
assert.Equal(t, val, "123")
// delete "key_1"
err = etcdKV.Remove("key_1")
assert.Nil(t, err)
val, err = etcdKV.Load("key_1")
assert.Error(t, err)
assert.Empty(t, val)
val, err = etcdKV.Load("key_2")
assert.Nil(t, err)
assert.Equal(t, val, "456")
keys, vals, err := etcdKV.LoadWithPrefix("key")
assert.Nil(t, err)
assert.Equal(t, len(keys), len(vals))
assert.Equal(t, len(keys), 1)
assert.Equal(t, keys[0], etcdKV.GetPath("key_2"))
assert.Equal(t, vals[0], "456")
// MultiRemove
err = etcdKV.Save("key_1", "111")
assert.Nil(t, err)
kvs := map[string]string{
"key_1": "123",
"key_2": "456",
"key_3": "789",
"key_4": "012",
}
err = etcdKV.MultiSave(kvs)
assert.Nil(t, err)
val, err = etcdKV.Load("key_1")
assert.Nil(t, err)
assert.Equal(t, val, "123")
val, err = etcdKV.Load("key_3")
assert.Nil(t, err)
assert.Equal(t, val, "789")
keys = []string{"key_1", "key_2", "key_3"}
err = etcdKV.MultiRemove(keys)
assert.Nil(t, err)
val, err = etcdKV.Load("key_1")
assert.Error(t, err)
assert.Empty(t, val)
}
func TestEtcdKV_MultiSaveAndRemove(t *testing.T) {
cli, err := newEtcdClient()
assert.Nil(t, err)
rootPath := "/etcd/test/root"
etcdKV := etcdkv.NewEtcdKV(cli, rootPath)
defer etcdKV.Close()
defer etcdKV.RemoveWithPrefix("")
err = etcdKV.Save("key_1", "123")
assert.Nil(t, err)
err = etcdKV.Save("key_2", "456")
assert.Nil(t, err)
err = etcdKV.Save("key_3", "789")
assert.Nil(t, err)
kvs := map[string]string{
"key_1": "111",
"key_2": "444",
}
keys := []string{"key_3"}
err = etcdKV.MultiSaveAndRemove(kvs, keys)
assert.Nil(t, err)
val, err := etcdKV.Load("key_1")
assert.Nil(t, err)
assert.Equal(t, val, "111")
val, err = etcdKV.Load("key_2")
assert.Nil(t, err)
assert.Equal(t, val, "444")
val, err = etcdKV.Load("key_3")
assert.Error(t, err)
assert.Empty(t, val)
}
func TestEtcdKV_MultiRemoveWithPrefix(t *testing.T) {
cli, err := newEtcdClient()
assert.Nil(t, err)
rootPath := "/etcd/test/root"
etcdKV := etcdkv.NewEtcdKV(cli, rootPath)
defer etcdKV.Close()
defer etcdKV.RemoveWithPrefix("")
err = etcdKV.Save("x/abc/1", "1")
assert.Nil(t, err)
err = etcdKV.Save("x/abc/2", "2")
assert.Nil(t, err)
err = etcdKV.Save("x/def/1", "10")
assert.Nil(t, err)
err = etcdKV.Save("x/def/2", "20")
assert.Nil(t, err)
err = etcdKV.MultiRemoveWithPrefix([]string{"x/abc", "x/def", "not-exist"})
assert.Nil(t, err)
k, v, err := etcdKV.LoadWithPrefix("x")
assert.Nil(t, err)
assert.Zero(t, len(k))
assert.Zero(t, len(v))
}
func TestEtcdKV_MultiSaveAndRemoveWithPrefix(t *testing.T) {
cli, err := newEtcdClient()
assert.Nil(t, err)
rootPath := "/etcd/test/root"
etcdKV := etcdkv.NewEtcdKV(cli, rootPath)
defer etcdKV.Close()
defer etcdKV.RemoveWithPrefix("")
err = etcdKV.Save("x/abc/1", "1")
assert.Nil(t, err)
err = etcdKV.Save("x/abc/2", "2")
assert.Nil(t, err)
err = etcdKV.Save("x/def/1", "10")
assert.Nil(t, err)
err = etcdKV.Save("x/def/2", "20")
assert.Nil(t, err)
err = etcdKV.MultiSaveAndRemoveWithPrefix(map[string]string{"y/k1": "v1", "y/k2": "v2"}, []string{"x/abc", "x/def", "not-exist"})
assert.Nil(t, err)
k, v, err := etcdKV.LoadWithPrefix("x")
assert.Nil(t, err)
assert.Zero(t, len(k))
assert.Zero(t, len(v))
k, v, err = etcdKV.LoadWithPrefix("y")
assert.Nil(t, err)
assert.Equal(t, len(k), len(v))
assert.Equal(t, len(k), 2)
assert.Equal(t, k[0], rootPath+"/y/k1")
assert.Equal(t, k[1], rootPath+"/y/k2")
assert.Equal(t, v[0], "v1")
assert.Equal(t, v[1], "v2")
}
func TestEtcdKV_Watch(t *testing.T) {
cli, err := newEtcdClient()
assert.Nil(t, err)
rootPath := "/etcd/test/root"
etcdKV := etcdkv.NewEtcdKV(cli, rootPath)
defer etcdKV.Close()
defer etcdKV.RemoveWithPrefix("")
ch := etcdKV.Watch("x/abc/1")
resp := <-ch
assert.True(t, resp.Created)
}
func TestEtcdKV_WatchPrefix(t *testing.T) {
cli, err := newEtcdClient()
assert.Nil(t, err)
rootPath := "/etcd/test/root"
etcdKV := etcdkv.NewEtcdKV(cli, rootPath)
defer etcdKV.Close()
defer etcdKV.RemoveWithPrefix("")
ch := etcdKV.WatchWithPrefix("x")
resp := <-ch
assert.True(t, resp.Created)
}
func TestEtcdKV_LoadWithVersion(t *testing.T) {
cli, err := newEtcdClient()
assert.Nil(t, err)
rootPath := "/etcd/test/root"
etcdKV := etcdkv.NewEtcdKV(cli, rootPath)
defer etcdKV.Close()
defer etcdKV.RemoveWithPrefix("")
key := "test-load-version"
err = etcdKV.Save(key, "v1")
assert.Nil(t, err)
_, values, _, err := etcdKV.LoadWithVersion(key)
assert.Nil(t, err)
assert.Equal(t, "v1", values[0])
}
func TestEtcdKV_WatchWithVersion(t *testing.T) {
cli, err := newEtcdClient()
assert.Nil(t, err)
rootPath := "/etcd/test/root"
etcdKV := etcdkv.NewEtcdKV(cli, rootPath)
defer etcdKV.Close()
defer etcdKV.RemoveWithPrefix("")
key := "test-version"
err = etcdKV.Save(key, "v")
assert.Nil(t, err)
_, _, revision, _ := etcdKV.LoadWithVersion(key)
ch := etcdKV.WatchWithVersion(key, revision+1)
err = etcdKV.Save(key, "v2")
assert.Nil(t, err)
resp2 := <-ch
assert.Equal(t, 1, len(resp2.Events))
assert.Equal(t, "v2", string(resp2.Events[0].Kv.Value))
assert.Equal(t, revision+1, resp2.Header.Revision)
}
func TestEtcdKV_CompareAndSwap(t *testing.T) {
cli, err := newEtcdClient()
assert.Nil(t, err)
rootPath := "/etcd/test/root"
etcdKV := etcdkv.NewEtcdKV(cli, rootPath)
defer etcdKV.Close()
defer etcdKV.RemoveWithPrefix("")
err = etcdKV.CompareVersionAndSwap("a/b/c", 0, "1")
assert.Nil(t, err)
value, err := etcdKV.Load("a/b/c")
assert.Nil(t, err)
assert.Equal(t, value, "1")
err = etcdKV.CompareVersionAndSwap("a/b/c", 0, "1")
assert.NotNil(t, err)
err = etcdKV.CompareValueAndSwap("a/b/c", "1", "2")
assert.Nil(t, err)
err = etcdKV.CompareValueAndSwap("a/b/c", "1", "2")
assert.NotNil(t, err)
etcdEndPoints := strings.Split(endpoints, ",")
te.Run("EtcdKV SaveAndLoad", func(t *testing.T) {
rootPath := "/etcd/test/root/saveandload"
etcdKV, err := etcdkv.NewEtcdKV(etcdEndPoints, rootPath)
require.NoError(t, err)
err = etcdKV.RemoveWithPrefix("")
require.NoError(t, err)
defer etcdKV.Close()
defer etcdKV.RemoveWithPrefix("")
saveAndLoadTests := []struct {
key string
value string
}{
{"test1", "value1"},
{"test2", "value2"},
{"test1/a", "value_a"},
{"test1/b", "value_b"},
}
for i, test := range saveAndLoadTests {
if i < 4 {
err = etcdKV.Save(test.key, test.value)
assert.NoError(t, err)
}
val, err := etcdKV.Load(test.key)
assert.NoError(t, err)
assert.Equal(t, test.value, val)
}
invalidLoadTests := []struct {
invalidKey string
}{
{"t"},
{"a"},
{"test1a"},
}
for _, test := range invalidLoadTests {
val, err := etcdKV.Load(test.invalidKey)
assert.Error(t, err)
assert.Zero(t, val)
}
loadPrefixTests := []struct {
prefix string
expectedKeys []string
expectedValues []string
expectedError error
}{
{"test", []string{
etcdKV.GetPath("test1"),
etcdKV.GetPath("test2"),
etcdKV.GetPath("test1/a"),
etcdKV.GetPath("test1/b")}, []string{"value1", "value2", "value_a", "value_b"}, nil},
{"test1", []string{
etcdKV.GetPath("test1"),
etcdKV.GetPath("test1/a"),
etcdKV.GetPath("test1/b")}, []string{"value1", "value_a", "value_b"}, nil},
{"test2", []string{etcdKV.GetPath("test2")}, []string{"value2"}, nil},
{"", []string{
etcdKV.GetPath("test1"),
etcdKV.GetPath("test2"),
etcdKV.GetPath("test1/a"),
etcdKV.GetPath("test1/b")}, []string{"value1", "value2", "value_a", "value_b"}, nil},
{"test1/a", []string{etcdKV.GetPath("test1/a")}, []string{"value_a"}, nil},
{"a", []string{}, []string{}, nil},
{"root", []string{}, []string{}, nil},
{"/etcd/test/root", []string{}, []string{}, nil},
}
for _, test := range loadPrefixTests {
actualKeys, actualValues, err := etcdKV.LoadWithPrefix(test.prefix)
assert.ElementsMatch(t, test.expectedKeys, actualKeys)
assert.ElementsMatch(t, test.expectedValues, actualValues)
assert.Equal(t, test.expectedError, err)
actualKeys, actualValues, versions, err := etcdKV.LoadWithPrefix2(test.prefix)
assert.ElementsMatch(t, test.expectedKeys, actualKeys)
assert.ElementsMatch(t, test.expectedValues, actualValues)
assert.NotZero(t, versions)
assert.Equal(t, test.expectedError, err)
}
removeTests := []struct {
validKey string
invalidKey string
}{
{"test1", "abc"},
{"test1/a", "test1/lskfjal"},
{"test1/b", "test1/b"},
{"test2", "-"},
}
for _, test := range removeTests {
err = etcdKV.Remove(test.validKey)
assert.NoError(t, err)
_, err = etcdKV.Load(test.validKey)
assert.Error(t, err)
err = etcdKV.Remove(test.validKey)
assert.NoError(t, err)
err = etcdKV.Remove(test.invalidKey)
assert.NoError(t, err)
}
})
te.Run("EtcdKV LoadWithRevision", func(t *testing.T) {
rootPath := "/etcd/test/root/LoadWithRevision"
etcdKV, err := etcdkv.NewEtcdKV(etcdEndPoints, rootPath)
assert.Nil(t, err)
defer etcdKV.Close()
defer etcdKV.RemoveWithPrefix("")
prepareKV := []struct {
inKey string
inValue string
}{
{"a", "a_version1"},
{"b", "b_version2"},
{"a", "a_version3"},
{"c", "c_version4"},
{"a/suba", "a_version5"},
}
for _, test := range prepareKV {
err = etcdKV.Save(test.inKey, test.inValue)
require.NoError(t, err)
}
loadWithRevisionTests := []struct {
inKey string
expectedKeyNo int
expectedValues []string
}{
{"a", 2, []string{"a_version3", "a_version5"}},
{"b", 1, []string{"b_version2"}},
{"c", 1, []string{"c_version4"}},
}
for _, test := range loadWithRevisionTests {
keys, values, revision, err := etcdKV.LoadWithRevision(test.inKey)
assert.NoError(t, err)
assert.Equal(t, test.expectedKeyNo, len(keys))
assert.ElementsMatch(t, test.expectedValues, values)
assert.NotZero(t, revision)
}
})
te.Run("EtcdKV MultiSaveAndMultiLoad", func(t *testing.T) {
rootPath := "/etcd/test/root/multi_save_and_multi_load"
etcdKV, err := etcdkv.NewEtcdKV(etcdEndPoints, rootPath)
assert.Nil(t, err)
defer etcdKV.Close()
defer etcdKV.RemoveWithPrefix("")
multiSaveTests := map[string]string{
"key_1": "value_1",
"key_2": "value_2",
"key_3/a": "value_3a",
"multikey_1": "multivalue_1",
"multikey_2": "multivalue_2",
"_": "other",
}
err = etcdKV.MultiSave(multiSaveTests)
assert.NoError(t, err)
for k, v := range multiSaveTests {
actualV, err := etcdKV.Load(k)
assert.NoError(t, err)
assert.Equal(t, v, actualV)
}
multiLoadTests := []struct {
inputKeys []string
expectedValues []string
}{
{[]string{"key_1"}, []string{"value_1"}},
{[]string{"key_1", "key_2", "key_3/a"}, []string{"value_1", "value_2", "value_3a"}},
{[]string{"multikey_1", "multikey_2"}, []string{"multivalue_1", "multivalue_2"}},
{[]string{"_"}, []string{"other"}},
}
for _, test := range multiLoadTests {
vs, err := etcdKV.MultiLoad(test.inputKeys)
assert.NoError(t, err)
assert.Equal(t, test.expectedValues, vs)
}
invalidMultiLoad := []struct {
invalidKeys []string
expectedValues []string
}{
{[]string{"a", "key_1"}, []string{"", "value_1"}},
{[]string{".....", "key_1"}, []string{"", "value_1"}},
{[]string{"*********"}, []string{""}},
{[]string{"key_1", "1"}, []string{"value_1", ""}},
}
for _, test := range invalidMultiLoad {
vs, err := etcdKV.MultiLoad(test.invalidKeys)
assert.Error(t, err)
assert.Equal(t, test.expectedValues, vs)
}
removeWithPrefixTests := []string{
"key_1",
"multi",
}
for _, k := range removeWithPrefixTests {
err = etcdKV.RemoveWithPrefix(k)
assert.NoError(t, err)
ks, vs, err := etcdKV.LoadWithPrefix(k)
assert.Empty(t, ks)
assert.Empty(t, vs)
assert.NoError(t, err)
}
multiRemoveTests := []string{
"key_2",
"key_3/a",
"multikey_2",
"_",
}
err = etcdKV.MultiRemove(multiRemoveTests)
assert.NoError(t, err)
ks, vs, err := etcdKV.LoadWithPrefix("")
assert.NoError(t, err)
assert.Empty(t, ks)
assert.Empty(t, vs)
multiSaveAndRemoveTests := []struct {
multiSaves map[string]string
multiRemoves []string
}{
{map[string]string{"key_1": "value_1"}, []string{}},
{map[string]string{"key_2": "value_2"}, []string{"key_1"}},
{map[string]string{"key_3/a": "value_3a"}, []string{"key_2"}},
{map[string]string{"multikey_1": "multivalue_1"}, []string{}},
{map[string]string{"multikey_2": "multivalue_2"}, []string{"multikey_1", "key_3/a"}},
{make(map[string]string), []string{"multikey_2"}},
}
for _, test := range multiSaveAndRemoveTests {
err = etcdKV.MultiSaveAndRemove(test.multiSaves, test.multiRemoves)
assert.NoError(t, err)
}
ks, vs, err = etcdKV.LoadWithPrefix("")
assert.NoError(t, err)
assert.Empty(t, ks)
assert.Empty(t, vs)
})
te.Run("EtcdKV MultiRemoveWithPrefix", func(t *testing.T) {
rootPath := "/etcd/test/root/multi_remove_with_prefix"
etcdKV, err := etcdkv.NewEtcdKV(etcdEndPoints, rootPath)
require.NoError(t, err)
defer etcdKV.Close()
defer etcdKV.RemoveWithPrefix("")
prepareTests := map[string]string{
"x/abc/1": "1",
"x/abc/2": "2",
"x/def/1": "10",
"x/def/2": "20",
"x/den/1": "100",
"x/den/2": "200",
}
err = etcdKV.MultiSave(prepareTests)
require.NoError(t, err)
multiRemoveWithPrefixTests := []struct {
prefix []string
testKey string
expectedValue string
}{
{[]string{"x/abc"}, "x/abc/1", ""},
{[]string{}, "x/abc/2", ""},
{[]string{}, "x/def/1", "10"},
{[]string{}, "x/def/2", "20"},
{[]string{}, "x/den/1", "100"},
{[]string{}, "x/den/2", "200"},
{[]string{}, "not-exist", ""},
{[]string{"x/def", "x/den"}, "x/def/1", ""},
{[]string{}, "x/def/1", ""},
{[]string{}, "x/def/2", ""},
{[]string{}, "x/den/1", ""},
{[]string{}, "x/den/2", ""},
{[]string{}, "not-exist", ""},
}
for _, test := range multiRemoveWithPrefixTests {
if len(test.prefix) > 0 {
err = etcdKV.MultiRemoveWithPrefix(test.prefix)
assert.NoError(t, err)
}
v, _ := etcdKV.Load(test.testKey)
assert.Equal(t, test.expectedValue, v)
}
k, v, err := etcdKV.LoadWithPrefix("/")
assert.NoError(t, err)
assert.Zero(t, len(k))
assert.Zero(t, len(v))
// MultiSaveAndRemoveWithPrefix
err = etcdKV.MultiSave(prepareTests)
require.NoError(t, err)
multiSaveAndRemoveWithPrefixTests := []struct {
multiSave map[string]string
prefix []string
loadPrefix string
lengthBeforeRemove int
lengthAfterRemove int
}{
{map[string]string{}, []string{"x/abc", "x/def", "x/den"}, "x", 6, 0},
{map[string]string{"y/a": "vvv", "y/b": "vvv"}, []string{}, "y", 0, 2},
{map[string]string{"y/c": "vvv"}, []string{}, "y", 2, 3},
{map[string]string{"p/a": "vvv"}, []string{"y/a", "y"}, "y", 3, 0},
{map[string]string{}, []string{"p"}, "p", 1, 0},
}
for _, test := range multiSaveAndRemoveWithPrefixTests {
k, _, err = etcdKV.LoadWithPrefix(test.loadPrefix)
assert.NoError(t, err)
assert.Equal(t, test.lengthBeforeRemove, len(k))
err = etcdKV.MultiSaveAndRemoveWithPrefix(test.multiSave, test.prefix)
assert.NoError(t, err)
k, _, err = etcdKV.LoadWithPrefix(test.loadPrefix)
assert.NoError(t, err)
assert.Equal(t, test.lengthAfterRemove, len(k))
}
})
te.Run("EtcdKV Watch", func(t *testing.T) {
rootPath := "/etcd/test/root/watch"
etcdKV, err := etcdkv.NewEtcdKV(etcdEndPoints, rootPath)
assert.Nil(t, err)
defer etcdKV.Close()
defer etcdKV.RemoveWithPrefix("")
ch := etcdKV.Watch("x")
resp := <-ch
assert.True(t, resp.Created)
ch = etcdKV.WatchWithPrefix("x")
resp = <-ch
assert.True(t, resp.Created)
})
te.Run("Etcd Revision", func(t *testing.T) {
rootPath := "/etcd/test/root/watch"
etcdKV, err := etcdkv.NewEtcdKV(etcdEndPoints, rootPath)
assert.Nil(t, err)
defer etcdKV.Close()
defer etcdKV.RemoveWithPrefix("")
revisionTests := []struct {
inKey string
fistValue string
secondValue string
}{
{"a", "v1", "v11"},
{"y", "v2", "v22"},
{"z", "v3", "v33"},
}
for _, test := range revisionTests {
err = etcdKV.Save(test.inKey, test.fistValue)
require.NoError(t, err)
_, _, revision, _ := etcdKV.LoadWithRevision(test.inKey)
ch := etcdKV.WatchWithRevision(test.inKey, revision+1)
err = etcdKV.Save(test.inKey, test.secondValue)
require.NoError(t, err)
resp := <-ch
assert.Equal(t, 1, len(resp.Events))
assert.Equal(t, test.secondValue, string(resp.Events[0].Kv.Value))
assert.Equal(t, revision+1, resp.Header.Revision)
}
err = etcdKV.CompareVersionAndSwap("a/b/c", 0, "1")
assert.NoError(t, err)
value, err := etcdKV.Load("a/b/c")
assert.NoError(t, err)
assert.Equal(t, value, "1")
err = etcdKV.CompareVersionAndSwap("a/b/c", 0, "1")
assert.Error(t, err)
err = etcdKV.CompareValueAndSwap("a/b/c", "1", "2")
assert.NoError(t, err)
err = etcdKV.CompareValueAndSwap("a/b/c", "1", "2")
assert.Error(t, err)
})
te.Run("Etcd Lease", func(t *testing.T) {
rootPath := "/etcd/test/root/lease"
etcdKV, err := etcdkv.NewEtcdKV(etcdEndPoints, rootPath)
assert.Nil(t, err)
defer etcdKV.Close()
defer etcdKV.RemoveWithPrefix("")
leaseID, err := etcdKV.Grant(10)
assert.NoError(t, err)
etcdKV.KeepAlive(leaseID)
tests := map[string]string{
"a/b": "v1",
"a/b/c": "v2",
"x": "v3",
}
for k, v := range tests {
err = etcdKV.SaveWithLease(k, v, leaseID)
assert.NoError(t, err)
err = etcdKV.SaveWithLease(k, v, clientv3.LeaseID(999))
assert.Error(t, err)
}
})
}

4
internal/msgstream/mq_msgstream_test.go
View File

@ -22,7 +22,6 @@ import (
"github.com/apache/pulsar-client-go/pulsar"
"github.com/stretchr/testify/assert"
"go.etcd.io/etcd/clientv3"
"github.com/milvus-io/milvus/internal/allocator"
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
@ -1028,11 +1027,10 @@ func initRmq(name string) *etcdkv.EtcdKV {
endpoints = "localhost:2379"
}
etcdEndpoints := strings.Split(endpoints, ",")
cli, err := clientv3.New(clientv3.Config{Endpoints: etcdEndpoints})
etcdKV, err := etcdkv.NewEtcdKV(etcdEndpoints, "/etcd/test/root")
if err != nil {
log.Fatalf("New clientv3 error = %v", err)
}
etcdKV := etcdkv.NewEtcdKV(cli, "/etcd/test/root")
idAllocator := allocator.NewGlobalIDAllocator("dummy", etcdKV)
_ = idAllocator.Initialize()

4
internal/querycoord/meta_test.go
View File

@ -16,15 +16,13 @@ import (
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.etcd.io/etcd/clientv3"
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
)
func TestReplica_Release(t *testing.T) {
etcdClient, err := clientv3.New(clientv3.Config{Endpoints: Params.EtcdEndpoints})
etcdKV, err := etcdkv.NewEtcdKV(Params.EtcdEndpoints, Params.MetaRootPath)
assert.Nil(t, err)
etcdKV := etcdkv.NewEtcdKV(etcdClient, Params.MetaRootPath)
meta, err := newMeta(etcdKV)
assert.Nil(t, err)
err = meta.addCollection(1, nil)

4
internal/querycoord/query_coord.go
View File

@ -22,7 +22,6 @@ import (
"github.com/coreos/etcd/mvcc/mvccpb"
"github.com/golang/protobuf/proto"
"go.etcd.io/etcd/clientv3"
"go.uber.org/zap"
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
@ -79,11 +78,10 @@ func (qc *QueryCoord) Register() error {
func (qc *QueryCoord) Init() error {
connectEtcdFn := func() error {
etcdClient, err := clientv3.New(clientv3.Config{Endpoints: Params.EtcdEndpoints})
etcdKV, err := etcdkv.NewEtcdKV(Params.EtcdEndpoints, Params.MetaRootPath)
if err != nil {
return err
}
etcdKV := etcdkv.NewEtcdKV(etcdClient, Params.MetaRootPath)
qc.kvClient = etcdKV
return nil
}

8
internal/querycoord/task_scheduler.go
View File

@ -148,7 +148,11 @@ func NewTaskScheduler(ctx context.Context, meta Meta, cluster *queryNodeCluster,
dataCoord: dataCoord,
}
s.triggerTaskQueue = NewTaskQueue()
idAllocator := allocator.NewGlobalIDAllocator("idTimestamp", tsoutil.NewTSOKVBase(Params.EtcdEndpoints, Params.KvRootPath, "queryCoordTaskID"))
etcdKV, err := tsoutil.NewTSOKVBase(Params.EtcdEndpoints, Params.KvRootPath, "queryCoordTaskID")
if err != nil {
return nil, err
}
idAllocator := allocator.NewGlobalIDAllocator("idTimestamp", etcdKV)
if err := idAllocator.Initialize(); err != nil {
log.Debug("query coordinator idAllocator initialize failed", zap.Error(err))
return nil, err
@ -156,7 +160,7 @@ func NewTaskScheduler(ctx context.Context, meta Meta, cluster *queryNodeCluster,
s.taskIDAllocator = func() (UniqueID, error) {
return idAllocator.AllocOne()
}
err := s.reloadFromKV()
err = s.reloadFromKV()
if err != nil {
return nil, err
}

5
internal/querycoord/task_scheduler_test.go
View File

@ -5,8 +5,6 @@ import (
"testing"
"time"
"go.etcd.io/etcd/clientv3"
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
"github.com/milvus-io/milvus/internal/log"
"github.com/milvus-io/milvus/internal/proto/commonpb"
@ -160,9 +158,8 @@ func TestWatchQueryChannel_ClearEtcdInfoAfterAssignedNodeDown(t *testing.T) {
}
func TestUnMarshalTask_LoadCollection(t *testing.T) {
etcdClient, err := clientv3.New(clientv3.Config{Endpoints: Params.EtcdEndpoints})
kv, err := etcdkv.NewEtcdKV(Params.EtcdEndpoints, Params.MetaRootPath)
assert.Nil(t, err)
kv := etcdkv.NewEtcdKV(etcdClient, Params.MetaRootPath)
loadTask := &LoadCollectionTask{
LoadCollectionRequest: &querypb.LoadCollectionRequest{

4
internal/querynode/query_node.go
View File

@ -30,7 +30,6 @@ import (
"strconv"
"sync/atomic"
"go.etcd.io/etcd/clientv3"
"go.uber.org/zap"
"github.com/milvus-io/milvus/internal/kv"
@ -102,11 +101,10 @@ func (node *QueryNode) Register() error {
func (node *QueryNode) Init() error {
//ctx := context.Background()
connectEtcdFn := func() error {
etcdClient, err := clientv3.New(clientv3.Config{Endpoints: Params.EtcdEndpoints})
etcdKV, err := etcdkv.NewEtcdKV(Params.EtcdEndpoints, Params.MetaRootPath)
if err != nil {
return err
}
etcdKV := etcdkv.NewEtcdKV(etcdClient, Params.MetaRootPath)
node.etcdKV = etcdKV
return err
}

4
internal/querynode/query_node_test.go
View File

@ -20,7 +20,6 @@ import (
"time"
"github.com/stretchr/testify/assert"
"go.etcd.io/etcd/clientv3"
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
"github.com/milvus-io/milvus/internal/msgstream"
@ -164,11 +163,10 @@ func newQueryNodeMock() *QueryNode {
}()
}
etcdClient, err := clientv3.New(clientv3.Config{Endpoints: Params.EtcdEndpoints})
etcdKV, err := etcdkv.NewEtcdKV(Params.EtcdEndpoints, Params.MetaRootPath)
if err != nil {
panic(err)
}
etcdKV := etcdkv.NewEtcdKV(etcdClient, Params.MetaRootPath)
msFactory, err := newMessageStreamFactory()
if err != nil {

23
internal/rootcoord/root_coord.go
View File

@ -902,7 +902,10 @@ func (c *Core) Init() error {
if c.MetaTable, initError = NewMetaTable(ms); initError != nil {
return initError
}
c.kvBase = etcdkv.NewEtcdKV(c.etcdCli, Params.KvRootPath)
if c.kvBase, initError = etcdkv.NewEtcdKV(Params.EtcdEndpoints, Params.KvRootPath); initError != nil {
return initError
}
return nil
}
log.Debug("RootCoord, Connect to Etcd")
@ -912,7 +915,11 @@ func (c *Core) Init() error {
}
log.Debug("RootCoord, Set TSO and ID Allocator")
idAllocator := allocator.NewGlobalIDAllocator("idTimestamp", tsoutil.NewTSOKVBase(Params.EtcdEndpoints, Params.KvRootPath, "gid"))
kv, initError := tsoutil.NewTSOKVBase(Params.EtcdEndpoints, Params.KvRootPath, "gid")
if initError != nil {
return
}
idAllocator := allocator.NewGlobalIDAllocator("idTimestamp", kv)
if initError = idAllocator.Initialize(); initError != nil {
return
}
@ -923,7 +930,11 @@ func (c *Core) Init() error {
return idAllocator.UpdateID()
}
tsoAllocator := tso.NewGlobalTSOAllocator("timestamp", tsoutil.NewTSOKVBase(Params.EtcdEndpoints, Params.KvRootPath, "tso"))
kv, initError = tsoutil.NewTSOKVBase(Params.EtcdEndpoints, Params.KvRootPath, "tso")
if initError != nil {
return
}
tsoAllocator := tso.NewGlobalTSOAllocator("timestamp", kv)
if initError = tsoAllocator.Initialize(); initError != nil {
return
}
@ -957,10 +968,16 @@ func (c *Core) Init() error {
c.chanTimeTick.GetProxy,
c.proxyClientManager.GetProxyClients,
)
if initError != nil {
return
}
c.proxyManager.AddSession(c.chanTimeTick.AddProxy, c.proxyClientManager.AddProxyClient)
c.proxyManager.DelSession(c.chanTimeTick.DelProxy, c.proxyClientManager.DelProxyClient)
initError = c.setMsgStreams()
if initError != nil {
return
}
})
if initError == nil {
log.Debug(typeutil.RootCoordRole, zap.String("State Code", internalpb.StateCode_name[int32(internalpb.StateCode_Initializing)]))

4
internal/tso/global_allocator_test.go
View File

@ -29,7 +29,9 @@ func TestGlobalTSOAllocator_All(t *testing.T) {
endpoints = "localhost:2379"
}
etcdEndpoints := strings.Split(endpoints, ",")
gTestTsoAllocator = NewGlobalTSOAllocator("timestamp", tsoutil.NewTSOKVBase(etcdEndpoints, "/test/root/kv", "tsoTest"))
etcdKV, err := tsoutil.NewTSOKVBase(etcdEndpoints, "/test/root/kv", "tsoTest")
assert.NoError(t, err)
gTestTsoAllocator = NewGlobalTSOAllocator("timestamp", etcdKV)
t.Run("Initialize", func(t *testing.T) {
err := gTestTsoAllocator.Initialize()
assert.Nil(t, err)

25
internal/util/rocksmq/server/rocksmq/rocksmq_impl_test.go
View File

@ -24,7 +24,6 @@ import (
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
"github.com/stretchr/testify/assert"
"go.etcd.io/etcd/clientv3"
)
func TestFixChannelName(t *testing.T) {
@ -34,19 +33,19 @@ func TestFixChannelName(t *testing.T) {
assert.Equal(t, len(fixName), FixedChannelNameLen)
}
func newEtcdClient() (*clientv3.Client, error) {
func etcdEndpoints() []string {
endpoints := os.Getenv("ETCD_ENDPOINTS")
if endpoints == "" {
endpoints = "localhost:2379"
}
etcdEndpoints := strings.Split(endpoints, ",")
return clientv3.New(clientv3.Config{Endpoints: etcdEndpoints})
return etcdEndpoints
}
func TestRocksMQ(t *testing.T) {
cli, err := newEtcdClient()
ep := etcdEndpoints()
etcdKV, err := etcdkv.NewEtcdKV(ep, "/etcd/test/root")
assert.Nil(t, err)
etcdKV := etcdkv.NewEtcdKV(cli, "/etcd/test/root")
defer etcdKV.Close()
idAllocator := allocator.NewGlobalIDAllocator("dummy", etcdKV)
_ = idAllocator.Initialize()
@ -97,9 +96,9 @@ func TestRocksMQ(t *testing.T) {
}
func TestRocksMQ_Loop(t *testing.T) {
cli, err := newEtcdClient()
ep := etcdEndpoints()
etcdKV, err := etcdkv.NewEtcdKV(ep, "/etcd/test/root")
assert.Nil(t, err)
etcdKV := etcdkv.NewEtcdKV(cli, "/etcd/test/root")
defer etcdKV.Close()
idAllocator := allocator.NewGlobalIDAllocator("dummy", etcdKV)
_ = idAllocator.Initialize()
@ -161,9 +160,9 @@ func TestRocksMQ_Loop(t *testing.T) {
}
func TestRocksMQ_Goroutines(t *testing.T) {
cli, err := newEtcdClient()
ep := etcdEndpoints()
etcdKV, err := etcdkv.NewEtcdKV(ep, "/etcd/test/root")
assert.Nil(t, err)
etcdKV := etcdkv.NewEtcdKV(cli, "/etcd/test/root")
defer etcdKV.Close()
idAllocator := allocator.NewGlobalIDAllocator("dummy", etcdKV)
_ = idAllocator.Initialize()
@ -228,9 +227,9 @@ func TestRocksMQ_Goroutines(t *testing.T) {
Consume: 90000 message / s
*/
func TestRocksMQ_Throughout(t *testing.T) {
cli, err := newEtcdClient()
ep := etcdEndpoints()
etcdKV, err := etcdkv.NewEtcdKV(ep, "/etcd/test/root")
assert.Nil(t, err)
etcdKV := etcdkv.NewEtcdKV(cli, "/etcd/test/root")
defer etcdKV.Close()
idAllocator := allocator.NewGlobalIDAllocator("dummy", etcdKV)
_ = idAllocator.Initialize()
@ -278,9 +277,9 @@ func TestRocksMQ_Throughout(t *testing.T) {
}
func TestRocksMQ_MultiChan(t *testing.T) {
cli, err := newEtcdClient()
ep := etcdEndpoints()
etcdKV, err := etcdkv.NewEtcdKV(ep, "/etcd/test/root")
assert.Nil(t, err)
etcdKV := etcdkv.NewEtcdKV(cli, "/etcd/test/root")
defer etcdKV.Close()
idAllocator := allocator.NewGlobalIDAllocator("dummy", etcdKV)
_ = idAllocator.Initialize()

28
internal/util/sessionutil/session_util_test.go
View File

@ -12,7 +12,6 @@ import (
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
"github.com/milvus-io/milvus/internal/util/paramtable"
"github.com/stretchr/testify/assert"
"go.etcd.io/etcd/clientv3"
"golang.org/x/net/context"
)
@ -29,11 +28,10 @@ func TestGetServerIDConcurrently(t *testing.T) {
}
etcdEndpoints := strings.Split(endpoints, ",")
cli, err := clientv3.New(clientv3.Config{Endpoints: etcdEndpoints})
assert.Nil(t, err)
etcdKV := etcdkv.NewEtcdKV(cli, "")
_, err = cli.Delete(ctx, metaRoot, clientv3.WithPrefix())
assert.Nil(t, err)
etcdKV, err := etcdkv.NewEtcdKV(etcdEndpoints, metaRoot)
assert.NoError(t, err)
err = etcdKV.RemoveWithPrefix("")
assert.NoError(t, err)
defer etcdKV.Close()
defer etcdKV.RemoveWithPrefix("")
@ -72,14 +70,13 @@ func TestInit(t *testing.T) {
if err != nil {
panic(err)
}
metaRoot := fmt.Sprintf("%d/%s", rand.Int(), DefaultServiceRoot)
etcdEndpoints := strings.Split(endpoints, ",")
cli, err := clientv3.New(clientv3.Config{Endpoints: etcdEndpoints})
assert.Nil(t, err)
etcdKV := etcdkv.NewEtcdKV(cli, "")
metaRoot := fmt.Sprintf("%d/%s", rand.Int(), DefaultServiceRoot)
_, err = cli.Delete(ctx, metaRoot, clientv3.WithPrefix())
assert.Nil(t, err)
etcdKV, err := etcdkv.NewEtcdKV(etcdEndpoints, metaRoot)
assert.NoError(t, err)
err = etcdKV.RemoveWithPrefix("")
assert.NoError(t, err)
defer etcdKV.Close()
defer etcdKV.RemoveWithPrefix("")
@ -103,12 +100,9 @@ func TestUpdateSessions(t *testing.T) {
}
etcdEndpoints := strings.Split(endpoints, ",")
cli, err := clientv3.New(clientv3.Config{Endpoints: etcdEndpoints})
assert.Nil(t, err)
etcdKV := etcdkv.NewEtcdKV(cli, "")
metaRoot := fmt.Sprintf("%d/%s", rand.Int(), DefaultServiceRoot)
_, err = cli.Delete(ctx, metaRoot, clientv3.WithPrefix())
assert.Nil(t, err)
etcdKV, err := etcdkv.NewEtcdKV(etcdEndpoints, "")
assert.NoError(t, err)
defer etcdKV.Close()
defer etcdKV.RemoveWithPrefix("")

9
internal/util/tsoutil/tso.go
View File

@ -16,7 +16,6 @@ import (
"time"
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
"go.etcd.io/etcd/clientv3"
)
const (
@ -51,10 +50,6 @@ func Mod24H(ts uint64) uint64 {
return (physical << logicalBits) | logical
}
func NewTSOKVBase(etcdEndpoints []string, tsoRoot, subPath string) *etcdkv.EtcdKV {
client, _ := clientv3.New(clientv3.Config{
Endpoints: etcdEndpoints,
DialTimeout: 5 * time.Second,
})
return etcdkv.NewEtcdKV(client, path.Join(tsoRoot, subPath))
func NewTSOKVBase(etcdEndpoints []string, tsoRoot, subPath string) (*etcdkv.EtcdKV, error) {
return etcdkv.NewEtcdKV(etcdEndpoints, path.Join(tsoRoot, subPath))
}