enhance: Pre-allocate ids for compaction (#34187)

This PR removes the dependency of compaction on the ID allocator by
pre-allocating the logID and segmentID.

issue: #33957

---------

Signed-off-by: bigsheeper <yihao.dai@zilliz.com>

internal/datacoord/compaction.go

@@ -577,18 +577,21 @@ func (c *compactionPlanHandler) createCompactTask(t *datapb.CompactionTask) (Com
 	case datapb.CompactionType_MixCompaction:
 		task = &mixCompactionTask{
 			CompactionTask: t,
+			allocator:      c.allocator,
 			meta:           c.meta,
 			sessions:       c.sessions,
 		}
 	case datapb.CompactionType_Level0DeleteCompaction:
 		task = &l0CompactionTask{
 			CompactionTask: t,
+			allocator:      c.allocator,
 			meta:           c.meta,
 			sessions:       c.sessions,
 		}
 	case datapb.CompactionType_ClusteringCompaction:
 		task = &clusteringCompactionTask{
 			CompactionTask:   t,
+			allocator:        c.allocator,
 			meta:             c.meta,
 			sessions:         c.sessions,
 			handler:          c.handler,

internal/datacoord/compaction_task_clustering.go

@@ -48,8 +48,9 @@ type clusteringCompactionTask struct {
 	*datapb.CompactionTask
 	plan   *datapb.CompactionPlan
 	result *datapb.CompactionPlanResult
-	span   trace.Span
 
+	span             trace.Span
+	allocator        allocator
 	meta             CompactionMeta
 	sessions         SessionManager
 	handler          Handler
@@ -133,6 +134,10 @@ func (t *clusteringCompactionTask) retryableProcess() error {
 }
 
 func (t *clusteringCompactionTask) BuildCompactionRequest() (*datapb.CompactionPlan, error) {
+	beginLogID, _, err := t.allocator.allocN(1)
+	if err != nil {
+		return nil, err
+	}
 	plan := &datapb.CompactionPlan{
 		PlanID:             t.GetPlanID(),
 		StartTime:          t.GetStartTime(),
@@ -147,6 +152,11 @@ func (t *clusteringCompactionTask) BuildCompactionRequest() (*datapb.CompactionP
 		PreferSegmentRows:  t.GetPreferSegmentRows(),
 		AnalyzeResultPath:  path.Join(t.meta.(*meta).chunkManager.RootPath(), common.AnalyzeStatsPath, metautil.JoinIDPath(t.AnalyzeTaskID, t.AnalyzeVersion)),
 		AnalyzeSegmentIds:  t.GetInputSegments(), // todo: if need
+		BeginLogID:         beginLogID,
+		PreAllocatedSegments: &datapb.IDRange{
+			Begin: t.GetResultSegments()[0],
+			End:   t.GetResultSegments()[1],
+		},
 	}
 	log := log.With(zap.Int64("taskID", t.GetTriggerID()), zap.Int64("planID", plan.GetPlanID()))
 

internal/datacoord/compaction_task_clustering_test.go

@@ -20,6 +20,8 @@ import (
 	"context"
 	"fmt"
 
+	"github.com/stretchr/testify/mock"
+
 	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
 	"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
 	"github.com/milvus-io/milvus/internal/metastore/kv/datacoord"
@@ -50,6 +52,8 @@ func (s *CompactionTaskSuite) TestClusteringCompactionSegmentMetaChange() {
 		},
 	})
 	session := NewSessionManagerImpl()
+	alloc := NewNMockAllocator(s.T())
+	alloc.EXPECT().allocN(mock.Anything).Return(100, 200, nil)
 
 	schema := ConstructScalarClusteringSchema("TestClusteringCompactionTask", 32, true)
 	pk := &schemapb.FieldSchema{
@@ -77,9 +81,11 @@ func (s *CompactionTaskSuite) TestClusteringCompactionSegmentMetaChange() {
 			Schema:             schema,
 			ClusteringKeyField: pk,
 			InputSegments:      []int64{101, 102},
+			ResultSegments:     []int64{1000, 1100},
 		},
-		meta:     meta,
-		sessions: session,
+		meta:      meta,
+		sessions:  session,
+		allocator: alloc,
 	}
 
 	task.processPipelining()

internal/datacoord/compaction_task_l0.go

@@ -37,11 +37,13 @@ var _ CompactionTask = (*l0CompactionTask)(nil)
 
 type l0CompactionTask struct {
 	*datapb.CompactionTask
-	plan     *datapb.CompactionPlan
-	result   *datapb.CompactionPlanResult
-	span     trace.Span
-	sessions SessionManager
-	meta     CompactionMeta
+	plan   *datapb.CompactionPlan
+	result *datapb.CompactionPlanResult
+
+	span      trace.Span
+	allocator allocator
+	sessions  SessionManager
+	meta      CompactionMeta
 }
 
 // Note: return True means exit this state machine.
@@ -230,6 +232,10 @@ func (t *l0CompactionTask) CleanLogPath() {
 }
 
 func (t *l0CompactionTask) BuildCompactionRequest() (*datapb.CompactionPlan, error) {
+	beginLogID, _, err := t.allocator.allocN(1)
+	if err != nil {
+		return nil, err
+	}
 	plan := &datapb.CompactionPlan{
 		PlanID:           t.GetPlanID(),
 		StartTime:        t.GetStartTime(),
@@ -239,6 +245,7 @@ func (t *l0CompactionTask) BuildCompactionRequest() (*datapb.CompactionPlan, err
 		CollectionTtl:    t.GetCollectionTtl(),
 		TotalRows:        t.GetTotalRows(),
 		Schema:           t.GetSchema(),
+		BeginLogID:       beginLogID,
 	}
 
 	log := log.With(zap.Int64("taskID", t.GetTriggerID()), zap.Int64("planID", plan.GetPlanID()))

internal/datacoord/compaction_task_l0_test.go

@@ -1,3 +1,19 @@
+// Licensed to the LF AI & Data foundation under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
 package datacoord
 
 import (
@@ -58,6 +74,9 @@ func (s *CompactionTaskSuite) TestProcessRefreshPlan_NormalL0() {
 		},
 		meta: s.mockMeta,
 	}
+	alloc := NewNMockAllocator(s.T())
+	alloc.EXPECT().allocN(mock.Anything).Return(100, 200, nil)
+	task.allocator = alloc
 	plan, err := task.BuildCompactionRequest()
 	s.Require().NoError(err)
 
@@ -88,6 +107,9 @@ func (s *CompactionTaskSuite) TestProcessRefreshPlan_SegmentNotFoundL0() {
 		},
 		meta: s.mockMeta,
 	}
+	alloc := NewNMockAllocator(s.T())
+	alloc.EXPECT().allocN(mock.Anything).Return(100, 200, nil)
+	task.allocator = alloc
 
 	_, err := task.BuildCompactionRequest()
 	s.Error(err)
@@ -121,10 +143,41 @@ func (s *CompactionTaskSuite) TestProcessRefreshPlan_SelectZeroSegmentsL0() {
 		},
 		meta: s.mockMeta,
 	}
+	alloc := NewNMockAllocator(s.T())
+	alloc.EXPECT().allocN(mock.Anything).Return(100, 200, nil)
+	task.allocator = alloc
 	_, err := task.BuildCompactionRequest()
 	s.Error(err)
 }
 
+func (s *CompactionTaskSuite) TestBuildCompactionRequestFailed_AllocFailed() {
+	var task CompactionTask
+
+	alloc := NewNMockAllocator(s.T())
+	alloc.EXPECT().allocN(mock.Anything).Return(100, 200, errors.New("mock alloc err"))
+
+	task = &l0CompactionTask{
+		allocator: alloc,
+	}
+	_, err := task.BuildCompactionRequest()
+	s.T().Logf("err=%v", err)
+	s.Error(err)
+
+	task = &mixCompactionTask{
+		allocator: alloc,
+	}
+	_, err = task.BuildCompactionRequest()
+	s.T().Logf("err=%v", err)
+	s.Error(err)
+
+	task = &clusteringCompactionTask{
+		allocator: alloc,
+	}
+	_, err = task.BuildCompactionRequest()
+	s.T().Logf("err=%v", err)
+	s.Error(err)
+}
+
 func generateTestL0Task(state datapb.CompactionTaskState) *l0CompactionTask {
 	return &l0CompactionTask{
 		CompactionTask: &datapb.CompactionTask{
@@ -145,9 +198,13 @@ func (s *CompactionTaskSuite) SetupSubTest() {
 }
 
 func (s *CompactionTaskSuite) TestProcessStateTrans() {
+	alloc := NewNMockAllocator(s.T())
+	alloc.EXPECT().allocN(mock.Anything).Return(100, 200, nil)
+
 	s.Run("test pipelining needReassignNodeID", func() {
 		t := generateTestL0Task(datapb.CompactionTaskState_pipelining)
 		t.NodeID = NullNodeID
+		t.allocator = alloc
 		got := t.Process()
 		s.False(got)
 		s.Equal(datapb.CompactionTaskState_pipelining, t.State)
@@ -157,6 +214,7 @@ func (s *CompactionTaskSuite) TestProcessStateTrans() {
 	s.Run("test pipelining BuildCompactionRequest failed", func() {
 		t := generateTestL0Task(datapb.CompactionTaskState_pipelining)
 		t.NodeID = 100
+		t.allocator = alloc
 		channel := "ch-1"
 		deltaLogs := []*datapb.FieldBinlog{getFieldBinlogIDs(101, 3)}
 
@@ -194,6 +252,7 @@ func (s *CompactionTaskSuite) TestProcessStateTrans() {
 	s.Run("test pipelining Compaction failed", func() {
 		t := generateTestL0Task(datapb.CompactionTaskState_pipelining)
 		t.NodeID = 100
+		t.allocator = alloc
 		channel := "ch-1"
 		deltaLogs := []*datapb.FieldBinlog{getFieldBinlogIDs(101, 3)}
 
@@ -234,6 +293,7 @@ func (s *CompactionTaskSuite) TestProcessStateTrans() {
 	s.Run("test pipelining success", func() {
 		t := generateTestL0Task(datapb.CompactionTaskState_pipelining)
 		t.NodeID = 100
+		t.allocator = alloc
 		channel := "ch-1"
 		deltaLogs := []*datapb.FieldBinlog{getFieldBinlogIDs(101, 3)}
 

internal/datacoord/compaction_task_mix.go

@@ -18,9 +18,11 @@ var _ CompactionTask = (*mixCompactionTask)(nil)
 
 type mixCompactionTask struct {
 	*datapb.CompactionTask
-	plan       *datapb.CompactionPlan
-	result     *datapb.CompactionPlanResult
+	plan   *datapb.CompactionPlan
+	result *datapb.CompactionPlanResult
+
 	span       trace.Span
+	allocator  allocator
 	sessions   SessionManager
 	meta       CompactionMeta
 	newSegment *SegmentInfo
@@ -328,6 +330,10 @@ func (t *mixCompactionTask) CleanLogPath() {
 }
 
 func (t *mixCompactionTask) BuildCompactionRequest() (*datapb.CompactionPlan, error) {
+	beginLogID, _, err := t.allocator.allocN(1)
+	if err != nil {
+		return nil, err
+	}
 	plan := &datapb.CompactionPlan{
 		PlanID:           t.GetPlanID(),
 		StartTime:        t.GetStartTime(),
@@ -337,6 +343,10 @@ func (t *mixCompactionTask) BuildCompactionRequest() (*datapb.CompactionPlan, er
 		CollectionTtl:    t.GetCollectionTtl(),
 		TotalRows:        t.GetTotalRows(),
 		Schema:           t.GetSchema(),
+		BeginLogID:       beginLogID,
+		PreAllocatedSegments: &datapb.IDRange{
+			Begin: t.GetResultSegments()[0],
+		},
 	}
 	log := log.With(zap.Int64("taskID", t.GetTriggerID()), zap.Int64("planID", plan.GetPlanID()))
 

internal/datacoord/compaction_task_mix_test.go

@@ -23,18 +23,22 @@ func (s *CompactionTaskSuite) TestProcessRefreshPlan_NormalMix() {
 	}).Times(2)
 	task := &mixCompactionTask{
 		CompactionTask: &datapb.CompactionTask{
-			PlanID:        1,
-			TriggerID:     19530,
-			CollectionID:  1,
-			PartitionID:   10,
-			Type:          datapb.CompactionType_MixCompaction,
-			NodeID:        1,
-			State:         datapb.CompactionTaskState_executing,
-			InputSegments: []int64{200, 201},
+			PlanID:         1,
+			TriggerID:      19530,
+			CollectionID:   1,
+			PartitionID:    10,
+			Type:           datapb.CompactionType_MixCompaction,
+			NodeID:         1,
+			State:          datapb.CompactionTaskState_executing,
+			InputSegments:  []int64{200, 201},
+			ResultSegments: []int64{100},
 		},
 		// plan: plan,
 		meta: s.mockMeta,
 	}
+	alloc := NewNMockAllocator(s.T())
+	alloc.EXPECT().allocN(mock.Anything).Return(100, 200, nil)
+	task.allocator = alloc
 	plan, err := task.BuildCompactionRequest()
 	s.Require().NoError(err)
 
@@ -53,18 +57,22 @@ func (s *CompactionTaskSuite) TestProcessRefreshPlan_MixSegmentNotFound() {
 		}).Once()
 		task := &mixCompactionTask{
 			CompactionTask: &datapb.CompactionTask{
-				PlanID:        1,
-				TriggerID:     19530,
-				CollectionID:  1,
-				PartitionID:   10,
-				Channel:       channel,
-				Type:          datapb.CompactionType_MixCompaction,
-				State:         datapb.CompactionTaskState_executing,
-				NodeID:        1,
-				InputSegments: []int64{200, 201},
+				PlanID:         1,
+				TriggerID:      19530,
+				CollectionID:   1,
+				PartitionID:    10,
+				Channel:        channel,
+				Type:           datapb.CompactionType_MixCompaction,
+				State:          datapb.CompactionTaskState_executing,
+				NodeID:         1,
+				InputSegments:  []int64{200, 201},
+				ResultSegments: []int64{100},
 			},
 			meta: s.mockMeta,
 		}
+		alloc := NewNMockAllocator(s.T())
+		alloc.EXPECT().allocN(mock.Anything).Return(100, 200, nil)
+		task.allocator = alloc
 		_, err := task.BuildCompactionRequest()
 		s.Error(err)
 		s.ErrorIs(err, merr.ErrSegmentNotFound)

internal/datacoord/compaction_trigger.go

@@ -399,7 +399,7 @@ func (t *compactionTrigger) handleGlobalSignal(signal *compactionSignal) error {
 		}
 
 		plans := t.generatePlans(group.segments, signal, ct)
-		currentID, _, err := t.allocator.allocN(int64(len(plans)))
+		currentID, _, err := t.allocator.allocN(int64(len(plans) * 2))
 		if err != nil {
 			return err
 		}
@@ -413,6 +413,8 @@ func (t *compactionTrigger) handleGlobalSignal(signal *compactionSignal) error {
 			start := time.Now()
 			planID := currentID
 			currentID++
+			targetSegmentID := currentID
+			currentID++
 			pts, _ := tsoutil.ParseTS(ct.startTime)
 			task := &datapb.CompactionTask{
 				PlanID:           planID,
@@ -426,6 +428,7 @@ func (t *compactionTrigger) handleGlobalSignal(signal *compactionSignal) error {
 				PartitionID:      group.partitionID,
 				Channel:          group.channelName,
 				InputSegments:    segIDs,
+				ResultSegments:   []int64{targetSegmentID}, // pre-allocated target segment
 				TotalRows:        totalRows,
 				Schema:           coll.Schema,
 			}
@@ -503,7 +506,7 @@ func (t *compactionTrigger) handleSignal(signal *compactionSignal) {
 	}
 
 	plans := t.generatePlans(segments, signal, ct)
-	currentID, _, err := t.allocator.allocN(int64(len(plans)))
+	currentID, _, err := t.allocator.allocN(int64(len(plans) * 2))
 	if err != nil {
 		log.Warn("fail to allocate id", zap.Error(err))
 		return
@@ -518,6 +521,8 @@ func (t *compactionTrigger) handleSignal(signal *compactionSignal) {
 		start := time.Now()
 		planID := currentID
 		currentID++
+		targetSegmentID := currentID
+		currentID++
 		pts, _ := tsoutil.ParseTS(ct.startTime)
 		if err := t.compactionHandler.enqueueCompaction(&datapb.CompactionTask{
 			PlanID:           planID,
@@ -531,6 +536,7 @@ func (t *compactionTrigger) handleSignal(signal *compactionSignal) {
 			PartitionID:      partitionID,
 			Channel:          channel,
 			InputSegments:    segmentIDS,
+			ResultSegments:   []int64{targetSegmentID}, // pre-allocated target segment
 			TotalRows:        totalRows,
 			Schema:           coll.Schema,
 		}); err != nil {