-
Notifications
You must be signed in to change notification settings - Fork 280
/
interceptor.go
157 lines (131 loc) · 5.61 KB
/
interceptor.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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
/*
* Copyright 2022 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.
*/
package rpc
import (
"context"
"sync"
"github.com/juju/ratelimit"
"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
"d7y.io/dragonfly/v2/internal/dferrors"
)
var (
// otelUnaryInterceptor is the unary interceptor for tracing.
otelUnaryInterceptor grpc.UnaryClientInterceptor
// otelStreamInterceptor is the stream interceptor for tracing.
otelStreamInterceptor grpc.StreamClientInterceptor
// interceptorsInitialized is used to ensure that otel interceptors are initialized only once.
interceptorsInitialized = sync.Once{}
)
// OTEL interceptors must be created once to avoid memory leak,
// refer to https://github.com/open-telemetry/opentelemetry-go-contrib/issues/4226 and
// https://github.com/argoproj/argo-cd/pull/15174.
func ensureOTELInterceptorInitialized() {
interceptorsInitialized.Do(func() {
otelUnaryInterceptor = otelgrpc.UnaryClientInterceptor()
otelStreamInterceptor = otelgrpc.StreamClientInterceptor()
})
}
// OTELUnaryClientInterceptor returns a new unary client interceptor that traces gRPC requests.
func OTELUnaryClientInterceptor() grpc.UnaryClientInterceptor {
ensureOTELInterceptorInitialized()
return otelUnaryInterceptor
}
// OTELStreamClientInterceptor returns a new stream client interceptor that traces gRPC requests.
func OTELStreamClientInterceptor() grpc.StreamClientInterceptor {
ensureOTELInterceptorInitialized()
return otelStreamInterceptor
}
// Refresher is the interface for refreshing dynconfig.
type Refresher interface {
Refresh() error
}
// RefresherUnaryClientInterceptor returns a new unary client interceptor that refresh dynconfig addresses when calling error.
func RefresherUnaryClientInterceptor(r Refresher) grpc.UnaryClientInterceptor {
return func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
err := invoker(ctx, method, req, reply, cc, opts...)
if s, ok := status.FromError(err); ok {
if s.Code() == codes.ResourceExhausted || s.Code() == codes.Unavailable {
// nolint
r.Refresh()
}
}
return err
}
}
// RefresherStreamClientInterceptor returns a new stream client interceptor that refresh dynconfig addresses when calling error.
func RefresherStreamClientInterceptor(r Refresher) grpc.StreamClientInterceptor {
return func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
clientStream, err := streamer(ctx, desc, cc, method, opts...)
if s, ok := status.FromError(err); ok {
if s.Code() == codes.ResourceExhausted || s.Code() == codes.Unavailable {
// nolint
r.Refresh()
}
}
return clientStream, err
}
}
// RateLimiterInterceptor is the interface for ratelimit interceptor.
type RateLimiterInterceptor struct {
// tokenBucket is token bucket of ratelimit.
tokenBucket *ratelimit.Bucket
}
// NewRateLimiterInterceptor returns a RateLimiterInterceptor instance.
func NewRateLimiterInterceptor(qps float64, burst int64) *RateLimiterInterceptor {
return &RateLimiterInterceptor{
tokenBucket: ratelimit.NewBucketWithRate(qps, burst),
}
}
// Limit is the predicate which limits the requests.
func (r *RateLimiterInterceptor) Limit() bool {
if r.tokenBucket.TakeAvailable(1) == 0 {
return true
}
return false
}
// ConvertErrorUnaryServerInterceptor returns a new unary server interceptor that convert error when trigger custom error.
func ConvertErrorUnaryServerInterceptor(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
h, err := handler(ctx, req)
if err != nil {
return h, dferrors.ConvertDfErrorToGRPCError(err)
}
return h, nil
}
// ConvertErrorStreamServerInterceptor returns a new stream server interceptor that convert error when trigger custom error.
func ConvertErrorStreamServerInterceptor(srv any, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
if err := handler(srv, ss); err != nil {
return dferrors.ConvertDfErrorToGRPCError(err)
}
return nil
}
// ConvertErrorUnaryClientInterceptor returns a new unary client interceptor that convert error when trigger custom error.
func ConvertErrorUnaryClientInterceptor(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
if err := invoker(ctx, method, req, reply, cc, opts...); err != nil {
return dferrors.ConvertGRPCErrorToDfError(err)
}
return nil
}
// ConvertErrorStreamClientInterceptor returns a new stream client interceptor that convert error when trigger custom error.
func ConvertErrorStreamClientInterceptor(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
s, err := streamer(ctx, desc, cc, method, opts...)
if err != nil {
return nil, dferrors.ConvertGRPCErrorToDfError(err)
}
return s, nil
}