Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
milvus
milvus
提交
49443e8a
M
milvus
项目概览
milvus
/
milvus
11 个月 前同步成功
通知
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,发现更多精彩内容 >>
未验证
提交
49443e8a
编写于
5月 24, 2021
作者:
Y
yukun
提交者:
GitHub
5月 24, 2021
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
Add retrieve taskscheduler implementation (#5353)
See also: #5257 Signed-off-by:
N
fishpenguin
<
kun.yu@zilliz.com
>
上级
bfc057d5
变更
2
隐藏空白更改
内联
并排
Showing
2 changed file
with
123 addition
and
29 deletion
+123
-29
internal/proxynode/impl.go
internal/proxynode/impl.go
+57
-1
internal/proxynode/task_scheduler.go
internal/proxynode/task_scheduler.go
+66
-28
未找到文件。
internal/proxynode/impl.go
浏览文件 @
49443e8a
...
...
@@ -28,6 +28,7 @@ import (
"github.com/milvus-io/milvus/internal/proto/milvuspb"
"github.com/milvus-io/milvus/internal/proto/proxypb"
"github.com/milvus-io/milvus/internal/proto/querypb"
"github.com/milvus-io/milvus/internal/proto/schemapb"
"github.com/milvus-io/milvus/internal/util/typeutil"
)
...
...
@@ -1170,7 +1171,62 @@ func (node *ProxyNode) Search(ctx context.Context, request *milvuspb.SearchReque
}
func
(
node
*
ProxyNode
)
Retrieve
(
ctx
context
.
Context
,
request
*
milvuspb
.
RetrieveRequest
)
(
*
milvuspb
.
RetrieveResults
,
error
)
{
return
nil
,
nil
rt
:=
&
RetrieveTask
{
ctx
:
ctx
,
Condition
:
NewTaskCondition
(
ctx
),
RetrieveRequest
:
&
internalpb
.
RetrieveRequest
{
Base
:
&
commonpb
.
MsgBase
{
MsgType
:
commonpb
.
MsgType_Retrieve
,
SourceID
:
Params
.
ProxyID
,
},
ResultChannelID
:
strconv
.
FormatInt
(
Params
.
ProxyID
,
10
),
},
queryMsgStream
:
node
.
queryMsgStream
,
resultBuf
:
make
(
chan
[]
*
internalpb
.
RetrieveResults
),
retrieve
:
request
,
}
err
:=
node
.
sched
.
DqQueue
.
Enqueue
(
rt
)
if
err
!=
nil
{
return
&
milvuspb
.
RetrieveResults
{
Status
:
&
commonpb
.
Status
{
ErrorCode
:
commonpb
.
ErrorCode_UnexpectedError
,
Reason
:
err
.
Error
(),
},
},
nil
}
log
.
Debug
(
"Retrieve"
,
zap
.
String
(
"role"
,
Params
.
RoleName
),
zap
.
Int64
(
"msgID"
,
rt
.
Base
.
MsgID
),
zap
.
Uint64
(
"timestamp"
,
rt
.
Base
.
Timestamp
),
zap
.
String
(
"db"
,
request
.
DbName
),
zap
.
String
(
"collection"
,
request
.
CollectionName
),
zap
.
Any
(
"partitions"
,
request
.
PartitionNames
),
zap
.
Any
(
"len(Ids)"
,
len
(
request
.
Ids
.
IdField
.
(
*
schemapb
.
IDs_IntId
)
.
IntId
.
Data
)))
defer
func
()
{
log
.
Debug
(
"Retrieve Done"
,
zap
.
Error
(
err
),
zap
.
String
(
"role"
,
Params
.
RoleName
),
zap
.
Int64
(
"msgID"
,
rt
.
Base
.
MsgID
),
zap
.
Uint64
(
"timestamp"
,
rt
.
Base
.
Timestamp
),
zap
.
String
(
"db"
,
request
.
DbName
),
zap
.
String
(
"collection"
,
request
.
CollectionName
),
zap
.
Any
(
"partitions"
,
request
.
PartitionNames
),
zap
.
Any
(
"len(Ids)"
,
len
(
request
.
Ids
.
IdField
.
(
*
schemapb
.
IDs_IntId
)
.
IntId
.
Data
)))
}()
err
=
rt
.
WaitToFinish
()
if
err
!=
nil
{
return
&
milvuspb
.
RetrieveResults
{
Status
:
&
commonpb
.
Status
{
ErrorCode
:
commonpb
.
ErrorCode_UnexpectedError
,
Reason
:
err
.
Error
(),
},
},
nil
}
return
rt
.
result
,
nil
}
func
(
node
*
ProxyNode
)
Flush
(
ctx
context
.
Context
,
request
*
milvuspb
.
FlushRequest
)
(
*
commonpb
.
Status
,
error
)
{
...
...
internal/proxynode/task_scheduler.go
浏览文件 @
49443e8a
...
...
@@ -401,13 +401,17 @@ func (sched *TaskScheduler) queryResultLoop() {
queryResultMsgStream
,
_
:=
sched
.
msFactory
.
NewQueryMsgStream
(
sched
.
ctx
)
queryResultMsgStream
.
AsConsumer
(
Params
.
SearchResultChannelNames
,
Params
.
ProxySubName
)
log
.
Debug
(
"proxynode"
,
zap
.
Strings
(
"search result channel names"
,
Params
.
SearchResultChannelNames
))
queryResultMsgStream
.
AsConsumer
(
Params
.
RetrieveResultChannelNames
,
Params
.
ProxySubName
)
log
.
Debug
(
"proxynode"
,
zap
.
Strings
(
"Retrieve result channel names"
,
Params
.
RetrieveResultChannelNames
))
log
.
Debug
(
"proxynode"
,
zap
.
String
(
"proxySubName"
,
Params
.
ProxySubName
))
queryNodeNum
:=
Params
.
QueryNodeNum
queryResultMsgStream
.
Start
()
defer
queryResultMsgStream
.
Close
()
queryResultBuf
:=
make
(
map
[
UniqueID
][]
*
internalpb
.
SearchResults
)
retrieveResultBuf
:=
make
(
map
[
UniqueID
][]
*
internalpb
.
RetrieveResults
)
for
{
select
{
...
...
@@ -422,41 +426,75 @@ func (sched *TaskScheduler) queryResultLoop() {
for
_
,
tsMsg
:=
range
msgPack
.
Msgs
{
sp
,
ctx
:=
trace
.
StartSpanFromContext
(
tsMsg
.
TraceCtx
())
tsMsg
.
SetTraceCtx
(
ctx
)
searchResultMsg
,
_
:=
tsMsg
.
(
*
msgstream
.
SearchResultMsg
)
reqID
:=
searchResultMsg
.
Base
.
MsgID
reqIDStr
:=
strconv
.
FormatInt
(
reqID
,
10
)
t
:=
sched
.
getTaskByReqID
(
reqID
)
if
t
==
nil
{
log
.
Debug
(
"proxynode"
,
zap
.
String
(
"QueryResult GetTaskByReqID failed, reqID = "
,
reqIDStr
))
delete
(
queryResultBuf
,
reqID
)
continue
}
if
searchResultMsg
,
srOk
:=
tsMsg
.
(
*
msgstream
.
SearchResultMsg
);
srOk
{
reqID
:=
searchResultMsg
.
Base
.
MsgID
reqIDStr
:=
strconv
.
FormatInt
(
reqID
,
10
)
t
:=
sched
.
getTaskByReqID
(
reqID
)
if
t
==
nil
{
log
.
Debug
(
"proxynode"
,
zap
.
String
(
"QueryResult GetTaskByReqID failed, reqID = "
,
reqIDStr
))
delete
(
queryResultBuf
,
reqID
)
continue
}
_
,
ok
=
queryResultBuf
[
reqID
]
if
!
ok
{
queryResultBuf
[
reqID
]
=
make
([]
*
internalpb
.
SearchResults
,
0
)
}
queryResultBuf
[
reqID
]
=
append
(
queryResultBuf
[
reqID
],
&
searchResultMsg
.
SearchResults
)
_
,
ok
=
queryResultBuf
[
reqID
]
if
!
ok
{
queryResultBuf
[
reqID
]
=
make
([]
*
internalpb
.
SearchResults
,
0
)
}
queryResultBuf
[
reqID
]
=
append
(
queryResultBuf
[
reqID
],
&
searchResultMsg
.
SearchResults
)
//t := sched.getTaskByReqID(reqID)
{
colName
:=
t
.
(
*
SearchTask
)
.
query
.
CollectionName
log
.
Debug
(
"Getcollection"
,
zap
.
String
(
"collection name"
,
colName
),
zap
.
String
(
"reqID"
,
reqIDStr
),
zap
.
Int
(
"answer cnt"
,
len
(
queryResultBuf
[
reqID
])))
//t := sched.getTaskByReqID(reqID)
{
colName
:=
t
.
(
*
SearchTask
)
.
query
.
CollectionName
log
.
Debug
(
"Getcollection"
,
zap
.
String
(
"collection name"
,
colName
),
zap
.
String
(
"reqID"
,
reqIDStr
),
zap
.
Int
(
"answer cnt"
,
len
(
queryResultBuf
[
reqID
])))
}
if
len
(
queryResultBuf
[
reqID
])
==
queryNodeNum
{
t
:=
sched
.
getTaskByReqID
(
reqID
)
if
t
!=
nil
{
qt
,
ok
:=
t
.
(
*
SearchTask
)
if
ok
{
qt
.
resultBuf
<-
queryResultBuf
[
reqID
]
delete
(
queryResultBuf
,
reqID
)
}
}
else
{
// log.Printf("task with reqID %v is nil", reqID)
}
}
sp
.
Finish
()
}
if
len
(
queryResultBuf
[
reqID
])
==
queryNodeNum
{
if
retrieveResultMsg
,
rtOk
:=
tsMsg
.
(
*
msgstream
.
RetrieveResultMsg
);
rtOk
{
reqID
:=
retrieveResultMsg
.
Base
.
MsgID
reqIDStr
:=
strconv
.
FormatInt
(
reqID
,
10
)
t
:=
sched
.
getTaskByReqID
(
reqID
)
if
t
!=
nil
{
qt
,
ok
:=
t
.
(
*
SearchTask
)
if
ok
{
qt
.
resultBuf
<-
queryResultBuf
[
reqID
]
delete
(
queryResultBuf
,
reqID
)
}
}
else
{
if
t
==
nil
{
log
.
Debug
(
"proxynode"
,
zap
.
String
(
"RetrieveResult GetTaskByReqID failed, reqID = "
,
reqIDStr
))
delete
(
retrieveResultBuf
,
reqID
)
continue
}
// log.Printf("task with reqID %v is nil", reqID)
_
,
ok
=
retrieveResultBuf
[
reqID
]
if
!
ok
{
retrieveResultBuf
[
reqID
]
=
make
([]
*
internalpb
.
RetrieveResults
,
0
)
}
retrieveResultBuf
[
reqID
]
=
append
(
retrieveResultBuf
[
reqID
],
&
retrieveResultMsg
.
RetrieveResults
)
{
colName
:=
t
.
(
*
RetrieveTask
)
.
retrieve
.
CollectionName
log
.
Debug
(
"Getcollection"
,
zap
.
String
(
"collection name"
,
colName
),
zap
.
String
(
"reqID"
,
reqIDStr
),
zap
.
Int
(
"answer cnt"
,
len
(
retrieveResultBuf
[
reqID
])))
}
if
len
(
retrieveResultBuf
[
reqID
])
==
queryNodeNum
{
t
:=
sched
.
getTaskByReqID
(
reqID
)
if
t
!=
nil
{
rt
,
ok
:=
t
.
(
*
RetrieveTask
)
if
ok
{
rt
.
resultBuf
<-
retrieveResultBuf
[
reqID
]
delete
(
retrieveResultBuf
,
reqID
)
}
}
else
{
}
}
sp
.
Finish
()
}
sp
.
Finish
()
}
case
<-
sched
.
ctx
.
Done
()
:
log
.
Debug
(
"proxynode server is closed ..."
)
...
...
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录