Skip to content

Commit

Permalink
enhance: Add ut for l0CompactionTask processExecuting
Browse files Browse the repository at this point in the history
See also: milvus-io#34796

Signed-off-by: yangxuan <xuan.yang@zilliz.com>
  • Loading branch information
XuanYang-cn committed Jul 18, 2024
1 parent c8bf6c8 commit a5994bc
Show file tree
Hide file tree
Showing 2 changed files with 151 additions and 25 deletions.
14 changes: 10 additions & 4 deletions internal/datacoord/compaction_task_l0.go
Original file line number Diff line number Diff line change
Expand Up @@ -108,10 +108,11 @@ func (t *l0CompactionTask) processExecuting() bool {
}
switch result.GetState() {
case datapb.CompactionTaskState_executing:
// will L0Compaction be timeouted?
if t.checkTimeout() {
err := t.updateAndSaveTaskMeta(setState(datapb.CompactionTaskState_timeout))
if err != nil {
log.Warn("l0CompactionTask failed to updateAndSaveTaskMeta", zap.Error(err))
log.Warn("l0CompactionTask failed to set task timeout state", zap.Error(err))
return false
}
return t.processTimeout()
Expand All @@ -124,12 +125,13 @@ func (t *l0CompactionTask) processExecuting() bool {
}

if err := t.updateAndSaveTaskMeta(setState(datapb.CompactionTaskState_meta_saved)); err != nil {
log.Warn("l0CompactionTask failed to save task meta_saved state", zap.Error(err))
return false
}
return t.processMetaSaved()
case datapb.CompactionTaskState_failed:
if err := t.updateAndSaveTaskMeta(setState(datapb.CompactionTaskState_failed)); err != nil {
log.Warn("l0CompactionTask failed to updateAndSaveTaskMeta", zap.Error(err))
log.Warn("l0CompactionTask failed to set task failed state", zap.Error(err))
return false
}
return t.processFailed()
Expand Down Expand Up @@ -356,11 +358,15 @@ func (t *l0CompactionTask) processFailed() bool {

func (t *l0CompactionTask) checkTimeout() bool {
if t.GetTimeoutInSeconds() > 0 {
diff := time.Since(time.Unix(t.GetStartTime(), 0)).Seconds()
start := time.Unix(t.GetStartTime(), 0)
diff := time.Since(start).Seconds()
if diff > float64(t.GetTimeoutInSeconds()) {
log.Warn("compaction timeout",
zap.Int64("taskID", t.GetTriggerID()),
zap.Int64("planID", t.GetPlanID()),
zap.Int64("nodeID", t.GetNodeID()),
zap.Int32("timeout in seconds", t.GetTimeoutInSeconds()),
zap.Int64("startTime", t.GetStartTime()),
zap.Time("startTime", start),
)
return true
}
Expand Down
162 changes: 141 additions & 21 deletions internal/datacoord/compaction_task_l0_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,17 +18,42 @@ package datacoord

import (
"context"
"testing"
"time"

"github.com/cockroachdb/errors"
"github.com/samber/lo"
"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/proto/datapb"
"github.com/milvus-io/milvus/pkg/util/merr"
)

func (s *CompactionTaskSuite) TestProcessRefreshPlan_NormalL0() {
func TestL0CompactionTaskSuite(t *testing.T) {
suite.Run(t, new(L0CompactionTaskSuite))
}

type L0CompactionTaskSuite struct {
suite.Suite

mockAlloc *NMockAllocator
mockMeta *MockCompactionMeta
mockSessMgr *MockSessionManager
}

func (s *L0CompactionTaskSuite) SetupTest() {
s.mockMeta = NewMockCompactionMeta(s.T())
s.mockSessMgr = NewMockSessionManager(s.T())
s.mockAlloc = NewNMockAllocator(s.T())
}

func (s *L0CompactionTaskSuite) SetupSubTest() {
s.SetupTest()
}

func (s *L0CompactionTaskSuite) TestProcessRefreshPlan_NormalL0() {
channel := "Ch-1"
deltaLogs := []*datapb.FieldBinlog{getFieldBinlogIDs(101, 3)}

Expand Down Expand Up @@ -88,7 +113,7 @@ func (s *CompactionTaskSuite) TestProcessRefreshPlan_NormalL0() {
s.ElementsMatch([]int64{200, 201, 202, 100, 101}, segIDs)
}

func (s *CompactionTaskSuite) TestProcessRefreshPlan_SegmentNotFoundL0() {
func (s *L0CompactionTaskSuite) TestProcessRefreshPlan_SegmentNotFoundL0() {
channel := "Ch-1"
s.mockMeta.EXPECT().GetHealthySegment(mock.Anything).RunAndReturn(func(segID int64) *SegmentInfo {
return nil
Expand Down Expand Up @@ -116,7 +141,7 @@ func (s *CompactionTaskSuite) TestProcessRefreshPlan_SegmentNotFoundL0() {
s.ErrorIs(err, merr.ErrSegmentNotFound)
}

func (s *CompactionTaskSuite) TestProcessRefreshPlan_SelectZeroSegmentsL0() {
func (s *L0CompactionTaskSuite) TestProcessRefreshPlan_SelectZeroSegmentsL0() {
channel := "Ch-1"
deltaLogs := []*datapb.FieldBinlog{getFieldBinlogIDs(101, 3)}
s.mockMeta.EXPECT().GetHealthySegment(mock.Anything).RunAndReturn(func(segID int64) *SegmentInfo {
Expand Down Expand Up @@ -150,7 +175,7 @@ func (s *CompactionTaskSuite) TestProcessRefreshPlan_SelectZeroSegmentsL0() {
s.Error(err)
}

func (s *CompactionTaskSuite) TestBuildCompactionRequestFailed_AllocFailed() {
func (s *L0CompactionTaskSuite) TestBuildCompactionRequestFailed_AllocFailed() {
var task CompactionTask

alloc := NewNMockAllocator(s.T())
Expand Down Expand Up @@ -178,7 +203,7 @@ func (s *CompactionTaskSuite) TestBuildCompactionRequestFailed_AllocFailed() {
s.Error(err)
}

func generateTestL0Task(state datapb.CompactionTaskState) *l0CompactionTask {
func (s *L0CompactionTaskSuite) generateTestL0Task(state datapb.CompactionTaskState) *l0CompactionTask {
return &l0CompactionTask{
CompactionTask: &datapb.CompactionTask{
PlanID: 1,
Expand All @@ -190,19 +215,20 @@ func generateTestL0Task(state datapb.CompactionTaskState) *l0CompactionTask {
State: state,
InputSegments: []int64{100, 101},
},
meta: s.mockMeta,
sessions: s.mockSessMgr,
allocator: s.mockAlloc,
}
}

func (s *CompactionTaskSuite) SetupSubTest() {
func (s *L0CompactionTaskSuite) SetupSubTest() {
s.SetupTest()
}

func (s *CompactionTaskSuite) TestProcessStateTrans() {
alloc := NewNMockAllocator(s.T())
alloc.EXPECT().allocN(mock.Anything).Return(100, 200, nil)

func (s *L0CompactionTaskSuite) TestPorcessStateTrans() {
s.mockAlloc.EXPECT().allocN(mock.Anything).Return(100, 200, nil)
s.Run("test pipelining needReassignNodeID", func() {
t := generateTestL0Task(datapb.CompactionTaskState_pipelining)
t := s.generateTestL0Task(datapb.CompactionTaskState_pipelining)
t.NodeID = NullNodeID
t.allocator = alloc
got := t.Process()
Expand All @@ -212,13 +238,12 @@ func (s *CompactionTaskSuite) TestProcessStateTrans() {
})

s.Run("test pipelining BuildCompactionRequest failed", func() {
t := generateTestL0Task(datapb.CompactionTaskState_pipelining)
t := s.generateTestL0Task(datapb.CompactionTaskState_pipelining)
t.NodeID = 100
t.allocator = alloc
channel := "ch-1"
deltaLogs := []*datapb.FieldBinlog{getFieldBinlogIDs(101, 3)}

t.meta = s.mockMeta
s.mockMeta.EXPECT().SelectSegments(mock.Anything, mock.Anything).Return(
[]*SegmentInfo{
{SegmentInfo: &datapb.SegmentInfo{
Expand All @@ -241,7 +266,6 @@ func (s *CompactionTaskSuite) TestProcessStateTrans() {
s.mockMeta.EXPECT().SaveCompactionTask(mock.Anything).Return(nil).Once()
s.mockMeta.EXPECT().SetSegmentsCompacting(mock.Anything, false).Return()

t.sessions = s.mockSessMgr
s.mockSessMgr.EXPECT().DropCompactionPlan(mock.Anything, mock.Anything).Return(nil).Once()

got := t.Process()
Expand All @@ -250,13 +274,12 @@ func (s *CompactionTaskSuite) TestProcessStateTrans() {
})

s.Run("test pipelining Compaction failed", func() {
t := generateTestL0Task(datapb.CompactionTaskState_pipelining)
t := s.generateTestL0Task(datapb.CompactionTaskState_pipelining)
t.NodeID = 100
t.allocator = alloc
channel := "ch-1"
deltaLogs := []*datapb.FieldBinlog{getFieldBinlogIDs(101, 3)}

t.meta = s.mockMeta
s.mockMeta.EXPECT().SelectSegments(mock.Anything, mock.Anything).Return(
[]*SegmentInfo{
{SegmentInfo: &datapb.SegmentInfo{
Expand All @@ -278,7 +301,6 @@ func (s *CompactionTaskSuite) TestProcessStateTrans() {
}).Twice()
s.mockMeta.EXPECT().SaveCompactionTask(mock.Anything).Return(nil)

t.sessions = s.mockSessMgr
s.mockSessMgr.EXPECT().Compaction(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, nodeID int64, plan *datapb.CompactionPlan) error {
s.Require().EqualValues(t.NodeID, nodeID)
return errors.New("mock error")
Expand All @@ -291,13 +313,12 @@ func (s *CompactionTaskSuite) TestProcessStateTrans() {
})

s.Run("test pipelining success", func() {
t := generateTestL0Task(datapb.CompactionTaskState_pipelining)
t := s.generateTestL0Task(datapb.CompactionTaskState_pipelining)
t.NodeID = 100
t.allocator = alloc
channel := "ch-1"
deltaLogs := []*datapb.FieldBinlog{getFieldBinlogIDs(101, 3)}

t.meta = s.mockMeta
s.mockMeta.EXPECT().SelectSegments(mock.Anything, mock.Anything).Return(
[]*SegmentInfo{
{SegmentInfo: &datapb.SegmentInfo{
Expand All @@ -319,14 +340,113 @@ func (s *CompactionTaskSuite) TestProcessStateTrans() {
}).Twice()
s.mockMeta.EXPECT().SaveCompactionTask(mock.Anything).Return(nil).Once()

t.sessions = s.mockSessMgr
s.mockSessMgr.EXPECT().Compaction(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, nodeID int64, plan *datapb.CompactionPlan) error {
s.Require().EqualValues(t.NodeID, nodeID)
return nil
})

got := t.Process()
s.False(got)
s.Equal(datapb.CompactionTaskState_executing, t.State)
s.Equal(datapb.CompactionTaskState_executing, t.GetState())
})

// stay in executing state when GetCompactionPlanResults error except ErrNodeNotFound
s.Run("test executing GetCompactionPlanResult fail NodeNotFound", func() {
t := s.generateTestL0Task(datapb.CompactionTaskState_executing)
t.NodeID = 100
s.Require().True(t.GetNodeID() > 0)

s.mockSessMgr.EXPECT().GetCompactionPlanResult(t.NodeID, mock.Anything).Return(nil, merr.WrapErrNodeNotFound(t.NodeID)).Once()
s.mockMeta.EXPECT().SaveCompactionTask(mock.Anything).Return(nil).Once()

got := t.Process()
s.False(got)
s.Equal(datapb.CompactionTaskState_pipelining, t.GetState())
s.EqualValues(NullNodeID, t.GetNodeID())
})

// stay in executing state when GetCompactionPlanResults error except ErrNodeNotFound
s.Run("test executing GetCompactionPlanResult fail mock error", func() {
t := s.generateTestL0Task(datapb.CompactionTaskState_executing)
t.NodeID = 100
s.Require().True(t.GetNodeID() > 0)

s.mockSessMgr.EXPECT().GetCompactionPlanResult(t.NodeID, mock.Anything).Return(nil, errors.New("mock error")).Times(12)
for i := 0; i < 12; i++ {
got := t.Process()
s.False(got)
s.Equal(datapb.CompactionTaskState_executing, t.GetState())
s.EqualValues(100, t.GetNodeID())
}
})

s.Run("test executing with result executing", func() {
t := s.generateTestL0Task(datapb.CompactionTaskState_executing)
t.NodeID = 100
s.Require().True(t.GetNodeID() > 0)

s.mockSessMgr.EXPECT().GetCompactionPlanResult(t.NodeID, mock.Anything).
Return(&datapb.CompactionPlanResult{
PlanID: t.GetPlanID(),
State: datapb.CompactionTaskState_executing,
}, nil).Twice()

got := t.Process()
s.False(got)

// test timeout
t.StartTime = time.Now().Add(-time.Hour).Unix()
t.TimeoutInSeconds = 10

s.mockMeta.EXPECT().SaveCompactionTask(mock.Anything).Return(nil)
s.mockMeta.EXPECT().SetSegmentsCompacting(mock.Anything, false).
RunAndReturn(func(inputs []int64, compacting bool) {
s.ElementsMatch(inputs, t.GetInputSegments())
s.False(compacting)
}).Once()

got = t.Process()
s.True(got)
s.Equal(datapb.CompactionTaskState_timeout, t.GetState())
})

s.Run("test executing with result executing timeout and updataAndSaveTaskMeta failed", func() {
t := s.generateTestL0Task(datapb.CompactionTaskState_executing)
t.NodeID = 100
s.Require().True(t.GetNodeID() > 0)

s.mockSessMgr.EXPECT().GetCompactionPlanResult(t.NodeID, mock.Anything).
Return(&datapb.CompactionPlanResult{
PlanID: t.GetPlanID(),
State: datapb.CompactionTaskState_executing,
}, nil).Once()
s.mockMeta.EXPECT().SaveCompactionTask(mock.Anything).Return(errors.New("mock error")).Once()

t.StartTime = time.Now().Add(-time.Hour).Unix()
t.TimeoutInSeconds = 10

got := t.Process()
s.False(got)
s.Equal(datapb.CompactionTaskState_executing, t.GetState())
})

s.Run("test executing with result completed", func() {
t := s.generateTestL0Task(datapb.CompactionTaskState_executing)
t.NodeID = 100
s.Require().True(t.GetNodeID() > 0)

s.mockSessMgr.EXPECT().GetCompactionPlanResult(t.NodeID, mock.Anything).
Return(&datapb.CompactionPlanResult{
PlanID: t.GetPlanID(),
State: datapb.CompactionTaskState_completed,
}, nil).Once()
s.mockMeta.EXPECT().SaveCompactionTask(mock.Anything).Return(errors.New("mock error")).Once()

t.StartTime = time.Now().Add(-time.Hour).Unix()
t.TimeoutInSeconds = 10

got := t.Process()
s.False(got)
s.Equal(datapb.CompactionTaskState_executing, t.GetState())
})
}

0 comments on commit a5994bc

Please sign in to comment.