Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
milvus
milvus
提交
f1daef22
M
milvus
项目概览
milvus
/
milvus
接近 1 年 前同步成功
通知
261
Star
22476
Fork
2472
代码
文件
提交
分支
Tags
贡献者
分支图
Diff
Issue
0
列表
看板
标记
里程碑
合并请求
0
DevOps
流水线
流水线任务
计划
Wiki
0
Wiki
分析
仓库
DevOps
项目成员
Pages
M
milvus
项目概览
项目概览
详情
发布
仓库
仓库
文件
提交
分支
标签
贡献者
分支图
比较
Issue
0
Issue
0
列表
看板
标记
里程碑
合并请求
0
合并请求
0
Pages
DevOps
DevOps
流水线
流水线任务
计划
分析
分析
仓库分析
DevOps
Wiki
0
Wiki
成员
成员
收起侧边栏
关闭侧边栏
动态
分支图
创建新Issue
流水线任务
提交
Issue看板
体验新版 GitCode,发现更多精彩内容 >>
未验证
提交
f1daef22
编写于
1月 31, 2023
作者:
X
Xiaofan
提交者:
GitHub
1月 31, 2023
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
Fix failed bulkload segment marked as sealed (#21799)
Signed-off-by:
N
xiaofan-luan
<
xiaofan.luan@zilliz.com
>
上级
a04238c2
变更
14
展开全部
隐藏空白更改
内联
并排
Showing
14 changed file
with
450 addition
and
450 deletion
+450
-450
internal/datacoord/meta.go
internal/datacoord/meta.go
+5
-2
internal/datacoord/segment_manager.go
internal/datacoord/segment_manager.go
+25
-16
internal/datacoord/segment_manager_test.go
internal/datacoord/segment_manager_test.go
+34
-3
internal/datacoord/server_test.go
internal/datacoord/server_test.go
+45
-24
internal/datacoord/services.go
internal/datacoord/services.go
+15
-7
internal/proto/data_coord.proto
internal/proto/data_coord.proto
+1
-0
internal/proto/datapb/data_coord.pb.go
internal/proto/datapb/data_coord.pb.go
+293
-285
internal/proxy/task.go
internal/proxy/task.go
+1
-0
internal/rootcoord/broker.go
internal/rootcoord/broker.go
+1
-5
internal/rootcoord/import_helper.go
internal/rootcoord/import_helper.go
+0
-14
internal/rootcoord/import_manager.go
internal/rootcoord/import_manager.go
+3
-20
internal/rootcoord/import_manager_test.go
internal/rootcoord/import_manager_test.go
+23
-64
internal/rootcoord/root_coord.go
internal/rootcoord/root_coord.go
+0
-1
internal/rootcoord/root_coord_test.go
internal/rootcoord/root_coord_test.go
+4
-9
未找到文件。
internal/datacoord/meta.go
浏览文件 @
f1daef22
...
...
@@ -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
()
...
...
internal/datacoord/segment_manager.go
浏览文件 @
f1daef22
...
...
@@ -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
.
dropEmpty
SealedSegment
(
t
,
channel
)
s
.
cleanup
SealedSegment
(
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
)
dropEmpty
SealedSegment
(
ts
Timestamp
,
channel
string
)
{
func
(
s
*
SegmentManager
)
cleanup
SealedSegment
(
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
{
...
...
internal/datacoord/segment_manager_test.go
浏览文件 @
f1daef22
...
...
@@ -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
])
...
...
internal/datacoord/server_test.go
浏览文件 @
f1daef22
...
...
@@ -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"
...
...
internal/datacoord/services.go
浏览文件 @
f1daef22
...
...
@@ -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
{
...
...
internal/proto/data_coord.proto
浏览文件 @
f1daef22
...
...
@@ -109,6 +109,7 @@ message FlushRequest {
int64
dbID
=
2
;
repeated
int64
segmentIDs
=
3
;
int64
collectionID
=
4
;
bool
isImport
=
5
;
}
message
FlushResponse
{
...
...
internal/proto/datapb/data_coord.pb.go
浏览文件 @
f1daef22
此差异已折叠。
点击以展开。
internal/proxy/task.go
浏览文件 @
f1daef22
...
...
@@ -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
{
...
...
internal/rootcoord/broker.go
浏览文件 @
f1daef22
...
...
@@ -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
)
}
...
...
internal/rootcoord/import_helper.go
浏览文件 @
f1daef22
...
...
@@ -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
)
...
...
internal/rootcoord/import_manager.go
浏览文件 @
f1daef22
...
...
@@ -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
))
}
}
}
...
...
internal/rootcoord/import_manager_test.go
浏览文件 @
f1daef22
...
...
@@ -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
,
call
MarkSegmentsDropped
,
call
GetSegmentStates
,
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
,
call
MarkSegmentsDropped
,
call
GetSegmentStates
,
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
,
call
MarkSegmentsDropped
,
call
GetSegmentStates
,
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
,
call
MarkSegmentsDropped
,
call
GetSegmentStates
,
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
,
call
MarkSegmentsDropped
,
call
GetSegmentStates
,
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
,
call
MarkSegmentsDropped
,
call
GetSegmentStates
,
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
,
call
MarkSegmentsDropped
,
call
GetSegmentStates
,
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
,
call
MarkSegmentsDropped
,
call
GetSegmentStates
,
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
,
call
MarkSegmentsDropped
,
call
GetSegmentStates
,
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
,
call
MarkSegmentsDropped
,
call
GetSegmentStates
,
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
,
call
MarkSegmentsDropped
,
call
GetSegmentStates
,
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
,
call
MarkSegmentsDropped
,
call
GetSegmentStates
,
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
,
call
MarkSegmentsDropped
,
call
GetSegmentStates
,
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
,
call
MarkSegmentsDropped
,
call
GetSegmentStates
,
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
,
call
MarkSegmentsDropped
,
call
GetSegmentStates
,
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
,
call
MarkSegmentsDropped
,
call
GetSegmentStates
,
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
,
call
MarkSegmentsDropped
,
call
GetSegmentStates
,
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
,
call
MarkSegmentsDropped
,
call
GetSegmentStates
,
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
,
call
MarkSegmentsDropped
,
call
GetSegmentStates
,
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"
...
...
internal/rootcoord/root_coord.go
浏览文件 @
f1daef22
...
...
@@ -394,7 +394,6 @@ func (c *Core) initImportManager() error {
impTaskKv
,
f
.
NewIDAllocator
(),
f
.
NewImportFunc
(),
f
.
NewMarkSegmentsDroppedFunc
(),
f
.
NewGetSegmentStatesFunc
(),
f
.
NewGetCollectionNameFunc
(),
f
.
NewUnsetIsImportingStateFunc
(),
...
...
internal/rootcoord/root_coord_test.go
浏览文件 @
f1daef22
...
...
@@ -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
,
call
MarkSegmentsDropped
,
call
GetSegmentStates
,
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
,
call
MarkSegmentsDropped
,
call
GetSegmentStates
,
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.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录