未验证 提交 f1daef22 编写于 作者: X Xiaofan 提交者: GitHub

Fix failed bulkload segment marked as sealed (#21799)

Signed-off-by: Nxiaofan-luan <xiaofan.luan@zilliz.com>
上级 a04238c2
......@@ -352,8 +352,11 @@ func (m *meta) SetState(segmentID UniqueID, targetState commonpb.SegmentState) e
log.Warn("meta update: setting segment state - segment not found",
zap.Int64("segment ID", segmentID),
zap.Any("target state", targetState))
// TODO: Should probably return error instead.
return nil
// idempotent drop
if targetState == commonpb.SegmentState_Dropped {
return nil
}
return fmt.Errorf("segment is not exist with ID = %d", segmentID)
}
// Persist segment updates first.
clonedSegment := curSegInfo.Clone()
......
......@@ -78,7 +78,7 @@ type Manager interface {
DropSegment(ctx context.Context, segmentID UniqueID)
// SealAllSegments seals all segments of collection with collectionID and return sealed segments.
// If segIDs is not empty, also seals segments in segIDs.
SealAllSegments(ctx context.Context, collectionID UniqueID, segIDs []UniqueID) ([]UniqueID, error)
SealAllSegments(ctx context.Context, collectionID UniqueID, segIDs []UniqueID, isImporting bool) ([]UniqueID, error)
// GetFlushableSegments returns flushable segment ids
GetFlushableSegments(ctx context.Context, channel string, ts Timestamp) ([]UniqueID, error)
// ExpireAllocations notifies segment status to expire old allocations
......@@ -261,7 +261,7 @@ func (s *SegmentManager) AllocSegment(ctx context.Context, collectionID UniqueID
requestRows, int64(maxCountPerSegment))
// create new segments and add allocations
expireTs, err := s.genExpireTs(ctx)
expireTs, err := s.genExpireTs(ctx, false)
if err != nil {
return nil, err
}
......@@ -296,7 +296,7 @@ func (s *SegmentManager) allocSegmentForImport(ctx context.Context, collectionID
// create new segments and add allocations to meta
// to avoid mixing up with growing segments, the segment state is "Importing"
expireTs, err := s.genExpireTs(ctx)
expireTs, err := s.genExpireTs(ctx, true)
if err != nil {
return nil, err
}
......@@ -329,13 +329,17 @@ func isGrowing(segment *SegmentInfo) bool {
return segment.GetState() == commonpb.SegmentState_Growing
}
func (s *SegmentManager) genExpireTs(ctx context.Context) (Timestamp, error) {
func (s *SegmentManager) genExpireTs(ctx context.Context, isImported bool) (Timestamp, error) {
ts, err := s.allocator.allocTimestamp(ctx)
if err != nil {
return 0, err
}
physicalTs, logicalTs := tsoutil.ParseTS(ts)
expirePhysicalTs := physicalTs.Add(time.Duration(Params.DataCoordCfg.SegAssignmentExpiration.GetAsFloat()) * time.Millisecond)
// for imported segment, clean up ImportTaskExpiration
if isImported {
expirePhysicalTs = physicalTs.Add(time.Duration(Params.RootCoordCfg.ImportTaskExpiration.GetAsFloat()) * time.Second)
}
expireTs := tsoutil.ComposeTS(expirePhysicalTs.UnixNano()/int64(time.Millisecond), int64(logicalTs))
return expireTs, nil
}
......@@ -416,7 +420,7 @@ func (s *SegmentManager) DropSegment(ctx context.Context, segmentID UniqueID) {
}
// SealAllSegments seals all segments of collection with collectionID and return sealed segments
func (s *SegmentManager) SealAllSegments(ctx context.Context, collectionID UniqueID, segIDs []UniqueID) ([]UniqueID, error) {
func (s *SegmentManager) SealAllSegments(ctx context.Context, collectionID UniqueID, segIDs []UniqueID, isImport bool) ([]UniqueID, error) {
_, sp := otel.Tracer(typeutil.DataCoordRole).Start(ctx, "Seal-Segments")
defer sp.End()
s.mu.Lock()
......@@ -435,12 +439,13 @@ func (s *SegmentManager) SealAllSegments(ctx context.Context, collectionID Uniqu
if info.CollectionID != collectionID {
continue
}
// idempotent sealed
if info.State == commonpb.SegmentState_Sealed {
ret = append(ret, id)
continue
}
if info.State == commonpb.SegmentState_Flushing ||
info.State == commonpb.SegmentState_Flushed {
// segment can be sealed only if it is growing or if it's importing
if (!isImport && info.State != commonpb.SegmentState_Growing) || (isImport && info.State != commonpb.SegmentState_Importing) {
continue
}
if err := s.meta.SetState(id, commonpb.SegmentState_Sealed); err != nil {
......@@ -462,7 +467,7 @@ func (s *SegmentManager) GetFlushableSegments(ctx context.Context, channel strin
return nil, err
}
s.dropEmptySealedSegment(t, channel)
s.cleanupSealedSegment(t, channel)
ret := make([]UniqueID, 0, len(s.segments))
for _, id := range s.segments {
......@@ -501,7 +506,7 @@ func (s *SegmentManager) ExpireAllocations(channel string, ts Timestamp) error {
return nil
}
func (s *SegmentManager) dropEmptySealedSegment(ts Timestamp, channel string) {
func (s *SegmentManager) cleanupSealedSegment(ts Timestamp, channel string) {
valids := make([]int64, 0, len(s.segments))
for _, id := range s.segments {
segment := s.meta.GetSegment(id)
......@@ -511,10 +516,18 @@ func (s *SegmentManager) dropEmptySealedSegment(ts Timestamp, channel string) {
}
if isEmptySealedSegment(segment, ts) {
log.Info("remove empty sealed segment", zap.Any("segment", id))
log.Info("remove empty sealed segment", zap.Int64("collection", segment.CollectionID), zap.Any("segment", id))
s.meta.SetState(id, commonpb.SegmentState_Dropped)
continue
}
// clean up importing segment since the task failed.
if segment.GetState() == commonpb.SegmentState_Importing && segment.GetLastExpireTime() < ts {
log.Info("cleanup staled importing segment", zap.Int64("collection", segment.CollectionID), zap.Int64("segment", id))
s.meta.SetState(id, commonpb.SegmentState_Dropped)
continue
}
valids = append(valids, id)
}
s.segments = valids
......@@ -533,9 +546,7 @@ func (s *SegmentManager) tryToSealSegment(ts Timestamp, channel string) error {
continue
}
channelInfo[info.InsertChannel] = append(channelInfo[info.InsertChannel], info)
if info.State == commonpb.SegmentState_Sealed ||
info.State == commonpb.SegmentState_Flushing ||
info.State == commonpb.SegmentState_Flushed {
if info.State != commonpb.SegmentState_Growing {
continue
}
// change shouldSeal to segment seal policy logic
......@@ -552,9 +563,7 @@ func (s *SegmentManager) tryToSealSegment(ts Timestamp, channel string) error {
for _, policy := range s.channelSealPolicies {
vs := policy(channel, segmentInfos, ts)
for _, info := range vs {
if info.State == commonpb.SegmentState_Sealed ||
info.State == commonpb.SegmentState_Flushing ||
info.State == commonpb.SegmentState_Flushed {
if info.State != commonpb.SegmentState_Growing {
continue
}
if err := s.meta.SetState(info.GetID(), commonpb.SegmentState_Sealed); err != nil {
......
......@@ -247,7 +247,7 @@ func TestSaveSegmentsToMeta(t *testing.T) {
allocations, err := segmentManager.AllocSegment(context.Background(), collID, 0, "c1", 1000)
assert.Nil(t, err)
assert.EqualValues(t, 1, len(allocations))
_, err = segmentManager.SealAllSegments(context.Background(), collID, nil)
_, err = segmentManager.SealAllSegments(context.Background(), collID, nil, false)
assert.Nil(t, err)
segment := meta.GetSegment(allocations[0].SegmentID)
assert.NotNil(t, segment)
......@@ -269,7 +269,7 @@ func TestSaveSegmentsToMetaWithSpecificSegments(t *testing.T) {
allocations, err := segmentManager.AllocSegment(context.Background(), collID, 0, "c1", 1000)
assert.Nil(t, err)
assert.EqualValues(t, 1, len(allocations))
_, err = segmentManager.SealAllSegments(context.Background(), collID, []int64{allocations[0].SegmentID})
_, err = segmentManager.SealAllSegments(context.Background(), collID, []int64{allocations[0].SegmentID}, false)
assert.Nil(t, err)
segment := meta.GetSegment(allocations[0].SegmentID)
assert.NotNil(t, segment)
......@@ -364,6 +364,37 @@ func TestExpireAllocation(t *testing.T) {
assert.EqualValues(t, 0, len(segment.allocations))
}
func TestCleanExpiredBulkloadSegment(t *testing.T) {
t.Run("expiredBulkloadSegment", func(t *testing.T) {
Params.Init()
mockAllocator := newMockAllocator()
meta, err := newMemoryMeta()
assert.Nil(t, err)
schema := newTestSchema()
collID, err := mockAllocator.allocID(context.Background())
assert.Nil(t, err)
meta.AddCollection(&collectionInfo{ID: collID, Schema: schema})
ms := newMockRootCoordService()
segmentManager := newSegmentManager(meta, mockAllocator, ms)
allocation, err := segmentManager.allocSegmentForImport(context.TODO(), collID, 0, "c1", 2, 1)
assert.Nil(t, err)
ids, err := segmentManager.GetFlushableSegments(context.TODO(), "c1", allocation.ExpireTime)
assert.Nil(t, err)
assert.EqualValues(t, len(ids), 0)
assert.EqualValues(t, len(segmentManager.segments), 1)
ids, err = segmentManager.GetFlushableSegments(context.TODO(), "c1", allocation.ExpireTime+1)
assert.Nil(t, err)
assert.Empty(t, ids)
assert.EqualValues(t, len(ids), 0)
assert.EqualValues(t, len(segmentManager.segments), 0)
})
}
func TestGetFlushableSegments(t *testing.T) {
t.Run("get flushable segments between small interval", func(t *testing.T) {
Params.Init()
......@@ -380,7 +411,7 @@ func TestGetFlushableSegments(t *testing.T) {
assert.Nil(t, err)
assert.EqualValues(t, 1, len(allocations))
ids, err := segmentManager.SealAllSegments(context.TODO(), collID, nil)
ids, err := segmentManager.SealAllSegments(context.TODO(), collID, nil, false)
assert.Nil(t, err)
assert.EqualValues(t, 1, len(ids))
assert.EqualValues(t, allocations[0].SegmentID, ids[0])
......
......@@ -252,6 +252,50 @@ func TestFlush(t *testing.T) {
assert.EqualValues(t, segID, ids[0])
})
t.Run("bulkload segment", func(t *testing.T) {
svr := newTestServer(t, nil)
defer closeTestServer(t, svr)
schema := newTestSchema()
svr.meta.AddCollection(&collectionInfo{ID: 0, Schema: schema, Partitions: []int64{}})
allocations, err := svr.segmentManager.allocSegmentForImport(context.TODO(), 0, 1, "channel-1", 1, 100)
assert.Nil(t, err)
expireTs := allocations.ExpireTime
segID := allocations.SegmentID
resp, err := svr.Flush(context.TODO(), req)
assert.Nil(t, err)
assert.EqualValues(t, commonpb.ErrorCode_Success, resp.Status.ErrorCode)
assert.EqualValues(t, 0, len(resp.SegmentIDs))
// should not flush anything since this is a normal flush
svr.meta.SetCurrentRows(segID, 1)
ids, err := svr.segmentManager.GetFlushableSegments(context.TODO(), "channel-1", expireTs)
assert.Nil(t, err)
assert.EqualValues(t, 0, len(ids))
req := &datapb.FlushRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_Flush,
MsgID: 0,
Timestamp: 0,
SourceID: 0,
},
DbID: 0,
CollectionID: 0,
IsImport: true,
}
resp, err = svr.Flush(context.TODO(), req)
assert.Nil(t, err)
assert.EqualValues(t, commonpb.ErrorCode_Success, resp.Status.ErrorCode)
assert.EqualValues(t, 1, len(resp.SegmentIDs))
ids, err = svr.segmentManager.GetFlushableSegments(context.TODO(), "channel-1", expireTs)
assert.Nil(t, err)
assert.EqualValues(t, 1, len(ids))
assert.EqualValues(t, segID, ids[0])
})
t.Run("closed server", func(t *testing.T) {
svr := newTestServer(t, nil)
closeTestServer(t, svr)
......@@ -1164,7 +1208,7 @@ func (s *spySegmentManager) DropSegment(ctx context.Context, segmentID UniqueID)
}
// SealAllSegments seals all segments of collection with collectionID and return sealed segments
func (s *spySegmentManager) SealAllSegments(ctx context.Context, collectionID UniqueID, segIDs []UniqueID) ([]UniqueID, error) {
func (s *spySegmentManager) SealAllSegments(ctx context.Context, collectionID UniqueID, segIDs []UniqueID, isImport bool) ([]UniqueID, error) {
panic("not implemented") // TODO: Implement
}
......@@ -3391,29 +3435,6 @@ func TestDataCoord_UnsetIsImportingState(t *testing.T) {
})
}
func TestDataCoord_MarkSegmentsDropped(t *testing.T) {
t.Run("normal case", func(t *testing.T) {
svr := newTestServer(t, nil)
defer closeTestServer(t, svr)
seg := buildSegment(100, 100, 100, "ch1", false)
svr.meta.AddSegment(seg)
status, err := svr.MarkSegmentsDropped(context.Background(), &datapb.MarkSegmentsDroppedRequest{
SegmentIds: []int64{100},
})
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, status.GetErrorCode())
// Trying to mark dropped of a segment that does not exist.
status, err = svr.MarkSegmentsDropped(context.Background(), &datapb.MarkSegmentsDroppedRequest{
SegmentIds: []int64{999},
})
assert.NoError(t, err)
// Returning success as SetState will succeed if segment does not exist. This should probably get fixed.
assert.Equal(t, commonpb.ErrorCode_Success, status.GetErrorCode())
})
}
func TestDataCoordServer_UpdateChannelCheckpoint(t *testing.T) {
mockVChannel := "fake-by-dev-rootcoord-dml-1-testchannelcp-v0"
mockPChannel := "fake-by-dev-rootcoord-dml-1"
......
......@@ -73,7 +73,10 @@ func (s *Server) GetStatisticsChannel(ctx context.Context) (*milvuspb.StringResp
// this api only guarantees all the segments requested is sealed
// these segments will be flushed only after the Flush policy is fulfilled
func (s *Server) Flush(ctx context.Context, req *datapb.FlushRequest) (*datapb.FlushResponse, error) {
log.Info("receive flush request", zap.Int64("dbID", req.GetDbID()), zap.Int64("collectionID", req.GetCollectionID()))
log.Info("receive flush request",
zap.Int64("dbID", req.GetDbID()),
zap.Int64("collectionID", req.GetCollectionID()),
zap.Bool("isImporting", req.GetIsImport()))
ctx, sp := otel.Tracer(typeutil.DataCoordRole).Start(ctx, "DataCoord-Flush")
defer sp.End()
resp := &datapb.FlushResponse{
......@@ -97,7 +100,7 @@ func (s *Server) Flush(ctx context.Context, req *datapb.FlushRequest) (*datapb.F
}
timeOfSeal, _ := tsoutil.ParseTS(ts)
sealedSegmentIDs, err := s.segmentManager.SealAllSegments(ctx, req.GetCollectionID(), req.GetSegmentIDs())
sealedSegmentIDs, err := s.segmentManager.SealAllSegments(ctx, req.GetCollectionID(), req.GetSegmentIDs(), req.GetIsImport())
if err != nil {
resp.Status.Reason = fmt.Sprintf("failed to flush %d, %s", req.CollectionID, err)
return resp, nil
......@@ -1165,6 +1168,7 @@ func (s *Server) Import(ctx context.Context, itr *datapb.ImportTaskRequest) (*da
nodes := s.sessionManager.getLiveNodeIDs()
if len(nodes) == 0 {
log.Error("import failed as all DataNodes are offline")
resp.Status.Reason = "no data node available"
return resp, nil
}
log.Info("available DataNodes are", zap.Int64s("node ID", nodes))
......@@ -1277,6 +1281,7 @@ func (s *Server) SaveImportSegment(ctx context.Context, req *datapb.SaveImportSe
zap.Error(err))
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: err.Error(),
}, nil
}
resp, err := cli.AddImportSegment(ctx,
......@@ -1296,6 +1301,7 @@ func (s *Server) SaveImportSegment(ctx context.Context, req *datapb.SaveImportSe
log.Error("failed to add segment", zap.Int64("DataNode ID", nodeID), zap.Error(err))
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: err.Error(),
}, nil
}
log.Info("succeed to add segment", zap.Int64("DataNode ID", nodeID), zap.Any("add segment req", req))
......@@ -1308,6 +1314,7 @@ func (s *Server) SaveImportSegment(ctx context.Context, req *datapb.SaveImportSe
log.Error("failed to SaveBinlogPaths", zap.Error(err))
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: err.Error(),
}, nil
}
return &commonpb.Status{
......@@ -1320,17 +1327,18 @@ func (s *Server) SaveImportSegment(ctx context.Context, req *datapb.SaveImportSe
func (s *Server) UnsetIsImportingState(ctx context.Context, req *datapb.UnsetIsImportingStateRequest) (*commonpb.Status, error) {
log.Info("unsetting isImport state of segments",
zap.Int64s("segments", req.GetSegmentIds()))
failure := false
var reportErr error
for _, segID := range req.GetSegmentIds() {
if err := s.meta.UnsetIsImporting(segID); err != nil {
// Fail-open.
log.Error("failed to unset segment is importing state", zap.Int64("segment ID", segID))
failure = true
reportErr = err
}
}
if failure {
if reportErr != nil {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: reportErr.Error(),
}, nil
}
return &commonpb.Status{
......@@ -1340,9 +1348,9 @@ func (s *Server) UnsetIsImportingState(ctx context.Context, req *datapb.UnsetIsI
// MarkSegmentsDropped marks the given segments as `Dropped`.
// An error status will be returned and error will be logged, if we failed to mark *all* segments.
// Deprecated, do not use it
func (s *Server) MarkSegmentsDropped(ctx context.Context, req *datapb.MarkSegmentsDroppedRequest) (*commonpb.Status, error) {
log.Info("marking segments dropped",
zap.Int64s("segments", req.GetSegmentIds()))
log.Info("marking segments dropped", zap.Int64s("segments", req.GetSegmentIds()))
failure := false
for _, segID := range req.GetSegmentIds() {
if err := s.meta.SetState(segID, commonpb.SegmentState_Dropped); err != nil {
......
......@@ -109,6 +109,7 @@ message FlushRequest {
int64 dbID = 2;
repeated int64 segmentIDs = 3;
int64 collectionID = 4;
bool isImport = 5;
}
message FlushResponse {
......
......@@ -1217,6 +1217,7 @@ func (ft *flushTask) Execute(ctx context.Context) error {
),
DbID: 0,
CollectionID: collID,
IsImport: false,
}
resp, err := ft.dataCoord.Flush(ctx, flushReq)
if err != nil {
......
......@@ -38,7 +38,6 @@ type Broker interface {
Flush(ctx context.Context, cID int64, segIDs []int64) error
Import(ctx context.Context, req *datapb.ImportTaskRequest) (*datapb.ImportTaskResponse, error)
UnsetIsImportingState(context.Context, *datapb.UnsetIsImportingStateRequest) (*commonpb.Status, error)
MarkSegmentsDropped(context.Context, *datapb.MarkSegmentsDroppedRequest) (*commonpb.Status, error)
GetSegmentStates(context.Context, *datapb.GetSegmentStatesRequest) (*datapb.GetSegmentStatesResponse, error)
GcConfirm(ctx context.Context, collectionID, partitionID UniqueID) bool
......@@ -135,6 +134,7 @@ func (b *ServerBroker) Flush(ctx context.Context, cID int64, segIDs []int64) err
DbID: 0,
SegmentIDs: segIDs,
CollectionID: cID,
IsImport: true,
})
if err != nil {
return errors.New("failed to call flush to data coordinator: " + err.Error())
......@@ -154,10 +154,6 @@ func (b *ServerBroker) UnsetIsImportingState(ctx context.Context, req *datapb.Un
return b.s.dataCoord.UnsetIsImportingState(ctx, req)
}
func (b *ServerBroker) MarkSegmentsDropped(ctx context.Context, req *datapb.MarkSegmentsDroppedRequest) (*commonpb.Status, error) {
return b.s.dataCoord.MarkSegmentsDropped(ctx, req)
}
func (b *ServerBroker) GetSegmentStates(ctx context.Context, req *datapb.GetSegmentStatesRequest) (*datapb.GetSegmentStatesResponse, error) {
return b.s.dataCoord.GetSegmentStates(ctx, req)
}
......
......@@ -15,7 +15,6 @@ import (
type GetCollectionNameFunc func(collID, partitionID UniqueID) (string, string, error)
type IDAllocator func(count uint32) (UniqueID, UniqueID, error)
type ImportFunc func(ctx context.Context, req *datapb.ImportTaskRequest) (*datapb.ImportTaskResponse, error)
type MarkSegmentsDroppedFunc func(ctx context.Context, segIDs []int64) (*commonpb.Status, error)
type GetSegmentStatesFunc func(ctx context.Context, req *datapb.GetSegmentStatesRequest) (*datapb.GetSegmentStatesResponse, error)
type DescribeIndexFunc func(ctx context.Context, colID UniqueID) (*indexpb.DescribeIndexResponse, error)
type GetSegmentIndexStateFunc func(ctx context.Context, collID UniqueID, indexName string, segIDs []UniqueID) ([]*indexpb.SegmentIndexState, error)
......@@ -25,7 +24,6 @@ type ImportFactory interface {
NewGetCollectionNameFunc() GetCollectionNameFunc
NewIDAllocator() IDAllocator
NewImportFunc() ImportFunc
NewMarkSegmentsDroppedFunc() MarkSegmentsDroppedFunc
NewGetSegmentStatesFunc() GetSegmentStatesFunc
NewDescribeIndexFunc() DescribeIndexFunc
NewGetSegmentIndexStateFunc() GetSegmentIndexStateFunc
......@@ -48,10 +46,6 @@ func (f ImportFactoryImpl) NewImportFunc() ImportFunc {
return ImportFuncWithCore(f.c)
}
func (f ImportFactoryImpl) NewMarkSegmentsDroppedFunc() MarkSegmentsDroppedFunc {
return MarkSegmentsDroppedWithCore(f.c)
}
func (f ImportFactoryImpl) NewGetSegmentStatesFunc() GetSegmentStatesFunc {
return GetSegmentStatesWithCore(f.c)
}
......@@ -101,14 +95,6 @@ func ImportFuncWithCore(c *Core) ImportFunc {
}
}
func MarkSegmentsDroppedWithCore(c *Core) MarkSegmentsDroppedFunc {
return func(ctx context.Context, segIDs []int64) (*commonpb.Status, error) {
return c.broker.MarkSegmentsDropped(ctx, &datapb.MarkSegmentsDroppedRequest{
SegmentIds: segIDs,
})
}
}
func GetSegmentStatesWithCore(c *Core) GetSegmentStatesFunc {
return func(ctx context.Context, req *datapb.GetSegmentStatesRequest) (*datapb.GetSegmentStatesResponse, error) {
return c.broker.GetSegmentStates(ctx, req)
......
......@@ -77,7 +77,6 @@ type importManager struct {
idAllocator func(count uint32) (typeutil.UniqueID, typeutil.UniqueID, error)
callImportService func(ctx context.Context, req *datapb.ImportTaskRequest) (*datapb.ImportTaskResponse, error)
getCollectionName func(collID, partitionID typeutil.UniqueID) (string, string, error)
callMarkSegmentsDropped func(ctx context.Context, segIDs []typeutil.UniqueID) (*commonpb.Status, error)
callGetSegmentStates func(ctx context.Context, req *datapb.GetSegmentStatesRequest) (*datapb.GetSegmentStatesResponse, error)
callUnsetIsImportingState func(context.Context, *datapb.UnsetIsImportingStateRequest) (*commonpb.Status, error)
}
......@@ -86,7 +85,6 @@ type importManager struct {
func newImportManager(ctx context.Context, client kv.TxnKV,
idAlloc func(count uint32) (typeutil.UniqueID, typeutil.UniqueID, error),
importService func(ctx context.Context, req *datapb.ImportTaskRequest) (*datapb.ImportTaskResponse, error),
markSegmentsDropped func(ctx context.Context, segIDs []typeutil.UniqueID) (*commonpb.Status, error),
getSegmentStates func(ctx context.Context, req *datapb.GetSegmentStatesRequest) (*datapb.GetSegmentStatesResponse, error),
getCollectionName func(collID, partitionID typeutil.UniqueID) (string, string, error),
unsetIsImportingState func(context.Context, *datapb.UnsetIsImportingStateRequest) (*commonpb.Status, error)) *importManager {
......@@ -102,7 +100,6 @@ func newImportManager(ctx context.Context, client kv.TxnKV,
lastReqID: 0,
idAllocator: idAlloc,
callImportService: importService,
callMarkSegmentsDropped: markSegmentsDropped,
callGetSegmentStates: getSegmentStates,
getCollectionName: getCollectionName,
callUnsetIsImportingState: unsetIsImportingState,
......@@ -228,7 +225,7 @@ func (m *importManager) sendOutTasks(ctx context.Context) error {
break
}
if err != nil {
log.Error("import task get error", zap.Error(err))
log.Warn("import task get error", zap.Error(err))
break
}
......@@ -1014,23 +1011,9 @@ func (m *importManager) removeBadImportSegments(ctx context.Context) {
log.Info("trying to mark segments as dropped",
zap.Int64("task ID", t.GetId()),
zap.Int64s("segment IDs", t.GetState().GetSegments()))
status, err := m.callMarkSegmentsDropped(ctx, t.GetState().GetSegments())
errMsg := "failed to mark all segments dropped, some segments might already have been dropped"
if err != nil {
log.Error(errMsg,
zap.Int64("task ID", t.GetId()),
zap.Int64s("segments", t.GetState().GetSegments()),
zap.Error(err))
}
if status.GetErrorCode() != commonpb.ErrorCode_Success {
log.Error(errMsg,
zap.Int64("task ID", t.GetId()),
zap.Int64s("segments", t.GetState().GetSegments()),
zap.Error(errors.New(status.GetReason())))
}
if err = m.setImportTaskState(t.GetId(), commonpb.ImportState_ImportFailedAndCleaned); err != nil {
log.Error(errMsg,
zap.Int64("task ID", t.GetId()))
log.Warn("failed to set ", zap.Int64("task ID", t.GetId()), zap.Error(err))
}
}
}
......
......@@ -101,11 +101,6 @@ func TestImportManager_NewImportManager(t *testing.T) {
},
}, nil
}
callMarkSegmentsDropped := func(ctx context.Context, segIDs []typeutil.UniqueID) (*commonpb.Status, error) {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
}, nil
}
callGetSegmentStates := func(ctx context.Context, req *datapb.GetSegmentStatesRequest) (*datapb.GetSegmentStatesResponse, error) {
return &datapb.GetSegmentStatesResponse{
Status: &commonpb.Status{
......@@ -119,7 +114,7 @@ func TestImportManager_NewImportManager(t *testing.T) {
defer wg.Done()
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
defer cancel()
mgr := newImportManager(ctx, mockKv, idAlloc, callImportServiceFn, callMarkSegmentsDropped, callGetSegmentStates, nil, nil)
mgr := newImportManager(ctx, mockKv, idAlloc, callImportServiceFn, callGetSegmentStates, nil, nil)
assert.NotNil(t, mgr)
// there are 2 tasks read from store, one is pending, the other is persisted.
......@@ -151,7 +146,7 @@ func TestImportManager_NewImportManager(t *testing.T) {
defer wg.Done()
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Nanosecond)
defer cancel()
mgr := newImportManager(ctx, mockKv, idAlloc, callImportServiceFn, callMarkSegmentsDropped, callGetSegmentStates, nil, nil)
mgr := newImportManager(ctx, mockKv, idAlloc, callImportServiceFn, callGetSegmentStates, nil, nil)
assert.NotNil(t, mgr)
mgr.init(context.TODO())
var wgLoop sync.WaitGroup
......@@ -170,7 +165,7 @@ func TestImportManager_NewImportManager(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Nanosecond)
defer cancel()
mgr := newImportManager(ctx, mockTxnKV, idAlloc, callImportServiceFn, callMarkSegmentsDropped, callGetSegmentStates, nil, nil)
mgr := newImportManager(ctx, mockTxnKV, idAlloc, callImportServiceFn, callGetSegmentStates, nil, nil)
assert.NotNil(t, mgr)
assert.Panics(t, func() {
mgr.init(context.TODO())
......@@ -187,7 +182,7 @@ func TestImportManager_NewImportManager(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Nanosecond)
defer cancel()
mgr := newImportManager(ctx, mockTxnKV, idAlloc, callImportServiceFn, callMarkSegmentsDropped, callGetSegmentStates, nil, nil)
mgr := newImportManager(ctx, mockTxnKV, idAlloc, callImportServiceFn, callGetSegmentStates, nil, nil)
assert.NotNil(t, mgr)
mgr.init(context.TODO())
})
......@@ -201,7 +196,7 @@ func TestImportManager_NewImportManager(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Nanosecond)
defer cancel()
mgr := newImportManager(ctx, mockTxnKV, idAlloc, callImportServiceFn, callMarkSegmentsDropped, callGetSegmentStates, nil, nil)
mgr := newImportManager(ctx, mockTxnKV, idAlloc, callImportServiceFn, callGetSegmentStates, nil, nil)
assert.NotNil(t, mgr)
mgr.init(context.TODO())
func() {
......@@ -221,7 +216,7 @@ func TestImportManager_NewImportManager(t *testing.T) {
defer wg.Done()
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
mgr := newImportManager(ctx, mockKv, idAlloc, callImportServiceFn, callMarkSegmentsDropped, callGetSegmentStates, nil, nil)
mgr := newImportManager(ctx, mockKv, idAlloc, callImportServiceFn, callGetSegmentStates, nil, nil)
assert.NotNil(t, mgr)
mgr.init(ctx)
var wgLoop sync.WaitGroup
......@@ -279,7 +274,7 @@ func TestImportManager_TestSetImportTaskState(t *testing.T) {
defer wg.Done()
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
defer cancel()
mgr := newImportManager(ctx, mockKv, idAlloc, nil, nil, nil, nil, nil)
mgr := newImportManager(ctx, mockKv, idAlloc, nil, nil, nil, nil)
assert.NotNil(t, mgr)
_, err := mgr.loadFromTaskStore(true)
assert.NoError(t, err)
......@@ -367,11 +362,6 @@ func TestImportManager_TestEtcdCleanUp(t *testing.T) {
}, nil
}
callMarkSegmentsDropped := func(ctx context.Context, segIDs []typeutil.UniqueID) (*commonpb.Status, error) {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
}, nil
}
callGetSegmentStates := func(ctx context.Context, req *datapb.GetSegmentStatesRequest) (*datapb.GetSegmentStatesResponse, error) {
return &datapb.GetSegmentStatesResponse{
Status: &commonpb.Status{
......@@ -381,7 +371,7 @@ func TestImportManager_TestEtcdCleanUp(t *testing.T) {
}
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
defer cancel()
mgr := newImportManager(ctx, mockKv, idAlloc, callImportServiceFn, callMarkSegmentsDropped, callGetSegmentStates, nil, nil)
mgr := newImportManager(ctx, mockKv, idAlloc, callImportServiceFn, callGetSegmentStates, nil, nil)
assert.NotNil(t, mgr)
_, err = mgr.loadFromTaskStore(true)
assert.NoError(t, err)
......@@ -463,11 +453,6 @@ func TestImportManager_TestFlipTaskStateLoop(t *testing.T) {
},
}, nil
}
callMarkSegmentsDropped := func(ctx context.Context, segIDs []typeutil.UniqueID) (*commonpb.Status, error) {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
}, nil
}
callGetSegmentStates := func(ctx context.Context, req *datapb.GetSegmentStatesRequest) (*datapb.GetSegmentStatesResponse, error) {
return &datapb.GetSegmentStatesResponse{
......@@ -490,7 +475,7 @@ func TestImportManager_TestFlipTaskStateLoop(t *testing.T) {
defer wg.Done()
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
mgr := newImportManager(ctx, mockKv, idAlloc, callImportServiceFn, callMarkSegmentsDropped,
mgr := newImportManager(ctx, mockKv, idAlloc, callImportServiceFn,
callGetSegmentStates, nil, callUnsetIsImportingState)
assert.NotNil(t, mgr)
var wgLoop sync.WaitGroup
......@@ -505,7 +490,7 @@ func TestImportManager_TestFlipTaskStateLoop(t *testing.T) {
defer wg.Done()
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
mgr := newImportManager(ctx, mockKv, idAlloc, callImportServiceFn, callMarkSegmentsDropped,
mgr := newImportManager(ctx, mockKv, idAlloc, callImportServiceFn,
callGetSegmentStates, nil, callUnsetIsImportingState)
assert.NotNil(t, mgr)
var wgLoop sync.WaitGroup
......@@ -520,7 +505,7 @@ func TestImportManager_TestFlipTaskStateLoop(t *testing.T) {
defer wg.Done()
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
mgr := newImportManager(ctx, mockKv, idAlloc, callImportServiceFn, callMarkSegmentsDropped,
mgr := newImportManager(ctx, mockKv, idAlloc, callImportServiceFn,
callGetSegmentStates, nil, callUnsetIsImportingState)
assert.NotNil(t, mgr)
var wgLoop sync.WaitGroup
......@@ -548,11 +533,6 @@ func TestImportManager_ImportJob(t *testing.T) {
defer paramtable.Get().Remove(Params.RootCoordCfg.ImportMaxPendingTaskCount.Key)
colID := int64(100)
mockKv := memkv.NewMemoryKV()
callMarkSegmentsDropped := func(ctx context.Context, segIDs []typeutil.UniqueID) (*commonpb.Status, error) {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
}, nil
}
callGetSegmentStates := func(ctx context.Context, req *datapb.GetSegmentStatesRequest) (*datapb.GetSegmentStatesResponse, error) {
return &datapb.GetSegmentStatesResponse{
Status: &commonpb.Status{
......@@ -561,7 +541,7 @@ func TestImportManager_ImportJob(t *testing.T) {
}, nil
}
// nil request
mgr := newImportManager(context.TODO(), mockKv, idAlloc, nil, callMarkSegmentsDropped, callGetSegmentStates, nil, nil)
mgr := newImportManager(context.TODO(), mockKv, idAlloc, nil, callGetSegmentStates, nil, nil)
resp := mgr.importJob(context.TODO(), nil, colID, 0)
assert.NotEqual(t, commonpb.ErrorCode_Success, resp.Status.ErrorCode)
......@@ -590,7 +570,7 @@ func TestImportManager_ImportJob(t *testing.T) {
// row-based case, task count equal to file count
// since the importServiceFunc return error, tasks will be kept in pending list
rowReq.Files = []string{"f1.json"}
mgr = newImportManager(context.TODO(), mockKv, idAlloc, importServiceFunc, callMarkSegmentsDropped, callGetSegmentStates, nil, nil)
mgr = newImportManager(context.TODO(), mockKv, idAlloc, importServiceFunc, callGetSegmentStates, nil, nil)
resp = mgr.importJob(context.TODO(), rowReq, colID, 0)
assert.Equal(t, len(rowReq.Files), len(mgr.pendingTasks))
assert.Equal(t, 0, len(mgr.workingTasks))
......@@ -603,7 +583,7 @@ func TestImportManager_ImportJob(t *testing.T) {
// column-based case, one quest one task
// since the importServiceFunc return error, tasks will be kept in pending list
mgr = newImportManager(context.TODO(), mockKv, idAlloc, importServiceFunc, callMarkSegmentsDropped, callGetSegmentStates, nil, nil)
mgr = newImportManager(context.TODO(), mockKv, idAlloc, importServiceFunc, callGetSegmentStates, nil, nil)
resp = mgr.importJob(context.TODO(), colReq, colID, 0)
assert.Equal(t, 1, len(mgr.pendingTasks))
assert.Equal(t, 0, len(mgr.workingTasks))
......@@ -617,13 +597,13 @@ func TestImportManager_ImportJob(t *testing.T) {
}
// row-based case, since the importServiceFunc return success, tasks will be sent to working list
mgr = newImportManager(context.TODO(), mockKv, idAlloc, importServiceFunc, callMarkSegmentsDropped, callGetSegmentStates, nil, nil)
mgr = newImportManager(context.TODO(), mockKv, idAlloc, importServiceFunc, callGetSegmentStates, nil, nil)
resp = mgr.importJob(context.TODO(), rowReq, colID, 0)
assert.Equal(t, 0, len(mgr.pendingTasks))
assert.Equal(t, len(rowReq.Files), len(mgr.workingTasks))
// column-based case, since the importServiceFunc return success, tasks will be sent to working list
mgr = newImportManager(context.TODO(), mockKv, idAlloc, importServiceFunc, callMarkSegmentsDropped, callGetSegmentStates, nil, nil)
mgr = newImportManager(context.TODO(), mockKv, idAlloc, importServiceFunc, callGetSegmentStates, nil, nil)
resp = mgr.importJob(context.TODO(), colReq, colID, 0)
assert.Equal(t, 0, len(mgr.pendingTasks))
assert.Equal(t, 1, len(mgr.workingTasks))
......@@ -647,7 +627,7 @@ func TestImportManager_ImportJob(t *testing.T) {
// row-based case, since the importServiceFunc return success for 1 task
// the first task is sent to working list, and 1 task left in pending list
mgr = newImportManager(context.TODO(), mockKv, idAlloc, importServiceFunc, callMarkSegmentsDropped, callGetSegmentStates, nil, nil)
mgr = newImportManager(context.TODO(), mockKv, idAlloc, importServiceFunc, callGetSegmentStates, nil, nil)
resp = mgr.importJob(context.TODO(), rowReq, colID, 0)
assert.Equal(t, 0, len(mgr.pendingTasks))
assert.Equal(t, 1, len(mgr.workingTasks))
......@@ -716,11 +696,6 @@ func TestImportManager_AllDataNodesBusy(t *testing.T) {
}, nil
}
callMarkSegmentsDropped := func(ctx context.Context, segIDs []typeutil.UniqueID) (*commonpb.Status, error) {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
}, nil
}
callGetSegmentStates := func(ctx context.Context, req *datapb.GetSegmentStatesRequest) (*datapb.GetSegmentStatesResponse, error) {
return &datapb.GetSegmentStatesResponse{
Status: &commonpb.Status{
......@@ -730,7 +705,7 @@ func TestImportManager_AllDataNodesBusy(t *testing.T) {
}
// each data node owns one task
mgr := newImportManager(context.TODO(), mockKv, idAlloc, importServiceFunc, callMarkSegmentsDropped, callGetSegmentStates, nil, nil)
mgr := newImportManager(context.TODO(), mockKv, idAlloc, importServiceFunc, callGetSegmentStates, nil, nil)
for i := 0; i < len(dnList); i++ {
resp := mgr.importJob(context.TODO(), rowReq, colID, 0)
assert.Equal(t, commonpb.ErrorCode_Success, resp.Status.ErrorCode)
......@@ -739,7 +714,7 @@ func TestImportManager_AllDataNodesBusy(t *testing.T) {
}
// all data nodes are busy, new task waiting in pending list
mgr = newImportManager(context.TODO(), mockKv, idAlloc, importServiceFunc, callMarkSegmentsDropped, callGetSegmentStates, nil, nil)
mgr = newImportManager(context.TODO(), mockKv, idAlloc, importServiceFunc, callGetSegmentStates, nil, nil)
resp := mgr.importJob(context.TODO(), rowReq, colID, 0)
assert.Equal(t, commonpb.ErrorCode_Success, resp.Status.ErrorCode)
assert.Equal(t, len(rowReq.Files), len(mgr.pendingTasks))
......@@ -747,7 +722,7 @@ func TestImportManager_AllDataNodesBusy(t *testing.T) {
// now all data nodes are free again, new task is executed instantly
count = 0
mgr = newImportManager(context.TODO(), mockKv, idAlloc, importServiceFunc, callMarkSegmentsDropped, callGetSegmentStates, nil, nil)
mgr = newImportManager(context.TODO(), mockKv, idAlloc, importServiceFunc, callGetSegmentStates, nil, nil)
resp = mgr.importJob(context.TODO(), colReq, colID, 0)
assert.Equal(t, commonpb.ErrorCode_Success, resp.Status.ErrorCode)
assert.Equal(t, 0, len(mgr.pendingTasks))
......@@ -796,12 +771,6 @@ func TestImportManager_TaskState(t *testing.T) {
PartitionName: "p1",
Files: []string{"f1.json"},
}
callMarkSegmentsDropped := func(ctx context.Context, segIDs []typeutil.UniqueID) (*commonpb.Status, error) {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
}, nil
}
callGetSegmentStates := func(ctx context.Context, req *datapb.GetSegmentStatesRequest) (*datapb.GetSegmentStatesResponse, error) {
return &datapb.GetSegmentStatesResponse{
Status: &commonpb.Status{
......@@ -811,7 +780,7 @@ func TestImportManager_TaskState(t *testing.T) {
}
// add 3 tasks, their ID is 10000, 10001, 10002, make sure updateTaskInfo() works correctly
mgr := newImportManager(context.TODO(), mockKv, idAlloc, importServiceFunc, callMarkSegmentsDropped, callGetSegmentStates, nil, nil)
mgr := newImportManager(context.TODO(), mockKv, idAlloc, importServiceFunc, callGetSegmentStates, nil, nil)
mgr.importJob(context.TODO(), rowReq, colID, 0)
rowReq.Files = []string{"f2.json"}
mgr.importJob(context.TODO(), rowReq, colID, 0)
......@@ -913,11 +882,6 @@ func TestImportManager_AllocFail(t *testing.T) {
Files: []string{"f1.json"},
}
callMarkSegmentsDropped := func(ctx context.Context, segIDs []typeutil.UniqueID) (*commonpb.Status, error) {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
}, nil
}
callGetSegmentStates := func(ctx context.Context, req *datapb.GetSegmentStatesRequest) (*datapb.GetSegmentStatesResponse, error) {
return &datapb.GetSegmentStatesResponse{
Status: &commonpb.Status{
......@@ -925,7 +889,7 @@ func TestImportManager_AllocFail(t *testing.T) {
},
}, nil
}
mgr := newImportManager(context.TODO(), mockKv, idAlloc, importServiceFunc, callMarkSegmentsDropped, callGetSegmentStates, nil, nil)
mgr := newImportManager(context.TODO(), mockKv, idAlloc, importServiceFunc, callGetSegmentStates, nil, nil)
resp := mgr.importJob(context.TODO(), rowReq, colID, 0)
assert.NotEqual(t, commonpb.ErrorCode_Success, resp.Status.ErrorCode)
assert.Equal(t, 0, len(mgr.pendingTasks))
......@@ -953,11 +917,6 @@ func TestImportManager_ListAllTasks(t *testing.T) {
}, nil
}
callMarkSegmentsDropped := func(ctx context.Context, segIDs []typeutil.UniqueID) (*commonpb.Status, error) {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
}, nil
}
callGetSegmentStates := func(ctx context.Context, req *datapb.GetSegmentStatesRequest) (*datapb.GetSegmentStatesResponse, error) {
return &datapb.GetSegmentStatesResponse{
Status: &commonpb.Status{
......@@ -993,7 +952,7 @@ func TestImportManager_ListAllTasks(t *testing.T) {
}
mockKv := memkv.NewMemoryKV()
mgr := newImportManager(context.TODO(), mockKv, idAlloc, fn, callMarkSegmentsDropped, callGetSegmentStates, getCollectionName, nil)
mgr := newImportManager(context.TODO(), mockKv, idAlloc, fn, callGetSegmentStates, getCollectionName, nil)
// add 10 tasks for collection1, id from 1 to 10
file1 := "f1.json"
......
......@@ -394,7 +394,6 @@ func (c *Core) initImportManager() error {
impTaskKv,
f.NewIDAllocator(),
f.NewImportFunc(),
f.NewMarkSegmentsDroppedFunc(),
f.NewGetSegmentStatesFunc(),
f.NewGetCollectionNameFunc(),
f.NewUnsetIsImportingStateFunc(),
......
......@@ -974,7 +974,7 @@ func TestCore_GetImportState(t *testing.T) {
t.Run("normal case", func(t *testing.T) {
ctx := context.Background()
c := newTestCore(withHealthyCode())
c.importManager = newImportManager(ctx, mockKv, nil, nil, nil, nil, nil, nil)
c.importManager = newImportManager(ctx, mockKv, nil, nil, nil, nil, nil)
resp, err := c.GetImportState(ctx, &milvuspb.GetImportStateRequest{
Task: 100,
})
......@@ -1058,7 +1058,7 @@ func TestCore_ListImportTasks(t *testing.T) {
ctx := context.Background()
c := newTestCore(withHealthyCode(), withMeta(meta))
c.importManager = newImportManager(ctx, mockKv, nil, nil, nil, nil, nil, nil)
c.importManager = newImportManager(ctx, mockKv, nil, nil, nil, nil, nil)
// list all tasks
resp, err := c.ListImportTasks(ctx, &milvuspb.ListImportTasksRequest{})
......@@ -1192,11 +1192,6 @@ func TestCore_ReportImport(t *testing.T) {
},
}, nil
}
callMarkSegmentsDropped := func(ctx context.Context, segIDs []typeutil.UniqueID) (*commonpb.Status, error) {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
}, nil
}
callGetSegmentStates := func(ctx context.Context, req *datapb.GetSegmentStatesRequest) (*datapb.GetSegmentStatesResponse, error) {
return &datapb.GetSegmentStatesResponse{
......@@ -1223,7 +1218,7 @@ func TestCore_ReportImport(t *testing.T) {
t.Run("report complete import with task not found", func(t *testing.T) {
ctx := context.Background()
c := newTestCore(withHealthyCode())
c.importManager = newImportManager(ctx, mockKv, idAlloc, callImportServiceFn, callMarkSegmentsDropped, callGetSegmentStates, nil, nil)
c.importManager = newImportManager(ctx, mockKv, idAlloc, callImportServiceFn, callGetSegmentStates, nil, nil)
resp, err := c.ReportImport(ctx, &rootcoordpb.ImportResult{
TaskId: 101,
State: commonpb.ImportState_ImportCompleted,
......@@ -1241,7 +1236,7 @@ func TestCore_ReportImport(t *testing.T) {
withTtSynchronizer(ticker),
withDataCoord(dc))
c.broker = newServerBroker(c)
c.importManager = newImportManager(ctx, mockKv, idAlloc, callImportServiceFn, callMarkSegmentsDropped, callGetSegmentStates, nil, callUnsetIsImportingState)
c.importManager = newImportManager(ctx, mockKv, idAlloc, callImportServiceFn, callGetSegmentStates, nil, callUnsetIsImportingState)
c.importManager.loadFromTaskStore(true)
c.importManager.sendOutTasks(ctx)
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册