feat: add context key to SyncProbes (#2485)

Signed-off-by: Gaius <gaius.qi@gmail.com>
This commit is contained in:
Gaius 2023-06-27 14:17:29 +08:00
parent 205e990c1b
commit 71e323f59d
No known key found for this signature in database
GPG Key ID: 8B4E5D1290FA2FFB
2 changed files with 8 additions and 2 deletions

View File

@ -308,7 +308,10 @@ func (v *v1) LeaveHost(ctx context.Context, req *schedulerv1.LeaveHostRequest, o
// SyncProbes sync probes of the host. // SyncProbes sync probes of the host.
func (v *v1) SyncProbes(ctx context.Context, req *schedulerv1.SyncProbesRequest, opts ...grpc.CallOption) (schedulerv1.Scheduler_SyncProbesClient, error) { 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 { if err != nil {
return nil, err return nil, err
} }

View File

@ -277,7 +277,10 @@ func (v *v2) LeaveHost(ctx context.Context, req *schedulerv2.LeaveHostRequest, o
// SyncProbes sync probes of the host. // SyncProbes sync probes of the host.
func (v *v2) SyncProbes(ctx context.Context, req *schedulerv2.SyncProbesRequest, opts ...grpc.CallOption) (schedulerv2.Scheduler_SyncProbesClient, error) { 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 { if err != nil {
return nil, err return nil, err
} }