fix: Make GetShardLeaders skip retry on not loaded collection (#37684)

issue: #37532

Signed-off-by: Wei Liu <wei.liu@zilliz.com>
pull/37699/head
wei liu 2024-11-14 21:14:40 +08:00 committed by GitHub
parent 1b4f7e3ac1
commit c5485bb1b1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 40 additions and 2 deletions

View File

@ -106,9 +106,8 @@ func (lb *LBPolicyImpl) GetShardLeaders(ctx context.Context, dbName string, coll
var err error
shardLeaders, err = globalMetaCache.GetShards(ctx, withCache, dbName, collName, collectionID)
if err != nil {
return !errors.Is(err, merr.ErrCollectionLoaded), err
return !errors.Is(err, merr.ErrCollectionNotLoaded), err
}
return false, nil
})

View File

@ -21,9 +21,12 @@ import (
"testing"
"github.com/cockroachdb/errors"
"github.com/pingcap/log"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/suite"
"go.uber.org/atomic"
"go.uber.org/zap"
"google.golang.org/grpc"
"google.golang.org/protobuf/proto"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
@ -456,6 +459,42 @@ func (s *LBPolicySuite) TestNewLBPolicy() {
policy.Close()
}
func (s *LBPolicySuite) TestGetShardLeaders() {
ctx := context.Background()
// ErrCollectionNotFullyLoaded is retriable, expected to retry until ctx done or success
counter := atomic.NewInt64(0)
globalMetaCache.DeprecateShardCache(dbName, s.collectionName)
s.qc.ExpectedCalls = nil
s.qc.EXPECT().GetShardLeaders(mock.Anything, mock.Anything).RunAndReturn(
func(ctx context.Context, req *querypb.GetShardLeadersRequest, opts ...grpc.CallOption) (*querypb.GetShardLeadersResponse, error) {
counter.Inc()
return nil, merr.ErrCollectionNotFullyLoaded
}).Times(5)
s.qc.EXPECT().GetShardLeaders(mock.Anything, mock.Anything).RunAndReturn(
func(ctx context.Context, req *querypb.GetShardLeadersRequest, opts ...grpc.CallOption) (*querypb.GetShardLeadersResponse, error) {
log.Info("return rpc success")
return nil, nil
}).Times(5)
_, err := s.lbPolicy.GetShardLeaders(ctx, dbName, s.collectionName, s.collectionID, true)
s.NoError(err)
s.Equal(int64(5), counter.Load())
// ErrServiceUnavailable is not retriable, expected to fail fast
counter.Store(0)
globalMetaCache.DeprecateShardCache(dbName, s.collectionName)
s.qc.ExpectedCalls = nil
s.qc.EXPECT().GetShardLeaders(mock.Anything, mock.Anything).RunAndReturn(
func(ctx context.Context, req *querypb.GetShardLeadersRequest, opts ...grpc.CallOption) (*querypb.GetShardLeadersResponse, error) {
counter.Inc()
return nil, merr.ErrCollectionNotLoaded
})
_, err = s.lbPolicy.GetShardLeaders(ctx, dbName, s.collectionName, s.collectionID, true)
log.Info("check err", zap.Error(err))
s.Error(err)
s.Equal(int64(1), counter.Load())
}
func TestLBPolicySuite(t *testing.T) {
suite.Run(t, new(LBPolicySuite))
}