fix: Deep copy ImportTask.segmentsInfo to prevent data race (milvus-i…

…o#34090) See also milvus-io#34089 Signed-off-by: Congqi Xia <congqi.xia@zilliz.com> Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
yellow-shine · Jul 2, 2024 · 62f71c6 · 62f71c6
1 parent 0b6c6ab
commit 62f71c6
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 10 deletions.
diff --git a/internal/datanode/importv2/task_import.go b/internal/datanode/importv2/task_import.go
@@ -22,7 +22,6 @@ import (
 	"time"
 
 	"github.com/cockroachdb/errors"
-	"github.com/golang/protobuf/proto"
 	"github.com/samber/lo"
 	"go.uber.org/zap"
 
@@ -36,6 +35,7 @@ import (
 	"github.com/milvus-io/milvus/pkg/log"
 	"github.com/milvus-io/milvus/pkg/util/conc"
 	"github.com/milvus-io/milvus/pkg/util/paramtable"
+	"github.com/milvus-io/milvus/pkg/util/typeutil"
 )
 
 type ImportTask struct {
@@ -107,11 +107,15 @@ func (t *ImportTask) GetSegmentsInfo() []*datapb.ImportSegmentInfo {
 
 func (t *ImportTask) Clone() Task {
 	ctx, cancel := context.WithCancel(t.ctx)
+	infos := make(map[int64]*datapb.ImportSegmentInfo)
+	for id, info := range t.segmentsInfo {
+		infos[id] = typeutil.Clone(info)
+	}
 	return &ImportTask{
-		ImportTaskV2: proto.Clone(t.ImportTaskV2).(*datapb.ImportTaskV2),
+		ImportTaskV2: typeutil.Clone(t.ImportTaskV2),
 		ctx:          ctx,
 		cancel:       cancel,
-		segmentsInfo: t.segmentsInfo,
+		segmentsInfo: infos,
 		req:          t.req,
 		metaCaches:   t.metaCaches,
 	}

diff --git a/internal/datanode/importv2/task_l0_import.go b/internal/datanode/importv2/task_l0_import.go
@@ -23,7 +23,6 @@ import (
 	"time"
 
 	"github.com/cockroachdb/errors"
-	"github.com/golang/protobuf/proto"
 	"github.com/samber/lo"
 	"go.uber.org/zap"
 
@@ -104,11 +103,15 @@ func (t *L0ImportTask) GetSegmentsInfo() []*datapb.ImportSegmentInfo {
 
 func (t *L0ImportTask) Clone() Task {
 	ctx, cancel := context.WithCancel(t.ctx)
+	infos := make(map[int64]*datapb.ImportSegmentInfo)
+	for id, info := range t.segmentsInfo {
+		infos[id] = typeutil.Clone(info)
+	}
 	return &L0ImportTask{
-		ImportTaskV2: proto.Clone(t.ImportTaskV2).(*datapb.ImportTaskV2),
+		ImportTaskV2: typeutil.Clone(t.ImportTaskV2),
 		ctx:          ctx,
 		cancel:       cancel,
-		segmentsInfo: t.segmentsInfo,
+		segmentsInfo: infos,
 		req:          t.req,
 		metaCaches:   t.metaCaches,
 	}

diff --git a/internal/datanode/importv2/task_l0_preimport.go b/internal/datanode/importv2/task_l0_preimport.go
@@ -23,7 +23,6 @@ import (
 	"time"
 
 	"github.com/cockroachdb/errors"
-	"github.com/golang/protobuf/proto"
 	"github.com/samber/lo"
 	"go.uber.org/zap"
 
@@ -102,7 +101,7 @@ func (t *L0PreImportTask) Cancel() {
 func (t *L0PreImportTask) Clone() Task {
 	ctx, cancel := context.WithCancel(t.ctx)
 	return &L0PreImportTask{
-		PreImportTask: proto.Clone(t.PreImportTask).(*datapb.PreImportTask),
+		PreImportTask: typeutil.Clone(t.PreImportTask),
 		ctx:           ctx,
 		cancel:        cancel,
 		partitionIDs:  t.GetPartitionIDs(),

diff --git a/internal/datanode/importv2/task_preimport.go b/internal/datanode/importv2/task_preimport.go
@@ -23,7 +23,6 @@ import (
 	"time"
 
 	"github.com/cockroachdb/errors"
-	"github.com/golang/protobuf/proto"
 	"github.com/samber/lo"
 	"go.uber.org/zap"
 
@@ -36,6 +35,7 @@ import (
 	"github.com/milvus-io/milvus/pkg/log"
 	"github.com/milvus-io/milvus/pkg/util/conc"
 	"github.com/milvus-io/milvus/pkg/util/paramtable"
+	"github.com/milvus-io/milvus/pkg/util/typeutil"
 )
 
 type PreImportTask struct {
@@ -108,7 +108,7 @@ func (p *PreImportTask) Cancel() {
 func (p *PreImportTask) Clone() Task {
 	ctx, cancel := context.WithCancel(p.ctx)
 	return &PreImportTask{
-		PreImportTask: proto.Clone(p.PreImportTask).(*datapb.PreImportTask),
+		PreImportTask: typeutil.Clone(p.PreImportTask),
 		ctx:           ctx,
 		cancel:        cancel,
 		partitionIDs:  p.GetPartitionIDs(),