fix: Retain collection early to prevent it from being released before query completion (#45413)

issue: #45314

This PR only ensures that no panic occurs. However, we still need to
provide protection for the delegator handling ongoing query tasks.

Signed-off-by: Cai Zhang <cai.zhang@zilliz.com>
This commit is contained in:
cai.zhang 2025-11-11 20:29:37 +08:00 committed by GitHub
parent 404797dd46
commit 216c576da2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 12 additions and 0 deletions

View File

@ -992,6 +992,16 @@ func (node *QueryNode) Query(ctx context.Context, req *querypb.QueryRequest) (*i
toMergeResults := make([]*internalpb.RetrieveResults, len(req.GetDmlChannels()))
runningGp, runningCtx := errgroup.WithContext(ctx)
if !node.manager.Collection.Ref(req.GetReq().GetCollectionID(), 1) {
err := merr.WrapErrCollectionNotLoaded(req.GetReq().GetCollectionID())
log.Warn("failed to query collection", zap.Error(err))
return &internalpb.RetrieveResults{
Status: merr.Status(err),
}, nil
}
defer func() {
node.manager.Collection.Unref(req.GetReq().GetCollectionID(), 1)
}()
for i, ch := range req.GetDmlChannels() {
ch := ch

View File

@ -1579,6 +1579,8 @@ func (suite *ServiceSuite) TestQuery_Failed() {
// data
schema := mock_segcore.GenTestCollectionSchema(suite.collectionName, schemapb.DataType_Int64, false)
suite.node.manager.Collection.PutOrRef(suite.collectionID, schema, nil, nil)
defer suite.node.manager.Collection.Unref(suite.collectionID, 1)
creq, err := suite.genCQueryRequest(10, IndexFaissIDMap, schema)
suite.NoError(err)
req := &querypb.QueryRequest{