mirror of
https://gitee.com/milvus-io/milvus.git
synced 2026-01-07 19:31:51 +08:00
975c91df16
issue: #44358 Implement complete snapshot management system including creation, deletion, listing, description, and restoration capabilities across all system components. Key features: - Create snapshots for entire collections - Drop snapshots by name with proper cleanup - List snapshots with collection filtering - Describe snapshot details and metadata Components added/modified: - Client SDK with full snapshot API support and options - DataCoord snapshot service with metadata management - Proxy layer with task-based snapshot operations - Protocol buffer definitions for snapshot RPCs - Comprehensive unit tests with mockey framework - Integration tests for end-to-end validation Technical implementation: - Snapshot metadata storage in etcd with proper indexing - File-based snapshot data persistence in object storage - Garbage collection integration for snapshot cleanup - Error handling and validation across all operations - Thread-safe operations with proper locking mechanisms <!-- This is an auto-generated comment: release notes by coderabbit.ai --> - Core invariant/assumption: snapshots are immutable point‑in‑time captures identified by (collection, snapshot name/ID); etcd snapshot metadata is authoritative for lifecycle (PENDING → COMMITTED → DELETING) and per‑segment manifests live in object storage (Avro / StorageV2). GC and restore logic must see snapshotRefIndex loaded (snapshotMeta.IsRefIndexLoaded) before reclaiming or relying on segment/index files. - New capability added: full end‑to‑end snapshot subsystem — client SDK APIs (Create/Drop/List/Describe/Restore + restore job queries), DataCoord SnapshotWriter/Reader (Avro + StorageV2 manifests), snapshotMeta in meta, SnapshotManager orchestration (create/drop/describe/list/restore), copy‑segment restore tasks/inspector/checker, proxy & RPC surface, GC integration, and docs/tests — enabling point‑in‑time collection snapshots persisted to object storage and restorations orchestrated across components. - Logic removed/simplified and why: duplicated recursive compaction/delta‑log traversal and ad‑hoc lookup code were consolidated behind two focused APIs/owners (Handler.GetDeltaLogFromCompactTo for delta traversal and SnapshotManager/SnapshotReader for snapshot I/O). MixCoord/coordinator broker paths were converted to thin RPC proxies. This eliminates multiple implementations of the same traversal/lookup, reducing divergence and simplifying responsibility boundaries. - Why this does NOT introduce data loss or regressions: snapshot create/drop use explicit two‑phase semantics (PENDING → COMMIT/DELETING) with SnapshotWriter writing manifests and metadata before commit; GC uses snapshotRefIndex guards and IsRefIndexLoaded/GetSnapshotBySegment/GetSnapshotByIndex checks to avoid removing referenced files; restore flow pre‑allocates job IDs, validates resources (partitions/indexes), performs rollback on failure (rollbackRestoreSnapshot), and converts/updates segment/index metadata only after successful copy tasks. Extensive unit and integration tests exercise pending/deleting/GC/restore/error paths to ensure idempotence and protection against premature deletion. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Wei Liu <wei.liu@zilliz.com>
548 lines
17 KiB
Go
548 lines
17 KiB
Go
// 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 (
|
|
"context"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/mock"
|
|
"github.com/stretchr/testify/suite"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
|
|
"github.com/milvus-io/milvus/internal/datacoord/broker"
|
|
task2 "github.com/milvus-io/milvus/internal/datacoord/task"
|
|
"github.com/milvus-io/milvus/internal/metastore/mocks"
|
|
"github.com/milvus-io/milvus/pkg/v2/proto/datapb"
|
|
"github.com/milvus-io/milvus/pkg/v2/taskcommon"
|
|
"github.com/milvus-io/milvus/pkg/v2/util/timerecord"
|
|
)
|
|
|
|
type CopySegmentInspectorSuite struct {
|
|
suite.Suite
|
|
|
|
collectionID int64
|
|
jobID int64
|
|
|
|
catalog *mocks.DataCoordCatalog
|
|
broker *broker.MockBroker
|
|
meta *meta
|
|
copyMeta CopySegmentMeta
|
|
scheduler *task2.MockGlobalScheduler
|
|
inspector *copySegmentInspector
|
|
}
|
|
|
|
func (s *CopySegmentInspectorSuite) SetupTest() {
|
|
var err error
|
|
|
|
s.collectionID = 1
|
|
s.jobID = 100
|
|
|
|
s.catalog = mocks.NewDataCoordCatalog(s.T())
|
|
s.catalog.EXPECT().ListCopySegmentJobs(mock.Anything).Return(nil, nil)
|
|
s.catalog.EXPECT().ListCopySegmentTasks(mock.Anything).Return(nil, nil)
|
|
s.catalog.EXPECT().ListChannelCheckpoint(mock.Anything).Return(nil, nil)
|
|
s.catalog.EXPECT().ListIndexes(mock.Anything).Return(nil, nil)
|
|
s.catalog.EXPECT().ListSegmentIndexes(mock.Anything).Return(nil, nil)
|
|
s.catalog.EXPECT().ListAnalyzeTasks(mock.Anything).Return(nil, nil)
|
|
s.catalog.EXPECT().ListCompactionTask(mock.Anything).Return(nil, nil)
|
|
s.catalog.EXPECT().ListPartitionStatsInfos(mock.Anything).Return(nil, nil)
|
|
s.catalog.EXPECT().ListStatsTasks(mock.Anything).Return(nil, nil)
|
|
s.catalog.EXPECT().ListFileResource(mock.Anything).Return(nil, 0, nil)
|
|
s.catalog.EXPECT().ListSnapshots(mock.Anything).Return(nil, nil)
|
|
|
|
s.broker = broker.NewMockBroker(s.T())
|
|
s.broker.EXPECT().ShowCollectionIDs(mock.Anything).Return(nil, nil)
|
|
|
|
s.meta, err = newMeta(context.TODO(), s.catalog, nil, s.broker)
|
|
s.NoError(err)
|
|
s.meta.AddCollection(&collectionInfo{
|
|
ID: s.collectionID,
|
|
Schema: newTestSchema(),
|
|
})
|
|
|
|
s.copyMeta, err = NewCopySegmentMeta(context.TODO(), s.catalog, s.meta, nil)
|
|
s.NoError(err)
|
|
|
|
s.scheduler = task2.NewMockGlobalScheduler(s.T())
|
|
|
|
s.inspector = NewCopySegmentInspector(
|
|
context.TODO(),
|
|
s.meta,
|
|
s.copyMeta,
|
|
s.scheduler,
|
|
).(*copySegmentInspector)
|
|
}
|
|
|
|
func (s *CopySegmentInspectorSuite) TearDownTest() {
|
|
s.inspector.Close()
|
|
}
|
|
|
|
func TestCopySegmentInspector(t *testing.T) {
|
|
suite.Run(t, new(CopySegmentInspectorSuite))
|
|
}
|
|
|
|
func (s *CopySegmentInspectorSuite) TestReloadFromMeta_NoPendingTasks() {
|
|
s.catalog.EXPECT().SaveCopySegmentJob(mock.Anything, mock.Anything).Return(nil)
|
|
s.catalog.EXPECT().SaveCopySegmentTask(mock.Anything, mock.Anything).Return(nil)
|
|
|
|
// Create a job
|
|
job := ©SegmentJob{
|
|
CopySegmentJob: &datapb.CopySegmentJob{
|
|
JobId: s.jobID,
|
|
CollectionId: s.collectionID,
|
|
State: datapb.CopySegmentJobState_CopySegmentJobExecuting,
|
|
},
|
|
tr: timerecord.NewTimeRecorder("test job"),
|
|
}
|
|
err := s.copyMeta.AddJob(context.TODO(), job)
|
|
s.NoError(err)
|
|
|
|
// Create a completed task (should not be enqueued)
|
|
task := ©SegmentTask{
|
|
copyMeta: s.copyMeta,
|
|
tr: timerecord.NewTimeRecorder("task"),
|
|
times: taskcommon.NewTimes(),
|
|
}
|
|
task.task.Store(&datapb.CopySegmentTask{
|
|
TaskId: 1001,
|
|
JobId: s.jobID,
|
|
CollectionId: s.collectionID,
|
|
State: datapb.CopySegmentTaskState_CopySegmentTaskCompleted,
|
|
})
|
|
err = s.copyMeta.AddTask(context.TODO(), task)
|
|
s.NoError(err)
|
|
|
|
// Reload should not enqueue completed tasks
|
|
s.inspector.reloadFromMeta()
|
|
}
|
|
|
|
func (s *CopySegmentInspectorSuite) TestReloadFromMeta_WithInProgressTasks() {
|
|
s.catalog.EXPECT().SaveCopySegmentJob(mock.Anything, mock.Anything).Return(nil)
|
|
s.catalog.EXPECT().SaveCopySegmentTask(mock.Anything, mock.Anything).Return(nil).Times(2)
|
|
|
|
// Create a job
|
|
job := ©SegmentJob{
|
|
CopySegmentJob: &datapb.CopySegmentJob{
|
|
JobId: s.jobID,
|
|
CollectionId: s.collectionID,
|
|
State: datapb.CopySegmentJobState_CopySegmentJobExecuting,
|
|
},
|
|
tr: timerecord.NewTimeRecorder("test job"),
|
|
}
|
|
err := s.copyMeta.AddJob(context.TODO(), job)
|
|
s.NoError(err)
|
|
|
|
// Create an in-progress task
|
|
task1 := ©SegmentTask{
|
|
copyMeta: s.copyMeta,
|
|
tr: timerecord.NewTimeRecorder("task1"),
|
|
times: taskcommon.NewTimes(),
|
|
}
|
|
task1.task.Store(&datapb.CopySegmentTask{
|
|
TaskId: 1001,
|
|
JobId: s.jobID,
|
|
CollectionId: s.collectionID,
|
|
State: datapb.CopySegmentTaskState_CopySegmentTaskInProgress,
|
|
})
|
|
err = s.copyMeta.AddTask(context.TODO(), task1)
|
|
s.NoError(err)
|
|
|
|
// Create a pending task (should not be enqueued by reload)
|
|
task2 := ©SegmentTask{
|
|
copyMeta: s.copyMeta,
|
|
tr: timerecord.NewTimeRecorder("task2"),
|
|
times: taskcommon.NewTimes(),
|
|
}
|
|
task2.task.Store(&datapb.CopySegmentTask{
|
|
TaskId: 1002,
|
|
JobId: s.jobID,
|
|
CollectionId: s.collectionID,
|
|
State: datapb.CopySegmentTaskState_CopySegmentTaskPending,
|
|
})
|
|
err = s.copyMeta.AddTask(context.TODO(), task2)
|
|
s.NoError(err)
|
|
|
|
// Expect only the in-progress task to be enqueued
|
|
s.scheduler.EXPECT().Enqueue(mock.MatchedBy(func(t any) bool {
|
|
copyTask, ok := t.(CopySegmentTask)
|
|
return ok && copyTask.GetTaskId() == 1001
|
|
})).Once()
|
|
|
|
s.inspector.reloadFromMeta()
|
|
}
|
|
|
|
func (s *CopySegmentInspectorSuite) TestProcessPending() {
|
|
s.catalog.EXPECT().SaveCopySegmentJob(mock.Anything, mock.Anything).Return(nil)
|
|
s.catalog.EXPECT().SaveCopySegmentTask(mock.Anything, mock.Anything).Return(nil)
|
|
|
|
// Create a job
|
|
job := ©SegmentJob{
|
|
CopySegmentJob: &datapb.CopySegmentJob{
|
|
JobId: s.jobID,
|
|
CollectionId: s.collectionID,
|
|
State: datapb.CopySegmentJobState_CopySegmentJobExecuting,
|
|
},
|
|
tr: timerecord.NewTimeRecorder("test job"),
|
|
}
|
|
err := s.copyMeta.AddJob(context.TODO(), job)
|
|
s.NoError(err)
|
|
|
|
// Create a pending task
|
|
task := ©SegmentTask{
|
|
copyMeta: s.copyMeta,
|
|
tr: timerecord.NewTimeRecorder("task"),
|
|
times: taskcommon.NewTimes(),
|
|
}
|
|
task.task.Store(&datapb.CopySegmentTask{
|
|
TaskId: 1001,
|
|
JobId: s.jobID,
|
|
CollectionId: s.collectionID,
|
|
State: datapb.CopySegmentTaskState_CopySegmentTaskPending,
|
|
})
|
|
err = s.copyMeta.AddTask(context.TODO(), task)
|
|
s.NoError(err)
|
|
|
|
// Expect task to be enqueued
|
|
s.scheduler.EXPECT().Enqueue(mock.MatchedBy(func(t any) bool {
|
|
copyTask, ok := t.(CopySegmentTask)
|
|
return ok && copyTask.GetTaskId() == 1001
|
|
})).Once()
|
|
|
|
s.inspector.processPending(task)
|
|
}
|
|
|
|
func (s *CopySegmentInspectorSuite) TestProcessFailed_DropTargetSegments() {
|
|
s.catalog.EXPECT().SaveCopySegmentJob(mock.Anything, mock.Anything).Return(nil)
|
|
s.catalog.EXPECT().SaveCopySegmentTask(mock.Anything, mock.Anything).Return(nil)
|
|
s.catalog.EXPECT().AddSegment(mock.Anything, mock.Anything).Return(nil)
|
|
s.catalog.EXPECT().AlterSegments(mock.Anything, mock.Anything).Return(nil)
|
|
|
|
// Create target segments
|
|
seg1 := NewSegmentInfo(&datapb.SegmentInfo{
|
|
ID: 101,
|
|
CollectionID: s.collectionID,
|
|
PartitionID: 10,
|
|
State: commonpb.SegmentState_Importing,
|
|
NumOfRows: 100,
|
|
InsertChannel: "ch1",
|
|
})
|
|
err := s.meta.AddSegment(context.TODO(), seg1)
|
|
s.NoError(err)
|
|
|
|
// Create a job
|
|
job := ©SegmentJob{
|
|
CopySegmentJob: &datapb.CopySegmentJob{
|
|
JobId: s.jobID,
|
|
CollectionId: s.collectionID,
|
|
State: datapb.CopySegmentJobState_CopySegmentJobFailed,
|
|
},
|
|
tr: timerecord.NewTimeRecorder("test job"),
|
|
}
|
|
err = s.copyMeta.AddJob(context.TODO(), job)
|
|
s.NoError(err)
|
|
|
|
// Create a failed task with target segment
|
|
idMappings := []*datapb.CopySegmentIDMapping{
|
|
{SourceSegmentId: 1, TargetSegmentId: 101, PartitionId: 10},
|
|
}
|
|
|
|
task := ©SegmentTask{
|
|
copyMeta: s.copyMeta,
|
|
tr: timerecord.NewTimeRecorder("task"),
|
|
times: taskcommon.NewTimes(),
|
|
}
|
|
task.task.Store(&datapb.CopySegmentTask{
|
|
TaskId: 1001,
|
|
JobId: s.jobID,
|
|
CollectionId: s.collectionID,
|
|
State: datapb.CopySegmentTaskState_CopySegmentTaskFailed,
|
|
IdMappings: idMappings,
|
|
Reason: "test failure",
|
|
})
|
|
err = s.copyMeta.AddTask(context.TODO(), task)
|
|
s.NoError(err)
|
|
|
|
// Process failed task
|
|
s.inspector.processFailed(task)
|
|
|
|
// Target segment should be marked as Dropped
|
|
segment := s.meta.GetSegment(context.TODO(), 101)
|
|
s.Equal(commonpb.SegmentState_Dropped, segment.GetState())
|
|
}
|
|
|
|
func (s *CopySegmentInspectorSuite) TestProcessFailed_NoTargetSegment() {
|
|
s.catalog.EXPECT().SaveCopySegmentJob(mock.Anything, mock.Anything).Return(nil)
|
|
s.catalog.EXPECT().SaveCopySegmentTask(mock.Anything, mock.Anything).Return(nil)
|
|
|
|
// Create a job
|
|
job := ©SegmentJob{
|
|
CopySegmentJob: &datapb.CopySegmentJob{
|
|
JobId: s.jobID,
|
|
CollectionId: s.collectionID,
|
|
State: datapb.CopySegmentJobState_CopySegmentJobFailed,
|
|
},
|
|
tr: timerecord.NewTimeRecorder("test job"),
|
|
}
|
|
err := s.copyMeta.AddJob(context.TODO(), job)
|
|
s.NoError(err)
|
|
|
|
// Create a failed task with non-existent target segment
|
|
idMappings := []*datapb.CopySegmentIDMapping{
|
|
{SourceSegmentId: 1, TargetSegmentId: 999, PartitionId: 10},
|
|
}
|
|
|
|
task := ©SegmentTask{
|
|
copyMeta: s.copyMeta,
|
|
tr: timerecord.NewTimeRecorder("task"),
|
|
times: taskcommon.NewTimes(),
|
|
}
|
|
task.task.Store(&datapb.CopySegmentTask{
|
|
TaskId: 1001,
|
|
JobId: s.jobID,
|
|
CollectionId: s.collectionID,
|
|
State: datapb.CopySegmentTaskState_CopySegmentTaskFailed,
|
|
IdMappings: idMappings,
|
|
Reason: "test failure",
|
|
})
|
|
err = s.copyMeta.AddTask(context.TODO(), task)
|
|
s.NoError(err)
|
|
|
|
// Process failed task should not panic when segment doesn't exist
|
|
s.NotPanics(func() {
|
|
s.inspector.processFailed(task)
|
|
})
|
|
}
|
|
|
|
func (s *CopySegmentInspectorSuite) TestInspect_ProcessPendingAndFailedTasks() {
|
|
s.catalog.EXPECT().SaveCopySegmentJob(mock.Anything, mock.Anything).Return(nil)
|
|
s.catalog.EXPECT().SaveCopySegmentTask(mock.Anything, mock.Anything).Return(nil).Times(3)
|
|
s.catalog.EXPECT().AddSegment(mock.Anything, mock.Anything).Return(nil)
|
|
s.catalog.EXPECT().AlterSegments(mock.Anything, mock.Anything).Return(nil)
|
|
|
|
// Create target segment
|
|
seg1 := NewSegmentInfo(&datapb.SegmentInfo{
|
|
ID: 101,
|
|
CollectionID: s.collectionID,
|
|
PartitionID: 10,
|
|
State: commonpb.SegmentState_Importing,
|
|
NumOfRows: 100,
|
|
InsertChannel: "ch1",
|
|
})
|
|
err := s.meta.AddSegment(context.TODO(), seg1)
|
|
s.NoError(err)
|
|
|
|
// Create a job
|
|
job := ©SegmentJob{
|
|
CopySegmentJob: &datapb.CopySegmentJob{
|
|
JobId: s.jobID,
|
|
CollectionId: s.collectionID,
|
|
State: datapb.CopySegmentJobState_CopySegmentJobExecuting,
|
|
},
|
|
tr: timerecord.NewTimeRecorder("test job"),
|
|
}
|
|
err = s.copyMeta.AddJob(context.TODO(), job)
|
|
s.NoError(err)
|
|
|
|
// Create a pending task
|
|
task1 := ©SegmentTask{
|
|
copyMeta: s.copyMeta,
|
|
tr: timerecord.NewTimeRecorder("task1"),
|
|
times: taskcommon.NewTimes(),
|
|
}
|
|
task1.task.Store(&datapb.CopySegmentTask{
|
|
TaskId: 1001,
|
|
JobId: s.jobID,
|
|
CollectionId: s.collectionID,
|
|
State: datapb.CopySegmentTaskState_CopySegmentTaskPending,
|
|
})
|
|
err = s.copyMeta.AddTask(context.TODO(), task1)
|
|
s.NoError(err)
|
|
|
|
// Create an in-progress task (should not be processed by inspect)
|
|
task2 := ©SegmentTask{
|
|
copyMeta: s.copyMeta,
|
|
tr: timerecord.NewTimeRecorder("task2"),
|
|
times: taskcommon.NewTimes(),
|
|
}
|
|
task2.task.Store(&datapb.CopySegmentTask{
|
|
TaskId: 1002,
|
|
JobId: s.jobID,
|
|
CollectionId: s.collectionID,
|
|
State: datapb.CopySegmentTaskState_CopySegmentTaskInProgress,
|
|
})
|
|
err = s.copyMeta.AddTask(context.TODO(), task2)
|
|
s.NoError(err)
|
|
|
|
// Create a failed task
|
|
idMappings := []*datapb.CopySegmentIDMapping{
|
|
{SourceSegmentId: 1, TargetSegmentId: 101, PartitionId: 10},
|
|
}
|
|
task3 := ©SegmentTask{
|
|
copyMeta: s.copyMeta,
|
|
tr: timerecord.NewTimeRecorder("task3"),
|
|
times: taskcommon.NewTimes(),
|
|
}
|
|
task3.task.Store(&datapb.CopySegmentTask{
|
|
TaskId: 1003,
|
|
JobId: s.jobID,
|
|
CollectionId: s.collectionID,
|
|
State: datapb.CopySegmentTaskState_CopySegmentTaskFailed,
|
|
IdMappings: idMappings,
|
|
})
|
|
err = s.copyMeta.AddTask(context.TODO(), task3)
|
|
s.NoError(err)
|
|
|
|
// Expect only the pending task to be enqueued
|
|
s.scheduler.EXPECT().Enqueue(mock.MatchedBy(func(t any) bool {
|
|
copyTask, ok := t.(CopySegmentTask)
|
|
return ok && copyTask.GetTaskId() == 1001
|
|
})).Once()
|
|
|
|
// Inspect should process pending and failed tasks
|
|
s.inspector.inspect()
|
|
|
|
// Verify failed task's target segment is dropped
|
|
segment := s.meta.GetSegment(context.TODO(), 101)
|
|
s.Equal(commonpb.SegmentState_Dropped, segment.GetState())
|
|
}
|
|
|
|
func (s *CopySegmentInspectorSuite) TestInspect_MultipleJobs() {
|
|
s.catalog.EXPECT().SaveCopySegmentJob(mock.Anything, mock.Anything).Return(nil).Times(2)
|
|
s.catalog.EXPECT().SaveCopySegmentTask(mock.Anything, mock.Anything).Return(nil).Times(2)
|
|
|
|
// Create two jobs
|
|
job1 := ©SegmentJob{
|
|
CopySegmentJob: &datapb.CopySegmentJob{
|
|
JobId: 100,
|
|
CollectionId: s.collectionID,
|
|
State: datapb.CopySegmentJobState_CopySegmentJobExecuting,
|
|
},
|
|
tr: timerecord.NewTimeRecorder("job1"),
|
|
}
|
|
err := s.copyMeta.AddJob(context.TODO(), job1)
|
|
s.NoError(err)
|
|
|
|
job2 := ©SegmentJob{
|
|
CopySegmentJob: &datapb.CopySegmentJob{
|
|
JobId: 200,
|
|
CollectionId: s.collectionID,
|
|
State: datapb.CopySegmentJobState_CopySegmentJobExecuting,
|
|
},
|
|
tr: timerecord.NewTimeRecorder("job2"),
|
|
}
|
|
err = s.copyMeta.AddJob(context.TODO(), job2)
|
|
s.NoError(err)
|
|
|
|
// Create pending tasks for both jobs
|
|
task1 := ©SegmentTask{
|
|
copyMeta: s.copyMeta,
|
|
tr: timerecord.NewTimeRecorder("task1"),
|
|
times: taskcommon.NewTimes(),
|
|
}
|
|
task1.task.Store(&datapb.CopySegmentTask{
|
|
TaskId: 1001,
|
|
JobId: 100,
|
|
CollectionId: s.collectionID,
|
|
State: datapb.CopySegmentTaskState_CopySegmentTaskPending,
|
|
})
|
|
err = s.copyMeta.AddTask(context.TODO(), task1)
|
|
s.NoError(err)
|
|
|
|
task2 := ©SegmentTask{
|
|
copyMeta: s.copyMeta,
|
|
tr: timerecord.NewTimeRecorder("task2"),
|
|
times: taskcommon.NewTimes(),
|
|
}
|
|
task2.task.Store(&datapb.CopySegmentTask{
|
|
TaskId: 2001,
|
|
JobId: 200,
|
|
CollectionId: s.collectionID,
|
|
State: datapb.CopySegmentTaskState_CopySegmentTaskPending,
|
|
})
|
|
err = s.copyMeta.AddTask(context.TODO(), task2)
|
|
s.NoError(err)
|
|
|
|
// Expect both tasks to be enqueued
|
|
s.scheduler.EXPECT().Enqueue(mock.Anything).Times(2)
|
|
|
|
s.inspector.inspect()
|
|
}
|
|
|
|
func (s *CopySegmentInspectorSuite) TestClose() {
|
|
// Close should be idempotent
|
|
s.NotPanics(func() {
|
|
s.inspector.Close()
|
|
s.inspector.Close()
|
|
})
|
|
}
|
|
|
|
func (s *CopySegmentInspectorSuite) TestInspect_EmptyJobs() {
|
|
// Inspect with no jobs should not panic
|
|
s.NotPanics(func() {
|
|
s.inspector.inspect()
|
|
})
|
|
}
|
|
|
|
func (s *CopySegmentInspectorSuite) TestInspect_JobsAreSorted() {
|
|
s.catalog.EXPECT().SaveCopySegmentJob(mock.Anything, mock.Anything).Return(nil).Times(3)
|
|
s.catalog.EXPECT().SaveCopySegmentTask(mock.Anything, mock.Anything).Return(nil).Times(3)
|
|
|
|
// Create jobs with non-sequential IDs to test sorting
|
|
jobs := []int64{300, 100, 200}
|
|
for _, jobID := range jobs {
|
|
job := ©SegmentJob{
|
|
CopySegmentJob: &datapb.CopySegmentJob{
|
|
JobId: jobID,
|
|
CollectionId: s.collectionID,
|
|
State: datapb.CopySegmentJobState_CopySegmentJobExecuting,
|
|
},
|
|
tr: timerecord.NewTimeRecorder("job"),
|
|
}
|
|
err := s.copyMeta.AddJob(context.TODO(), job)
|
|
s.NoError(err)
|
|
|
|
// Create a pending task for each job
|
|
task := ©SegmentTask{
|
|
copyMeta: s.copyMeta,
|
|
tr: timerecord.NewTimeRecorder("task"),
|
|
times: taskcommon.NewTimes(),
|
|
}
|
|
task.task.Store(&datapb.CopySegmentTask{
|
|
TaskId: jobID*10 + 1,
|
|
JobId: jobID,
|
|
CollectionId: s.collectionID,
|
|
State: datapb.CopySegmentTaskState_CopySegmentTaskPending,
|
|
})
|
|
err = s.copyMeta.AddTask(context.TODO(), task)
|
|
s.NoError(err)
|
|
}
|
|
|
|
// Track the order of task IDs being enqueued
|
|
var enqueuedTaskIDs []int64
|
|
s.scheduler.EXPECT().Enqueue(mock.Anything).Run(func(t task2.Task) {
|
|
copyTask, ok := t.(CopySegmentTask)
|
|
s.True(ok)
|
|
enqueuedTaskIDs = append(enqueuedTaskIDs, copyTask.GetTaskId())
|
|
}).Times(3)
|
|
|
|
s.inspector.inspect()
|
|
|
|
// Verify tasks were processed in order of job IDs (100, 200, 300)
|
|
s.Equal([]int64{1001, 2001, 3001}, enqueuedTaskIDs)
|
|
}
|