Skip to content

Commit

Permalink
feat: record the duration waiting in the proxy queue (milvus-io#34744)
Browse files Browse the repository at this point in the history 
Signed-off-by: longjiquan <[email protected]>
  • Loading branch information
@longjiquan
longjiquan committed Jul 18, 2024
1 parent 3d9ec2a commit 0d4ef5c
Show file tree
Hide file tree
Showing 13 changed files with 140 additions and 72 deletions.
1 change: 1 addition & 0 deletions internal/proxy/mock_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,7 @@ func newMockIDAllocatorInterface() allocator.Interface {
}

type mockTask struct {
baseTask
*TaskCondition
id UniqueID
name string
Expand Down
103 changes: 61 additions & 42 deletions internal/proxy/task.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ import (
"context"
"fmt"
"math"
"time"

"github.com/cockroachdb/errors"
"github.com/golang/protobuf/proto"
Expand Down Expand Up @@ -107,6 +108,20 @@ type task interface {
PostExecute(ctx context.Context) error
WaitToFinish() error
Notify(err error)
SetOnEnqueueTime()
GetDurationInQueue() time.Duration
}

type baseTask struct {
onEnqueueTime time.Time
}

func (bt *baseTask) SetOnEnqueueTime() {
bt.onEnqueueTime = time.Now()
}

func (bt *baseTask) GetDurationInQueue() time.Duration {
return time.Since(bt.onEnqueueTime)
}

type dmlTask interface {
Expand Down Expand Up @@ -362,12 +377,12 @@ func (t *dropCollectionTask) OnEnqueue() error {
if t.Base == nil {
t.Base = commonpbutil.NewMsgBase()
}
t.Base.MsgType = commonpb.MsgType_DropCollection
t.Base.SourceID = paramtable.GetNodeID()
return nil
}

func (t *dropCollectionTask) PreExecute(ctx context.Context) error {
t.Base.MsgType = commonpb.MsgType_DropCollection
t.Base.SourceID = paramtable.GetNodeID()

if err := validateCollectionName(t.CollectionName); err != nil {
return err
Expand Down Expand Up @@ -427,12 +442,12 @@ func (t *hasCollectionTask) SetTs(ts Timestamp) {

func (t *hasCollectionTask) OnEnqueue() error {
t.Base = commonpbutil.NewMsgBase()
t.Base.MsgType = commonpb.MsgType_HasCollection
t.Base.SourceID = paramtable.GetNodeID()
return nil
}

func (t *hasCollectionTask) PreExecute(ctx context.Context) error {
t.Base.MsgType = commonpb.MsgType_HasCollection
t.Base.SourceID = paramtable.GetNodeID()

if err := validateCollectionName(t.CollectionName); err != nil {
return err
Expand Down Expand Up @@ -501,12 +516,12 @@ func (t *describeCollectionTask) SetTs(ts Timestamp) {

func (t *describeCollectionTask) OnEnqueue() error {
t.Base = commonpbutil.NewMsgBase()
t.Base.MsgType = commonpb.MsgType_DescribeCollection
t.Base.SourceID = paramtable.GetNodeID()
return nil
}

func (t *describeCollectionTask) PreExecute(ctx context.Context) error {
t.Base.MsgType = commonpb.MsgType_DescribeCollection
t.Base.SourceID = paramtable.GetNodeID()

if t.CollectionID != 0 && len(t.CollectionName) == 0 {
return nil
Expand Down Expand Up @@ -636,12 +651,12 @@ func (t *showCollectionsTask) SetTs(ts Timestamp) {

func (t *showCollectionsTask) OnEnqueue() error {
t.Base = commonpbutil.NewMsgBase()
t.Base.MsgType = commonpb.MsgType_ShowCollections
t.Base.SourceID = paramtable.GetNodeID()
return nil
}

func (t *showCollectionsTask) PreExecute(ctx context.Context) error {
t.Base.MsgType = commonpb.MsgType_ShowCollections
t.Base.SourceID = paramtable.GetNodeID()
if t.GetType() == milvuspb.ShowType_InMemory {
for _, collectionName := range t.CollectionNames {
if err := validateCollectionName(collectionName); err != nil {
Expand Down Expand Up @@ -796,12 +811,12 @@ func (t *alterCollectionTask) OnEnqueue() error {
if t.Base == nil {
t.Base = commonpbutil.NewMsgBase()
}
t.Base.MsgType = commonpb.MsgType_AlterCollection
t.Base.SourceID = paramtable.GetNodeID()
return nil
}

func (t *alterCollectionTask) PreExecute(ctx context.Context) error {
t.Base.MsgType = commonpb.MsgType_AlterCollection
t.Base.SourceID = paramtable.GetNodeID()

return nil
}
Expand Down Expand Up @@ -860,12 +875,12 @@ func (t *createPartitionTask) OnEnqueue() error {
if t.Base == nil {
t.Base = commonpbutil.NewMsgBase()
}
t.Base.MsgType = commonpb.MsgType_CreatePartition
t.Base.SourceID = paramtable.GetNodeID()
return nil
}

func (t *createPartitionTask) PreExecute(ctx context.Context) error {
t.Base.MsgType = commonpb.MsgType_CreatePartition
t.Base.SourceID = paramtable.GetNodeID()

collName, partitionTag := t.CollectionName, t.PartitionName

Expand Down Expand Up @@ -948,12 +963,12 @@ func (t *dropPartitionTask) OnEnqueue() error {
if t.Base == nil {
t.Base = commonpbutil.NewMsgBase()
}
t.Base.MsgType = commonpb.MsgType_DropPartition
t.Base.SourceID = paramtable.GetNodeID()
return nil
}

func (t *dropPartitionTask) PreExecute(ctx context.Context) error {
t.Base.MsgType = commonpb.MsgType_DropPartition
t.Base.SourceID = paramtable.GetNodeID()

collName, partitionTag := t.CollectionName, t.PartitionName

Expand Down Expand Up @@ -1059,12 +1074,12 @@ func (t *hasPartitionTask) SetTs(ts Timestamp) {

func (t *hasPartitionTask) OnEnqueue() error {
t.Base = commonpbutil.NewMsgBase()
t.Base.MsgType = commonpb.MsgType_HasPartition
t.Base.SourceID = paramtable.GetNodeID()
return nil
}

func (t *hasPartitionTask) PreExecute(ctx context.Context) error {
t.Base.MsgType = commonpb.MsgType_HasPartition
t.Base.SourceID = paramtable.GetNodeID()

collName, partitionTag := t.CollectionName, t.PartitionName

Expand Down Expand Up @@ -1136,12 +1151,12 @@ func (t *showPartitionsTask) SetTs(ts Timestamp) {

func (t *showPartitionsTask) OnEnqueue() error {
t.Base = commonpbutil.NewMsgBase()
t.Base.MsgType = commonpb.MsgType_ShowPartitions
t.Base.SourceID = paramtable.GetNodeID()
return nil
}

func (t *showPartitionsTask) PreExecute(ctx context.Context) error {
t.Base.MsgType = commonpb.MsgType_ShowPartitions
t.Base.SourceID = paramtable.GetNodeID()

if err := validateCollectionName(t.CollectionName); err != nil {
return err
Expand Down Expand Up @@ -1301,12 +1316,12 @@ func (t *flushTask) OnEnqueue() error {
if t.Base == nil {
t.Base = commonpbutil.NewMsgBase()
}
t.Base.MsgType = commonpb.MsgType_Flush
t.Base.SourceID = paramtable.GetNodeID()
return nil
}

func (t *flushTask) PreExecute(ctx context.Context) error {
t.Base.MsgType = commonpb.MsgType_Flush
t.Base.SourceID = paramtable.GetNodeID()
return nil
}

Expand Down Expand Up @@ -1407,14 +1422,14 @@ func (t *loadCollectionTask) OnEnqueue() error {
if t.Base == nil {
t.Base = commonpbutil.NewMsgBase()
}
t.Base.MsgType = commonpb.MsgType_LoadCollection
t.Base.SourceID = paramtable.GetNodeID()
return nil
}

func (t *loadCollectionTask) PreExecute(ctx context.Context) error {
log.Ctx(ctx).Debug("loadCollectionTask PreExecute",
zap.String("role", typeutil.ProxyRole))
t.Base.MsgType = commonpb.MsgType_LoadCollection
t.Base.SourceID = paramtable.GetNodeID()

collName := t.CollectionName

Expand Down Expand Up @@ -1558,12 +1573,12 @@ func (t *releaseCollectionTask) OnEnqueue() error {
if t.Base == nil {
t.Base = commonpbutil.NewMsgBase()
}
t.Base.MsgType = commonpb.MsgType_ReleaseCollection
t.Base.SourceID = paramtable.GetNodeID()
return nil
}

func (t *releaseCollectionTask) PreExecute(ctx context.Context) error {
t.Base.MsgType = commonpb.MsgType_ReleaseCollection
t.Base.SourceID = paramtable.GetNodeID()

collName := t.CollectionName

Expand Down Expand Up @@ -1652,12 +1667,12 @@ func (t *loadPartitionsTask) OnEnqueue() error {
if t.Base == nil {
t.Base = commonpbutil.NewMsgBase()
}
t.Base.MsgType = commonpb.MsgType_LoadPartitions
t.Base.SourceID = paramtable.GetNodeID()
return nil
}

func (t *loadPartitionsTask) PreExecute(ctx context.Context) error {
t.Base.MsgType = commonpb.MsgType_LoadPartitions
t.Base.SourceID = paramtable.GetNodeID()

collName := t.CollectionName

Expand Down Expand Up @@ -1801,12 +1816,12 @@ func (t *releasePartitionsTask) OnEnqueue() error {
if t.Base == nil {
t.Base = commonpbutil.NewMsgBase()
}
t.Base.MsgType = commonpb.MsgType_ReleasePartitions
t.Base.SourceID = paramtable.GetNodeID()
return nil
}

func (t *releasePartitionsTask) PreExecute(ctx context.Context) error {
t.Base.MsgType = commonpb.MsgType_ReleasePartitions
t.Base.SourceID = paramtable.GetNodeID()

collName := t.CollectionName

Expand Down Expand Up @@ -1905,12 +1920,12 @@ func (t *CreateResourceGroupTask) OnEnqueue() error {
if t.Base == nil {
t.Base = commonpbutil.NewMsgBase()
}
t.Base.MsgType = commonpb.MsgType_CreateResourceGroup
t.Base.SourceID = paramtable.GetNodeID()
return nil
}

func (t *CreateResourceGroupTask) PreExecute(ctx context.Context) error {
t.Base.MsgType = commonpb.MsgType_CreateResourceGroup
t.Base.SourceID = paramtable.GetNodeID()

return nil
}
Expand Down Expand Up @@ -1969,12 +1984,12 @@ func (t *DropResourceGroupTask) OnEnqueue() error {
if t.Base == nil {
t.Base = commonpbutil.NewMsgBase()
}
t.Base.MsgType = commonpb.MsgType_DropResourceGroup
t.Base.SourceID = paramtable.GetNodeID()
return nil
}

func (t *DropResourceGroupTask) PreExecute(ctx context.Context) error {
t.Base.MsgType = commonpb.MsgType_DropResourceGroup
t.Base.SourceID = paramtable.GetNodeID()

return nil
}
Expand Down Expand Up @@ -2030,13 +2045,15 @@ func (t *DescribeResourceGroupTask) SetTs(ts Timestamp) {
}

func (t *DescribeResourceGroupTask) OnEnqueue() error {
t.Base = commonpbutil.NewMsgBase()
if t.Base == nil {
t.Base = commonpbutil.NewMsgBase()
}
t.Base.MsgType = commonpb.MsgType_DescribeResourceGroup
t.Base.SourceID = paramtable.GetNodeID()
return nil
}

func (t *DescribeResourceGroupTask) PreExecute(ctx context.Context) error {
t.Base.MsgType = commonpb.MsgType_DescribeResourceGroup
t.Base.SourceID = paramtable.GetNodeID()

return nil
}
Expand Down Expand Up @@ -2154,12 +2171,12 @@ func (t *TransferNodeTask) OnEnqueue() error {
if t.Base == nil {
t.Base = commonpbutil.NewMsgBase()
}
t.Base.MsgType = commonpb.MsgType_TransferNode
t.Base.SourceID = paramtable.GetNodeID()
return nil
}

func (t *TransferNodeTask) PreExecute(ctx context.Context) error {
t.Base.MsgType = commonpb.MsgType_TransferNode
t.Base.SourceID = paramtable.GetNodeID()

return nil
}
Expand Down Expand Up @@ -2218,12 +2235,12 @@ func (t *TransferReplicaTask) OnEnqueue() error {
if t.Base == nil {
t.Base = commonpbutil.NewMsgBase()
}
t.Base.MsgType = commonpb.MsgType_TransferReplica
t.Base.SourceID = paramtable.GetNodeID()
return nil
}

func (t *TransferReplicaTask) PreExecute(ctx context.Context) error {
t.Base.MsgType = commonpb.MsgType_TransferReplica
t.Base.SourceID = paramtable.GetNodeID()

return nil
}
Expand Down Expand Up @@ -2288,13 +2305,15 @@ func (t *ListResourceGroupsTask) SetTs(ts Timestamp) {
}

func (t *ListResourceGroupsTask) OnEnqueue() error {
t.Base = commonpbutil.NewMsgBase()
if t.Base == nil {
t.Base = commonpbutil.NewMsgBase()
}
t.Base.MsgType = commonpb.MsgType_ListResourceGroups
t.Base.SourceID = paramtable.GetNodeID()
return nil
}

func (t *ListResourceGroupsTask) PreExecute(ctx context.Context) error {
t.Base.MsgType = commonpb.MsgType_ListResourceGroups
t.Base.SourceID = paramtable.GetNodeID()

return nil
}
Expand Down
Loading

0 comments on commit 0d4ef5c

Please sign in to comment.