From 71e323f59d0ca9c63d1d9d1b607db08431a3986c Mon Sep 17 00:00:00 2001 From: Gaius Date: Tue, 27 Jun 2023 14:17:29 +0800 Subject: [PATCH] feat: add context key to SyncProbes (#2485) Signed-off-by: Gaius --- pkg/rpc/scheduler/client/client_v1.go | 5 ++++- pkg/rpc/scheduler/client/client_v2.go | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/pkg/rpc/scheduler/client/client_v1.go b/pkg/rpc/scheduler/client/client_v1.go index 3488d34dd..e732678f3 100644 --- a/pkg/rpc/scheduler/client/client_v1.go +++ b/pkg/rpc/scheduler/client/client_v1.go @@ -308,7 +308,10 @@ func (v *v1) LeaveHost(ctx context.Context, req *schedulerv1.LeaveHostRequest, o // SyncProbes sync probes of the host. func (v *v1) SyncProbes(ctx context.Context, req *schedulerv1.SyncProbesRequest, opts ...grpc.CallOption) (schedulerv1.Scheduler_SyncProbesClient, error) { - stream, err := v.SchedulerClient.SyncProbes(ctx, opts...) + stream, err := v.SchedulerClient.SyncProbes( + context.WithValue(ctx, pkgbalancer.ContextKey, req.Host.Id), + opts..., + ) if err != nil { return nil, err } diff --git a/pkg/rpc/scheduler/client/client_v2.go b/pkg/rpc/scheduler/client/client_v2.go index 4601cced2..82a0b11e8 100644 --- a/pkg/rpc/scheduler/client/client_v2.go +++ b/pkg/rpc/scheduler/client/client_v2.go @@ -277,7 +277,10 @@ func (v *v2) LeaveHost(ctx context.Context, req *schedulerv2.LeaveHostRequest, o // SyncProbes sync probes of the host. func (v *v2) SyncProbes(ctx context.Context, req *schedulerv2.SyncProbesRequest, opts ...grpc.CallOption) (schedulerv2.Scheduler_SyncProbesClient, error) { - stream, err := v.SchedulerClient.SyncProbes(ctx, opts...) + stream, err := v.SchedulerClient.SyncProbes( + context.WithValue(ctx, pkgbalancer.ContextKey, req.Host.Id), + opts..., + ) if err != nil { return nil, err }