From a7e6f08183bed139fd11f22e35c19fea47e392d5 Mon Sep 17 00:00:00 2001
From: congqixia <congqi.xia@zilliz.com>
Date: Fri, 14 Jul 2023 17:04:31 +0800
Subject: [PATCH] Fix DeleteCollection not protected (#25599)

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
---
 internal/querynodev2/segments/collection.go | 7 ++++++-
 internal/querynodev2/segments/plan.go       | 5 +++++
 internal/querynodev2/segments/plan_test.go  | 6 ++++++
 3 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/internal/querynodev2/segments/collection.go b/internal/querynodev2/segments/collection.go
index abb3d9f139..f14ddc8943 100644
--- a/internal/querynodev2/segments/collection.go
+++ b/internal/querynodev2/segments/collection.go
@@ -232,8 +232,13 @@ func DeleteCollection(collection *Collection) {
 		void
 		deleteCollection(CCollection collection);
 	*/
+	collection.mu.Lock()
+	defer collection.mu.Unlock()
+
 	cPtr := collection.collectionPtr
-	C.DeleteCollection(cPtr)
+	if cPtr != nil {
+		C.DeleteCollection(cPtr)
+	}
 
 	collection.collectionPtr = nil
 }
diff --git a/internal/querynodev2/segments/plan.go b/internal/querynodev2/segments/plan.go
index 969d7c0d7a..ed2a41058c 100644
--- a/internal/querynodev2/segments/plan.go
+++ b/internal/querynodev2/segments/plan.go
@@ -32,6 +32,7 @@ import (
 	"github.com/cockroachdb/errors"
 
 	"github.com/milvus-io/milvus/internal/proto/querypb"
+	"github.com/milvus-io/milvus/pkg/util/merr"
 	. "github.com/milvus-io/milvus/pkg/util/typeutil"
 )
 
@@ -178,6 +179,10 @@ func NewRetrievePlan(col *Collection, expr []byte, timestamp Timestamp, msgID Un
 	col.mu.RLock()
 	defer col.mu.RUnlock()
 
+	if col.collectionPtr == nil {
+		return nil, merr.WrapErrCollectionNotFound(col.id, "collection released")
+	}
+
 	var cPlan C.CRetrievePlan
 	status := C.CreateRetrievePlanByExpr(col.collectionPtr, unsafe.Pointer(&expr[0]), (C.int64_t)(len(expr)), &cPlan)
 
diff --git a/internal/querynodev2/segments/plan_test.go b/internal/querynodev2/segments/plan_test.go
index f05d0809d5..abd41f3636 100644
--- a/internal/querynodev2/segments/plan_test.go
+++ b/internal/querynodev2/segments/plan_test.go
@@ -70,6 +70,12 @@ func (suite *PlanSuite) TestPlanFail() {
 	suite.Error(err)
 }
 
+func (suite *PlanSuite) TestQueryPlanCollectionReleased() {
+	collection := &Collection{id: suite.collectionID}
+	_, err := NewRetrievePlan(collection, nil, 0, 0)
+	suite.Error(err)
+}
+
 func TestPlan(t *testing.T) {
 	suite.Run(t, new(PlanSuite))
 }