-
Notifications
You must be signed in to change notification settings - Fork 278
/
client_v1.go
96 lines (82 loc) · 2.96 KB
/
client_v1.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
/*
* Copyright 2023 The Dragonfly Authors
*
* Licensed 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.
*/
//go:generate mockgen -destination mocks/client_v1_mock.go -source client_v1.go -package mocks
package client
import (
"context"
"time"
grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
grpc_zap "github.com/grpc-ecosystem/go-grpc-middleware/logging/zap"
grpc_retry "github.com/grpc-ecosystem/go-grpc-middleware/retry"
grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
"google.golang.org/grpc"
trainerv1 "d7y.io/api/v2/pkg/apis/trainer/v1"
logger "d7y.io/dragonfly/v2/internal/dflog"
)
const (
// maxRetries is maximum number of retries.
maxRetries = 3
// backoffWaitBetween is waiting for a fixed period of
// time between calls in backoff linear.
backoffWaitBetween = 500 * time.Millisecond
)
// GetV1ByAddr returns v1 version of the trainer client by address.
func GetV1ByAddr(ctx context.Context, target string, opts ...grpc.DialOption) (V1, error) {
conn, err := grpc.DialContext(
ctx,
target,
append([]grpc.DialOption{
grpc.WithUnaryInterceptor(grpc_middleware.ChainUnaryClient(
otelgrpc.UnaryClientInterceptor(),
grpc_prometheus.UnaryClientInterceptor,
grpc_zap.UnaryClientInterceptor(logger.GrpcLogger.Desugar()),
grpc_retry.UnaryClientInterceptor(
grpc_retry.WithMax(maxRetries),
grpc_retry.WithBackoff(grpc_retry.BackoffLinear(backoffWaitBetween)),
),
)),
grpc.WithStreamInterceptor(grpc_middleware.ChainStreamClient(
otelgrpc.StreamClientInterceptor(),
grpc_prometheus.StreamClientInterceptor,
grpc_zap.StreamClientInterceptor(logger.GrpcLogger.Desugar()),
)),
}, opts...)...,
)
if err != nil {
return nil, err
}
return &v1{
TrainerClient: trainerv1.NewTrainerClient(conn),
ClientConn: conn,
}, nil
}
// V1 is the interface for v1 version of the grpc client.
type V1 interface {
// Train models of scheduler using dataset.
Train(context.Context, ...grpc.CallOption) (trainerv1.Trainer_TrainClient, error)
// Close tears down the ClientConn and all underlying connections.
Close() error
}
// v1 provides v1 version of the trainer grpc function.
type v1 struct {
trainerv1.TrainerClient
*grpc.ClientConn
}
// Train models of scheduler using dataset.
func (v *v1) Train(ctx context.Context, opts ...grpc.CallOption) (trainerv1.Trainer_TrainClient, error) {
return v.TrainerClient.Train(ctx, opts...)
}