enhance: [cp 3.0] record and expose DataNode task cost_time / cost_cpu_num via QueryTask (#51459)

issue: #49156
pr: #51443

## Summary

Cherry-pick of #51443 to 3.0 (clean pick, cumulative diff byte-identical
to the master PR):

- Record `cost_time` / `cost_cpu_num` on DataNode index build tasks and
surface them through `QueryTask` response `Properties`.
- Reserve the unified `cost_time` / `cost_cpu_num` keys under
`pkg/taskcommon` so future PRs can wire up the other
`CreateTask`-dispatched task types.
- Introduce small `internal/datanode/taskcost` helper (`NowMs`,
`CalcCostTimeMs`, `EstimateIndexBuildCPUNum`) for index build task cost
accounting.
- Final state/failReason and execution-end cost are written in a single
`TaskManager` critical section, and `QueryTask` reads them from a single
snapshot, so state and cost in one response are always consistent.

## Notes

Producer-side only; DataCoord / proxy / metrics exporter do not consume
these properties yet. `cost_cpu_num` is a coarse, statistics-only
estimate for observability and is not used for coordinator-side
scheduling.

Includes both master commits: the feature (`8a4ab16f1f`) and the
adversarial-review fixes (`055842e1c9`).

---------

Signed-off-by: cai.zhang <cai.zhang@zilliz.com>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
cai.zhang
2026-07-18 21:52:40 +08:00
committed by GitHub
co-authored by Claude Opus 4.7
parent f6af132935
commit af5dee46de
11 changed files with 448 additions and 15 deletions
+40 -5
View File
@@ -21,10 +21,12 @@ import (
"context"
"runtime/debug"
"sync"
"time"
"github.com/cockroachdb/errors"
"go.uber.org/atomic"
"github.com/milvus-io/milvus/internal/datanode/taskcost"
"github.com/milvus-io/milvus/internal/storagev2"
"github.com/milvus-io/milvus/pkg/v3/mlog"
"github.com/milvus-io/milvus/pkg/v3/proto/indexpb"
@@ -242,19 +244,52 @@ func (sched *TaskScheduler) processTask(t Task) {
}()
sched.TaskQueue.AddActiveTask(t)
defer sched.TaskQueue.PopActiveTask(t.Name())
mlog.Debug(t.Ctx(), "process task", mlog.String("task", t.Name()))
var (
indexTask *indexBuildTask
costCPUNum int64
execStart time.Time
)
if ibt, ok := t.(*indexBuildTask); ok {
indexTask = ibt
costCPUNum = taskcost.EstimateIndexBuildCPUNum(indexTask.IsVectorIndex())
// execStart carries a monotonic clock reading; CostTimeMs derived from
// it is immune to wall-clock steps. ExecStartMs/ExecEndMs stay wall-clock
// timestamps for external exposure.
execStart = time.Now()
indexTask.manager.StoreIndexTaskExecutionStart(indexTask.req.GetClusterID(), indexTask.req.GetBuildID(), taskcost.NowMs(), costCPUNum)
mlog.Debug(t.Ctx(), "process task", mlog.String("task", t.Name()), mlog.Int64("costCPUNum", costCPUNum))
} else {
mlog.Debug(t.Ctx(), "process task", mlog.String("task", t.Name()))
}
pipelines := []func(context.Context) error{t.PreExecute, t.Execute, t.PostExecute}
for _, fn := range pipelines {
if err := wrap(fn); err != nil {
mlog.Warn(t.Ctx(), "process task failed", mlog.Err(err))
t.SetState(getStateFromError(err), err.Error())
if indexTask != nil {
costTimeMs := taskcost.ElapsedMs(execStart)
// End bookkeeping and final state must land in one critical
// section, so a concurrent QueryTask never sees a final cost
// paired with an in-progress state.
indexTask.SetStateWithCost(getStateFromError(err), err.Error(), taskcost.NowMs(), costTimeMs)
mlog.Warn(t.Ctx(), "process task failed", mlog.Err(err), mlog.Int64("costTimeMs", costTimeMs), mlog.Int64("costCPUNum", costCPUNum))
} else {
t.SetState(getStateFromError(err), err.Error())
mlog.Warn(t.Ctx(), "process task failed", mlog.Err(err))
}
return
}
}
t.SetState(indexpb.JobState_JobStateFinished, "")
if indexTask != nil {
costTimeMs := taskcost.ElapsedMs(execStart)
indexTask.SetStateWithCost(indexpb.JobState_JobStateFinished, "", taskcost.NowMs(), costTimeMs)
mlog.Debug(t.Ctx(), "process task completed", mlog.String("task", t.Name()), mlog.Int64("costTimeMs", costTimeMs), mlog.Int64("costCPUNum", costCPUNum))
} else {
t.SetState(indexpb.JobState_JobStateFinished, "")
mlog.Debug(t.Ctx(), "process task completed", mlog.String("task", t.Name()))
}
// Publish filesystem metrics after index task completion
if indexTask, ok := t.(*indexBuildTask); ok {
if indexTask != nil {
if indexTask.req != nil && indexTask.req.GetStorageConfig() != nil {
storagev2.PublishFilesystemMetricsWithConfig(indexTask.req.GetStorageConfig())
}
+100
View File
@@ -23,10 +23,16 @@ import (
"testing"
"time"
"github.com/bytedance/mockey"
"github.com/cockroachdb/errors"
"github.com/stretchr/testify/assert"
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
"github.com/milvus-io/milvus/pkg/v3/common"
"github.com/milvus-io/milvus/pkg/v3/proto/indexpb"
"github.com/milvus-io/milvus/pkg/v3/proto/workerpb"
"github.com/milvus-io/milvus/pkg/v3/util/hardware"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
)
@@ -230,3 +236,97 @@ func TestIndexTaskScheduler(t *testing.T) {
assert.Equal(t, task.GetState(), indexpb.JobState_JobStateFinished)
}
}
func newSchedulerIndexBuildTask(t *testing.T, manager *TaskManager, buildID int64) *indexBuildTask {
ctx, cancel := context.WithCancel(context.Background())
t.Cleanup(cancel)
req := &workerpb.CreateJobRequest{
ClusterID: "test-cluster",
BuildID: buildID,
IndexParams: []*commonpb.KeyValuePair{
{Key: common.IndexTypeKey, Value: "STL_SORT"},
},
Field: &schemapb.FieldSchema{
FieldID: 100,
DataType: schemapb.DataType_Int64,
},
}
manager.LoadOrStoreIndexTask(req.GetClusterID(), req.GetBuildID(), &IndexTaskInfo{
State: commonpb.IndexState_InProgress,
})
return NewIndexBuildTask(ctx, cancel, req, nil, manager, nil)
}
func TestIndexTaskSchedulerRecordsIndexTaskCost(t *testing.T) {
paramtable.Init()
t.Run("success records execution cost", func(t *testing.T) {
manager := NewTaskManager(context.Background())
task := newSchedulerIndexBuildTask(t, manager, 1001)
preMock := mockey.Mock((*indexBuildTask).PreExecute).Return(nil).Build()
defer preMock.UnPatch()
executeMock := mockey.Mock((*indexBuildTask).Execute).Return(nil).Build()
defer executeMock.UnPatch()
postMock := mockey.Mock((*indexBuildTask).PostExecute).Return(nil).Build()
defer postMock.UnPatch()
scheduler := NewTaskScheduler(context.Background())
scheduler.processTask(task)
info := manager.GetIndexTaskInfo("test-cluster", 1001)
assert.NotNil(t, info)
assert.Equal(t, commonpb.IndexState_Finished, info.State)
assert.Greater(t, info.ExecStartMs, int64(0))
assert.GreaterOrEqual(t, info.ExecEndMs, info.ExecStartMs)
assert.GreaterOrEqual(t, info.CostTimeMs, int64(0))
assert.Equal(t, int64(1), info.CostCPUNum)
})
t.Run("pre execute failure still records execution end", func(t *testing.T) {
manager := NewTaskManager(context.Background())
task := newSchedulerIndexBuildTask(t, manager, 1002)
expectedErr := errors.New("pre execute failed")
preMock := mockey.Mock((*indexBuildTask).PreExecute).Return(expectedErr).Build()
defer preMock.UnPatch()
scheduler := NewTaskScheduler(context.Background())
scheduler.processTask(task)
info := manager.GetIndexTaskInfo("test-cluster", 1002)
assert.NotNil(t, info)
assert.Equal(t, commonpb.IndexState_Retry, info.State)
assert.Equal(t, expectedErr.Error(), info.FailReason)
assert.Greater(t, info.ExecStartMs, int64(0))
assert.GreaterOrEqual(t, info.ExecEndMs, info.ExecStartMs)
assert.GreaterOrEqual(t, info.CostTimeMs, int64(0))
assert.Equal(t, int64(1), info.CostCPUNum)
})
t.Run("vector index records build pool cpu num", func(t *testing.T) {
manager := NewTaskManager(context.Background())
task := newSchedulerIndexBuildTask(t, manager, 1003)
vecMock := mockey.Mock((*indexBuildTask).IsVectorIndex).Return(true).Build()
defer vecMock.UnPatch()
preMock := mockey.Mock((*indexBuildTask).PreExecute).Return(nil).Build()
defer preMock.UnPatch()
executeMock := mockey.Mock((*indexBuildTask).Execute).Return(nil).Build()
defer executeMock.UnPatch()
postMock := mockey.Mock((*indexBuildTask).PostExecute).Return(nil).Build()
defer postMock.UnPatch()
scheduler := NewTaskScheduler(context.Background())
scheduler.processTask(task)
info := manager.GetIndexTaskInfo("test-cluster", 1003)
assert.NotNil(t, info)
assert.Equal(t, commonpb.IndexState_Finished, info.State)
assert.Greater(t, info.ExecStartMs, int64(0))
assert.GreaterOrEqual(t, info.ExecEndMs, info.ExecStartMs)
assert.GreaterOrEqual(t, info.CostTimeMs, int64(0))
assert.Equal(t, int64(hardware.GetCPUNum()), info.CostCPUNum)
})
}
+13
View File
@@ -118,6 +118,19 @@ func (it *indexBuildTask) Name() string {
func (it *indexBuildTask) SetState(state indexpb.JobState, failReason string) {
it.manager.StoreIndexTaskState(it.req.GetClusterID(), it.req.GetBuildID(), commonpb.IndexState(state), failReason)
it.observeStateMetrics(state)
}
// SetStateWithCost stores the final state/failReason together with the
// execution-end cost bookkeeping in one TaskManager critical section, so a
// concurrent QueryTask can never observe a final cost paired with a stale
// in-progress state.
func (it *indexBuildTask) SetStateWithCost(state indexpb.JobState, failReason string, endMs, costTimeMs int64) {
it.manager.StoreIndexTaskExecutionEndWithState(it.req.GetClusterID(), it.req.GetBuildID(), endMs, costTimeMs, commonpb.IndexState(state), failReason)
it.observeStateMetrics(state)
}
func (it *indexBuildTask) observeStateMetrics(state indexpb.JobState) {
if state == indexpb.JobState_JobStateFinished {
metrics.DataNodeBuildIndexLatency.WithLabelValues(paramtable.GetStringNodeID()).Observe(it.tr.ElapseSpan().Seconds())
metrics.DataNodeIndexTaskLatencyInQueue.WithLabelValues(paramtable.GetStringNodeID()).Observe(float64(it.queueDur.Milliseconds()))
+49
View File
@@ -44,6 +44,10 @@ type IndexTaskInfo struct {
CurrentIndexVersion int32
CurrentScalarIndexVersion int32
IndexStorePathVersion indexpb.IndexStorePathVersion
ExecStartMs int64
ExecEndMs int64
CostTimeMs int64
CostCPUNum int64
// task statistics
statistic *indexpb.JobInfo
@@ -60,6 +64,10 @@ func (i *IndexTaskInfo) Clone() *IndexTaskInfo {
CurrentIndexVersion: i.CurrentIndexVersion,
CurrentScalarIndexVersion: i.CurrentScalarIndexVersion,
IndexStorePathVersion: i.IndexStorePathVersion,
ExecStartMs: i.ExecStartMs,
ExecEndMs: i.ExecEndMs,
CostTimeMs: i.CostTimeMs,
CostCPUNum: i.CostCPUNum,
statistic: typeutil.Clone(i.statistic),
}
}
@@ -130,6 +138,37 @@ func (m *TaskManager) StoreIndexTaskState(ClusterID string, buildID typeutil.Uni
}
}
func (m *TaskManager) StoreIndexTaskExecutionStart(clusterID string, buildID typeutil.UniqueID, startMs int64, costCPUNum int64) {
key := Key{ClusterID: clusterID, TaskID: buildID}
m.stateLock.Lock()
defer m.stateLock.Unlock()
if task, ok := m.indexTasks[key]; ok {
task.ExecStartMs = startMs
task.CostCPUNum = costCPUNum
task.ExecEndMs = 0
task.CostTimeMs = 0
}
}
// StoreIndexTaskExecutionEndWithState records the execution-end cost bookkeeping
// together with the final task state/failReason in a single critical section, so
// a concurrent reader can never observe a final cost paired with a stale state
// (or vice versa).
func (m *TaskManager) StoreIndexTaskExecutionEndWithState(clusterID string, buildID typeutil.UniqueID, endMs int64, costTimeMs int64, state commonpb.IndexState, failReason string) {
key := Key{ClusterID: clusterID, TaskID: buildID}
m.stateLock.Lock()
defer m.stateLock.Unlock()
if task, ok := m.indexTasks[key]; ok {
mlog.Debug(m.ctx, "store task execution end with state", mlog.String("clusterID", clusterID), mlog.FieldBuildID(buildID),
mlog.String("state", state.String()), mlog.String("fail reason", failReason),
mlog.Int64("costTimeMs", costTimeMs))
task.ExecEndMs = endMs
task.CostTimeMs = costTimeMs
task.State = state
task.FailReason = failReason
}
}
func (m *TaskManager) ForeachIndexTaskInfo(fn func(ClusterID string, buildID typeutil.UniqueID, info *IndexTaskInfo)) {
m.stateLock.Lock()
defer m.stateLock.Unlock()
@@ -138,6 +177,16 @@ func (m *TaskManager) ForeachIndexTaskInfo(fn func(ClusterID string, buildID typ
}
}
func (m *TaskManager) GetIndexTaskInfo(clusterID string, buildID typeutil.UniqueID) *IndexTaskInfo {
m.stateLock.Lock()
defer m.stateLock.Unlock()
if info, ok := m.indexTasks[Key{ClusterID: clusterID, TaskID: buildID}]; ok {
return info.Clone()
}
return nil
}
func (m *TaskManager) StoreIndexFilesAndStatistic(
ClusterID string,
buildID typeutil.UniqueID,
@@ -156,3 +156,32 @@ func (s *statsTaskInfoSuite) TestIndexTaskInfoReturnsIndexStorePathVersion() {
cloned := info.Clone()
s.Equal(indexpb.IndexStorePathVersion_INDEX_STORE_PATH_VERSION_COLLECTION_ROOTED, cloned.IndexStorePathVersion)
}
func (s *statsTaskInfoSuite) Test_IndexTaskCostMethods() {
buildID := int64(200)
s.manager.LoadOrStoreIndexTask(s.cluster, buildID, &IndexTaskInfo{State: commonpb.IndexState_InProgress})
s.manager.StoreIndexTaskExecutionStart(s.cluster, buildID, 111, 4)
s.manager.StoreIndexTaskExecutionEndWithState(s.cluster, buildID, 222, 111, commonpb.IndexState_Failed, "mock failure")
info := s.manager.GetIndexTaskInfo(s.cluster, buildID)
s.Require().NotNil(info)
s.Equal(int64(111), info.ExecStartMs)
s.Equal(int64(222), info.ExecEndMs)
s.Equal(int64(111), info.CostTimeMs)
s.Equal(int64(4), info.CostCPUNum)
// state/failReason land together with the cost in one critical section
s.Equal(commonpb.IndexState_Failed, info.State)
s.Equal("mock failure", info.FailReason)
info.CostCPUNum = 999
reloaded := s.manager.GetIndexTaskInfo(s.cluster, buildID)
s.Equal(int64(4), reloaded.CostCPUNum)
s.manager.StoreIndexTaskExecutionStart(s.cluster, buildID, 333, 8)
reloaded = s.manager.GetIndexTaskInfo(s.cluster, buildID)
s.Equal(int64(333), reloaded.ExecStartMs)
s.Equal(int64(8), reloaded.CostCPUNum)
s.Equal(int64(0), reloaded.ExecEndMs)
s.Equal(int64(0), reloaded.CostTimeMs)
}
+24 -8
View File
@@ -814,15 +814,31 @@ func (node *DataNode) QueryTask(ctx context.Context, request *workerpb.QueryTask
}
return wrapQueryTaskResult(resp, resProperties)
case taskcommon.Index:
resp, err := node.queryIndexTask(ctx, &workerpb.QueryJobsRequest{ClusterID: clusterID, TaskIDs: []int64{taskID}})
if err != nil {
return nil, err
}
// State/reason and cost must come from one snapshot cloned under one
// lock, so a concurrently completing task can never yield a final cost
// paired with an in-progress state (or vice versa).
resProperties := taskcommon.NewProperties(nil)
results := resp.GetIndexJobResults().GetResults()
if len(results) > 0 {
resProperties.AppendTaskState(taskcommon.State(results[0].GetState()))
resProperties.AppendReason(results[0].GetFailReason())
info := node.taskManager.GetIndexTaskInfo(clusterID, taskID)
if info == nil {
resProperties.AppendCostTime(0)
resProperties.AppendCostCPUNum(0)
resp := &workerpb.QueryJobsV2Response{
Status: merr.Status(merr.WrapErrServiceInternalMsg("tasks '%v' not found", []int64{taskID})),
}
return wrapQueryTaskResult(resp, resProperties)
}
resProperties.AppendTaskState(taskcommon.State(info.State))
resProperties.AppendReason(info.FailReason)
resProperties.AppendCostTime(info.CostTimeMs)
resProperties.AppendCostCPUNum(info.CostCPUNum)
resp := &workerpb.QueryJobsV2Response{
Status: merr.Success(),
ClusterID: clusterID,
Result: &workerpb.QueryJobsV2Response_IndexJobResults{
IndexJobResults: &workerpb.IndexJobResults{
Results: []*workerpb.IndexTaskInfo{info.ToIndexTaskInfo(taskID)},
},
},
}
return wrapQueryTaskResult(resp, resProperties)
case taskcommon.Stats:
+24
View File
@@ -34,6 +34,7 @@ import (
"github.com/milvus-io/milvus/internal/compaction"
"github.com/milvus-io/milvus/internal/datanode/compactor"
"github.com/milvus-io/milvus/internal/datanode/external"
"github.com/milvus-io/milvus/internal/datanode/index"
"github.com/milvus-io/milvus/internal/storage"
"github.com/milvus-io/milvus/internal/util/sessionutil"
"github.com/milvus-io/milvus/pkg/v3/common"
@@ -678,6 +679,29 @@ func (s *DataNodeServicesSuite) TestQueryTask() {
s.True(strings.Contains(resp.GetStatus().GetReason(), "not found"))
})
s.Run("query index task with cost", func() {
s.node.taskManager.LoadOrStoreIndexTask("cluster-0", 101, &index.IndexTaskInfo{State: commonpb.IndexState_InProgress})
s.node.taskManager.StoreIndexTaskExecutionStart("cluster-0", 101, 100, 3)
s.node.taskManager.StoreIndexTaskExecutionEndWithState("cluster-0", 101, 180, 80, commonpb.IndexState_Finished, "")
req := &workerpb.QueryTaskRequest{
Properties: map[string]string{
taskcommon.ClusterIDKey: "cluster-0",
taskcommon.TypeKey: taskcommon.Index,
taskcommon.TaskIDKey: "101",
},
}
resp, err := s.node.QueryTask(s.ctx, req)
s.NoError(merr.CheckRPCCall(resp, err))
props := taskcommon.NewProperties(resp.GetProperties())
// state and cost come from the same snapshot
state, err := props.GetTaskState()
s.NoError(err)
s.Equal(taskcommon.State(commonpb.IndexState_Finished), state)
s.Equal(int64(80), props.GetCostTime())
s.Equal(int64(3), props.GetCostCPUNum())
})
s.Run("invalid task type", func() {
req := &workerpb.QueryTaskRequest{
Properties: map[string]string{
+60
View File
@@ -0,0 +1,60 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package taskcost
import (
"time"
"github.com/milvus-io/milvus/pkg/v3/util/hardware"
)
// NowMs returns the current wall-clock time in milliseconds. Use it only for
// externally exposed timestamps (ExecStartMs / ExecEndMs); never subtract two
// NowMs readings to derive a duration — use ElapsedMs instead.
func NowMs() int64 {
return time.Now().UnixMilli()
}
// ElapsedMs returns the duration since start in milliseconds. It relies on the
// monotonic clock reading embedded in start (time.Since), so the result is
// immune to wall-clock steps such as NTP corrections and is never negative.
func ElapsedMs(start time.Time) int64 {
return time.Since(start).Milliseconds()
}
// EstimateIndexBuildCPUNum returns the approximate number of CPU threads an
// index build task consumes during execution.
//
// The value is task-level statistics for observability only and is never used
// for coordinator-side scheduling, so approximation error is acceptable (e.g.
// standalone sizes the build pool to NumCPU * buildIndexThreadPoolRatio, and
// concurrent builds share one pool).
//
// Vector indexes go through knowhere's build thread pool (sized to NumCPU
// in cluster mode, see internal/datanode/index/init_segcore.go:74), so they
// report hardware.GetCPUNum(). All current scalar indexes build
// single-threaded: the tantivy wrapper hardcodes 1 thread
// (internal/core/thirdparty/tantivy/tantivy-wrapper.h:23 — covers INVERTED /
// NGRAM / TEXT_MATCH / JSON_INVERTED), and the rest (BITMAP / STL_SORT /
// STRING_SORT / MARISA / HYBRID / RTREE) are plain loops with no thread
// pool / OpenMP / std::async. They therefore report 1.
func EstimateIndexBuildCPUNum(isVectorIndex bool) int64 {
if isVectorIndex {
return int64(hardware.GetCPUNum())
}
return 1
}
+49
View File
@@ -0,0 +1,49 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package taskcost
import (
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/milvus-io/milvus/pkg/v3/util/hardware"
)
func TestNowMs(t *testing.T) {
before := time.Now().UnixMilli()
got := NowMs()
after := time.Now().UnixMilli()
assert.GreaterOrEqual(t, got, before)
assert.LessOrEqual(t, got, after)
}
func TestElapsedMs(t *testing.T) {
start := time.Now()
time.Sleep(10 * time.Millisecond)
got := ElapsedMs(start)
// Sleep guarantees at least 10ms elapsed on the monotonic clock.
assert.GreaterOrEqual(t, got, int64(10))
}
func TestEstimateIndexBuildCPUNum(t *testing.T) {
assert.Equal(t, int64(hardware.GetCPUNum()), EstimateIndexBuildCPUNum(true),
"vector index saturates knowhere build pool")
assert.Equal(t, int64(1), EstimateIndexBuildCPUNum(false),
"all scalar indexes currently build single-threaded")
}
+34 -2
View File
@@ -37,8 +37,10 @@ const (
CollectionIDKey = "collection_id"
// result
StateKey = "task_state"
ReasonKey = "task_reason"
StateKey = "task_state"
ReasonKey = "task_reason"
CostTimeKey = "cost_time"
CostCPUNumKey = "cost_cpu_num"
)
type Properties map[string]string
@@ -102,6 +104,14 @@ func (p Properties) AppendTaskState(state State) {
p[StateKey] = state.String()
}
func (p Properties) AppendCostTime(costTime int64) {
p[CostTimeKey] = fmt.Sprintf("%d", costTime)
}
func (p Properties) AppendCostCPUNum(costCPUNum int64) {
p[CostCPUNumKey] = fmt.Sprintf("%d", costCPUNum)
}
func (p Properties) GetTaskType() (Type, error) {
if _, ok := p[TypeKey]; !ok {
return "", WrapErrTaskPropertyLack(TypeKey, p[TaskIDKey])
@@ -206,3 +216,25 @@ func (p Properties) GetCollectionID() (int64, error) {
}
return collectionID, nil
}
func (p Properties) GetCostTime() int64 {
if _, ok := p[CostTimeKey]; !ok {
return 0
}
costTime, err := strconv.ParseInt(p[CostTimeKey], 10, 64)
if err != nil {
return 0
}
return costTime
}
func (p Properties) GetCostCPUNum() int64 {
if _, ok := p[CostCPUNumKey]; !ok {
return 0
}
costCPUNum, err := strconv.ParseInt(p[CostCPUNumKey], 10, 64)
if err != nil {
return 0
}
return costCPUNum
}
+26
View File
@@ -100,3 +100,29 @@ func TestProperties_CollectionID(t *testing.T) {
assert.Equal(t, int64(-1), collectionID)
})
}
func TestProperties_CostFields(t *testing.T) {
t.Run("append and get", func(t *testing.T) {
props := NewProperties(nil)
props.AppendCostTime(123)
props.AppendCostCPUNum(4)
assert.Equal(t, int64(123), props.GetCostTime())
assert.Equal(t, int64(4), props.GetCostCPUNum())
})
t.Run("missing keys", func(t *testing.T) {
props := NewProperties(nil)
assert.Equal(t, int64(0), props.GetCostTime())
assert.Equal(t, int64(0), props.GetCostCPUNum())
})
t.Run("invalid values", func(t *testing.T) {
props := NewProperties(map[string]string{
CostTimeKey: "abc",
CostCPUNumKey: "xyz",
})
assert.Equal(t, int64(0), props.GetCostTime())
assert.Equal(t, int64(0), props.GetCostCPUNum())
})
}