admin/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2026-01-07 19:31:51 +08:00
Code Issues Packages Projects Releases Wiki Activity
milvus/internal/datacoord/copy_segment_meta_test.go
wei liu 975c91df16
feat: Add comprehensive snapshot functionality for collections (#44361) 
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>
2026-01-06 10:15:24 +08:00

869 lines
26 KiB
Go
Raw Blame History

// 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/cockroachdb/errors"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/suite"
"github.com/milvus-io/milvus/internal/datacoord/broker"
"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 CopySegmentMetaSuite struct {
suite.Suite
collectionID int64
catalog *mocks.DataCoordCatalog
broker *broker.MockBroker
meta *meta
copyMeta CopySegmentMeta
}
func (s *CopySegmentMetaSuite) SetupTest() {
var err error
s.collectionID = 1
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)
}
func TestCopySegmentMeta(t *testing.T) {
suite.Run(t, new(CopySegmentMetaSuite))
}
func (s *CopySegmentMetaSuite) TestNewCopySegmentMeta_Success() {
catalog := mocks.NewDataCoordCatalog(s.T())
catalog.EXPECT().ListCopySegmentJobs(mock.Anything).Return(nil, nil)
catalog.EXPECT().ListCopySegmentTasks(mock.Anything).Return(nil, nil)
catalog.EXPECT().ListChannelCheckpoint(mock.Anything).Return(nil, nil)
catalog.EXPECT().ListIndexes(mock.Anything).Return(nil, nil)
catalog.EXPECT().ListSegmentIndexes(mock.Anything).Return(nil, nil)
catalog.EXPECT().ListAnalyzeTasks(mock.Anything).Return(nil, nil)
catalog.EXPECT().ListCompactionTask(mock.Anything).Return(nil, nil)
catalog.EXPECT().ListPartitionStatsInfos(mock.Anything).Return(nil, nil)
catalog.EXPECT().ListStatsTasks(mock.Anything).Return(nil, nil)
catalog.EXPECT().ListFileResource(mock.Anything).Return(nil, 0, nil)
catalog.EXPECT().ListSnapshots(mock.Anything).Return(nil, nil)
broker := broker.NewMockBroker(s.T())
broker.EXPECT().ShowCollectionIDs(mock.Anything).Return(nil, nil)
meta, err := newMeta(context.TODO(), catalog, nil, broker)
s.NoError(err)
copyMeta, err := NewCopySegmentMeta(context.TODO(), catalog, meta, nil)
s.NoError(err)
s.NotNil(copyMeta)
}
func (s *CopySegmentMetaSuite) TestNewCopySegmentMeta_ListJobsError() {
catalog := mocks.NewDataCoordCatalog(s.T())
catalog.EXPECT().ListCopySegmentJobs(mock.Anything).Return(nil, errors.New("list jobs error"))
copyMeta, err := NewCopySegmentMeta(context.TODO(), catalog, nil, nil)
s.Error(err)
s.Nil(copyMeta)
s.Contains(err.Error(), "list jobs error")
}
func (s *CopySegmentMetaSuite) TestNewCopySegmentMeta_ListTasksError() {
catalog := mocks.NewDataCoordCatalog(s.T())
catalog.EXPECT().ListCopySegmentJobs(mock.Anything).Return(nil, nil)
catalog.EXPECT().ListCopySegmentTasks(mock.Anything).Return(nil, errors.New("list tasks error"))
copyMeta, err := NewCopySegmentMeta(context.TODO(), catalog, nil, nil)
s.Error(err)
s.Nil(copyMeta)
s.Contains(err.Error(), "list tasks error")
}
func (s *CopySegmentMetaSuite) TestNewCopySegmentMeta_RestoreJobs() {
catalog := mocks.NewDataCoordCatalog(s.T())
restoredJobs := []*datapb.CopySegmentJob{
{
JobId: 100,
CollectionId: 1,
State: datapb.CopySegmentJobState_CopySegmentJobPending,
},
{
JobId: 200,
CollectionId: 1,
State: datapb.CopySegmentJobState_CopySegmentJobExecuting,
},
}
catalog.EXPECT().ListCopySegmentJobs(mock.Anything).Return(restoredJobs, nil)
catalog.EXPECT().ListCopySegmentTasks(mock.Anything).Return(nil, nil)
catalog.EXPECT().ListChannelCheckpoint(mock.Anything).Return(nil, nil)
catalog.EXPECT().ListIndexes(mock.Anything).Return(nil, nil)
catalog.EXPECT().ListSegmentIndexes(mock.Anything).Return(nil, nil)
catalog.EXPECT().ListAnalyzeTasks(mock.Anything).Return(nil, nil)
catalog.EXPECT().ListCompactionTask(mock.Anything).Return(nil, nil)
catalog.EXPECT().ListPartitionStatsInfos(mock.Anything).Return(nil, nil)
catalog.EXPECT().ListStatsTasks(mock.Anything).Return(nil, nil)
catalog.EXPECT().ListFileResource(mock.Anything).Return(nil, 0, nil)
catalog.EXPECT().ListSnapshots(mock.Anything).Return(nil, nil)
broker := broker.NewMockBroker(s.T())
broker.EXPECT().ShowCollectionIDs(mock.Anything).Return(nil, nil)
meta, err := newMeta(context.TODO(), catalog, nil, broker)
s.NoError(err)
copyMeta, err := NewCopySegmentMeta(context.TODO(), catalog, meta, nil)
s.NoError(err)
// Verify jobs are restored
job1 := copyMeta.GetJob(context.TODO(), 100)
s.NotNil(job1)
s.Equal(int64(100), job1.GetJobId())
job2 := copyMeta.GetJob(context.TODO(), 200)
s.NotNil(job2)
s.Equal(int64(200), job2.GetJobId())
}
func (s *CopySegmentMetaSuite) TestNewCopySegmentMeta_RestoreTasks() {
catalog := mocks.NewDataCoordCatalog(s.T())
restoredTasks := []*datapb.CopySegmentTask{
{
TaskId: 1001,
JobId: 100,
CollectionId: 1,
State: datapb.CopySegmentTaskState_CopySegmentTaskPending,
},
{
TaskId: 1002,
JobId: 100,
CollectionId: 1,
State: datapb.CopySegmentTaskState_CopySegmentTaskInProgress,
},
}
catalog.EXPECT().ListCopySegmentJobs(mock.Anything).Return(nil, nil)
catalog.EXPECT().ListCopySegmentTasks(mock.Anything).Return(restoredTasks, nil)
catalog.EXPECT().ListChannelCheckpoint(mock.Anything).Return(nil, nil)
catalog.EXPECT().ListIndexes(mock.Anything).Return(nil, nil)
catalog.EXPECT().ListSegmentIndexes(mock.Anything).Return(nil, nil)
catalog.EXPECT().ListAnalyzeTasks(mock.Anything).Return(nil, nil)
catalog.EXPECT().ListCompactionTask(mock.Anything).Return(nil, nil)
catalog.EXPECT().ListPartitionStatsInfos(mock.Anything).Return(nil, nil)
catalog.EXPECT().ListStatsTasks(mock.Anything).Return(nil, nil)
catalog.EXPECT().ListFileResource(mock.Anything).Return(nil, 0, nil)
catalog.EXPECT().ListSnapshots(mock.Anything).Return(nil, nil)
broker := broker.NewMockBroker(s.T())
broker.EXPECT().ShowCollectionIDs(mock.Anything).Return(nil, nil)
meta, err := newMeta(context.TODO(), catalog, nil, broker)
s.NoError(err)
copyMeta, err := NewCopySegmentMeta(context.TODO(), catalog, meta, nil)
s.NoError(err)
// Verify tasks are restored
task1 := copyMeta.GetTask(context.TODO(), 1001)
s.NotNil(task1)
s.Equal(int64(1001), task1.GetTaskId())
task2 := copyMeta.GetTask(context.TODO(), 1002)
s.NotNil(task2)
s.Equal(int64(1002), task2.GetTaskId())
}
func (s *CopySegmentMetaSuite) TestAddJob_Success() {
s.catalog.EXPECT().SaveCopySegmentJob(mock.Anything, mock.Anything).Return(nil)
job := &copySegmentJob{
CopySegmentJob: &datapb.CopySegmentJob{
JobId: 100,
CollectionId: s.collectionID,
State: datapb.CopySegmentJobState_CopySegmentJobPending,
},
tr: timerecord.NewTimeRecorder("test job"),
}
err := s.copyMeta.AddJob(context.TODO(), job)
s.NoError(err)
// Verify job is added
retrievedJob := s.copyMeta.GetJob(context.TODO(), 100)
s.NotNil(retrievedJob)
s.Equal(int64(100), retrievedJob.GetJobId())
}
func (s *CopySegmentMetaSuite) TestAddJob_CatalogError() {
s.catalog.EXPECT().SaveCopySegmentJob(mock.Anything, mock.Anything).Return(errors.New("catalog error"))
job := &copySegmentJob{
CopySegmentJob: &datapb.CopySegmentJob{
JobId: 100,
CollectionId: s.collectionID,
},
tr: timerecord.NewTimeRecorder("test job"),
}
err := s.copyMeta.AddJob(context.TODO(), job)
s.Error(err)
s.Contains(err.Error(), "catalog error")
// Verify job is not added
retrievedJob := s.copyMeta.GetJob(context.TODO(), 100)
s.Nil(retrievedJob)
}
func (s *CopySegmentMetaSuite) TestUpdateJob_Success() {
s.catalog.EXPECT().SaveCopySegmentJob(mock.Anything, mock.Anything).Return(nil).Times(2)
// Add a job first
job := &copySegmentJob{
CopySegmentJob: &datapb.CopySegmentJob{
JobId: 100,
CollectionId: s.collectionID,
State: datapb.CopySegmentJobState_CopySegmentJobPending,
},
tr: timerecord.NewTimeRecorder("test job"),
}
err := s.copyMeta.AddJob(context.TODO(), job)
s.NoError(err)
// Update the job
err = s.copyMeta.UpdateJob(context.TODO(), 100,
UpdateCopyJobState(datapb.CopySegmentJobState_CopySegmentJobExecuting),
UpdateCopyJobReason("executing"))
s.NoError(err)
// Verify job is updated
updatedJob := s.copyMeta.GetJob(context.TODO(), 100)
s.Equal(datapb.CopySegmentJobState_CopySegmentJobExecuting, updatedJob.GetState())
s.Equal("executing", updatedJob.GetReason())
}
func (s *CopySegmentMetaSuite) TestUpdateJob_NotFound() {
// Try to update non-existent job (should not error, just no-op)
err := s.copyMeta.UpdateJob(context.TODO(), 999,
UpdateCopyJobState(datapb.CopySegmentJobState_CopySegmentJobExecuting))
s.NoError(err)
}
func (s *CopySegmentMetaSuite) TestUpdateJob_CatalogError() {
s.catalog.EXPECT().SaveCopySegmentJob(mock.Anything, mock.Anything).Return(nil).Once()
s.catalog.EXPECT().SaveCopySegmentJob(mock.Anything, mock.Anything).Return(errors.New("catalog error")).Once()
// Add a job first
job := &copySegmentJob{
CopySegmentJob: &datapb.CopySegmentJob{
JobId: 100,
CollectionId: s.collectionID,
State: datapb.CopySegmentJobState_CopySegmentJobPending,
},
tr: timerecord.NewTimeRecorder("test job"),
}
err := s.copyMeta.AddJob(context.TODO(), job)
s.NoError(err)
// Update should fail
err = s.copyMeta.UpdateJob(context.TODO(), 100,
UpdateCopyJobState(datapb.CopySegmentJobState_CopySegmentJobExecuting))
s.Error(err)
s.Contains(err.Error(), "catalog error")
}
func (s *CopySegmentMetaSuite) TestGetJob() {
s.catalog.EXPECT().SaveCopySegmentJob(mock.Anything, mock.Anything).Return(nil)
job := &copySegmentJob{
CopySegmentJob: &datapb.CopySegmentJob{
JobId: 100,
CollectionId: s.collectionID,
},
tr: timerecord.NewTimeRecorder("test job"),
}
s.copyMeta.AddJob(context.TODO(), job)
// Get existing job
retrievedJob := s.copyMeta.GetJob(context.TODO(), 100)
s.NotNil(retrievedJob)
s.Equal(int64(100), retrievedJob.GetJobId())
// Get non-existent job
nonExistent := s.copyMeta.GetJob(context.TODO(), 999)
s.Nil(nonExistent)
}
func (s *CopySegmentMetaSuite) TestGetJobBy() {
s.catalog.EXPECT().SaveCopySegmentJob(mock.Anything, mock.Anything).Return(nil).Times(3)
// Add multiple jobs
jobs := []*copySegmentJob{
{
CopySegmentJob: &datapb.CopySegmentJob{
JobId: 100,
CollectionId: 1,
State: datapb.CopySegmentJobState_CopySegmentJobPending,
},
tr: timerecord.NewTimeRecorder("job1"),
},
{
CopySegmentJob: &datapb.CopySegmentJob{
JobId: 200,
CollectionId: 1,
State: datapb.CopySegmentJobState_CopySegmentJobExecuting,
},
tr: timerecord.NewTimeRecorder("job2"),
},
{
CopySegmentJob: &datapb.CopySegmentJob{
JobId: 300,
CollectionId: 2,
State: datapb.CopySegmentJobState_CopySegmentJobPending,
},
tr: timerecord.NewTimeRecorder("job3"),
},
}
for _, job := range jobs {
s.copyMeta.AddJob(context.TODO(), job)
}
// Filter by collection ID
filtered := s.copyMeta.GetJobBy(context.TODO(), WithCopyJobCollectionID(1))
s.Len(filtered, 2)
// Filter by state
filtered = s.copyMeta.GetJobBy(context.TODO(),
WithCopyJobStates(datapb.CopySegmentJobState_CopySegmentJobPending))
s.Len(filtered, 2)
// Filter by collection ID and state
filtered = s.copyMeta.GetJobBy(context.TODO(),
WithCopyJobCollectionID(1),
WithCopyJobStates(datapb.CopySegmentJobState_CopySegmentJobPending))
s.Len(filtered, 1)
s.Equal(int64(100), filtered[0].GetJobId())
}
func (s *CopySegmentMetaSuite) TestCountJobBy() {
s.catalog.EXPECT().SaveCopySegmentJob(mock.Anything, mock.Anything).Return(nil).Times(3)
// Add multiple jobs
for i := 0; i < 3; i++ {
job := &copySegmentJob{
CopySegmentJob: &datapb.CopySegmentJob{
JobId: int64(100 + i),
CollectionId: s.collectionID,
State: datapb.CopySegmentJobState_CopySegmentJobPending,
},
tr: timerecord.NewTimeRecorder("job"),
}
s.copyMeta.AddJob(context.TODO(), job)
}
count := s.copyMeta.CountJobBy(context.TODO(),
WithCopyJobStates(datapb.CopySegmentJobState_CopySegmentJobPending))
s.Equal(3, count)
count = s.copyMeta.CountJobBy(context.TODO(),
WithCopyJobStates(datapb.CopySegmentJobState_CopySegmentJobExecuting))
s.Equal(0, count)
}
func (s *CopySegmentMetaSuite) TestRemoveJob_Success() {
s.catalog.EXPECT().SaveCopySegmentJob(mock.Anything, mock.Anything).Return(nil)
s.catalog.EXPECT().DropCopySegmentJob(mock.Anything, int64(100)).Return(nil)
// Add a job
job := &copySegmentJob{
CopySegmentJob: &datapb.CopySegmentJob{
JobId: 100,
CollectionId: s.collectionID,
},
tr: timerecord.NewTimeRecorder("test job"),
}
s.copyMeta.AddJob(context.TODO(), job)
// Remove the job
err := s.copyMeta.RemoveJob(context.TODO(), 100)
s.NoError(err)
// Verify job is removed
retrievedJob := s.copyMeta.GetJob(context.TODO(), 100)
s.Nil(retrievedJob)
}
func (s *CopySegmentMetaSuite) TestRemoveJob_NotFound() {
// Remove non-existent job (should not error)
err := s.copyMeta.RemoveJob(context.TODO(), 999)
s.NoError(err)
}
func (s *CopySegmentMetaSuite) TestRemoveJob_CatalogError() {
s.catalog.EXPECT().SaveCopySegmentJob(mock.Anything, mock.Anything).Return(nil)
s.catalog.EXPECT().DropCopySegmentJob(mock.Anything, int64(100)).Return(errors.New("catalog error"))
// Add a job
job := &copySegmentJob{
CopySegmentJob: &datapb.CopySegmentJob{
JobId: 100,
CollectionId: s.collectionID,
},
tr: timerecord.NewTimeRecorder("test job"),
}
s.copyMeta.AddJob(context.TODO(), job)
// Remove should fail
err := s.copyMeta.RemoveJob(context.TODO(), 100)
s.Error(err)
s.Contains(err.Error(), "catalog error")
// Job should still exist in memory
retrievedJob := s.copyMeta.GetJob(context.TODO(), 100)
s.NotNil(retrievedJob)
}
func (s *CopySegmentMetaSuite) TestAddTask_Success() {
s.catalog.EXPECT().SaveCopySegmentTask(mock.Anything, mock.Anything).Return(nil)
task := &copySegmentTask{
copyMeta: s.copyMeta,
tr: timerecord.NewTimeRecorder("task"),
times: taskcommon.NewTimes(),
}
task.task.Store(&datapb.CopySegmentTask{
TaskId: 1001,
JobId: 100,
CollectionId: s.collectionID,
State: datapb.CopySegmentTaskState_CopySegmentTaskPending,
})
err := s.copyMeta.AddTask(context.TODO(), task)
s.NoError(err)
// Verify task is added
retrievedTask := s.copyMeta.GetTask(context.TODO(), 1001)
s.NotNil(retrievedTask)
s.Equal(int64(1001), retrievedTask.GetTaskId())
}
func (s *CopySegmentMetaSuite) TestAddTask_CatalogError() {
s.catalog.EXPECT().SaveCopySegmentTask(mock.Anything, mock.Anything).Return(errors.New("catalog error"))
task := &copySegmentTask{
copyMeta: s.copyMeta,
tr: timerecord.NewTimeRecorder("task"),
times: taskcommon.NewTimes(),
}
task.task.Store(&datapb.CopySegmentTask{
TaskId: 1001,
JobId: 100,
CollectionId: s.collectionID,
})
err := s.copyMeta.AddTask(context.TODO(), task)
s.Error(err)
s.Contains(err.Error(), "catalog error")
// Verify task is not added
retrievedTask := s.copyMeta.GetTask(context.TODO(), 1001)
s.Nil(retrievedTask)
}
func (s *CopySegmentMetaSuite) TestUpdateTask_Success() {
s.catalog.EXPECT().SaveCopySegmentTask(mock.Anything, mock.Anything).Return(nil).Times(2)
// Add a task first
task := &copySegmentTask{
copyMeta: s.copyMeta,
tr: timerecord.NewTimeRecorder("task"),
times: taskcommon.NewTimes(),
}
task.task.Store(&datapb.CopySegmentTask{
TaskId: 1001,
JobId: 100,
CollectionId: s.collectionID,
State: datapb.CopySegmentTaskState_CopySegmentTaskPending,
})
err := s.copyMeta.AddTask(context.TODO(), task)
s.NoError(err)
// Update the task
err = s.copyMeta.UpdateTask(context.TODO(), 1001,
UpdateCopyTaskState(datapb.CopySegmentTaskState_CopySegmentTaskInProgress),
UpdateCopyTaskReason("executing"))
s.NoError(err)
// Verify task is updated
updatedTask := s.copyMeta.GetTask(context.TODO(), 1001)
s.Equal(datapb.CopySegmentTaskState_CopySegmentTaskInProgress, updatedTask.GetState())
s.Equal("executing", updatedTask.GetReason())
}
func (s *CopySegmentMetaSuite) TestUpdateTask_NotFound() {
// Try to update non-existent task (should not error, just no-op)
err := s.copyMeta.UpdateTask(context.TODO(), 9999,
UpdateCopyTaskState(datapb.CopySegmentTaskState_CopySegmentTaskInProgress))
s.NoError(err)
}
func (s *CopySegmentMetaSuite) TestGetTask() {
s.catalog.EXPECT().SaveCopySegmentTask(mock.Anything, mock.Anything).Return(nil)
task := &copySegmentTask{
copyMeta: s.copyMeta,
tr: timerecord.NewTimeRecorder("task"),
times: taskcommon.NewTimes(),
}
task.task.Store(&datapb.CopySegmentTask{
TaskId: 1001,
JobId: 100,
CollectionId: s.collectionID,
})
s.copyMeta.AddTask(context.TODO(), task)
// Get existing task
retrievedTask := s.copyMeta.GetTask(context.TODO(), 1001)
s.NotNil(retrievedTask)
s.Equal(int64(1001), retrievedTask.GetTaskId())
// Get non-existent task
nonExistent := s.copyMeta.GetTask(context.TODO(), 9999)
s.Nil(nonExistent)
}
func (s *CopySegmentMetaSuite) TestGetTaskBy() {
s.catalog.EXPECT().SaveCopySegmentTask(mock.Anything, mock.Anything).Return(nil).Times(3)
// Add multiple tasks
tasks := []struct {
taskID int64
jobID int64
state datapb.CopySegmentTaskState
}{
{1001, 100, datapb.CopySegmentTaskState_CopySegmentTaskPending},
{1002, 100, datapb.CopySegmentTaskState_CopySegmentTaskInProgress},
{1003, 200, datapb.CopySegmentTaskState_CopySegmentTaskPending},
}
for _, t := range tasks {
task := &copySegmentTask{
copyMeta: s.copyMeta,
tr: timerecord.NewTimeRecorder("task"),
times: taskcommon.NewTimes(),
}
task.task.Store(&datapb.CopySegmentTask{
TaskId: t.taskID,
JobId: t.jobID,
CollectionId: s.collectionID,
State: t.state,
})
s.copyMeta.AddTask(context.TODO(), task)
}
// Filter by job ID
filtered := s.copyMeta.GetTaskBy(context.TODO(), WithCopyTaskJob(100))
s.Len(filtered, 2)
// Filter by state
filtered = s.copyMeta.GetTaskBy(context.TODO(),
WithCopyTaskStates(datapb.CopySegmentTaskState_CopySegmentTaskPending))
s.Len(filtered, 2)
// Filter by job ID and state
filtered = s.copyMeta.GetTaskBy(context.TODO(),
WithCopyTaskJob(100),
WithCopyTaskStates(datapb.CopySegmentTaskState_CopySegmentTaskPending))
s.Len(filtered, 1)
s.Equal(int64(1001), filtered[0].GetTaskId())
}
func (s *CopySegmentMetaSuite) TestGetTasksByJobID() {
s.catalog.EXPECT().SaveCopySegmentTask(mock.Anything, mock.Anything).Return(nil).Times(3)
// Create tasks for different jobs
tasks := []struct {
taskID int64
jobID int64
collectionID int64
}{
{taskID: 1001, jobID: 100, collectionID: 1},
{taskID: 1002, jobID: 100, collectionID: 1},
{taskID: 1003, jobID: 200, collectionID: 2},
}
for _, t := range tasks {
task := &copySegmentTask{
copyMeta: s.copyMeta,
tr: timerecord.NewTimeRecorder("task"),
times: taskcommon.NewTimes(),
}
task.task.Store(&datapb.CopySegmentTask{
TaskId: t.taskID,
JobId: t.jobID,
CollectionId: t.collectionID,
})
s.copyMeta.AddTask(context.TODO(), task)
}
// GetTasksByJobID should return tasks for job 100
result := s.copyMeta.GetTasksByJobID(context.TODO(), 100)
s.Len(result, 2)
taskIDs := make([]int64, 0)
for _, t := range result {
taskIDs = append(taskIDs, t.GetTaskId())
}
s.ElementsMatch([]int64{1001, 1002}, taskIDs)
// GetTasksByJobID should return tasks for job 200
result = s.copyMeta.GetTasksByJobID(context.TODO(), 200)
s.Len(result, 1)
s.Equal(int64(1003), result[0].GetTaskId())
// GetTasksByJobID should return empty for non-existent job
result = s.copyMeta.GetTasksByJobID(context.TODO(), 999)
s.Len(result, 0)
}
func (s *CopySegmentMetaSuite) TestGetTasksByCollectionID() {
s.catalog.EXPECT().SaveCopySegmentTask(mock.Anything, mock.Anything).Return(nil).Times(3)
// Create tasks for different collections
tasks := []struct {
taskID int64
jobID int64
collectionID int64
}{
{taskID: 1001, jobID: 100, collectionID: 1},
{taskID: 1002, jobID: 100, collectionID: 1},
{taskID: 1003, jobID: 200, collectionID: 2},
}
for _, t := range tasks {
task := &copySegmentTask{
copyMeta: s.copyMeta,
tr: timerecord.NewTimeRecorder("task"),
times: taskcommon.NewTimes(),
}
task.task.Store(&datapb.CopySegmentTask{
TaskId: t.taskID,
JobId: t.jobID,
CollectionId: t.collectionID,
})
s.copyMeta.AddTask(context.TODO(), task)
}
// GetTasksByCollectionID should return tasks for collection 1
result := s.copyMeta.GetTasksByCollectionID(context.TODO(), 1)
s.Len(result, 2)
taskIDs := make([]int64, 0)
for _, t := range result {
taskIDs = append(taskIDs, t.GetTaskId())
}
s.ElementsMatch([]int64{1001, 1002}, taskIDs)
// GetTasksByCollectionID should return tasks for collection 2
result = s.copyMeta.GetTasksByCollectionID(context.TODO(), 2)
s.Len(result, 1)
s.Equal(int64(1003), result[0].GetTaskId())
// GetTasksByCollectionID should return empty for non-existent collection
result = s.copyMeta.GetTasksByCollectionID(context.TODO(), 999)
s.Len(result, 0)
}
func (s *CopySegmentMetaSuite) TestSecondaryIndexCleanup() {
s.catalog.EXPECT().SaveCopySegmentTask(mock.Anything, mock.Anything).Return(nil).Times(2)
s.catalog.EXPECT().DropCopySegmentTask(mock.Anything, mock.Anything).Return(nil).Times(2)
// Add tasks
task1 := &copySegmentTask{
copyMeta: s.copyMeta,
tr: timerecord.NewTimeRecorder("task"),
times: taskcommon.NewTimes(),
}
task1.task.Store(&datapb.CopySegmentTask{
TaskId: 1001,
JobId: 100,
CollectionId: 1,
})
task2 := &copySegmentTask{
copyMeta: s.copyMeta,
tr: timerecord.NewTimeRecorder("task"),
times: taskcommon.NewTimes(),
}
task2.task.Store(&datapb.CopySegmentTask{
TaskId: 1002,
JobId: 100,
CollectionId: 1,
})
s.copyMeta.AddTask(context.TODO(), task1)
s.copyMeta.AddTask(context.TODO(), task2)
// Both tasks exist in indexes
s.Len(s.copyMeta.GetTasksByJobID(context.TODO(), 100), 2)
s.Len(s.copyMeta.GetTasksByCollectionID(context.TODO(), 1), 2)
// Remove first task
err := s.copyMeta.RemoveTask(context.TODO(), 1001)
s.NoError(err)
// Index should be updated
s.Len(s.copyMeta.GetTasksByJobID(context.TODO(), 100), 1)
s.Len(s.copyMeta.GetTasksByCollectionID(context.TODO(), 1), 1)
// Remove second task
err = s.copyMeta.RemoveTask(context.TODO(), 1002)
s.NoError(err)
// Index should be empty
s.Len(s.copyMeta.GetTasksByJobID(context.TODO(), 100), 0)
s.Len(s.copyMeta.GetTasksByCollectionID(context.TODO(), 1), 0)
}
func (s *CopySegmentMetaSuite) TestRemoveTask_Success() {
s.catalog.EXPECT().SaveCopySegmentTask(mock.Anything, mock.Anything).Return(nil)
s.catalog.EXPECT().DropCopySegmentTask(mock.Anything, int64(1001)).Return(nil)
// Add a task
task := &copySegmentTask{
copyMeta: s.copyMeta,
tr: timerecord.NewTimeRecorder("task"),
times: taskcommon.NewTimes(),
}
task.task.Store(&datapb.CopySegmentTask{
TaskId: 1001,
JobId: 100,
CollectionId: s.collectionID,
})
s.copyMeta.AddTask(context.TODO(), task)
// Remove the task
err := s.copyMeta.RemoveTask(context.TODO(), 1001)
s.NoError(err)
// Verify task is removed
retrievedTask := s.copyMeta.GetTask(context.TODO(), 1001)
s.Nil(retrievedTask)
}
func (s *CopySegmentMetaSuite) TestRemoveTask_NotFound() {
// Remove non-existent task (should not error)
err := s.copyMeta.RemoveTask(context.TODO(), 9999)
s.NoError(err)
}
func (s *CopySegmentMetaSuite) TestRemoveTask_CatalogError() {
s.catalog.EXPECT().SaveCopySegmentTask(mock.Anything, mock.Anything).Return(nil)
s.catalog.EXPECT().DropCopySegmentTask(mock.Anything, int64(1001)).Return(errors.New("catalog error"))
// Add a task
task := &copySegmentTask{
copyMeta: s.copyMeta,
tr: timerecord.NewTimeRecorder("task"),
times: taskcommon.NewTimes(),
}
task.task.Store(&datapb.CopySegmentTask{
TaskId: 1001,
JobId: 100,
CollectionId: s.collectionID,
})
s.copyMeta.AddTask(context.TODO(), task)
// Remove should fail
err := s.copyMeta.RemoveTask(context.TODO(), 1001)
s.Error(err)
s.Contains(err.Error(), "catalog error")
// Task should still exist in memory
retrievedTask := s.copyMeta.GetTask(context.TODO(), 1001)
s.NotNil(retrievedTask)
}
func (s *CopySegmentMetaSuite) TestCopySegmentTasks_Operations() {
tasks := newCopySegmentTasks()
// Test empty tasks
s.Nil(tasks.get(1001))
s.Empty(tasks.listTasks())
// Test add
task1 := &copySegmentTask{
tr: timerecord.NewTimeRecorder("task1"),
times: taskcommon.NewTimes(),
}
task1.task.Store(&datapb.CopySegmentTask{TaskId: 1001})
tasks.add(task1)
task2 := &copySegmentTask{
tr: timerecord.NewTimeRecorder("task2"),
times: taskcommon.NewTimes(),
}
task2.task.Store(&datapb.CopySegmentTask{TaskId: 1002})
tasks.add(task2)
// Test get
s.NotNil(tasks.get(1001))
s.NotNil(tasks.get(1002))
s.Nil(tasks.get(9999))
// Test listTasks
allTasks := tasks.listTasks()
s.Len(allTasks, 2)
// Test remove
tasks.remove(1001)
s.Nil(tasks.get(1001))
s.NotNil(tasks.get(1002))
s.Len(tasks.listTasks(), 1)
}
Reference in New Issue View Git Blame Copy Permalink