mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-11-30 10:59:32 +08:00
797218a8ad
Signed-off-by: yah01 <yang.cen@zilliz.com> Co-authored-by: Congqi Xia <congqi.xia@zilliz.com> Co-authored-by: Congqi Xia <congqi.xia@zilliz.com>
2056 lines
46 KiB
Go
2056 lines
46 KiB
Go
// 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
|
|
// with the License. You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// 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.
|
|
|
|
package querynode
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/milvus-io/milvus/internal/proto/commonpb"
|
|
"github.com/milvus-io/milvus/internal/proto/internalpb"
|
|
"github.com/milvus-io/milvus/internal/proto/querypb"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
type mockNodeDetector struct {
|
|
initNodes []nodeEvent
|
|
evtCh chan nodeEvent
|
|
}
|
|
|
|
func (m *mockNodeDetector) watchNodes(collectionID int64, replicaID int64, vchannelName string) ([]nodeEvent, <-chan nodeEvent) {
|
|
return m.initNodes, m.evtCh
|
|
}
|
|
|
|
type mockSegmentDetector struct {
|
|
initSegments []segmentEvent
|
|
evtCh chan segmentEvent
|
|
}
|
|
|
|
func (m *mockSegmentDetector) watchSegments(collectionID int64, replicaID int64, vchannelName string) ([]segmentEvent, <-chan segmentEvent) {
|
|
return m.initSegments, m.evtCh
|
|
}
|
|
|
|
type mockShardQueryNode struct {
|
|
searchResult *internalpb.SearchResults
|
|
searchErr error
|
|
queryResult *internalpb.RetrieveResults
|
|
queryErr error
|
|
releaseSegmentsResult *commonpb.Status
|
|
releaseSegmentsErr error
|
|
}
|
|
|
|
func (m *mockShardQueryNode) Search(_ context.Context, _ *querypb.SearchRequest) (*internalpb.SearchResults, error) {
|
|
return m.searchResult, m.searchErr
|
|
}
|
|
|
|
func (m *mockShardQueryNode) Query(_ context.Context, _ *querypb.QueryRequest) (*internalpb.RetrieveResults, error) {
|
|
return m.queryResult, m.queryErr
|
|
}
|
|
|
|
func (m *mockShardQueryNode) ReleaseSegments(ctx context.Context, in *querypb.ReleaseSegmentsRequest) (*commonpb.Status, error) {
|
|
return m.releaseSegmentsResult, m.releaseSegmentsErr
|
|
}
|
|
|
|
func (m *mockShardQueryNode) Stop() error {
|
|
return nil
|
|
}
|
|
|
|
func buildMockQueryNode(nodeID int64, addr string) shardQueryNode {
|
|
return &mockShardQueryNode{
|
|
searchResult: &internalpb.SearchResults{},
|
|
queryResult: &internalpb.RetrieveResults{},
|
|
}
|
|
}
|
|
|
|
func segmentEventsToSyncInfo(events []segmentEvent) []*querypb.ReplicaSegmentsInfo {
|
|
infos := make([]*querypb.ReplicaSegmentsInfo, 0, len(events))
|
|
for _, event := range events {
|
|
for _, nodeID := range event.nodeIDs {
|
|
infos = append(infos, &querypb.ReplicaSegmentsInfo{
|
|
NodeId: nodeID,
|
|
SegmentIds: []int64{event.segmentID},
|
|
PartitionId: event.partitionID,
|
|
})
|
|
}
|
|
}
|
|
|
|
return infos
|
|
}
|
|
|
|
func TestShardCluster_Create(t *testing.T) {
|
|
collectionID := int64(1)
|
|
vchannelName := "dml_1_1_v0"
|
|
replicaID := int64(0)
|
|
|
|
t.Run("empty shard cluster", func(t *testing.T) {
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{}, &mockSegmentDetector{}, buildMockQueryNode)
|
|
assert.NotPanics(t, func() { sc.Close() })
|
|
// close twice
|
|
assert.NotPanics(t, func() { sc.Close() })
|
|
})
|
|
|
|
t.Run("init nodes", func(t *testing.T) {
|
|
nodeEvent := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
isLeader: true,
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
}
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{
|
|
initNodes: nodeEvent,
|
|
}, &mockSegmentDetector{}, buildMockQueryNode)
|
|
defer sc.Close()
|
|
|
|
for _, e := range nodeEvent {
|
|
node, has := sc.getNode(e.nodeID)
|
|
assert.True(t, has)
|
|
assert.Equal(t, e.nodeAddr, node.nodeAddr)
|
|
}
|
|
sc.mut.Lock()
|
|
defer sc.mut.Unlock()
|
|
|
|
require.NotNil(t, sc.leader)
|
|
assert.Equal(t, int64(1), sc.leader.nodeID)
|
|
})
|
|
|
|
t.Run("init segments", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
}
|
|
|
|
segmentEvents := []segmentEvent{
|
|
{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoading,
|
|
},
|
|
{
|
|
segmentID: 3,
|
|
nodeIDs: []int64{3},
|
|
state: segmentStateOffline,
|
|
},
|
|
}
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{
|
|
initNodes: nodeEvents,
|
|
}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
}, buildMockQueryNode)
|
|
defer sc.Close()
|
|
|
|
for _, e := range segmentEvents {
|
|
sc.mut.RLock()
|
|
segment, has := sc.segments[e.segmentID]
|
|
_, inCluster := sc.pickNode(e)
|
|
sc.mut.RUnlock()
|
|
if inCluster {
|
|
assert.True(t, has)
|
|
assert.Equal(t, e.segmentID, segment.segmentID)
|
|
assert.Contains(t, e.nodeIDs, segment.nodeID)
|
|
assert.Equal(t, e.state, segment.state)
|
|
} else {
|
|
assert.False(t, has)
|
|
}
|
|
}
|
|
assert.EqualValues(t, unavailable, sc.state.Load())
|
|
})
|
|
}
|
|
|
|
func TestShardCluster_nodeEvent(t *testing.T) {
|
|
collectionID := int64(1)
|
|
vchannelName := "dml_1_1_v0"
|
|
replicaID := int64(0)
|
|
|
|
t.Run("only nodes", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
}
|
|
evtCh := make(chan nodeEvent, 10)
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{
|
|
initNodes: nodeEvents,
|
|
evtCh: evtCh,
|
|
}, &mockSegmentDetector{}, buildMockQueryNode)
|
|
defer sc.Close()
|
|
|
|
evtCh <- nodeEvent{
|
|
nodeID: 3,
|
|
nodeAddr: "addr_3",
|
|
eventType: nodeAdd,
|
|
}
|
|
// same event
|
|
evtCh <- nodeEvent{
|
|
nodeID: 3,
|
|
nodeAddr: "addr_3",
|
|
eventType: nodeAdd,
|
|
}
|
|
|
|
assert.Eventually(t, func() bool {
|
|
node, has := sc.getNode(3)
|
|
return has && node.nodeAddr == "addr_3"
|
|
}, time.Second, time.Millisecond)
|
|
|
|
evtCh <- nodeEvent{
|
|
nodeID: 3,
|
|
nodeAddr: "addr_new",
|
|
eventType: nodeAdd,
|
|
}
|
|
assert.Eventually(t, func() bool {
|
|
node, has := sc.getNode(3)
|
|
return has && node.nodeAddr == "addr_new"
|
|
}, time.Second, time.Millisecond)
|
|
|
|
evtCh <- nodeEvent{
|
|
nodeID: 3,
|
|
nodeAddr: "addr_new",
|
|
eventType: nodeDel,
|
|
}
|
|
assert.Eventually(t, func() bool {
|
|
_, has := sc.getNode(3)
|
|
return !has
|
|
}, time.Second, time.Millisecond)
|
|
assert.Equal(t, int32(available), sc.state.Load())
|
|
|
|
evtCh <- nodeEvent{
|
|
nodeID: 4,
|
|
nodeAddr: "addr_new",
|
|
eventType: nodeDel,
|
|
}
|
|
|
|
close(evtCh)
|
|
})
|
|
|
|
t.Run("with segments", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
}
|
|
segmentEvents := []segmentEvent{
|
|
{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoading,
|
|
},
|
|
{
|
|
segmentID: 3,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateOffline,
|
|
},
|
|
}
|
|
|
|
evtCh := make(chan nodeEvent, 10)
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{
|
|
initNodes: nodeEvents,
|
|
evtCh: evtCh,
|
|
}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
}, buildMockQueryNode)
|
|
defer sc.Close()
|
|
|
|
evtCh <- nodeEvent{
|
|
nodeID: 3,
|
|
nodeAddr: "addr_3",
|
|
eventType: nodeAdd,
|
|
}
|
|
|
|
assert.Eventually(t, func() bool {
|
|
node, has := sc.getNode(3)
|
|
return has && node.nodeAddr == "addr_3"
|
|
}, time.Second, time.Millisecond)
|
|
|
|
evtCh <- nodeEvent{
|
|
nodeID: 3,
|
|
nodeAddr: "addr_new",
|
|
eventType: nodeAdd,
|
|
}
|
|
assert.Eventually(t, func() bool {
|
|
node, has := sc.getNode(3)
|
|
return has && node.nodeAddr == "addr_new"
|
|
}, time.Second, time.Millisecond)
|
|
|
|
// remove node 2
|
|
evtCh <- nodeEvent{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
eventType: nodeDel,
|
|
}
|
|
assert.Eventually(t, func() bool {
|
|
_, has := sc.getNode(2)
|
|
return !has
|
|
}, time.Second, time.Millisecond)
|
|
assert.Equal(t, int32(unavailable), sc.state.Load())
|
|
|
|
segment, has := sc.getSegment(2)
|
|
assert.True(t, has)
|
|
assert.Equal(t, segmentStateOffline, segment.state)
|
|
|
|
close(evtCh)
|
|
})
|
|
}
|
|
|
|
func TestShardCluster_segmentEvent(t *testing.T) {
|
|
collectionID := int64(1)
|
|
vchannelName := "dml_1_1_v0"
|
|
replicaID := int64(0)
|
|
|
|
t.Run("from loading", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
{
|
|
nodeID: 3,
|
|
nodeAddr: "addr_3",
|
|
},
|
|
}
|
|
|
|
segmentEvents := []segmentEvent{
|
|
{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoading,
|
|
},
|
|
{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoading,
|
|
},
|
|
{
|
|
segmentID: 3,
|
|
nodeIDs: []int64{3},
|
|
state: segmentStateLoading,
|
|
},
|
|
}
|
|
|
|
evtCh := make(chan segmentEvent, 10)
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{
|
|
initNodes: nodeEvents,
|
|
}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
evtCh: evtCh,
|
|
}, buildMockQueryNode)
|
|
defer sc.Close()
|
|
|
|
evtCh <- segmentEvent{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoading,
|
|
eventType: segmentAdd,
|
|
}
|
|
|
|
evtCh <- segmentEvent{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
eventType: segmentAdd,
|
|
}
|
|
|
|
evtCh <- segmentEvent{
|
|
segmentID: 3,
|
|
nodeIDs: []int64{3},
|
|
state: segmentStateOffline,
|
|
eventType: segmentAdd,
|
|
}
|
|
assert.Eventually(t, func() bool {
|
|
seg, has := sc.getSegment(2)
|
|
return has && seg.nodeID == 2 && seg.state == segmentStateLoaded
|
|
}, time.Second, time.Millisecond)
|
|
assert.Eventually(t, func() bool {
|
|
seg, has := sc.getSegment(3)
|
|
return has && seg.nodeID == 3 && seg.state == segmentStateOffline
|
|
}, time.Second, time.Millisecond)
|
|
// put this check behind other make sure event is processed
|
|
assert.Eventually(t, func() bool {
|
|
seg, has := sc.getSegment(1)
|
|
return has && seg.nodeID == 1 && seg.state == segmentStateLoading
|
|
}, time.Second, time.Millisecond)
|
|
|
|
// node id not match
|
|
evtCh <- segmentEvent{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
eventType: segmentAdd,
|
|
}
|
|
// will not change
|
|
assert.Eventually(t, func() bool {
|
|
seg, has := sc.getSegment(1)
|
|
return has && seg.nodeID == 1 && seg.state == segmentStateLoading
|
|
}, time.Second, time.Millisecond)
|
|
close(evtCh)
|
|
})
|
|
|
|
t.Run("from loaded", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
{
|
|
nodeID: 3,
|
|
nodeAddr: "addr_3",
|
|
},
|
|
}
|
|
segmentEvents := []segmentEvent{
|
|
{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 3,
|
|
nodeIDs: []int64{3},
|
|
state: segmentStateLoaded,
|
|
},
|
|
}
|
|
|
|
evtCh := make(chan segmentEvent, 10)
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{initNodes: nodeEvents}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
evtCh: evtCh,
|
|
}, buildMockQueryNode)
|
|
defer sc.Close()
|
|
|
|
sc.SyncSegments(segmentEventsToSyncInfo(segmentEvents), segmentStateLoaded)
|
|
|
|
// make reference greater than 0
|
|
_, versionID := sc.segmentAllocations(nil)
|
|
defer sc.finishUsage(versionID)
|
|
|
|
evtCh <- segmentEvent{
|
|
segmentID: 4,
|
|
nodeIDs: []int64{4},
|
|
state: segmentStateLoaded,
|
|
eventType: segmentAdd,
|
|
}
|
|
evtCh <- segmentEvent{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
eventType: segmentAdd,
|
|
}
|
|
|
|
evtCh <- segmentEvent{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoading,
|
|
eventType: segmentAdd,
|
|
}
|
|
|
|
evtCh <- segmentEvent{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
eventType: segmentAdd,
|
|
}
|
|
|
|
evtCh <- segmentEvent{
|
|
segmentID: 3,
|
|
nodeIDs: []int64{3},
|
|
state: segmentStateOffline,
|
|
eventType: segmentAdd,
|
|
}
|
|
|
|
assert.Eventually(t, func() bool {
|
|
seg, has := sc.getSegment(1)
|
|
return has && seg.nodeID == 1 && seg.state == segmentStateLoading
|
|
}, time.Second, time.Millisecond)
|
|
assert.Eventually(t, func() bool {
|
|
seg, has := sc.getSegment(3)
|
|
return has && seg.nodeID == 3 && seg.state == segmentStateOffline
|
|
}, time.Second, time.Millisecond)
|
|
// put this check behind other make sure event is processed
|
|
assert.Eventually(t, func() bool {
|
|
seg, has := sc.getSegment(2)
|
|
return has && seg.nodeID == 2 && seg.state == segmentStateLoaded
|
|
}, time.Second, time.Millisecond)
|
|
|
|
_, has := sc.getSegment(4)
|
|
assert.False(t, has)
|
|
|
|
})
|
|
|
|
t.Run("from loaded, node changed", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
}
|
|
|
|
segmentEvents := []segmentEvent{
|
|
{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
},
|
|
}
|
|
|
|
evtCh := make(chan segmentEvent, 10)
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{initNodes: nodeEvents}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
evtCh: evtCh,
|
|
}, buildMockQueryNode)
|
|
defer sc.Close()
|
|
|
|
sc.SyncSegments(segmentEventsToSyncInfo(segmentEvents), segmentStateLoaded)
|
|
|
|
// make reference greater than 0
|
|
_, versionID := sc.segmentAllocations(nil)
|
|
defer sc.finishUsage(versionID)
|
|
|
|
// bring segment online in the other querynode
|
|
evtCh <- segmentEvent{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
eventType: segmentAdd,
|
|
}
|
|
|
|
evtCh <- segmentEvent{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoaded,
|
|
eventType: segmentAdd,
|
|
}
|
|
|
|
assert.Eventually(t, func() bool {
|
|
seg, has := sc.getSegment(1)
|
|
return has && seg.nodeID == 2 && seg.state == segmentStateLoaded
|
|
}, time.Second, time.Millisecond)
|
|
|
|
assert.Eventually(t, func() bool {
|
|
seg, has := sc.getSegment(2)
|
|
return has && seg.nodeID == 1 && seg.state == segmentStateLoaded
|
|
}, time.Second, time.Millisecond)
|
|
|
|
})
|
|
|
|
t.Run("from offline", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
{
|
|
nodeID: 3,
|
|
nodeAddr: "addr_3",
|
|
},
|
|
}
|
|
|
|
segmentEvents := []segmentEvent{
|
|
{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateOffline,
|
|
},
|
|
{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateOffline,
|
|
},
|
|
{
|
|
segmentID: 3,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateOffline,
|
|
},
|
|
}
|
|
|
|
evtCh := make(chan segmentEvent, 10)
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{initNodes: nodeEvents}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
evtCh: evtCh,
|
|
}, buildMockQueryNode)
|
|
defer sc.Close()
|
|
|
|
sc.SyncSegments(segmentEventsToSyncInfo(nil), segmentStateLoaded)
|
|
|
|
evtCh <- segmentEvent{
|
|
segmentID: 3,
|
|
nodeIDs: []int64{3},
|
|
state: segmentStateOffline,
|
|
eventType: segmentAdd,
|
|
}
|
|
evtCh <- segmentEvent{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoading,
|
|
eventType: segmentAdd,
|
|
}
|
|
|
|
evtCh <- segmentEvent{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
eventType: segmentAdd,
|
|
}
|
|
|
|
assert.Eventually(t, func() bool {
|
|
seg, has := sc.getSegment(1)
|
|
return has && seg.nodeID == 1 && seg.state == segmentStateLoading
|
|
}, time.Second, time.Millisecond)
|
|
assert.Eventually(t, func() bool {
|
|
seg, has := sc.getSegment(2)
|
|
return has && seg.nodeID == 2 && seg.state == segmentStateLoaded
|
|
}, time.Second, time.Millisecond)
|
|
// put this check behind other make sure event is processed
|
|
assert.Eventually(t, func() bool {
|
|
seg, has := sc.getSegment(3)
|
|
return has && seg.nodeID == 3 && seg.state == segmentStateOffline
|
|
}, time.Second, time.Millisecond)
|
|
|
|
close(evtCh)
|
|
})
|
|
|
|
t.Run("remove segments", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
{
|
|
nodeID: 3,
|
|
nodeAddr: "addr_3",
|
|
},
|
|
}
|
|
|
|
segmentEvents := []segmentEvent{
|
|
{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoading,
|
|
},
|
|
{
|
|
segmentID: 3,
|
|
nodeIDs: []int64{3},
|
|
state: segmentStateOffline,
|
|
},
|
|
}
|
|
|
|
evtCh := make(chan segmentEvent, 10)
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{initNodes: nodeEvents}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
evtCh: evtCh,
|
|
}, buildMockQueryNode)
|
|
defer sc.Close()
|
|
|
|
evtCh <- segmentEvent{
|
|
segmentID: 3,
|
|
nodeIDs: []int64{3},
|
|
eventType: segmentDel,
|
|
}
|
|
evtCh <- segmentEvent{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
eventType: segmentDel,
|
|
}
|
|
evtCh <- segmentEvent{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
eventType: segmentDel,
|
|
}
|
|
|
|
assert.Eventually(t, func() bool {
|
|
_, has := sc.getSegment(1)
|
|
return !has
|
|
}, time.Second, time.Millisecond)
|
|
assert.Eventually(t, func() bool {
|
|
_, has := sc.getSegment(2)
|
|
return !has
|
|
}, time.Second, time.Millisecond)
|
|
// put this check behind other make sure event is processed
|
|
assert.Eventually(t, func() bool {
|
|
_, has := sc.getSegment(3)
|
|
return !has
|
|
}, time.Second, time.Millisecond)
|
|
})
|
|
|
|
t.Run("remove failed", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
{
|
|
nodeID: 3,
|
|
nodeAddr: "addr_3",
|
|
},
|
|
}
|
|
|
|
segmentEvents := []segmentEvent{
|
|
{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoading,
|
|
},
|
|
{
|
|
segmentID: 3,
|
|
nodeIDs: []int64{3},
|
|
state: segmentStateOffline,
|
|
},
|
|
}
|
|
|
|
evtCh := make(chan segmentEvent, 10)
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{initNodes: nodeEvents}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
evtCh: evtCh,
|
|
}, buildMockQueryNode)
|
|
defer sc.Close()
|
|
|
|
sc.SyncSegments(segmentEventsToSyncInfo(nil), segmentStateLoaded)
|
|
// non-exist segment
|
|
evtCh <- segmentEvent{
|
|
segmentID: 4,
|
|
nodeIDs: []int64{3},
|
|
eventType: segmentDel,
|
|
}
|
|
// segment node id not match
|
|
evtCh <- segmentEvent{
|
|
segmentID: 3,
|
|
nodeIDs: []int64{4},
|
|
eventType: segmentDel,
|
|
}
|
|
|
|
// use add segment as event process signal
|
|
evtCh <- segmentEvent{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
eventType: segmentAdd,
|
|
}
|
|
assert.Eventually(t, func() bool {
|
|
seg, has := sc.getSegment(2)
|
|
return has && seg.nodeID == 2 && seg.state == segmentStateLoaded
|
|
}, time.Second, time.Millisecond)
|
|
|
|
_, has := sc.getSegment(3)
|
|
assert.True(t, has)
|
|
})
|
|
}
|
|
|
|
func TestShardCluster_SyncSegments(t *testing.T) {
|
|
collectionID := int64(1)
|
|
vchannelName := "dml_1_1_v0"
|
|
replicaID := int64(0)
|
|
|
|
t.Run("sync new segments", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
{
|
|
nodeID: 3,
|
|
nodeAddr: "addr_3",
|
|
},
|
|
}
|
|
|
|
segmentEvents := []segmentEvent{}
|
|
|
|
evtCh := make(chan segmentEvent, 10)
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{initNodes: nodeEvents}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
evtCh: evtCh,
|
|
}, buildMockQueryNode)
|
|
defer sc.Close()
|
|
|
|
sc.SyncSegments([]*querypb.ReplicaSegmentsInfo{
|
|
{
|
|
NodeId: 1,
|
|
SegmentIds: []int64{1},
|
|
},
|
|
{
|
|
NodeId: 2,
|
|
SegmentIds: []int64{2},
|
|
},
|
|
{
|
|
NodeId: 3,
|
|
SegmentIds: []int64{3},
|
|
},
|
|
}, segmentStateLoaded)
|
|
assert.Eventually(t, func() bool {
|
|
seg, has := sc.getSegment(1)
|
|
return has && seg.nodeID == 1 && seg.state == segmentStateLoaded
|
|
}, time.Second, time.Millisecond)
|
|
assert.Eventually(t, func() bool {
|
|
seg, has := sc.getSegment(2)
|
|
return has && seg.nodeID == 2 && seg.state == segmentStateLoaded
|
|
}, time.Second, time.Millisecond)
|
|
assert.Eventually(t, func() bool {
|
|
seg, has := sc.getSegment(3)
|
|
return has && seg.nodeID == 3 && seg.state == segmentStateLoaded
|
|
}, time.Second, time.Millisecond)
|
|
|
|
})
|
|
|
|
t.Run("sync existing segments", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
{
|
|
nodeID: 3,
|
|
nodeAddr: "addr_3",
|
|
},
|
|
}
|
|
|
|
segmentEvents := []segmentEvent{
|
|
{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateOffline,
|
|
},
|
|
{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateOffline,
|
|
},
|
|
{
|
|
segmentID: 3,
|
|
nodeIDs: []int64{3},
|
|
state: segmentStateOffline,
|
|
},
|
|
}
|
|
|
|
evtCh := make(chan segmentEvent, 10)
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{initNodes: nodeEvents}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
evtCh: evtCh,
|
|
}, buildMockQueryNode)
|
|
defer sc.Close()
|
|
|
|
sc.SyncSegments([]*querypb.ReplicaSegmentsInfo{
|
|
{
|
|
NodeId: 1,
|
|
SegmentIds: []int64{1},
|
|
},
|
|
{
|
|
NodeId: 2,
|
|
SegmentIds: []int64{2},
|
|
},
|
|
{
|
|
NodeId: 3,
|
|
SegmentIds: []int64{3},
|
|
},
|
|
}, segmentStateLoaded)
|
|
assert.Eventually(t, func() bool {
|
|
seg, has := sc.getSegment(1)
|
|
return has && seg.nodeID == 1 && seg.state == segmentStateLoaded
|
|
}, time.Second, time.Millisecond)
|
|
assert.Eventually(t, func() bool {
|
|
seg, has := sc.getSegment(2)
|
|
return has && seg.nodeID == 2 && seg.state == segmentStateLoaded
|
|
}, time.Second, time.Millisecond)
|
|
assert.Eventually(t, func() bool {
|
|
seg, has := sc.getSegment(3)
|
|
return has && seg.nodeID == 3 && seg.state == segmentStateLoaded
|
|
}, time.Second, time.Millisecond)
|
|
})
|
|
|
|
}
|
|
|
|
var streamingDoNothing = func(context.Context) error { return nil }
|
|
var streamingError = func(context.Context) error { return errors.New("mock streaming error") }
|
|
|
|
func TestShardCluster_Search(t *testing.T) {
|
|
collectionID := int64(1)
|
|
vchannelName := "dml_1_1_v0"
|
|
replicaID := int64(0)
|
|
ctx := context.Background()
|
|
|
|
t.Run("search unavailable cluster", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
{
|
|
nodeID: 3,
|
|
nodeAddr: "addr_3",
|
|
},
|
|
}
|
|
|
|
segmentEvents := []segmentEvent{
|
|
{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateOffline,
|
|
},
|
|
{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateOffline,
|
|
},
|
|
{
|
|
segmentID: 3,
|
|
nodeIDs: []int64{3},
|
|
state: segmentStateOffline,
|
|
},
|
|
}
|
|
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{initNodes: nodeEvents}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
}, buildMockQueryNode)
|
|
|
|
defer sc.Close()
|
|
require.EqualValues(t, unavailable, sc.state.Load())
|
|
|
|
_, err := sc.Search(ctx, &querypb.SearchRequest{
|
|
DmlChannel: vchannelName,
|
|
}, streamingDoNothing)
|
|
assert.Error(t, err)
|
|
})
|
|
|
|
t.Run("search wrong channel", func(t *testing.T) {
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{}, &mockSegmentDetector{}, buildMockQueryNode)
|
|
|
|
defer sc.Close()
|
|
|
|
_, err := sc.Search(ctx, &querypb.SearchRequest{
|
|
DmlChannel: vchannelName + "_suffix",
|
|
}, streamingDoNothing)
|
|
assert.Error(t, err)
|
|
})
|
|
|
|
t.Run("normal search", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
}
|
|
|
|
segmentEvents := []segmentEvent{
|
|
{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 3,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
},
|
|
}
|
|
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{
|
|
initNodes: nodeEvents,
|
|
}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
}, buildMockQueryNode)
|
|
|
|
defer sc.Close()
|
|
// setup first version
|
|
sc.SyncSegments(nil, segmentStateLoaded)
|
|
|
|
require.EqualValues(t, available, sc.state.Load())
|
|
|
|
result, err := sc.Search(ctx, &querypb.SearchRequest{
|
|
DmlChannel: vchannelName,
|
|
}, streamingDoNothing)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, len(nodeEvents), len(result))
|
|
})
|
|
|
|
t.Run("with streaming fail", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
}
|
|
|
|
segmentEvents := []segmentEvent{
|
|
{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 3,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
},
|
|
}
|
|
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{
|
|
initNodes: nodeEvents,
|
|
}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
}, buildMockQueryNode)
|
|
|
|
defer sc.Close()
|
|
// setup first version
|
|
sc.SyncSegments(nil, segmentStateLoaded)
|
|
|
|
require.EqualValues(t, available, sc.state.Load())
|
|
|
|
_, err := sc.Search(ctx, &querypb.SearchRequest{
|
|
DmlChannel: vchannelName,
|
|
}, func(ctx context.Context) error { return errors.New("mocked") })
|
|
assert.Error(t, err)
|
|
})
|
|
|
|
t.Run("partial fail", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
}
|
|
|
|
segmentEvents := []segmentEvent{
|
|
{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 3,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
},
|
|
}
|
|
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{
|
|
initNodes: nodeEvents,
|
|
}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
}, func(nodeID int64, addr string) shardQueryNode {
|
|
if nodeID != 2 { // hard code error one
|
|
return buildMockQueryNode(nodeID, addr)
|
|
}
|
|
return &mockShardQueryNode{
|
|
searchErr: errors.New("mocked error"),
|
|
queryErr: errors.New("mocked error"),
|
|
}
|
|
})
|
|
|
|
defer sc.Close()
|
|
// setup first version
|
|
sc.SyncSegments(nil, segmentStateLoaded)
|
|
|
|
require.EqualValues(t, available, sc.state.Load())
|
|
|
|
_, err := sc.Search(ctx, &querypb.SearchRequest{
|
|
DmlChannel: vchannelName,
|
|
}, streamingDoNothing)
|
|
assert.Error(t, err)
|
|
})
|
|
|
|
t.Run("test meta error", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
}
|
|
segmentEvents := []segmentEvent{
|
|
{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
},
|
|
}
|
|
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{
|
|
initNodes: nodeEvents,
|
|
}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
}, buildMockQueryNode)
|
|
|
|
//mock meta error
|
|
sc.mut.Lock()
|
|
sc.segments[3] = shardSegmentInfo{
|
|
segmentID: 3,
|
|
nodeID: 3, // node does not exist
|
|
state: segmentStateLoaded,
|
|
}
|
|
sc.mut.Unlock()
|
|
|
|
defer sc.Close()
|
|
// setup first version
|
|
sc.SyncSegments(nil, segmentStateLoaded)
|
|
|
|
require.EqualValues(t, available, sc.state.Load())
|
|
|
|
_, err := sc.Search(ctx, &querypb.SearchRequest{
|
|
DmlChannel: vchannelName,
|
|
}, streamingDoNothing)
|
|
assert.Error(t, err)
|
|
})
|
|
}
|
|
|
|
func TestShardCluster_Query(t *testing.T) {
|
|
collectionID := int64(1)
|
|
vchannelName := "dml_1_1_v0"
|
|
replicaID := int64(0)
|
|
ctx := context.Background()
|
|
|
|
t.Run("query unavailable cluster", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
{
|
|
nodeID: 3,
|
|
nodeAddr: "addr_3",
|
|
},
|
|
}
|
|
|
|
segmentEvents := []segmentEvent{
|
|
{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateOffline,
|
|
},
|
|
{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateOffline,
|
|
},
|
|
{
|
|
segmentID: 3,
|
|
nodeIDs: []int64{3},
|
|
state: segmentStateOffline,
|
|
},
|
|
}
|
|
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{initNodes: nodeEvents}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
}, buildMockQueryNode)
|
|
|
|
defer sc.Close()
|
|
// setup first version
|
|
sc.SyncSegments(nil, segmentStateLoaded)
|
|
|
|
require.EqualValues(t, unavailable, sc.state.Load())
|
|
|
|
_, err := sc.Query(ctx, &querypb.QueryRequest{
|
|
DmlChannel: vchannelName,
|
|
}, streamingDoNothing)
|
|
assert.Error(t, err)
|
|
})
|
|
t.Run("query wrong channel", func(t *testing.T) {
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{}, &mockSegmentDetector{}, buildMockQueryNode)
|
|
|
|
defer sc.Close()
|
|
// setup first version
|
|
sc.SyncSegments(nil, segmentStateLoaded)
|
|
|
|
_, err := sc.Query(ctx, &querypb.QueryRequest{
|
|
DmlChannel: vchannelName + "_suffix",
|
|
}, streamingDoNothing)
|
|
assert.Error(t, err)
|
|
})
|
|
t.Run("normal query", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
}
|
|
|
|
segmentEvents := []segmentEvent{
|
|
{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 3,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
},
|
|
}
|
|
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{
|
|
initNodes: nodeEvents,
|
|
}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
}, buildMockQueryNode)
|
|
|
|
defer sc.Close()
|
|
// setup first version
|
|
sc.SyncSegments(nil, segmentStateLoaded)
|
|
|
|
require.EqualValues(t, available, sc.state.Load())
|
|
|
|
result, err := sc.Query(ctx, &querypb.QueryRequest{
|
|
DmlChannel: vchannelName,
|
|
}, streamingDoNothing)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, len(nodeEvents), len(result))
|
|
})
|
|
t.Run("with streaming fail", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
}
|
|
|
|
segmentEvents := []segmentEvent{
|
|
{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 3,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
},
|
|
}
|
|
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{
|
|
initNodes: nodeEvents,
|
|
}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
}, buildMockQueryNode)
|
|
|
|
defer sc.Close()
|
|
// setup first version
|
|
sc.SyncSegments(nil, segmentStateLoaded)
|
|
|
|
require.EqualValues(t, available, sc.state.Load())
|
|
|
|
_, err := sc.Query(ctx, &querypb.QueryRequest{
|
|
DmlChannel: vchannelName,
|
|
}, func(ctx context.Context) error { return errors.New("mocked") })
|
|
assert.Error(t, err)
|
|
})
|
|
|
|
t.Run("partial fail", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
}
|
|
|
|
segmentEvents := []segmentEvent{
|
|
{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 3,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
},
|
|
}
|
|
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{
|
|
initNodes: nodeEvents,
|
|
}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
}, func(nodeID int64, addr string) shardQueryNode {
|
|
if nodeID != 2 { // hard code error one
|
|
return buildMockQueryNode(nodeID, addr)
|
|
}
|
|
return &mockShardQueryNode{
|
|
searchErr: errors.New("mocked error"),
|
|
queryErr: errors.New("mocked error"),
|
|
}
|
|
})
|
|
|
|
defer sc.Close()
|
|
// setup first version
|
|
sc.SyncSegments(nil, segmentStateLoaded)
|
|
|
|
require.EqualValues(t, available, sc.state.Load())
|
|
|
|
_, err := sc.Query(ctx, &querypb.QueryRequest{
|
|
DmlChannel: vchannelName,
|
|
}, streamingDoNothing)
|
|
assert.Error(t, err)
|
|
})
|
|
t.Run("test meta error", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
}
|
|
|
|
segmentEvents := []segmentEvent{
|
|
{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
},
|
|
}
|
|
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{
|
|
initNodes: nodeEvents,
|
|
}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
}, buildMockQueryNode)
|
|
|
|
//mock meta error
|
|
sc.mut.Lock()
|
|
sc.segments[3] = shardSegmentInfo{
|
|
segmentID: 3,
|
|
nodeID: 3, // node does not exist
|
|
state: segmentStateLoaded,
|
|
}
|
|
sc.mut.Unlock()
|
|
|
|
defer sc.Close()
|
|
// setup first version
|
|
sc.SyncSegments(nil, segmentStateLoaded)
|
|
|
|
require.EqualValues(t, available, sc.state.Load())
|
|
|
|
_, err := sc.Query(ctx, &querypb.QueryRequest{
|
|
DmlChannel: vchannelName,
|
|
}, streamingDoNothing)
|
|
assert.Error(t, err)
|
|
})
|
|
|
|
}
|
|
|
|
func TestShardCluster_Version(t *testing.T) {
|
|
collectionID := int64(1)
|
|
vchannelName := "dml_1_1_v0"
|
|
replicaID := int64(0)
|
|
// ctx := context.Background()
|
|
t.Run("alloc with non-serviceable", func(t *testing.T) {
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{}, &mockSegmentDetector{}, buildMockQueryNode)
|
|
defer sc.Close()
|
|
|
|
_, v := sc.segmentAllocations(nil)
|
|
assert.Equal(t, int64(0), v)
|
|
})
|
|
|
|
t.Run("normal alloc & finish", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
}
|
|
|
|
segmentEvents := []segmentEvent{
|
|
{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
},
|
|
}
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{
|
|
initNodes: nodeEvents,
|
|
}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
}, buildMockQueryNode)
|
|
defer sc.Close()
|
|
|
|
sc.SyncSegments(nil, segmentStateLoaded)
|
|
_, version := sc.segmentAllocations(nil)
|
|
|
|
sc.mut.RLock()
|
|
assert.Equal(t, version, sc.currentVersion.versionID)
|
|
assert.Equal(t, int64(1), sc.currentVersion.inUse.Load())
|
|
sc.mut.RUnlock()
|
|
|
|
sc.finishUsage(version)
|
|
sc.mut.RLock()
|
|
|
|
assert.Equal(t, int64(0), sc.currentVersion.inUse.Load())
|
|
sc.mut.RUnlock()
|
|
})
|
|
|
|
t.Run("wait segments online", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
}
|
|
|
|
segmentEvents := []segmentEvent{
|
|
{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
},
|
|
}
|
|
evtCh := make(chan segmentEvent, 10)
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{
|
|
initNodes: nodeEvents,
|
|
}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
evtCh: evtCh,
|
|
}, buildMockQueryNode)
|
|
defer sc.Close()
|
|
|
|
sc.SyncSegments(nil, segmentStateLoaded)
|
|
|
|
assert.True(t, sc.segmentsOnline([]shardSegmentInfo{{nodeID: 1, segmentID: 1}, {nodeID: 2, segmentID: 2}}))
|
|
assert.False(t, sc.segmentsOnline([]shardSegmentInfo{{nodeID: 1, segmentID: 1}, {nodeID: 2, segmentID: 2}, {nodeID: 1, segmentID: 3}}))
|
|
|
|
sig := make(chan struct{})
|
|
go func() {
|
|
sc.waitSegmentsOnline([]shardSegmentInfo{{nodeID: 1, segmentID: 1}, {nodeID: 2, segmentID: 2}, {nodeID: 1, segmentID: 3}})
|
|
close(sig)
|
|
}()
|
|
|
|
evtCh <- segmentEvent{
|
|
eventType: segmentAdd,
|
|
segmentID: 3,
|
|
nodeIDs: []int64{1, 4},
|
|
state: segmentStateLoaded,
|
|
}
|
|
|
|
<-sig
|
|
assert.True(t, sc.segmentsOnline([]shardSegmentInfo{{nodeID: 1, segmentID: 1}, {nodeID: 2, segmentID: 2}, {nodeID: 1, segmentID: 3}}))
|
|
})
|
|
}
|
|
|
|
func TestShardCluster_HandoffSegments(t *testing.T) {
|
|
collectionID := int64(1)
|
|
otherCollectionID := int64(2)
|
|
vchannelName := "dml_1_1_v0"
|
|
otherVchannelName := "dml_1_2_v0"
|
|
replicaID := int64(0)
|
|
|
|
t.Run("handoff without using segments", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
isLeader: true,
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
}
|
|
|
|
segmentEvents := []segmentEvent{
|
|
{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
},
|
|
}
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{
|
|
initNodes: nodeEvents,
|
|
}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
}, buildMockQueryNode)
|
|
defer sc.Close()
|
|
|
|
sc.SyncSegments(segmentEventsToSyncInfo(nil), segmentStateLoaded)
|
|
|
|
err := sc.HandoffSegments(&querypb.SegmentChangeInfo{
|
|
OnlineSegments: []*querypb.SegmentInfo{
|
|
{SegmentID: 2, NodeID: 2, CollectionID: collectionID, DmChannel: vchannelName, NodeIds: []UniqueID{2}},
|
|
},
|
|
OfflineSegments: []*querypb.SegmentInfo{
|
|
{SegmentID: 1, NodeID: 1, CollectionID: collectionID, DmChannel: vchannelName, NodeIds: []UniqueID{1}},
|
|
},
|
|
})
|
|
if err != nil {
|
|
t.Log(err.Error())
|
|
}
|
|
assert.NoError(t, err)
|
|
|
|
sc.mut.RLock()
|
|
_, has := sc.segments[1]
|
|
sc.mut.RUnlock()
|
|
|
|
assert.False(t, has)
|
|
})
|
|
t.Run("handoff with growing segment(segment not recorded)", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
isLeader: true,
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
}
|
|
|
|
segmentEvents := []segmentEvent{
|
|
{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
},
|
|
}
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{
|
|
initNodes: nodeEvents,
|
|
}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
}, buildMockQueryNode)
|
|
defer sc.Close()
|
|
|
|
err := sc.HandoffSegments(&querypb.SegmentChangeInfo{
|
|
OnlineSegments: []*querypb.SegmentInfo{
|
|
{SegmentID: 2, NodeID: 2, CollectionID: collectionID, DmChannel: vchannelName, NodeIds: []UniqueID{2}},
|
|
{SegmentID: 4, NodeID: 2, CollectionID: otherCollectionID, DmChannel: otherVchannelName, NodeIds: []UniqueID{2}},
|
|
},
|
|
OfflineSegments: []*querypb.SegmentInfo{
|
|
{SegmentID: 3, NodeID: 1, CollectionID: collectionID, DmChannel: vchannelName, NodeIds: []UniqueID{1}},
|
|
{SegmentID: 5, NodeID: 2, CollectionID: otherCollectionID, DmChannel: otherVchannelName, NodeIds: []UniqueID{2}},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
sc.mut.RLock()
|
|
_, has := sc.segments[3]
|
|
sc.mut.RUnlock()
|
|
|
|
assert.False(t, has)
|
|
})
|
|
t.Run("handoff wait online and usage", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
isLeader: true,
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
}
|
|
|
|
segmentEvents := []segmentEvent{
|
|
{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
},
|
|
}
|
|
evtCh := make(chan segmentEvent, 10)
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{
|
|
initNodes: nodeEvents,
|
|
}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
evtCh: evtCh,
|
|
}, buildMockQueryNode)
|
|
defer sc.Close()
|
|
|
|
sc.SyncSegments(segmentEventsToSyncInfo(nil), segmentStateLoaded)
|
|
|
|
//add in-use count
|
|
_, versionID := sc.segmentAllocations(nil)
|
|
|
|
sig := make(chan struct{})
|
|
go func() {
|
|
err := sc.HandoffSegments(&querypb.SegmentChangeInfo{
|
|
OnlineSegments: []*querypb.SegmentInfo{
|
|
{SegmentID: 3, NodeID: 1, CollectionID: collectionID, DmChannel: vchannelName, NodeIds: []UniqueID{1}},
|
|
},
|
|
OfflineSegments: []*querypb.SegmentInfo{
|
|
{SegmentID: 1, NodeID: 1, CollectionID: collectionID, DmChannel: vchannelName, NodeIds: []UniqueID{1}},
|
|
},
|
|
})
|
|
|
|
assert.NoError(t, err)
|
|
close(sig)
|
|
}()
|
|
|
|
sc.finishUsage(versionID)
|
|
|
|
evtCh <- segmentEvent{
|
|
eventType: segmentAdd,
|
|
segmentID: 3,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoaded,
|
|
}
|
|
|
|
// wait for handoff appended into list
|
|
assert.Eventually(t, func() bool {
|
|
sc.mut.RLock()
|
|
defer sc.mut.RUnlock()
|
|
return sc.currentVersion.versionID != versionID
|
|
}, time.Second, time.Millisecond*10)
|
|
|
|
tmpAllocs, nVersionID := sc.segmentAllocations(nil)
|
|
found := false
|
|
for _, segments := range tmpAllocs {
|
|
if inList(segments, int64(1)) {
|
|
found = true
|
|
break
|
|
}
|
|
}
|
|
// segment 1 shall not be allocated again!
|
|
assert.False(t, found)
|
|
sc.finishUsage(nVersionID)
|
|
// rc shall be 0 now
|
|
sc.finishUsage(versionID)
|
|
|
|
// wait handoff finished
|
|
<-sig
|
|
|
|
sc.mut.RLock()
|
|
_, has := sc.segments[1]
|
|
sc.mut.RUnlock()
|
|
|
|
assert.False(t, has)
|
|
})
|
|
|
|
t.Run("load balance wait online and usage", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
isLeader: true,
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
}
|
|
|
|
segmentEvents := []segmentEvent{
|
|
{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
},
|
|
}
|
|
evtCh := make(chan segmentEvent, 10)
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{
|
|
initNodes: nodeEvents,
|
|
}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
evtCh: evtCh,
|
|
}, buildMockQueryNode)
|
|
defer sc.Close()
|
|
|
|
sc.SyncSegments(segmentEventsToSyncInfo(nil), segmentStateLoaded)
|
|
|
|
// add rc to all segments
|
|
_, versionID := sc.segmentAllocations(nil)
|
|
|
|
sig := make(chan struct{})
|
|
go func() {
|
|
err := sc.HandoffSegments(&querypb.SegmentChangeInfo{
|
|
OnlineSegments: []*querypb.SegmentInfo{
|
|
{SegmentID: 1, NodeID: 2, CollectionID: collectionID, DmChannel: vchannelName, NodeIds: []UniqueID{1}},
|
|
},
|
|
OfflineSegments: []*querypb.SegmentInfo{
|
|
{SegmentID: 1, NodeID: 1, CollectionID: collectionID, DmChannel: vchannelName, NodeIds: []UniqueID{1}},
|
|
},
|
|
})
|
|
|
|
assert.NoError(t, err)
|
|
close(sig)
|
|
}()
|
|
|
|
evtCh <- segmentEvent{
|
|
eventType: segmentAdd,
|
|
segmentID: 1,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
}
|
|
sc.finishUsage(versionID)
|
|
|
|
// wait for handoff appended into list
|
|
assert.Eventually(t, func() bool {
|
|
sc.mut.RLock()
|
|
defer sc.mut.RUnlock()
|
|
return sc.currentVersion.versionID != versionID
|
|
}, time.Second, time.Millisecond*10)
|
|
|
|
tmpAllocs, tmpVersionID := sc.segmentAllocations(nil)
|
|
for nodeID, segments := range tmpAllocs {
|
|
for _, segment := range segments {
|
|
if segment == int64(1) {
|
|
assert.Equal(t, int64(2), nodeID)
|
|
}
|
|
}
|
|
}
|
|
sc.finishUsage(tmpVersionID)
|
|
// rc shall be 0 now
|
|
sc.finishUsage(versionID)
|
|
|
|
// wait handoff finished
|
|
<-sig
|
|
|
|
sc.mut.RLock()
|
|
info, has := sc.segments[1]
|
|
sc.mut.RUnlock()
|
|
|
|
assert.True(t, has)
|
|
assert.Equal(t, int64(2), info.nodeID)
|
|
})
|
|
|
|
t.Run("handoff from non-exist node", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
isLeader: true,
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
}
|
|
|
|
segmentEvents := []segmentEvent{
|
|
{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
},
|
|
}
|
|
evtCh := make(chan segmentEvent, 10)
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{
|
|
initNodes: nodeEvents,
|
|
}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
evtCh: evtCh,
|
|
}, buildMockQueryNode)
|
|
defer sc.Close()
|
|
|
|
sc.SyncSegments(segmentEventsToSyncInfo(nil), segmentStateLoaded)
|
|
|
|
err := sc.HandoffSegments(&querypb.SegmentChangeInfo{
|
|
OnlineSegments: []*querypb.SegmentInfo{
|
|
{SegmentID: 2, NodeID: 2, CollectionID: collectionID, DmChannel: vchannelName, NodeIds: []UniqueID{2}},
|
|
},
|
|
OfflineSegments: []*querypb.SegmentInfo{
|
|
{SegmentID: 1, NodeID: 3, CollectionID: collectionID, DmChannel: vchannelName, NodeIds: []UniqueID{3}},
|
|
},
|
|
})
|
|
|
|
assert.NoError(t, err)
|
|
})
|
|
|
|
t.Run("release failed", func(t *testing.T) {
|
|
nodeEvents := []nodeEvent{
|
|
{
|
|
nodeID: 1,
|
|
nodeAddr: "addr_1",
|
|
isLeader: true,
|
|
},
|
|
{
|
|
nodeID: 2,
|
|
nodeAddr: "addr_2",
|
|
},
|
|
}
|
|
|
|
segmentEvents := []segmentEvent{
|
|
{
|
|
segmentID: 1,
|
|
nodeIDs: []int64{1},
|
|
state: segmentStateLoaded,
|
|
},
|
|
{
|
|
segmentID: 2,
|
|
nodeIDs: []int64{2},
|
|
state: segmentStateLoaded,
|
|
},
|
|
}
|
|
evtCh := make(chan segmentEvent, 10)
|
|
mqn := &mockShardQueryNode{}
|
|
sc := NewShardCluster(collectionID, replicaID, vchannelName,
|
|
&mockNodeDetector{
|
|
initNodes: nodeEvents,
|
|
}, &mockSegmentDetector{
|
|
initSegments: segmentEvents,
|
|
evtCh: evtCh,
|
|
}, func(nodeID int64, addr string) shardQueryNode {
|
|
return mqn
|
|
})
|
|
defer sc.Close()
|
|
|
|
mqn.releaseSegmentsErr = errors.New("mocked error")
|
|
|
|
err := sc.HandoffSegments(&querypb.SegmentChangeInfo{
|
|
OfflineSegments: []*querypb.SegmentInfo{
|
|
{SegmentID: 1, NodeID: 1, CollectionID: collectionID, DmChannel: vchannelName, NodeIds: []UniqueID{1}},
|
|
},
|
|
})
|
|
|
|
assert.Error(t, err)
|
|
|
|
mqn.releaseSegmentsErr = nil
|
|
mqn.releaseSegmentsResult = &commonpb.Status{
|
|
ErrorCode: commonpb.ErrorCode_UnexpectedError,
|
|
Reason: "mocked error",
|
|
}
|
|
|
|
err = sc.HandoffSegments(&querypb.SegmentChangeInfo{
|
|
OfflineSegments: []*querypb.SegmentInfo{
|
|
{SegmentID: 2, NodeID: 2, CollectionID: collectionID, DmChannel: vchannelName, NodeIds: []UniqueID{2}},
|
|
},
|
|
})
|
|
|
|
assert.Error(t, err)
|
|
|
|
})
|
|
}
|