Files
wehub-resource-sync 498b235461
Build and test / Build and test AMD64 Ubuntu 22.04 (push) Failing after 0s
Publish Builder / amazonlinux2023 (push) Failing after 1s
Build and test / UT for Go (push) Has been skipped
Publish KRTE Images / KRTE (push) Failing after 1s
Build and test / Integration Test (push) Has been skipped
Build and test / Upload Code Coverage (push) Has been skipped
Publish Builder / rockylinux9 (push) Failing after 1s
Publish Builder / ubuntu22.04 (push) Failing after 0s
Publish Builder / ubuntu24.04 (push) Failing after 0s
Publish Gpu Builder / publish-gpu-builder (push) Failing after 1s
Publish Test Images / PyTest (push) Failing after 0s
Build and test / UT for Cpp (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 12:31:17 +08:00

411 lines
14 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 utils
import (
"context"
"testing"
"time"
"github.com/blang/semver/v4"
"github.com/bytedance/mockey"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/suite"
"github.com/milvus-io/milvus/internal/querycoordv2/meta"
"github.com/milvus-io/milvus/internal/querycoordv2/session"
"github.com/milvus-io/milvus/internal/storagev2/packed"
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
)
type UtilTestSuite struct {
suite.Suite
nodeMgr *session.NodeManager
}
func (suite *UtilTestSuite) SetupTest() {
suite.nodeMgr = session.NewNodeManager()
}
func (suite *UtilTestSuite) setNodeAvailable(nodes ...int64) {
for _, node := range nodes {
nodeInfo := session.NewNodeInfo(session.ImmutableNodeInfo{
NodeID: node,
Address: "",
Hostname: "localhost",
})
nodeInfo.SetLastHeartbeat(time.Now())
suite.nodeMgr.Add(nodeInfo)
}
}
func (suite *UtilTestSuite) TestCheckLeaderAvaliable() {
leadview := &meta.LeaderView{
ID: 1,
Channel: "test",
Segments: map[int64]*querypb.SegmentDist{2: {NodeID: 2}},
TargetVersion: 1011,
}
mockTargetManager := meta.NewMockTargetManager(suite.T())
mockTargetManager.EXPECT().GetSealedSegmentsByChannel(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(map[int64]*datapb.SegmentInfo{
2: {
ID: 2,
InsertChannel: "test",
},
}).Maybe()
mockTargetManager.EXPECT().GetCollectionTargetVersion(mock.Anything, mock.Anything, mock.Anything).Return(1011).Maybe()
suite.setNodeAvailable(1, 2)
err := CheckDelegatorDataReady(suite.nodeMgr, mockTargetManager, leadview, meta.CurrentTarget)
suite.NoError(err)
}
func (suite *UtilTestSuite) TestCheckLeaderAvaliableFailed() {
suite.Run("leader not available", func() {
leadview := &meta.LeaderView{
ID: 1,
Channel: "test",
Segments: map[int64]*querypb.SegmentDist{2: {NodeID: 2}},
TargetVersion: 1011,
}
mockTargetManager := meta.NewMockTargetManager(suite.T())
mockTargetManager.EXPECT().GetSealedSegmentsByChannel(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(map[int64]*datapb.SegmentInfo{
2: {
ID: 2,
InsertChannel: "test",
},
}).Maybe()
mockTargetManager.EXPECT().GetCollectionTargetVersion(mock.Anything, mock.Anything, mock.Anything).Return(1011).Maybe()
// leader nodeID=1 not available
suite.setNodeAvailable(2)
err := CheckDelegatorDataReady(suite.nodeMgr, mockTargetManager, leadview, meta.CurrentTarget)
suite.Error(err)
})
suite.Run("shard worker not available", func() {
leadview := &meta.LeaderView{
ID: 11111,
Channel: "test",
Segments: map[int64]*querypb.SegmentDist{2: {NodeID: 2}},
TargetVersion: 1011,
}
mockTargetManager := meta.NewMockTargetManager(suite.T())
mockTargetManager.EXPECT().GetSealedSegmentsByChannel(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(map[int64]*datapb.SegmentInfo{
2: {
ID: 2,
InsertChannel: "test",
},
}).Maybe()
mockTargetManager.EXPECT().GetCollectionTargetVersion(mock.Anything, mock.Anything, mock.Anything).Return(1011).Maybe()
// leader nodeID=2 not available
suite.setNodeAvailable(1)
err := CheckDelegatorDataReady(suite.nodeMgr, mockTargetManager, leadview, meta.CurrentTarget)
suite.Error(err)
})
suite.Run("segment lacks", func() {
leadview := &meta.LeaderView{
ID: 1,
Channel: "test",
Segments: map[int64]*querypb.SegmentDist{2: {NodeID: 2}},
TargetVersion: 1011,
}
mockTargetManager := meta.NewMockTargetManager(suite.T())
mockTargetManager.EXPECT().GetSealedSegmentsByChannel(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(map[int64]*datapb.SegmentInfo{
// target segmentID=1 not in leadView
1: {
ID: 1,
InsertChannel: "test",
},
}).Maybe()
mockTargetManager.EXPECT().GetCollectionTargetVersion(mock.Anything, mock.Anything, mock.Anything).Return(1011).Maybe()
suite.setNodeAvailable(1, 2)
err := CheckDelegatorDataReady(suite.nodeMgr, mockTargetManager, leadview, meta.CurrentTarget)
suite.Error(err)
})
suite.Run("target version not synced", func() {
leadview := &meta.LeaderView{
ID: 1,
Channel: "test",
Segments: map[int64]*querypb.SegmentDist{2: {NodeID: 2}},
}
mockTargetManager := meta.NewMockTargetManager(suite.T())
mockTargetManager.EXPECT().GetSealedSegmentsByChannel(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(map[int64]*datapb.SegmentInfo{
// target segmentID=1 not in leadView
1: {
ID: 1,
InsertChannel: "test",
},
}).Maybe()
mockTargetManager.EXPECT().GetCollectionTargetVersion(mock.Anything, mock.Anything, mock.Anything).Return(1011).Maybe()
suite.setNodeAvailable(1, 2)
err := CheckDelegatorDataReady(suite.nodeMgr, mockTargetManager, leadview, meta.CurrentTarget)
suite.Error(err)
})
suite.Run("catching up streaming data", func() {
leadview := &meta.LeaderView{
ID: 1,
Channel: "test",
Segments: map[int64]*querypb.SegmentDist{2: {NodeID: 2}},
TargetVersion: 1011,
Status: &querypb.LeaderViewStatus{
Serviceable: true,
CatchingUpStreamingData: true, // still catching up
},
}
// When catching up, function returns early without calling targetMgr
// so we can pass nil as targetMgr
suite.setNodeAvailable(1, 2)
err := CheckDelegatorDataReady(suite.nodeMgr, nil, leadview, meta.CurrentTarget)
suite.Error(err)
suite.Contains(err.Error(), "catching up streaming data")
})
suite.Run("caught up streaming data", func() {
leadview := &meta.LeaderView{
ID: 1,
Channel: "test",
Segments: map[int64]*querypb.SegmentDist{2: {NodeID: 2}},
TargetVersion: 1011,
Status: &querypb.LeaderViewStatus{
Serviceable: true,
CatchingUpStreamingData: false, // already caught up
},
}
// Use mockey to mock TargetManager.GetSealedSegmentsByChannel
targetMgr := &meta.TargetManager{}
mockGetSealedSegments := mockey.Mock(mockey.GetMethod(targetMgr, "GetSealedSegmentsByChannel")).
Return(map[int64]*datapb.SegmentInfo{
2: {
ID: 2,
InsertChannel: "test",
},
}).Build()
defer mockGetSealedSegments.UnPatch()
suite.setNodeAvailable(1, 2)
err := CheckDelegatorDataReady(suite.nodeMgr, targetMgr, leadview, meta.CurrentTarget)
suite.NoError(err)
})
}
func (suite *UtilTestSuite) TestGetChannelRWAndRONodesFor260() {
nodes := []int64{1, 2, 3, 4, 5}
nodeManager := session.NewNodeManager()
r := meta.NewReplica(&querypb.Replica{
Nodes: nodes,
})
rwNodes, roNodes := GetChannelRWAndRONodesFor260(r, nodeManager)
suite.ElementsMatch(rwNodes, []int64{})
suite.ElementsMatch(roNodes, []int64{1, 2, 3, 4, 5})
nodeManager.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
NodeID: 1,
Address: "127.0.0.1:0",
Hostname: "localhost",
Version: semver.MustParse("2.5.0"),
}))
nodeManager.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
NodeID: 2,
Address: "127.0.0.1:0",
Hostname: "localhost",
Version: semver.MustParse("2.6.0-dev"),
}))
rwNodes, roNodes = GetChannelRWAndRONodesFor260(r, nodeManager)
suite.ElementsMatch(rwNodes, []int64{})
suite.ElementsMatch(roNodes, []int64{1, 3, 4, 5})
}
func (suite *UtilTestSuite) TestFilterOutNodeLessThan260() {
nodes := []int64{1, 2, 3, 4, 5}
nodeManager := session.NewNodeManager()
filteredNodes := filterNodeLessThan260(nodes, nodeManager)
suite.ElementsMatch(filteredNodes, []int64{1, 2, 3, 4, 5})
nodeManager.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
NodeID: 1,
Address: "127.0.0.1:0",
Hostname: "localhost",
Version: semver.MustParse("2.5.0"),
}))
filteredNodes = filterNodeLessThan260(nodes, nodeManager)
suite.ElementsMatch(filteredNodes, []int64{1, 2, 3, 4, 5})
nodeManager.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
NodeID: 2,
Address: "127.0.0.1:0",
Hostname: "localhost",
Version: semver.MustParse("2.6.0-dev"),
}))
filteredNodes = filterNodeLessThan260(nodes, nodeManager)
suite.ElementsMatch(filteredNodes, []int64{1, 3, 4, 5})
nodeManager.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
NodeID: 3,
Address: "127.0.0.1:0",
Hostname: "localhost",
Version: semver.MustParse("2.6.0"),
}))
filteredNodes = filterNodeLessThan260(nodes, nodeManager)
suite.ElementsMatch(filteredNodes, []int64{1, 4, 5})
}
func (suite *UtilTestSuite) TestCheckSegmentDataReady_ManifestComparison() {
basePath := "/data/insert_log/col/part/seg"
collectionID := int64(100)
segmentID := int64(200)
nodeID := int64(1)
newDistManager := func(manifestPath string) *meta.DistributionManager {
dm := meta.NewDistributionManager(session.NewNodeManager())
dm.SegmentDistManager.Update(nodeID, &meta.Segment{
SegmentInfo: &datapb.SegmentInfo{
ID: segmentID,
CollectionID: collectionID,
},
Node: nodeID,
ManifestPath: manifestPath,
})
return dm
}
newTargetMgr := func(manifestPath string) meta.TargetManagerInterface {
m := meta.NewMockTargetManager(suite.T())
m.EXPECT().GetSealedSegmentsByCollection(mock.Anything, collectionID, mock.Anything).
Return(map[int64]*datapb.SegmentInfo{
segmentID: {
ID: segmentID,
CollectionID: collectionID,
ManifestPath: manifestPath,
},
}).Maybe()
return m
}
suite.Run("same manifest version - ready", func() {
manifest := packed.MarshalManifestPath(basePath, 5)
err := CheckSegmentDataReady(context.Background(), collectionID, newDistManager(manifest), newTargetMgr(manifest), meta.NextTarget)
suite.NoError(err)
})
suite.Run("dist newer than target - ready", func() {
distManifest := packed.MarshalManifestPath(basePath, 10)
targetManifest := packed.MarshalManifestPath(basePath, 5)
err := CheckSegmentDataReady(context.Background(), collectionID, newDistManager(distManifest), newTargetMgr(targetManifest), meta.NextTarget)
suite.NoError(err)
})
suite.Run("dist older than target - not ready", func() {
distManifest := packed.MarshalManifestPath(basePath, 1)
targetManifest := packed.MarshalManifestPath(basePath, 5)
err := CheckSegmentDataReady(context.Background(), collectionID, newDistManager(distManifest), newTargetMgr(targetManifest), meta.NextTarget)
suite.Error(err)
})
suite.Run("both empty manifest - ready", func() {
err := CheckSegmentDataReady(context.Background(), collectionID, newDistManager(""), newTargetMgr(""), meta.NextTarget)
suite.NoError(err)
})
suite.Run("segment not in dist - not ready", func() {
dm := meta.NewDistributionManager(session.NewNodeManager())
targetManifest := packed.MarshalManifestPath(basePath, 5)
err := CheckSegmentDataReady(context.Background(), collectionID, dm, newTargetMgr(targetManifest), meta.NextTarget)
suite.Error(err)
})
}
func (suite *UtilTestSuite) TestCheckSegmentDataReady_DataVersion() {
basePath := "/data/insert_log/col/part/seg"
collectionID := int64(100)
segmentID := int64(200)
nodeID := int64(1)
// Use matching manifest in dist and target so DataVersion is the only gate.
manifest := packed.MarshalManifestPath(basePath, 5)
newDistManager := func(distDataVersion *int32) *meta.DistributionManager {
dm := meta.NewDistributionManager(session.NewNodeManager())
dm.SegmentDistManager.Update(nodeID, &meta.Segment{
SegmentInfo: &datapb.SegmentInfo{
ID: segmentID,
CollectionID: collectionID,
},
Node: nodeID,
ManifestPath: manifest,
DataVersion: distDataVersion,
})
return dm
}
newTargetMgr := func(targetDataVersion int32) meta.TargetManagerInterface {
m := meta.NewMockTargetManager(suite.T())
m.EXPECT().GetSealedSegmentsByCollection(mock.Anything, collectionID, mock.Anything).
Return(map[int64]*datapb.SegmentInfo{
segmentID: {
ID: segmentID,
CollectionID: collectionID,
ManifestPath: manifest,
DataVersion: targetDataVersion,
},
}).Maybe()
return m
}
int32Ptr := func(v int32) *int32 { return &v }
suite.Run("dist data version older than target - not ready", func() {
err := CheckSegmentDataReady(context.Background(), collectionID,
newDistManager(int32Ptr(1)), newTargetMgr(5), meta.NextTarget)
suite.Error(err)
})
suite.Run("dist data version equals target - ready", func() {
err := CheckSegmentDataReady(context.Background(), collectionID,
newDistManager(int32Ptr(5)), newTargetMgr(5), meta.NextTarget)
suite.NoError(err)
})
suite.Run("dist data version newer than target - ready", func() {
err := CheckSegmentDataReady(context.Background(), collectionID,
newDistManager(int32Ptr(10)), newTargetMgr(5), meta.NextTarget)
suite.NoError(err)
})
// mixed-version rollout: an old QueryNode does not report DataVersion,
// so the dist-side pointer is nil. Skip the check rather than loop Reopen forever.
suite.Run("dist data version nil (old QueryNode) - ready", func() {
err := CheckSegmentDataReady(context.Background(), collectionID,
newDistManager(nil), newTargetMgr(5), meta.NextTarget)
suite.NoError(err)
})
suite.Run("dist data version nil and target zero - ready", func() {
err := CheckSegmentDataReady(context.Background(), collectionID,
newDistManager(nil), newTargetMgr(0), meta.NextTarget)
suite.NoError(err)
})
}
func TestUtilSuite(t *testing.T) {
suite.Run(t, new(UtilTestSuite))
}