From 715b5153b84db24d23ae5a5af233b4607e1d3460 Mon Sep 17 00:00:00 2001 From: wei liu Date: Thu, 7 Aug 2025 12:13:40 +0800 Subject: [PATCH] enhance: Improve delegator serviceable check logic in PinReadableSegments (#43768) issue: #43767 - Enhance serviceable check logic to properly handle full vs partial result requirements - For full result (requiredLoadRatio >= 1.0): check queryView.Serviceable() - For partial result (requiredLoadRatio < 1.0): check load ratio satisfaction - Add comprehensive unit tests covering all serviceable check scenarios This enhancement ensures delegator correctly validates serviceability based on the requested result completeness, improving reliability of query operations. Signed-off-by: Wei Liu --- .../querynodev2/delegator/distribution.go | 21 ++- .../delegator/distribution_test.go | 173 ++++++++++++++++++ 2 files changed, 190 insertions(+), 4 deletions(-) diff --git a/internal/querynodev2/delegator/distribution.go b/internal/querynodev2/delegator/distribution.go index 574a2fed56..22705fef89 100644 --- a/internal/querynodev2/delegator/distribution.go +++ b/internal/querynodev2/delegator/distribution.go @@ -17,7 +17,6 @@ package delegator import ( - "fmt" "sync" "github.com/samber/lo" @@ -164,9 +163,23 @@ func (d *distribution) PinReadableSegments(requiredLoadRatio float64, partitions d.mut.RLock() defer d.mut.RUnlock() - if d.queryView.GetLoadedRatio() < requiredLoadRatio { - return nil, nil, nil, -1, merr.WrapErrChannelNotAvailable(d.channelName, - fmt.Sprintf("channel distribution is not serviceable, required load ratio is %f, current load ratio is %f", requiredLoadRatio, d.queryView.GetLoadedRatio())) + requireFullResult := requiredLoadRatio >= 1.0 + loadRatioSatisfy := d.queryView.GetLoadedRatio() >= requiredLoadRatio + var isServiceable bool + if requireFullResult { + isServiceable = d.queryView.Serviceable() + } else { + isServiceable = loadRatioSatisfy + } + + if !isServiceable { + log.Warn("channel distribution is not serviceable", + zap.String("channel", d.channelName), + zap.Float64("requiredLoadRatio", requiredLoadRatio), + zap.Float64("currentLoadRatio", d.queryView.GetLoadedRatio()), + zap.Bool("serviceable", d.queryView.Serviceable()), + ) + return nil, nil, nil, -1, merr.WrapErrChannelNotAvailable(d.channelName, "channel distribution is not serviceable") } current := d.current.Load() diff --git a/internal/querynodev2/delegator/distribution_test.go b/internal/querynodev2/delegator/distribution_test.go index c32ecdbc58..09ba7dd7d6 100644 --- a/internal/querynodev2/delegator/distribution_test.go +++ b/internal/querynodev2/delegator/distribution_test.go @@ -1410,3 +1410,176 @@ func (s *DistributionSuite) TestSyncTargetVersion_RedundantGrowingLogic() { }) } } + +func TestPinReadableSegments(t *testing.T) { + // Create test distribution + queryView := NewChannelQueryView(nil, nil, []int64{1}, initialTargetVersion) + dist := NewDistribution("test-channel", queryView) + + // Add some test segments + dist.AddDistributions([]SegmentEntry{ + {NodeID: 1, SegmentID: 1, PartitionID: 1, Version: 1}, + {NodeID: 1, SegmentID: 2, PartitionID: 1, Version: 1}, + }...) + + // Setup query view + dist.SyncTargetVersion(&querypb.SyncAction{ + TargetVersion: 1000, + SealedSegmentRowCount: map[int64]int64{1: 100, 2: 100}, + GrowingInTarget: []int64{}, + }, []int64{1}) + + // Test case 1: requireFullResult=true, Serviceable=false + mockServiceable := mockey.Mock((*channelQueryView).Serviceable).Return(false).Build() + mockGetLoadedRatio := mockey.Mock((*channelQueryView).GetLoadedRatio).Return(0.8).Build() + + sealed, growing, _, _, err := dist.PinReadableSegments(1.0, 1) + + assert.Error(t, err) + assert.Nil(t, sealed) + assert.Nil(t, growing) + assert.Contains(t, err.Error(), "channel distribution is not serviceable") + + // Test case 2: requireFullResult=true, Serviceable=true + mockServiceable.UnPatch() + mockServiceable = mockey.Mock((*channelQueryView).Serviceable).Return(true).Build() + + sealed, growing, _, _, err = dist.PinReadableSegments(1.0, 1) + + assert.NoError(t, err) + assert.NotNil(t, sealed) + assert.NotNil(t, growing) + + // Test case 3: requireFullResult=false, loadRatioSatisfy=false + mockServiceable.UnPatch() + mockGetLoadedRatio.UnPatch() + mockGetLoadedRatio = mockey.Mock((*channelQueryView).GetLoadedRatio).Return(0.5).Build() + + sealed, growing, _, _, err = dist.PinReadableSegments(0.8, 1) + + assert.Error(t, err) + assert.Nil(t, sealed) + assert.Nil(t, growing) + assert.Contains(t, err.Error(), "channel distribution is not serviceable") + + // Test case 4: requireFullResult=false, loadRatioSatisfy=true + mockGetLoadedRatio.UnPatch() + mockGetLoadedRatio = mockey.Mock((*channelQueryView).GetLoadedRatio).Return(0.9).Build() + defer mockGetLoadedRatio.UnPatch() + + sealed, growing, _, _, err = dist.PinReadableSegments(0.8, 1) + + assert.NoError(t, err) + assert.NotNil(t, sealed) + assert.NotNil(t, growing) +} + +func TestPinReadableSegments_ServiceableLogic(t *testing.T) { + // Create test distribution + queryView := NewChannelQueryView(nil, nil, []int64{1}, initialTargetVersion) + dist := NewDistribution("test-channel", queryView) + + // Add test segments + dist.AddDistributions([]SegmentEntry{ + {NodeID: 1, SegmentID: 1, PartitionID: 1, Version: 1}, + }...) + + // Setup query view + dist.SyncTargetVersion(&querypb.SyncAction{ + TargetVersion: 1000, + SealedSegmentRowCount: map[int64]int64{1: 100}, + GrowingInTarget: []int64{}, + }, []int64{1}) + + // Test case: requireFullResult=true, Serviceable=false, GetLoadedRatio=1.0 + // This tests the case where load ratio is satisfied but serviceable is false + mockServiceable := mockey.Mock((*channelQueryView).Serviceable).Return(false).Build() + mockGetLoadedRatio := mockey.Mock((*channelQueryView).GetLoadedRatio).Return(1.0).Build() + defer mockServiceable.UnPatch() + defer mockGetLoadedRatio.UnPatch() + + sealed, growing, _, _, err := dist.PinReadableSegments(1.0, 1) + + assert.Error(t, err) + assert.Nil(t, sealed) + assert.Nil(t, growing) + assert.Contains(t, err.Error(), "channel distribution is not serviceable") +} + +func TestPinReadableSegments_LoadRatioLogic(t *testing.T) { + // Create test distribution + queryView := NewChannelQueryView(nil, nil, []int64{1}, initialTargetVersion) + dist := NewDistribution("test-channel", queryView) + + // Add test segments + dist.AddDistributions([]SegmentEntry{ + {NodeID: 1, SegmentID: 1, PartitionID: 1, Version: 1}, + }...) + + // Setup query view + dist.SyncTargetVersion(&querypb.SyncAction{ + TargetVersion: 1000, + SealedSegmentRowCount: map[int64]int64{1: 100}, + GrowingInTarget: []int64{}, + }, []int64{1}) + + // Test case: requireFullResult=false, loadRatioSatisfy=false + // This tests the case where partial result is requested but load ratio is insufficient + mockGetLoadedRatio := mockey.Mock((*channelQueryView).GetLoadedRatio).Return(0.3).Build() + defer mockGetLoadedRatio.UnPatch() + + sealed, growing, _, _, err := dist.PinReadableSegments(0.5, 1) + + assert.Error(t, err) + assert.Nil(t, sealed) + assert.Nil(t, growing) + assert.Contains(t, err.Error(), "channel distribution is not serviceable") +} + +func TestPinReadableSegments_EdgeCases(t *testing.T) { + // Create test distribution + queryView := NewChannelQueryView(nil, nil, []int64{1}, initialTargetVersion) + dist := NewDistribution("test-channel", queryView) + + // Add test segments + dist.AddDistributions([]SegmentEntry{ + {NodeID: 1, SegmentID: 1, PartitionID: 1, Version: 1}, + }...) + + // Setup query view + dist.SyncTargetVersion(&querypb.SyncAction{ + TargetVersion: 1000, + SealedSegmentRowCount: map[int64]int64{1: 100}, + GrowingInTarget: []int64{}, + }, []int64{1}) + + // Test case 1: requiredLoadRatio = 0.0 (edge case) + mockGetLoadedRatio := mockey.Mock((*channelQueryView).GetLoadedRatio).Return(0.0).Build() + + sealed, growing, _, _, err := dist.PinReadableSegments(0.0, 1) + + assert.NoError(t, err) + assert.NotNil(t, sealed) + assert.NotNil(t, growing) + + // Test case 2: requiredLoadRatio = 0.0, GetLoadedRatio = 0.0 (exact match) + mockGetLoadedRatio.UnPatch() + mockGetLoadedRatio = mockey.Mock((*channelQueryView).GetLoadedRatio).Return(0.0).Build() + + sealed, growing, _, _, err = dist.PinReadableSegments(0.0, 1) + + assert.NoError(t, err) + assert.NotNil(t, sealed) + assert.NotNil(t, growing) + + // Test case 3: requiredLoadRatio = 0.0, GetLoadedRatio = 0.1 (satisfied) + mockGetLoadedRatio.UnPatch() + mockGetLoadedRatio = mockey.Mock((*channelQueryView).GetLoadedRatio).Return(0.1).Build() + defer mockGetLoadedRatio.UnPatch() + + sealed, growing, _, _, err = dist.PinReadableSegments(0.0, 1) + + assert.NoError(t, err) + assert.NotNil(t, sealed) + assert.NotNil(t, growing) +}