mirror of https://github.com/milvus-io/milvus.git
improve code coverage for data service (#4998)
* rename mock.go to mock_test.go Signed-off-by: yudong.cai <yudong.cai@zilliz.com> * add allocator_test.go to improve code coverage Signed-off-by: yudong.cai <yudong.cai@zilliz.com> * update meta_test.go Signed-off-by: yudong.cai <yudong.cai@zilliz.com> * update watcher_test.go Signed-off-by: yudong.cai <yudong.cai@zilliz.com> * fix race Signed-off-by: yudong.cai <yudong.cai@zilliz.com>pull/5002/head
Cai Yudong
GitHub
parent
4f7e393bfd
commit
111376ba40
|
|
@ -38,7 +38,7 @@ func (allocator *allocator) allocTimestamp() (Timestamp, error) {
|
|||
ctx := context.TODO()
|
||||
resp, err := allocator.masterClient.AllocTimestamp(ctx, &masterpb.AllocTimestampRequest{
|
||||
Base: &commonpb.MsgBase{
|
||||
MsgType: commonpb.MsgType_ShowCollections,
|
||||
MsgType: commonpb.MsgType_RequestTSO,
|
||||
MsgID: -1, // todo add msg id
|
||||
Timestamp: 0, // todo
|
||||
SourceID: Params.NodeID,
|
||||
|
|
@ -55,7 +55,7 @@ func (allocator *allocator) allocID() (UniqueID, error) {
|
|||
ctx := context.TODO()
|
||||
resp, err := allocator.masterClient.AllocID(ctx, &masterpb.AllocIDRequest{
|
||||
Base: &commonpb.MsgBase{
|
||||
MsgType: commonpb.MsgType_ShowCollections,
|
||||
MsgType: commonpb.MsgType_RequestID,
|
||||
MsgID: -1, // todo add msg id
|
||||
Timestamp: 0, // todo
|
||||
SourceID: Params.NodeID,
|
||||
|
|
|
|||
|
|
@ -0,0 +1,33 @@
|
|||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed 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 dataservice
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestAllocator_Basic(t *testing.T) {
|
||||
ms := newMockMasterService()
|
||||
allocator := newAllocator(ms)
|
||||
|
||||
t.Run("Test allocTimestamp", func(t *testing.T) {
|
||||
_, err := allocator.allocTimestamp()
|
||||
assert.NoError(t, err)
|
||||
})
|
||||
|
||||
t.Run("Test allocID", func(t *testing.T) {
|
||||
_, err := allocator.allocID()
|
||||
assert.NoError(t, err)
|
||||
})
|
||||
}
|
||||
|
|
@ -14,133 +14,293 @@ import (
|
|||
"testing"
|
||||
|
||||
"github.com/golang/protobuf/proto"
|
||||
|
||||
"github.com/milvus-io/milvus/internal/proto/datapb"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestCollection(t *testing.T) {
|
||||
func TestMeta_Basic(t *testing.T) {
|
||||
const collID = UniqueID(0)
|
||||
const partID0 = UniqueID(100)
|
||||
const partID1 = UniqueID(101)
|
||||
const channelName = "c1"
|
||||
|
||||
mockAllocator := newMockAllocator()
|
||||
meta, err := newMemoryMeta(mockAllocator)
|
||||
assert.Nil(t, err)
|
||||
|
||||
testSchema := newTestSchema()
|
||||
id, err := mockAllocator.allocID()
|
||||
assert.Nil(t, err)
|
||||
err = meta.AddCollection(&datapb.CollectionInfo{
|
||||
ID: id,
|
||||
collInfo := &datapb.CollectionInfo{
|
||||
ID: collID,
|
||||
Schema: testSchema,
|
||||
Partitions: []UniqueID{100},
|
||||
})
|
||||
assert.Nil(t, err)
|
||||
err = meta.AddCollection(&datapb.CollectionInfo{
|
||||
ID: id,
|
||||
Schema: testSchema,
|
||||
})
|
||||
assert.NotNil(t, err)
|
||||
has := meta.HasCollection(id)
|
||||
assert.True(t, has)
|
||||
collection, err := meta.GetCollection(id)
|
||||
assert.Nil(t, err)
|
||||
assert.EqualValues(t, id, collection.ID)
|
||||
assert.EqualValues(t, testSchema, collection.Schema)
|
||||
assert.EqualValues(t, 1, len(collection.Partitions))
|
||||
assert.EqualValues(t, 100, collection.Partitions[0])
|
||||
err = meta.DropCollection(id)
|
||||
assert.Nil(t, err)
|
||||
has = meta.HasCollection(id)
|
||||
assert.False(t, has)
|
||||
_, err = meta.GetCollection(id)
|
||||
assert.NotNil(t, err)
|
||||
}
|
||||
|
||||
func TestSegment(t *testing.T) {
|
||||
mockAllocator := newMockAllocator()
|
||||
meta, err := newMemoryMeta(mockAllocator)
|
||||
assert.Nil(t, err)
|
||||
id, err := mockAllocator.allocID()
|
||||
assert.Nil(t, err)
|
||||
segID, err := mockAllocator.allocID()
|
||||
assert.Nil(t, err)
|
||||
segmentInfo, err := BuildSegment(id, 100, segID, "c1")
|
||||
assert.Nil(t, err)
|
||||
err = meta.AddSegment(segmentInfo)
|
||||
assert.Nil(t, err)
|
||||
info, err := meta.GetSegment(segmentInfo.ID)
|
||||
assert.Nil(t, err)
|
||||
assert.True(t, proto.Equal(info, segmentInfo))
|
||||
ids := meta.GetSegmentsOfCollection(id)
|
||||
assert.EqualValues(t, 1, len(ids))
|
||||
assert.EqualValues(t, segmentInfo.ID, ids[0])
|
||||
ids = meta.GetSegmentsOfPartition(id, 100)
|
||||
assert.EqualValues(t, 1, len(ids))
|
||||
assert.EqualValues(t, segmentInfo.ID, ids[0])
|
||||
err = meta.SealSegment(segmentInfo.ID, 100)
|
||||
assert.Nil(t, err)
|
||||
err = meta.FlushSegment(segmentInfo.ID, 200)
|
||||
assert.Nil(t, err)
|
||||
info, err = meta.GetSegment(segmentInfo.ID)
|
||||
assert.Nil(t, err)
|
||||
assert.NotZero(t, info.SealedTime)
|
||||
assert.NotZero(t, info.FlushedTime)
|
||||
}
|
||||
|
||||
func TestPartition(t *testing.T) {
|
||||
mockAllocator := newMockAllocator()
|
||||
meta, err := newMemoryMeta(mockAllocator)
|
||||
assert.Nil(t, err)
|
||||
testSchema := newTestSchema()
|
||||
id, err := mockAllocator.allocID()
|
||||
assert.Nil(t, err)
|
||||
|
||||
err = meta.AddPartition(id, 10)
|
||||
assert.NotNil(t, err)
|
||||
err = meta.AddCollection(&datapb.CollectionInfo{
|
||||
ID: id,
|
||||
Partitions: []UniqueID{partID0, partID1},
|
||||
}
|
||||
collInfoWoPartition := &datapb.CollectionInfo{
|
||||
ID: collID,
|
||||
Schema: testSchema,
|
||||
Partitions: []UniqueID{},
|
||||
})
|
||||
assert.Nil(t, err)
|
||||
err = meta.AddPartition(id, 10)
|
||||
assert.Nil(t, err)
|
||||
err = meta.AddPartition(id, 10)
|
||||
assert.NotNil(t, err)
|
||||
collection, err := meta.GetCollection(id)
|
||||
assert.Nil(t, err)
|
||||
assert.EqualValues(t, 10, collection.Partitions[0])
|
||||
err = meta.DropPartition(id, 10)
|
||||
assert.Nil(t, err)
|
||||
collection, err = meta.GetCollection(id)
|
||||
assert.Nil(t, err)
|
||||
assert.EqualValues(t, 0, len(collection.Partitions))
|
||||
err = meta.DropPartition(id, 10)
|
||||
assert.NotNil(t, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetCount(t *testing.T) {
|
||||
mockAllocator := newMockAllocator()
|
||||
meta, err := newMemoryMeta(mockAllocator)
|
||||
assert.Nil(t, err)
|
||||
id, err := mockAllocator.allocID()
|
||||
assert.Nil(t, err)
|
||||
segID, err := mockAllocator.allocID()
|
||||
assert.Nil(t, err)
|
||||
nums, err := meta.GetNumRowsOfCollection(id)
|
||||
assert.Nil(t, err)
|
||||
assert.EqualValues(t, 0, nums)
|
||||
segment, err := BuildSegment(id, 100, segID, "c1")
|
||||
assert.Nil(t, err)
|
||||
segment.NumRows = 100
|
||||
err = meta.AddSegment(segment)
|
||||
assert.Nil(t, err)
|
||||
segID, err = mockAllocator.allocID()
|
||||
assert.Nil(t, err)
|
||||
segment, err = BuildSegment(id, 100, segID, "c1")
|
||||
assert.Nil(t, err)
|
||||
segment.NumRows = 300
|
||||
err = meta.AddSegment(segment)
|
||||
assert.Nil(t, err)
|
||||
nums, err = meta.GetNumRowsOfCollection(id)
|
||||
assert.Nil(t, err)
|
||||
assert.EqualValues(t, 400, nums)
|
||||
t.Run("Test Collection", func(t *testing.T) {
|
||||
// check add collection
|
||||
err = meta.AddCollection(collInfo)
|
||||
assert.Nil(t, err)
|
||||
|
||||
// check add existed collection
|
||||
err = meta.AddCollection(collInfo)
|
||||
assert.NotNil(t, err)
|
||||
|
||||
// check has collection
|
||||
has := meta.HasCollection(collID)
|
||||
assert.True(t, has)
|
||||
|
||||
// check partition info
|
||||
collInfo, err = meta.GetCollection(collID)
|
||||
assert.Nil(t, err)
|
||||
assert.EqualValues(t, collID, collInfo.ID)
|
||||
assert.EqualValues(t, testSchema, collInfo.Schema)
|
||||
assert.EqualValues(t, 2, len(collInfo.Partitions))
|
||||
assert.EqualValues(t, partID0, collInfo.Partitions[0])
|
||||
assert.EqualValues(t, partID1, collInfo.Partitions[1])
|
||||
|
||||
// check drop collection
|
||||
err = meta.DropCollection(collID)
|
||||
assert.Nil(t, err)
|
||||
has = meta.HasCollection(collID)
|
||||
assert.False(t, has)
|
||||
_, err = meta.GetCollection(collID)
|
||||
assert.NotNil(t, err)
|
||||
})
|
||||
|
||||
t.Run("Test Partition", func(t *testing.T) {
|
||||
err = meta.AddCollection(collInfoWoPartition)
|
||||
assert.Nil(t, err)
|
||||
|
||||
// check add partition
|
||||
err = meta.AddPartition(collID, partID0)
|
||||
assert.Nil(t, err)
|
||||
err = meta.AddPartition(collID, partID1)
|
||||
assert.Nil(t, err)
|
||||
exist0 := meta.HasPartition(collID, partID0)
|
||||
assert.True(t, exist0)
|
||||
exist1 := meta.HasPartition(collID, partID1)
|
||||
assert.True(t, exist1)
|
||||
|
||||
// check add existed partition
|
||||
err = meta.AddPartition(collID, partID0)
|
||||
assert.NotNil(t, err)
|
||||
|
||||
// check GetCollection
|
||||
collInfo, err = meta.GetCollection(collID)
|
||||
assert.Nil(t, err)
|
||||
assert.EqualValues(t, 2, len(collInfo.Partitions))
|
||||
assert.Contains(t, collInfo.Partitions, partID0)
|
||||
assert.Contains(t, collInfo.Partitions, partID1)
|
||||
|
||||
// check DropPartition
|
||||
err = meta.DropPartition(collID, partID0)
|
||||
assert.Nil(t, err)
|
||||
exist0 = meta.HasPartition(collID, partID0)
|
||||
assert.False(t, exist0)
|
||||
exist1 = meta.HasPartition(collID, partID1)
|
||||
assert.True(t, exist1)
|
||||
|
||||
// check DropPartition twice
|
||||
err = meta.DropPartition(collID, partID0)
|
||||
assert.NotNil(t, err)
|
||||
|
||||
err = meta.DropCollection(collID)
|
||||
assert.Nil(t, err)
|
||||
})
|
||||
|
||||
t.Run("Test Segment", func(t *testing.T) {
|
||||
err = meta.AddCollection(collInfoWoPartition)
|
||||
assert.Nil(t, err)
|
||||
err = meta.AddPartition(collID, partID0)
|
||||
assert.Nil(t, err)
|
||||
|
||||
// create seg0 for partition0, seg0/seg1 for partition1
|
||||
segID0_0, err := mockAllocator.allocID()
|
||||
assert.Nil(t, err)
|
||||
segInfo0_0, err := BuildSegment(collID, partID0, segID0_0, channelName)
|
||||
assert.Nil(t, err)
|
||||
segID1_0, err := mockAllocator.allocID()
|
||||
assert.Nil(t, err)
|
||||
segInfo1_0, err := BuildSegment(collID, partID1, segID1_0, channelName)
|
||||
assert.Nil(t, err)
|
||||
segID1_1, err := mockAllocator.allocID()
|
||||
assert.Nil(t, err)
|
||||
segInfo1_1, err := BuildSegment(collID, partID1, segID1_1, channelName)
|
||||
assert.Nil(t, err)
|
||||
|
||||
// check AddSegment
|
||||
err = meta.AddSegment(segInfo0_0)
|
||||
assert.Nil(t, err)
|
||||
err = meta.AddSegment(segInfo0_0)
|
||||
assert.NotNil(t, err)
|
||||
err = meta.AddSegment(segInfo1_0)
|
||||
assert.Nil(t, err)
|
||||
err = meta.AddSegment(segInfo1_1)
|
||||
assert.Nil(t, err)
|
||||
|
||||
// check GetSegment
|
||||
info0_0, err := meta.GetSegment(segID0_0)
|
||||
assert.Nil(t, err)
|
||||
assert.True(t, proto.Equal(info0_0, segInfo0_0))
|
||||
info1_0, err := meta.GetSegment(segID1_0)
|
||||
assert.Nil(t, err)
|
||||
assert.True(t, proto.Equal(info1_0, segInfo1_0))
|
||||
|
||||
// check GetSegmentsOfCollection
|
||||
segIDs := meta.GetSegmentsOfCollection(collID)
|
||||
assert.EqualValues(t, 3, len(segIDs))
|
||||
assert.Contains(t, segIDs, segID0_0)
|
||||
assert.Contains(t, segIDs, segID1_0)
|
||||
assert.Contains(t, segIDs, segID1_1)
|
||||
|
||||
// check GetSegmentsOfPartition
|
||||
segIDs = meta.GetSegmentsOfPartition(collID, partID0)
|
||||
assert.EqualValues(t, 1, len(segIDs))
|
||||
assert.Contains(t, segIDs, segID0_0)
|
||||
segIDs = meta.GetSegmentsOfPartition(collID, partID1)
|
||||
assert.EqualValues(t, 2, len(segIDs))
|
||||
assert.Contains(t, segIDs, segID1_0)
|
||||
assert.Contains(t, segIDs, segID1_1)
|
||||
|
||||
// check DropSegment
|
||||
err = meta.DropSegment(segID1_0)
|
||||
assert.Nil(t, err)
|
||||
segIDs = meta.GetSegmentsOfPartition(collID, partID1)
|
||||
assert.EqualValues(t, 1, len(segIDs))
|
||||
assert.Contains(t, segIDs, segID1_1)
|
||||
|
||||
// check OpenSegment/SealSegment/FlushSegment
|
||||
err = meta.OpenSegment(segID0_0, 100)
|
||||
assert.Nil(t, err)
|
||||
err = meta.SealSegment(segID0_0, 200)
|
||||
assert.Nil(t, err)
|
||||
err = meta.FlushSegment(segID0_0, 300)
|
||||
assert.Nil(t, err)
|
||||
|
||||
info0_0, err = meta.GetSegment(segID0_0)
|
||||
assert.Nil(t, err)
|
||||
assert.NotZero(t, info0_0.OpenTime)
|
||||
assert.NotZero(t, info0_0.SealedTime)
|
||||
assert.NotZero(t, info0_0.FlushedTime)
|
||||
|
||||
err = meta.DropPartition(collID, partID0)
|
||||
assert.Nil(t, err)
|
||||
err = meta.DropCollection(collID)
|
||||
assert.Nil(t, err)
|
||||
})
|
||||
|
||||
t.Run("Test GetCount", func(t *testing.T) {
|
||||
const rowCount0 = 100
|
||||
const rowCount1 = 300
|
||||
const dim = 1024
|
||||
|
||||
// no segment
|
||||
nums, err := meta.GetNumRowsOfCollection(collID)
|
||||
assert.Nil(t, err)
|
||||
assert.EqualValues(t, 0, nums)
|
||||
memSize, err := meta.GetMemSizeOfCollection(collID)
|
||||
assert.Nil(t, err)
|
||||
assert.EqualValues(t, 0, memSize)
|
||||
|
||||
// add seg1 with 100 rows
|
||||
segID0, err := mockAllocator.allocID()
|
||||
assert.Nil(t, err)
|
||||
segInfo0, err := BuildSegment(collID, partID0, segID0, channelName)
|
||||
assert.Nil(t, err)
|
||||
segInfo0.NumRows = rowCount0
|
||||
segInfo0.MemSize = rowCount0 * dim * 4
|
||||
err = meta.AddSegment(segInfo0)
|
||||
assert.Nil(t, err)
|
||||
|
||||
// update seg1 to 300 rows
|
||||
segInfo0.NumRows = rowCount1
|
||||
segInfo0.MemSize = rowCount1 * dim * 4
|
||||
err = meta.UpdateSegment(segInfo0)
|
||||
assert.Nil(t, err)
|
||||
|
||||
nums, err = meta.GetNumRowsOfCollection(collID)
|
||||
assert.Nil(t, err)
|
||||
assert.EqualValues(t, rowCount1, nums)
|
||||
memSize, err = meta.GetMemSizeOfCollection(collID)
|
||||
assert.Nil(t, err)
|
||||
assert.EqualValues(t, rowCount1*dim*4, memSize)
|
||||
|
||||
// check update non-exist segment
|
||||
segInfoNonExist := segInfo0
|
||||
segInfoNonExist.ID, err = mockAllocator.allocID()
|
||||
assert.Nil(t, err)
|
||||
err = meta.UpdateSegment(segInfo0)
|
||||
assert.NotNil(t, err)
|
||||
|
||||
// add seg2 with 300 rows
|
||||
segID1, err := mockAllocator.allocID()
|
||||
assert.Nil(t, err)
|
||||
segInfo1, err := BuildSegment(collID, partID0, segID1, channelName)
|
||||
assert.Nil(t, err)
|
||||
segInfo1.NumRows = rowCount1
|
||||
segInfo1.MemSize = rowCount1 * dim * 4
|
||||
err = meta.AddSegment(segInfo1)
|
||||
assert.Nil(t, err)
|
||||
|
||||
// check partition/collection statistics
|
||||
nums, err = meta.GetNumRowsOfPartition(collID, partID0)
|
||||
assert.Nil(t, err)
|
||||
assert.EqualValues(t, (rowCount1 + rowCount1), nums)
|
||||
nums, err = meta.GetNumRowsOfCollection(collID)
|
||||
assert.Nil(t, err)
|
||||
assert.EqualValues(t, (rowCount1 + rowCount1), nums)
|
||||
memSize, err = meta.GetMemSizeOfCollection(collID)
|
||||
assert.Nil(t, err)
|
||||
assert.EqualValues(t, (rowCount1+rowCount1)*dim*4, memSize)
|
||||
})
|
||||
|
||||
t.Run("Test Invalid", func(t *testing.T) {
|
||||
collIDInvalid := UniqueID(10000)
|
||||
partIDInvalid := UniqueID(10001)
|
||||
segIDInvalid := UniqueID(10002)
|
||||
|
||||
// check drop non-exist collection
|
||||
err = meta.DropCollection(collID)
|
||||
assert.NotNil(t, err)
|
||||
|
||||
// add partition wo collection
|
||||
err = meta.AddPartition(collID, partID0)
|
||||
assert.NotNil(t, err)
|
||||
|
||||
// has partition wo collection
|
||||
exist := meta.HasPartition(collID, partID0)
|
||||
assert.False(t, exist)
|
||||
|
||||
err = meta.AddCollection(collInfo)
|
||||
assert.Nil(t, err)
|
||||
|
||||
// check drop non-exist partition
|
||||
err = meta.DropPartition(collIDInvalid, partID0)
|
||||
assert.NotNil(t, err)
|
||||
err = meta.DropPartition(collID, partIDInvalid)
|
||||
assert.NotNil(t, err)
|
||||
|
||||
// check drop non-exist segment
|
||||
err = meta.DropSegment(segIDInvalid)
|
||||
assert.NotNil(t, err)
|
||||
|
||||
// check open non-exist segment
|
||||
err = meta.OpenSegment(segIDInvalid, 100)
|
||||
assert.NotNil(t, err)
|
||||
|
||||
// check seal non-exist segment
|
||||
err = meta.SealSegment(segIDInvalid, 200)
|
||||
assert.NotNil(t, err)
|
||||
|
||||
// check flush non-exist segment
|
||||
err = meta.FlushSegment(segIDInvalid, 300)
|
||||
assert.NotNil(t, err)
|
||||
|
||||
err = meta.DropCollection(collID)
|
||||
assert.Nil(t, err)
|
||||
})
|
||||
}
|
||||
|
|
|
|||
|
|
@ -321,6 +321,15 @@ func (allocator *segmentAllocator) SealSegment(ctx context.Context, segmentID Un
|
|||
return nil
|
||||
}
|
||||
|
||||
func (allocator *segmentAllocator) HasSegment(ctx context.Context, segmentID UniqueID) bool {
|
||||
sp, _ := trace.StartSpanFromContext(ctx)
|
||||
defer sp.Finish()
|
||||
allocator.mu.Lock()
|
||||
defer allocator.mu.Unlock()
|
||||
_, ok := allocator.segments[segmentID]
|
||||
return ok
|
||||
}
|
||||
|
||||
func (allocator *segmentAllocator) DropSegment(ctx context.Context, segmentID UniqueID) {
|
||||
sp, _ := trace.StartSpanFromContext(ctx)
|
||||
defer sp.Finish()
|
||||
|
|
|
|||
|
|
@ -81,7 +81,6 @@ func (watcher *dataNodeTimeTickWatcher) StartBackgroundLoop(ctx context.Context)
|
|||
case msg := <-watcher.msgQueue:
|
||||
if err := watcher.handleTimeTickMsg(msg); err != nil {
|
||||
log.Error("handle time tick error", zap.Error(err))
|
||||
continue
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -25,9 +25,12 @@ import (
|
|||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestDataNodeTTWatcher(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
func TestWatcher(t *testing.T) {
|
||||
const collID = UniqueID(0)
|
||||
const partID = UniqueID(100)
|
||||
|
||||
Params.Init()
|
||||
|
||||
cluster := newDataNodeCluster()
|
||||
defer cluster.ShutDownClients()
|
||||
schema := newTestSchema()
|
||||
|
|
@ -36,75 +39,103 @@ func TestDataNodeTTWatcher(t *testing.T) {
|
|||
assert.Nil(t, err)
|
||||
segAllocator := newSegmentAllocator(meta, allocator)
|
||||
assert.Nil(t, err)
|
||||
watcher := newDataNodeTimeTickWatcher(meta, segAllocator, cluster)
|
||||
|
||||
id, err := allocator.allocID()
|
||||
assert.Nil(t, err)
|
||||
err = meta.AddCollection(&datapb.CollectionInfo{
|
||||
collInfo := &datapb.CollectionInfo{
|
||||
Schema: schema,
|
||||
ID: id,
|
||||
})
|
||||
assert.Nil(t, err)
|
||||
|
||||
cases := []struct {
|
||||
sealed bool
|
||||
allocation bool
|
||||
expired bool
|
||||
expected bool
|
||||
}{
|
||||
{false, false, true, false},
|
||||
{false, true, true, false},
|
||||
{false, true, false, false},
|
||||
{true, false, true, true},
|
||||
{true, true, false, false},
|
||||
{true, true, true, true},
|
||||
ID: collID,
|
||||
}
|
||||
|
||||
segmentIDs := make([]UniqueID, len(cases))
|
||||
for i, c := range cases {
|
||||
segID, err := allocator.allocID()
|
||||
segmentIDs[i] = segID
|
||||
assert.Nil(t, err)
|
||||
segmentInfo, err := BuildSegment(id, 100, segID, "channel"+strconv.Itoa(i))
|
||||
assert.Nil(t, err)
|
||||
err = meta.AddSegment(segmentInfo)
|
||||
assert.Nil(t, err)
|
||||
err = segAllocator.OpenSegment(ctx, segmentInfo)
|
||||
assert.Nil(t, err)
|
||||
if c.allocation && c.expired {
|
||||
_, _, _, err := segAllocator.AllocSegment(ctx, id, 100, "channel"+strconv.Itoa(i), 100)
|
||||
assert.Nil(t, err)
|
||||
}
|
||||
}
|
||||
t.Run("Test ProxyTimeTickWatcher", func(t *testing.T) {
|
||||
proxyWatcher := newProxyTimeTickWatcher(segAllocator)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
go proxyWatcher.StartBackgroundLoop(ctx)
|
||||
|
||||
time.Sleep(time.Duration(Params.SegIDAssignExpiration+1000) * time.Millisecond)
|
||||
for i, c := range cases {
|
||||
if c.allocation && !c.expired {
|
||||
_, _, _, err := segAllocator.AllocSegment(ctx, id, 100, "channel"+strconv.Itoa(i), 100)
|
||||
assert.Nil(t, err)
|
||||
}
|
||||
if c.sealed {
|
||||
err := segAllocator.SealSegment(ctx, segmentIDs[i])
|
||||
assert.Nil(t, err)
|
||||
}
|
||||
}
|
||||
ts, err := allocator.allocTimestamp()
|
||||
assert.Nil(t, err)
|
||||
|
||||
err = watcher.handleTimeTickMsg(&msgstream.TimeTickMsg{
|
||||
BaseMsg: msgstream.BaseMsg{
|
||||
HashValues: []uint32{0},
|
||||
},
|
||||
TimeTickMsg: internalpb.TimeTickMsg{
|
||||
Base: &commonpb.MsgBase{
|
||||
MsgType: commonpb.MsgType_TimeTick,
|
||||
Timestamp: ts,
|
||||
msg := &msgstream.TimeTickMsg{
|
||||
TimeTickMsg: internalpb.TimeTickMsg{
|
||||
Base: &commonpb.MsgBase{
|
||||
MsgType: commonpb.MsgType_TimeTick,
|
||||
Timestamp: 100,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
proxyWatcher.Watch(msg)
|
||||
time.Sleep(time.Second)
|
||||
})
|
||||
|
||||
t.Run("Test DataNodeTimeTickWatcher", func(t *testing.T) {
|
||||
datanodeWatcher := newDataNodeTimeTickWatcher(meta, segAllocator, cluster)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
go datanodeWatcher.StartBackgroundLoop(ctx)
|
||||
|
||||
err = meta.AddCollection(collInfo)
|
||||
assert.Nil(t, err)
|
||||
|
||||
cases := []struct {
|
||||
sealed bool
|
||||
allocation bool
|
||||
expired bool
|
||||
expected bool
|
||||
}{
|
||||
{false, false, true, false},
|
||||
{false, true, true, false},
|
||||
{false, true, false, false},
|
||||
{true, false, true, true},
|
||||
{true, true, false, false},
|
||||
{true, true, true, true},
|
||||
}
|
||||
|
||||
segIDs := make([]UniqueID, len(cases))
|
||||
for i, c := range cases {
|
||||
segID, err := allocator.allocID()
|
||||
assert.Nil(t, err)
|
||||
segIDs[i] = segID
|
||||
segInfo, err := BuildSegment(collID, partID, segID, "channel"+strconv.Itoa(i))
|
||||
assert.Nil(t, err)
|
||||
err = meta.AddSegment(segInfo)
|
||||
assert.Nil(t, err)
|
||||
err = segAllocator.OpenSegment(ctx, segInfo)
|
||||
assert.Nil(t, err)
|
||||
if c.allocation && c.expired {
|
||||
_, _, _, err := segAllocator.AllocSegment(ctx, collID, partID, "channel"+strconv.Itoa(i), 100)
|
||||
assert.Nil(t, err)
|
||||
}
|
||||
}
|
||||
|
||||
time.Sleep(time.Duration(Params.SegIDAssignExpiration+1000) * time.Millisecond)
|
||||
for i, c := range cases {
|
||||
if c.allocation && !c.expired {
|
||||
_, _, _, err := segAllocator.AllocSegment(ctx, collID, partID, "channel"+strconv.Itoa(i), 100)
|
||||
assert.Nil(t, err)
|
||||
}
|
||||
if c.sealed {
|
||||
err := segAllocator.SealSegment(ctx, segIDs[i])
|
||||
assert.Nil(t, err)
|
||||
}
|
||||
}
|
||||
ts, err := allocator.allocTimestamp()
|
||||
assert.Nil(t, err)
|
||||
|
||||
ttMsg := &msgstream.TimeTickMsg{
|
||||
BaseMsg: msgstream.BaseMsg{
|
||||
HashValues: []uint32{0},
|
||||
},
|
||||
TimeTickMsg: internalpb.TimeTickMsg{
|
||||
Base: &commonpb.MsgBase{
|
||||
MsgType: commonpb.MsgType_TimeTick,
|
||||
Timestamp: ts,
|
||||
},
|
||||
},
|
||||
}
|
||||
datanodeWatcher.Watch(ttMsg)
|
||||
|
||||
time.Sleep(time.Second)
|
||||
|
||||
// check flushed segments been removed from segAllocator
|
||||
for i, c := range cases {
|
||||
ok := segAllocator.HasSegment(ctx, segIDs[i])
|
||||
assert.EqualValues(t, !c.expected, ok)
|
||||
}
|
||||
})
|
||||
assert.Nil(t, err)
|
||||
for i, c := range cases {
|
||||
_, ok := segAllocator.segments[segmentIDs[i]]
|
||||
assert.EqualValues(t, !c.expected, ok)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in New Issue