2021-10-25 20:23:16 +08:00
|
|
|
// Licensed to the LF AI & Data foundation under one
|
|
|
|
// or more contributor license agreements. See the NOTICE file
|
|
|
|
// distributed with this work for additional information
|
|
|
|
// regarding copyright ownership. The ASF licenses this file
|
|
|
|
// to you under the Apache License, Version 2.0 (the
|
|
|
|
// "License"); you may not use this file except in compliance
|
2021-04-19 11:35:38 +08:00
|
|
|
// with the License. You may obtain a copy of the License at
|
|
|
|
//
|
2021-10-25 20:23:16 +08:00
|
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
2021-04-19 11:35:38 +08:00
|
|
|
//
|
2021-10-25 20:23:16 +08:00
|
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
|
|
// See the License for the specific language governing permissions and
|
|
|
|
// limitations under the License.
|
2021-10-14 15:44:34 +08:00
|
|
|
|
2021-06-22 10:42:07 +08:00
|
|
|
package datacoord
|
2021-02-02 18:53:10 +08:00
|
|
|
|
|
|
|
import (
|
2021-05-26 19:06:56 +08:00
|
|
|
"context"
|
2021-02-02 18:53:10 +08:00
|
|
|
"testing"
|
2022-04-01 11:33:28 +08:00
|
|
|
"time"
|
2021-02-02 18:53:10 +08:00
|
|
|
|
2023-02-26 11:31:49 +08:00
|
|
|
"github.com/cockroachdb/errors"
|
2021-10-14 15:44:34 +08:00
|
|
|
"github.com/golang/protobuf/proto"
|
2023-04-06 19:14:32 +08:00
|
|
|
"github.com/stretchr/testify/assert"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"stathat.com/c/consistent"
|
|
|
|
|
2021-09-06 11:12:42 +08:00
|
|
|
"github.com/milvus-io/milvus/internal/kv"
|
2022-03-28 22:33:27 +08:00
|
|
|
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
|
2021-05-26 19:06:56 +08:00
|
|
|
"github.com/milvus-io/milvus/internal/proto/datapb"
|
2021-10-14 15:44:34 +08:00
|
|
|
"github.com/milvus-io/milvus/internal/types"
|
2021-02-02 18:53:10 +08:00
|
|
|
)
|
|
|
|
|
2022-03-28 22:33:27 +08:00
|
|
|
func getMetaKv(t *testing.T) kv.MetaKv {
|
|
|
|
rootPath := "/etcd/test/root/" + t.Name()
|
|
|
|
metakv, err := etcdkv.NewMetaKvFactory(rootPath, &Params.EtcdCfg)
|
|
|
|
require.NoError(t, err)
|
|
|
|
|
|
|
|
return metakv
|
|
|
|
}
|
|
|
|
|
2021-10-14 15:44:34 +08:00
|
|
|
func TestClusterCreate(t *testing.T) {
|
2022-03-28 22:33:27 +08:00
|
|
|
kv := getMetaKv(t)
|
|
|
|
defer func() {
|
|
|
|
kv.RemoveWithPrefix("")
|
|
|
|
kv.Close()
|
|
|
|
}()
|
|
|
|
|
2021-10-14 15:44:34 +08:00
|
|
|
t.Run("startup normally", func(t *testing.T) {
|
2022-03-28 22:33:27 +08:00
|
|
|
|
|
|
|
ctx, cancel := context.WithCancel(context.TODO())
|
|
|
|
defer cancel()
|
|
|
|
|
2021-10-14 15:44:34 +08:00
|
|
|
sessionManager := NewSessionManager()
|
2021-11-17 23:25:12 +08:00
|
|
|
channelManager, err := NewChannelManager(kv, newMockHandler())
|
2021-10-14 15:44:34 +08:00
|
|
|
assert.Nil(t, err)
|
|
|
|
cluster := NewCluster(sessionManager, channelManager)
|
|
|
|
defer cluster.Close()
|
|
|
|
addr := "localhost:8080"
|
|
|
|
info := &NodeInfo{
|
|
|
|
NodeID: 1,
|
|
|
|
Address: addr,
|
|
|
|
}
|
|
|
|
nodes := []*NodeInfo{info}
|
2022-03-28 22:33:27 +08:00
|
|
|
err = cluster.Startup(ctx, nodes)
|
2021-10-14 15:44:34 +08:00
|
|
|
assert.Nil(t, err)
|
|
|
|
dataNodes := sessionManager.GetSessions()
|
|
|
|
assert.EqualValues(t, 1, len(dataNodes))
|
|
|
|
assert.EqualValues(t, "localhost:8080", dataNodes[0].info.Address)
|
|
|
|
})
|
2021-07-12 11:03:52 +08:00
|
|
|
|
2021-10-14 15:44:34 +08:00
|
|
|
t.Run("startup with existed channel data", func(t *testing.T) {
|
2022-03-28 22:33:27 +08:00
|
|
|
defer kv.RemoveWithPrefix("")
|
|
|
|
|
|
|
|
ctx, cancel := context.WithCancel(context.TODO())
|
|
|
|
defer cancel()
|
|
|
|
|
2021-10-14 15:44:34 +08:00
|
|
|
var err error
|
|
|
|
info1 := &datapb.ChannelWatchInfo{
|
|
|
|
Vchan: &datapb.VchannelInfo{
|
|
|
|
CollectionID: 1,
|
|
|
|
ChannelName: "channel1",
|
|
|
|
},
|
|
|
|
}
|
|
|
|
info1Data, err := proto.Marshal(info1)
|
|
|
|
assert.Nil(t, err)
|
2022-12-07 18:01:19 +08:00
|
|
|
err = kv.Save(Params.CommonCfg.DataCoordWatchSubPath.GetValue()+"/1/channel1", string(info1Data))
|
2021-10-14 15:44:34 +08:00
|
|
|
assert.Nil(t, err)
|
2021-07-12 11:03:52 +08:00
|
|
|
|
2021-10-14 15:44:34 +08:00
|
|
|
sessionManager := NewSessionManager()
|
2021-11-17 23:25:12 +08:00
|
|
|
channelManager, err := NewChannelManager(kv, newMockHandler())
|
2021-10-14 15:44:34 +08:00
|
|
|
assert.Nil(t, err)
|
|
|
|
cluster := NewCluster(sessionManager, channelManager)
|
|
|
|
defer cluster.Close()
|
|
|
|
|
2022-03-28 22:33:27 +08:00
|
|
|
err = cluster.Startup(ctx, []*NodeInfo{{NodeID: 1, Address: "localhost:9999"}})
|
2021-10-14 15:44:34 +08:00
|
|
|
assert.Nil(t, err)
|
|
|
|
|
|
|
|
channels := channelManager.GetChannels()
|
2022-08-22 14:42:52 +08:00
|
|
|
assert.EqualValues(t, []*NodeChannelInfo{{1, []*channel{{Name: "channel1", CollectionID: 1}}}}, channels)
|
2021-10-14 15:44:34 +08:00
|
|
|
})
|
2021-07-12 11:03:52 +08:00
|
|
|
|
2021-10-14 15:44:34 +08:00
|
|
|
t.Run("remove all nodes and restart with other nodes", func(t *testing.T) {
|
2022-03-28 22:33:27 +08:00
|
|
|
defer kv.RemoveWithPrefix("")
|
|
|
|
|
|
|
|
ctx, cancel := context.WithCancel(context.TODO())
|
|
|
|
defer cancel()
|
|
|
|
|
2021-10-14 15:44:34 +08:00
|
|
|
sessionManager := NewSessionManager()
|
2021-11-17 23:25:12 +08:00
|
|
|
channelManager, err := NewChannelManager(kv, newMockHandler())
|
2021-10-14 15:44:34 +08:00
|
|
|
assert.Nil(t, err)
|
|
|
|
cluster := NewCluster(sessionManager, channelManager)
|
|
|
|
|
|
|
|
addr := "localhost:8080"
|
|
|
|
info := &NodeInfo{
|
|
|
|
NodeID: 1,
|
|
|
|
Address: addr,
|
2021-07-12 11:03:52 +08:00
|
|
|
}
|
2021-10-14 15:44:34 +08:00
|
|
|
nodes := []*NodeInfo{info}
|
2022-03-28 22:33:27 +08:00
|
|
|
err = cluster.Startup(ctx, nodes)
|
2021-10-14 15:44:34 +08:00
|
|
|
assert.Nil(t, err)
|
|
|
|
|
|
|
|
err = cluster.UnRegister(info)
|
|
|
|
assert.Nil(t, err)
|
|
|
|
sessions := sessionManager.GetSessions()
|
|
|
|
assert.Empty(t, sessions)
|
|
|
|
|
|
|
|
cluster.Close()
|
|
|
|
|
|
|
|
sessionManager2 := NewSessionManager()
|
2021-11-17 23:25:12 +08:00
|
|
|
channelManager2, err := NewChannelManager(kv, newMockHandler())
|
2021-10-14 15:44:34 +08:00
|
|
|
assert.Nil(t, err)
|
|
|
|
clusterReload := NewCluster(sessionManager2, channelManager2)
|
|
|
|
defer clusterReload.Close()
|
|
|
|
|
|
|
|
addr = "localhost:8081"
|
|
|
|
info = &NodeInfo{
|
|
|
|
NodeID: 2,
|
|
|
|
Address: addr,
|
2021-07-12 11:03:52 +08:00
|
|
|
}
|
2021-10-14 15:44:34 +08:00
|
|
|
nodes = []*NodeInfo{info}
|
2022-03-28 22:33:27 +08:00
|
|
|
err = clusterReload.Startup(ctx, nodes)
|
2021-10-14 15:44:34 +08:00
|
|
|
assert.Nil(t, err)
|
|
|
|
sessions = sessionManager2.GetSessions()
|
|
|
|
assert.EqualValues(t, 1, len(sessions))
|
|
|
|
assert.EqualValues(t, 2, sessions[0].info.NodeID)
|
|
|
|
assert.EqualValues(t, addr, sessions[0].info.Address)
|
|
|
|
channels := channelManager2.GetChannels()
|
|
|
|
assert.EqualValues(t, 1, len(channels))
|
|
|
|
assert.EqualValues(t, 2, channels[0].NodeID)
|
|
|
|
})
|
|
|
|
|
|
|
|
t.Run("loadKv Fails", func(t *testing.T) {
|
2022-03-28 22:33:27 +08:00
|
|
|
defer kv.RemoveWithPrefix("")
|
|
|
|
|
|
|
|
fkv := &loadPrefixFailKV{MetaKv: kv}
|
2021-11-17 23:25:12 +08:00
|
|
|
_, err := NewChannelManager(fkv, newMockHandler())
|
2021-10-14 15:44:34 +08:00
|
|
|
assert.NotNil(t, err)
|
|
|
|
})
|
2021-07-12 11:03:52 +08:00
|
|
|
}
|
|
|
|
|
2021-09-06 11:12:42 +08:00
|
|
|
// a mock kv that always fail when LoadWithPrefix
|
2021-09-10 11:26:00 +08:00
|
|
|
type loadPrefixFailKV struct {
|
2022-03-28 22:33:27 +08:00
|
|
|
kv.MetaKv
|
2021-09-06 11:12:42 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
// LoadWithPrefix override behavior
|
2021-09-10 11:26:00 +08:00
|
|
|
func (kv *loadPrefixFailKV) LoadWithPrefix(key string) ([]string, []string, error) {
|
2021-09-06 11:12:42 +08:00
|
|
|
return []string{}, []string{}, errors.New("mocked fail")
|
|
|
|
}
|
|
|
|
|
2021-10-14 15:44:34 +08:00
|
|
|
func TestRegister(t *testing.T) {
|
2022-03-28 22:33:27 +08:00
|
|
|
kv := getMetaKv(t)
|
|
|
|
defer func() {
|
|
|
|
kv.RemoveWithPrefix("")
|
|
|
|
kv.Close()
|
|
|
|
}()
|
|
|
|
|
2021-10-14 15:44:34 +08:00
|
|
|
t.Run("register to empty cluster", func(t *testing.T) {
|
2022-03-28 22:33:27 +08:00
|
|
|
defer kv.RemoveWithPrefix("")
|
|
|
|
|
|
|
|
ctx, cancel := context.WithCancel(context.TODO())
|
|
|
|
defer cancel()
|
|
|
|
|
2021-10-14 15:44:34 +08:00
|
|
|
sessionManager := NewSessionManager()
|
2021-11-17 23:25:12 +08:00
|
|
|
channelManager, err := NewChannelManager(kv, newMockHandler())
|
2021-10-14 15:44:34 +08:00
|
|
|
assert.Nil(t, err)
|
|
|
|
cluster := NewCluster(sessionManager, channelManager)
|
|
|
|
defer cluster.Close()
|
|
|
|
addr := "localhost:8080"
|
2022-03-28 22:33:27 +08:00
|
|
|
err = cluster.Startup(ctx, nil)
|
2021-10-14 15:44:34 +08:00
|
|
|
assert.Nil(t, err)
|
|
|
|
info := &NodeInfo{
|
|
|
|
NodeID: 1,
|
|
|
|
Address: addr,
|
|
|
|
}
|
|
|
|
err = cluster.Register(info)
|
|
|
|
assert.Nil(t, err)
|
|
|
|
sessions := sessionManager.GetSessions()
|
|
|
|
assert.EqualValues(t, 1, len(sessions))
|
|
|
|
assert.EqualValues(t, "localhost:8080", sessions[0].info.Address)
|
|
|
|
})
|
2021-09-06 11:12:42 +08:00
|
|
|
|
2021-10-14 15:44:34 +08:00
|
|
|
t.Run("register to empty cluster with buffer channels", func(t *testing.T) {
|
2022-03-28 22:33:27 +08:00
|
|
|
defer kv.RemoveWithPrefix("")
|
|
|
|
|
|
|
|
ctx, cancel := context.WithCancel(context.TODO())
|
|
|
|
defer cancel()
|
|
|
|
|
2021-10-14 15:44:34 +08:00
|
|
|
sessionManager := NewSessionManager()
|
2021-11-17 23:25:12 +08:00
|
|
|
channelManager, err := NewChannelManager(kv, newMockHandler())
|
2021-10-14 15:44:34 +08:00
|
|
|
assert.Nil(t, err)
|
|
|
|
err = channelManager.Watch(&channel{
|
2021-10-26 20:15:41 +08:00
|
|
|
Name: "ch1",
|
|
|
|
CollectionID: 0,
|
2021-10-14 15:44:34 +08:00
|
|
|
})
|
|
|
|
assert.Nil(t, err)
|
|
|
|
cluster := NewCluster(sessionManager, channelManager)
|
|
|
|
defer cluster.Close()
|
|
|
|
addr := "localhost:8080"
|
2022-03-28 22:33:27 +08:00
|
|
|
err = cluster.Startup(ctx, nil)
|
2021-10-14 15:44:34 +08:00
|
|
|
assert.Nil(t, err)
|
|
|
|
info := &NodeInfo{
|
|
|
|
NodeID: 1,
|
|
|
|
Address: addr,
|
|
|
|
}
|
|
|
|
err = cluster.Register(info)
|
|
|
|
assert.Nil(t, err)
|
2022-02-22 13:15:51 +08:00
|
|
|
bufferChannels := channelManager.GetBufferChannels()
|
2021-10-14 15:44:34 +08:00
|
|
|
assert.Empty(t, bufferChannels.Channels)
|
|
|
|
nodeChannels := channelManager.GetChannels()
|
|
|
|
assert.EqualValues(t, 1, len(nodeChannels))
|
|
|
|
assert.EqualValues(t, 1, nodeChannels[0].NodeID)
|
2021-10-26 20:15:41 +08:00
|
|
|
assert.EqualValues(t, "ch1", nodeChannels[0].Channels[0].Name)
|
2021-09-06 11:12:42 +08:00
|
|
|
})
|
2021-05-26 19:06:56 +08:00
|
|
|
|
2021-10-14 15:44:34 +08:00
|
|
|
t.Run("register and restart with no channel", func(t *testing.T) {
|
2022-03-28 22:33:27 +08:00
|
|
|
defer kv.RemoveWithPrefix("")
|
|
|
|
|
|
|
|
ctx, cancel := context.WithCancel(context.TODO())
|
|
|
|
defer cancel()
|
|
|
|
|
2021-10-14 15:44:34 +08:00
|
|
|
sessionManager := NewSessionManager()
|
2021-11-17 23:25:12 +08:00
|
|
|
channelManager, err := NewChannelManager(kv, newMockHandler())
|
2021-10-14 15:44:34 +08:00
|
|
|
assert.Nil(t, err)
|
|
|
|
cluster := NewCluster(sessionManager, channelManager)
|
|
|
|
addr := "localhost:8080"
|
2022-03-28 22:33:27 +08:00
|
|
|
err = cluster.Startup(ctx, nil)
|
2021-10-14 15:44:34 +08:00
|
|
|
assert.Nil(t, err)
|
|
|
|
info := &NodeInfo{
|
|
|
|
NodeID: 1,
|
|
|
|
Address: addr,
|
|
|
|
}
|
|
|
|
err = cluster.Register(info)
|
|
|
|
assert.Nil(t, err)
|
|
|
|
cluster.Close()
|
2021-05-26 19:06:56 +08:00
|
|
|
|
2021-10-14 15:44:34 +08:00
|
|
|
sessionManager2 := NewSessionManager()
|
2021-11-17 23:25:12 +08:00
|
|
|
channelManager2, err := NewChannelManager(kv, newMockHandler())
|
2021-10-14 15:44:34 +08:00
|
|
|
assert.Nil(t, err)
|
|
|
|
restartCluster := NewCluster(sessionManager2, channelManager2)
|
|
|
|
defer restartCluster.Close()
|
|
|
|
channels := channelManager2.GetChannels()
|
|
|
|
assert.Empty(t, channels)
|
|
|
|
})
|
2021-05-26 19:06:56 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
func TestUnregister(t *testing.T) {
|
2022-03-28 22:33:27 +08:00
|
|
|
kv := getMetaKv(t)
|
|
|
|
defer func() {
|
|
|
|
kv.RemoveWithPrefix("")
|
|
|
|
kv.Close()
|
|
|
|
}()
|
|
|
|
|
2021-07-15 16:38:31 +08:00
|
|
|
t.Run("remove node after unregister", func(t *testing.T) {
|
2022-03-28 22:33:27 +08:00
|
|
|
defer kv.RemoveWithPrefix("")
|
|
|
|
|
|
|
|
ctx, cancel := context.WithCancel(context.TODO())
|
|
|
|
defer cancel()
|
|
|
|
|
2021-10-14 15:44:34 +08:00
|
|
|
sessionManager := NewSessionManager()
|
2021-11-17 23:25:12 +08:00
|
|
|
channelManager, err := NewChannelManager(kv, newMockHandler())
|
2021-07-15 16:38:31 +08:00
|
|
|
assert.Nil(t, err)
|
2021-10-14 15:44:34 +08:00
|
|
|
cluster := NewCluster(sessionManager, channelManager)
|
2021-07-15 16:38:31 +08:00
|
|
|
defer cluster.Close()
|
|
|
|
addr := "localhost:8080"
|
2021-10-14 15:44:34 +08:00
|
|
|
info := &NodeInfo{
|
|
|
|
Address: addr,
|
|
|
|
NodeID: 1,
|
2021-07-15 16:38:31 +08:00
|
|
|
}
|
2021-10-14 15:44:34 +08:00
|
|
|
nodes := []*NodeInfo{info}
|
2022-03-28 22:33:27 +08:00
|
|
|
err = cluster.Startup(ctx, nodes)
|
2021-10-14 15:44:34 +08:00
|
|
|
assert.Nil(t, err)
|
|
|
|
err = cluster.UnRegister(nodes[0])
|
|
|
|
assert.Nil(t, err)
|
|
|
|
sessions := sessionManager.GetSessions()
|
|
|
|
assert.Empty(t, sessions)
|
2021-07-15 16:38:31 +08:00
|
|
|
})
|
|
|
|
|
|
|
|
t.Run("move channels to online nodes after unregister", func(t *testing.T) {
|
2022-03-28 22:33:27 +08:00
|
|
|
defer kv.RemoveWithPrefix("")
|
|
|
|
|
|
|
|
ctx, cancel := context.WithCancel(context.TODO())
|
|
|
|
defer cancel()
|
|
|
|
|
2021-10-14 15:44:34 +08:00
|
|
|
sessionManager := NewSessionManager()
|
2021-11-17 23:25:12 +08:00
|
|
|
channelManager, err := NewChannelManager(kv, newMockHandler())
|
2021-07-15 16:38:31 +08:00
|
|
|
assert.Nil(t, err)
|
2021-10-14 15:44:34 +08:00
|
|
|
cluster := NewCluster(sessionManager, channelManager)
|
2021-07-15 16:38:31 +08:00
|
|
|
defer cluster.Close()
|
2021-10-14 15:44:34 +08:00
|
|
|
|
|
|
|
nodeInfo1 := &NodeInfo{
|
|
|
|
Address: "localhost:8080",
|
|
|
|
NodeID: 1,
|
2021-07-15 16:38:31 +08:00
|
|
|
}
|
2021-10-14 15:44:34 +08:00
|
|
|
nodeInfo2 := &NodeInfo{
|
|
|
|
Address: "localhost:8081",
|
|
|
|
NodeID: 2,
|
2021-07-15 16:38:31 +08:00
|
|
|
}
|
2021-10-14 15:44:34 +08:00
|
|
|
nodes := []*NodeInfo{nodeInfo1, nodeInfo2}
|
2022-03-28 22:33:27 +08:00
|
|
|
err = cluster.Startup(ctx, nodes)
|
2021-10-14 15:44:34 +08:00
|
|
|
assert.Nil(t, err)
|
|
|
|
err = cluster.Watch("ch1", 1)
|
|
|
|
assert.Nil(t, err)
|
|
|
|
err = cluster.UnRegister(nodeInfo1)
|
|
|
|
assert.Nil(t, err)
|
|
|
|
|
|
|
|
channels := channelManager.GetChannels()
|
|
|
|
assert.EqualValues(t, 1, len(channels))
|
|
|
|
assert.EqualValues(t, 2, channels[0].NodeID)
|
|
|
|
assert.EqualValues(t, 1, len(channels[0].Channels))
|
2021-10-26 20:15:41 +08:00
|
|
|
assert.EqualValues(t, "ch1", channels[0].Channels[0].Name)
|
2021-07-15 16:38:31 +08:00
|
|
|
})
|
|
|
|
|
|
|
|
t.Run("remove all channels after unregsiter", func(t *testing.T) {
|
2022-03-28 22:33:27 +08:00
|
|
|
defer kv.RemoveWithPrefix("")
|
|
|
|
|
|
|
|
ctx, cancel := context.WithCancel(context.TODO())
|
|
|
|
defer cancel()
|
|
|
|
|
2021-10-14 15:44:34 +08:00
|
|
|
var mockSessionCreator = func(ctx context.Context, addr string) (types.DataNode, error) {
|
|
|
|
return newMockDataNodeClient(1, nil)
|
2021-07-15 16:38:31 +08:00
|
|
|
}
|
2021-10-14 15:44:34 +08:00
|
|
|
sessionManager := NewSessionManager(withSessionCreator(mockSessionCreator))
|
2021-11-17 23:25:12 +08:00
|
|
|
channelManager, err := NewChannelManager(kv, newMockHandler())
|
2021-07-15 16:38:31 +08:00
|
|
|
assert.Nil(t, err)
|
2021-10-14 15:44:34 +08:00
|
|
|
cluster := NewCluster(sessionManager, channelManager)
|
2021-07-15 16:38:31 +08:00
|
|
|
defer cluster.Close()
|
2021-10-14 15:44:34 +08:00
|
|
|
|
|
|
|
nodeInfo := &NodeInfo{
|
|
|
|
Address: "localhost:8080",
|
|
|
|
NodeID: 1,
|
2021-07-15 16:38:31 +08:00
|
|
|
}
|
2022-03-28 22:33:27 +08:00
|
|
|
err = cluster.Startup(ctx, []*NodeInfo{nodeInfo})
|
2021-10-14 15:44:34 +08:00
|
|
|
assert.Nil(t, err)
|
|
|
|
err = cluster.Watch("ch_1", 1)
|
|
|
|
assert.Nil(t, err)
|
|
|
|
err = cluster.UnRegister(nodeInfo)
|
|
|
|
assert.Nil(t, err)
|
|
|
|
channels := channelManager.GetChannels()
|
|
|
|
assert.Empty(t, channels)
|
2022-02-22 13:15:51 +08:00
|
|
|
channel := channelManager.GetBufferChannels()
|
2021-10-14 15:44:34 +08:00
|
|
|
assert.NotNil(t, channel)
|
|
|
|
assert.EqualValues(t, 1, len(channel.Channels))
|
2021-10-26 20:15:41 +08:00
|
|
|
assert.EqualValues(t, "ch_1", channel.Channels[0].Name)
|
2021-10-14 15:44:34 +08:00
|
|
|
})
|
|
|
|
}
|
|
|
|
|
|
|
|
func TestWatchIfNeeded(t *testing.T) {
|
2022-03-28 22:33:27 +08:00
|
|
|
kv := getMetaKv(t)
|
|
|
|
defer func() {
|
|
|
|
kv.RemoveWithPrefix("")
|
|
|
|
kv.Close()
|
|
|
|
}()
|
|
|
|
|
2021-10-14 15:44:34 +08:00
|
|
|
t.Run("add deplicated channel to cluster", func(t *testing.T) {
|
2022-03-28 22:33:27 +08:00
|
|
|
defer kv.RemoveWithPrefix("")
|
|
|
|
|
|
|
|
ctx, cancel := context.WithCancel(context.TODO())
|
|
|
|
defer cancel()
|
|
|
|
|
2021-10-14 15:44:34 +08:00
|
|
|
var mockSessionCreator = func(ctx context.Context, addr string) (types.DataNode, error) {
|
|
|
|
return newMockDataNodeClient(1, nil)
|
2021-07-15 16:38:31 +08:00
|
|
|
}
|
2021-10-14 15:44:34 +08:00
|
|
|
sessionManager := NewSessionManager(withSessionCreator(mockSessionCreator))
|
2021-11-17 23:25:12 +08:00
|
|
|
channelManager, err := NewChannelManager(kv, newMockHandler())
|
2021-10-14 15:44:34 +08:00
|
|
|
assert.Nil(t, err)
|
|
|
|
cluster := NewCluster(sessionManager, channelManager)
|
|
|
|
defer cluster.Close()
|
|
|
|
|
|
|
|
addr := "localhost:8080"
|
|
|
|
info := &NodeInfo{
|
|
|
|
Address: addr,
|
|
|
|
NodeID: 1,
|
2021-07-15 16:38:31 +08:00
|
|
|
}
|
2021-10-14 15:44:34 +08:00
|
|
|
|
2022-03-28 22:33:27 +08:00
|
|
|
err = cluster.Startup(ctx, []*NodeInfo{info})
|
2021-10-14 15:44:34 +08:00
|
|
|
assert.Nil(t, err)
|
|
|
|
err = cluster.Watch("ch1", 1)
|
|
|
|
assert.Nil(t, err)
|
|
|
|
channels := channelManager.GetChannels()
|
|
|
|
assert.EqualValues(t, 1, len(channels))
|
2021-10-26 20:15:41 +08:00
|
|
|
assert.EqualValues(t, "ch1", channels[0].Channels[0].Name)
|
2021-10-14 15:44:34 +08:00
|
|
|
})
|
|
|
|
|
|
|
|
t.Run("watch channel to empty cluster", func(t *testing.T) {
|
2022-03-28 22:33:27 +08:00
|
|
|
defer kv.RemoveWithPrefix("")
|
|
|
|
|
2021-10-14 15:44:34 +08:00
|
|
|
sessionManager := NewSessionManager()
|
2021-11-17 23:25:12 +08:00
|
|
|
channelManager, err := NewChannelManager(kv, newMockHandler())
|
2021-10-14 15:44:34 +08:00
|
|
|
assert.Nil(t, err)
|
|
|
|
cluster := NewCluster(sessionManager, channelManager)
|
|
|
|
defer cluster.Close()
|
|
|
|
|
|
|
|
err = cluster.Watch("ch1", 1)
|
2021-07-15 16:38:31 +08:00
|
|
|
assert.Nil(t, err)
|
2021-10-14 15:44:34 +08:00
|
|
|
|
|
|
|
channels := channelManager.GetChannels()
|
|
|
|
assert.Empty(t, channels)
|
2022-02-22 13:15:51 +08:00
|
|
|
channel := channelManager.GetBufferChannels()
|
2021-10-14 15:44:34 +08:00
|
|
|
assert.NotNil(t, channel)
|
2021-10-26 20:15:41 +08:00
|
|
|
assert.EqualValues(t, "ch1", channel.Channels[0].Name)
|
2021-07-15 16:38:31 +08:00
|
|
|
})
|
2021-04-08 15:08:34 +08:00
|
|
|
}
|
|
|
|
|
2021-10-14 15:44:34 +08:00
|
|
|
func TestConsistentHashPolicy(t *testing.T) {
|
2022-03-28 22:33:27 +08:00
|
|
|
kv := getMetaKv(t)
|
|
|
|
defer func() {
|
|
|
|
kv.RemoveWithPrefix("")
|
|
|
|
kv.Close()
|
|
|
|
}()
|
|
|
|
|
2021-10-14 15:44:34 +08:00
|
|
|
sessionManager := NewSessionManager()
|
|
|
|
chash := consistent.New()
|
|
|
|
factory := NewConsistentHashChannelPolicyFactory(chash)
|
2021-11-17 23:25:12 +08:00
|
|
|
channelManager, err := NewChannelManager(kv, newMockHandler(), withFactory(factory))
|
2021-07-12 11:03:52 +08:00
|
|
|
assert.Nil(t, err)
|
2021-10-14 15:44:34 +08:00
|
|
|
cluster := NewCluster(sessionManager, channelManager)
|
2021-07-12 11:03:52 +08:00
|
|
|
defer cluster.Close()
|
2021-10-14 15:44:34 +08:00
|
|
|
|
|
|
|
hash := consistent.New()
|
|
|
|
hash.Add("1")
|
|
|
|
hash.Add("2")
|
|
|
|
hash.Add("3")
|
|
|
|
|
|
|
|
nodeInfo1 := &NodeInfo{
|
|
|
|
NodeID: 1,
|
|
|
|
Address: "localhost:1111",
|
|
|
|
}
|
|
|
|
nodeInfo2 := &NodeInfo{
|
|
|
|
NodeID: 2,
|
|
|
|
Address: "localhost:2222",
|
2021-05-26 19:06:56 +08:00
|
|
|
}
|
2021-10-14 15:44:34 +08:00
|
|
|
nodeInfo3 := &NodeInfo{
|
|
|
|
NodeID: 3,
|
|
|
|
Address: "localhost:3333",
|
|
|
|
}
|
|
|
|
err = cluster.Register(nodeInfo1)
|
|
|
|
assert.Nil(t, err)
|
|
|
|
err = cluster.Register(nodeInfo2)
|
|
|
|
assert.Nil(t, err)
|
|
|
|
err = cluster.Register(nodeInfo3)
|
2021-05-26 19:06:56 +08:00
|
|
|
assert.Nil(t, err)
|
2021-10-14 15:44:34 +08:00
|
|
|
|
|
|
|
channels := []string{"ch1", "ch2", "ch3"}
|
|
|
|
for _, c := range channels {
|
|
|
|
err = cluster.Watch(c, 1)
|
|
|
|
assert.Nil(t, err)
|
|
|
|
idstr, err := hash.Get(c)
|
|
|
|
assert.Nil(t, err)
|
|
|
|
id, err := deformatNodeID(idstr)
|
|
|
|
assert.Nil(t, err)
|
|
|
|
match := channelManager.Match(id, c)
|
|
|
|
assert.True(t, match)
|
|
|
|
}
|
|
|
|
|
|
|
|
hash.Remove("1")
|
2021-10-21 10:57:03 +08:00
|
|
|
err = cluster.UnRegister(nodeInfo1)
|
|
|
|
assert.Nil(t, err)
|
2021-10-14 15:44:34 +08:00
|
|
|
for _, c := range channels {
|
|
|
|
idstr, err := hash.Get(c)
|
|
|
|
assert.Nil(t, err)
|
|
|
|
id, err := deformatNodeID(idstr)
|
|
|
|
assert.Nil(t, err)
|
|
|
|
match := channelManager.Match(id, c)
|
|
|
|
assert.True(t, match)
|
|
|
|
}
|
|
|
|
|
|
|
|
hash.Remove("2")
|
2021-10-21 10:57:03 +08:00
|
|
|
err = cluster.UnRegister(nodeInfo2)
|
|
|
|
assert.Nil(t, err)
|
2021-10-14 15:44:34 +08:00
|
|
|
for _, c := range channels {
|
|
|
|
idstr, err := hash.Get(c)
|
|
|
|
assert.Nil(t, err)
|
|
|
|
id, err := deformatNodeID(idstr)
|
|
|
|
assert.Nil(t, err)
|
|
|
|
match := channelManager.Match(id, c)
|
|
|
|
assert.True(t, match)
|
|
|
|
}
|
|
|
|
|
|
|
|
hash.Remove("3")
|
2021-10-21 10:57:03 +08:00
|
|
|
err = cluster.UnRegister(nodeInfo3)
|
|
|
|
assert.Nil(t, err)
|
2022-02-22 13:15:51 +08:00
|
|
|
bufferChannels := channelManager.GetBufferChannels()
|
2021-10-14 15:44:34 +08:00
|
|
|
assert.EqualValues(t, 3, len(bufferChannels.Channels))
|
2021-02-02 18:53:10 +08:00
|
|
|
}
|
2021-10-20 15:02:36 +08:00
|
|
|
|
|
|
|
func TestCluster_Flush(t *testing.T) {
|
2022-03-28 22:33:27 +08:00
|
|
|
kv := getMetaKv(t)
|
|
|
|
defer func() {
|
|
|
|
kv.RemoveWithPrefix("")
|
|
|
|
kv.Close()
|
|
|
|
}()
|
|
|
|
|
|
|
|
ctx, cancel := context.WithCancel(context.TODO())
|
|
|
|
defer cancel()
|
2021-10-20 15:02:36 +08:00
|
|
|
sessionManager := NewSessionManager()
|
2021-11-17 23:25:12 +08:00
|
|
|
channelManager, err := NewChannelManager(kv, newMockHandler())
|
2021-10-20 15:02:36 +08:00
|
|
|
assert.Nil(t, err)
|
|
|
|
cluster := NewCluster(sessionManager, channelManager)
|
|
|
|
defer cluster.Close()
|
|
|
|
addr := "localhost:8080"
|
|
|
|
info := &NodeInfo{
|
|
|
|
Address: addr,
|
|
|
|
NodeID: 1,
|
|
|
|
}
|
|
|
|
nodes := []*NodeInfo{info}
|
2022-03-28 22:33:27 +08:00
|
|
|
err = cluster.Startup(ctx, nodes)
|
2021-10-20 15:02:36 +08:00
|
|
|
assert.Nil(t, err)
|
|
|
|
|
|
|
|
err = cluster.Watch("chan-1", 1)
|
|
|
|
assert.NoError(t, err)
|
|
|
|
|
|
|
|
// flush empty should impact nothing
|
|
|
|
assert.NotPanics(t, func() {
|
2022-11-10 22:13:04 +08:00
|
|
|
err := cluster.Flush(context.Background(), 1, "chan-1", []*datapb.SegmentInfo{})
|
2022-10-22 12:09:28 +08:00
|
|
|
assert.NoError(t, err)
|
2021-10-20 15:02:36 +08:00
|
|
|
})
|
|
|
|
|
|
|
|
// flush not watched channel
|
|
|
|
assert.NotPanics(t, func() {
|
2022-11-10 22:13:04 +08:00
|
|
|
err := cluster.Flush(context.Background(), 1, "chan-2", []*datapb.SegmentInfo{{ID: 1}})
|
2022-10-22 12:09:28 +08:00
|
|
|
assert.Error(t, err)
|
2021-10-20 15:02:36 +08:00
|
|
|
})
|
2022-10-22 12:09:28 +08:00
|
|
|
|
|
|
|
// flush from wrong datanode
|
|
|
|
assert.NotPanics(t, func() {
|
2022-11-10 22:13:04 +08:00
|
|
|
err := cluster.Flush(context.Background(), 2, "chan-1", []*datapb.SegmentInfo{{ID: 1}})
|
2022-10-22 12:09:28 +08:00
|
|
|
assert.Error(t, err)
|
|
|
|
})
|
|
|
|
|
2021-10-20 15:02:36 +08:00
|
|
|
//TODO add a method to verify datanode has flush request after client injection is available
|
|
|
|
}
|
2022-04-01 11:33:28 +08:00
|
|
|
|
|
|
|
func TestCluster_Import(t *testing.T) {
|
|
|
|
kv := getMetaKv(t)
|
|
|
|
defer func() {
|
|
|
|
kv.RemoveWithPrefix("")
|
|
|
|
kv.Close()
|
|
|
|
}()
|
|
|
|
|
|
|
|
ctx, cancel := context.WithTimeout(context.TODO(), 100*time.Millisecond)
|
|
|
|
defer cancel()
|
|
|
|
sessionManager := NewSessionManager()
|
|
|
|
channelManager, err := NewChannelManager(kv, newMockHandler())
|
|
|
|
assert.Nil(t, err)
|
|
|
|
cluster := NewCluster(sessionManager, channelManager)
|
|
|
|
defer cluster.Close()
|
|
|
|
addr := "localhost:8080"
|
|
|
|
info := &NodeInfo{
|
|
|
|
Address: addr,
|
|
|
|
NodeID: 1,
|
|
|
|
}
|
|
|
|
nodes := []*NodeInfo{info}
|
|
|
|
err = cluster.Startup(ctx, nodes)
|
|
|
|
assert.Nil(t, err)
|
|
|
|
|
|
|
|
err = cluster.Watch("chan-1", 1)
|
|
|
|
assert.NoError(t, err)
|
|
|
|
|
|
|
|
assert.NotPanics(t, func() {
|
|
|
|
cluster.Import(ctx, 1, &datapb.ImportTaskRequest{})
|
|
|
|
})
|
|
|
|
time.Sleep(500 * time.Millisecond)
|
|
|
|
}
|
2022-05-25 14:34:00 +08:00
|
|
|
|
|
|
|
func TestCluster_ReCollectSegmentStats(t *testing.T) {
|
|
|
|
kv := getMetaKv(t)
|
|
|
|
defer func() {
|
|
|
|
kv.RemoveWithPrefix("")
|
|
|
|
kv.Close()
|
|
|
|
}()
|
|
|
|
|
|
|
|
t.Run("recollect succeed", func(t *testing.T) {
|
|
|
|
ctx, cancel := context.WithCancel(context.TODO())
|
|
|
|
defer cancel()
|
|
|
|
var mockSessionCreator = func(ctx context.Context, addr string) (types.DataNode, error) {
|
|
|
|
return newMockDataNodeClient(1, nil)
|
|
|
|
}
|
|
|
|
sessionManager := NewSessionManager(withSessionCreator(mockSessionCreator))
|
|
|
|
channelManager, err := NewChannelManager(kv, newMockHandler())
|
|
|
|
assert.Nil(t, err)
|
|
|
|
cluster := NewCluster(sessionManager, channelManager)
|
|
|
|
defer cluster.Close()
|
|
|
|
addr := "localhost:8080"
|
|
|
|
info := &NodeInfo{
|
|
|
|
Address: addr,
|
|
|
|
NodeID: 1,
|
|
|
|
}
|
|
|
|
nodes := []*NodeInfo{info}
|
|
|
|
err = cluster.Startup(ctx, nodes)
|
|
|
|
assert.Nil(t, err)
|
|
|
|
|
|
|
|
err = cluster.Watch("chan-1", 1)
|
|
|
|
assert.NoError(t, err)
|
|
|
|
|
|
|
|
assert.NotPanics(t, func() {
|
2023-02-27 10:41:46 +08:00
|
|
|
cluster.ReCollectSegmentStats(ctx)
|
2022-05-25 14:34:00 +08:00
|
|
|
})
|
|
|
|
time.Sleep(500 * time.Millisecond)
|
|
|
|
})
|
|
|
|
|
|
|
|
t.Run("recollect failed", func(t *testing.T) {
|
|
|
|
ctx, cancel := context.WithCancel(context.TODO())
|
|
|
|
defer cancel()
|
|
|
|
sessionManager := NewSessionManager()
|
|
|
|
channelManager, err := NewChannelManager(kv, newMockHandler())
|
|
|
|
assert.Nil(t, err)
|
|
|
|
cluster := NewCluster(sessionManager, channelManager)
|
|
|
|
defer cluster.Close()
|
|
|
|
addr := "localhost:8080"
|
|
|
|
info := &NodeInfo{
|
|
|
|
Address: addr,
|
|
|
|
NodeID: 1,
|
|
|
|
}
|
|
|
|
nodes := []*NodeInfo{info}
|
|
|
|
err = cluster.Startup(ctx, nodes)
|
|
|
|
assert.Nil(t, err)
|
|
|
|
|
|
|
|
err = cluster.Watch("chan-1", 1)
|
|
|
|
assert.NoError(t, err)
|
|
|
|
|
|
|
|
assert.NotPanics(t, func() {
|
2023-02-27 10:41:46 +08:00
|
|
|
cluster.ReCollectSegmentStats(ctx)
|
2022-05-25 14:34:00 +08:00
|
|
|
})
|
|
|
|
time.Sleep(500 * time.Millisecond)
|
|
|
|
})
|
|
|
|
}
|