/
retry.go
396 lines (361 loc) · 13 KB
/
retry.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
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
// Copyright (c) The go-grpc-middleware Authors.
// Licensed under the Apache License 2.0.
package retry
import (
"context"
"errors"
"fmt"
"io"
"strconv"
"sync"
"time"
"github.com/grpc-ecosystem/go-grpc-middleware/v2/metadata"
"golang.org/x/net/trace"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
grpcMetadata "google.golang.org/grpc/metadata"
"google.golang.org/grpc/status"
)
const (
AttemptMetadataKey = "x-retry-attempt"
)
// UnaryClientInterceptor returns a new retrying unary client interceptor.
//
// The default configuration of the interceptor is to not retry *at all*. This behaviour can be
// changed through options (e.g. WithMax) on creation of the interceptor or on call (through grpc.CallOptions).
func UnaryClientInterceptor(optFuncs ...CallOption) grpc.UnaryClientInterceptor {
intOpts := reuseOrNewWithCallOptions(defaultOptions, optFuncs)
return func(parentCtx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
grpcOpts, retryOpts := filterCallOptions(opts)
callOpts := reuseOrNewWithCallOptions(intOpts, retryOpts)
// short circuit for simplicity, and avoiding allocations.
if callOpts.max == 0 {
return invoker(parentCtx, method, req, reply, cc, grpcOpts...)
}
var lastErr error
for attempt := uint(0); attempt < callOpts.max; attempt++ {
if err := waitRetryBackoff(attempt, parentCtx, callOpts); err != nil {
return err
}
callCtx, cancel := perCallContext(parentCtx, callOpts, attempt)
defer cancel() // Clean up potential resources.
lastErr = invoker(callCtx, method, req, reply, cc, grpcOpts...)
if lastErr == nil {
return nil
}
callOpts.onRetryCallback(parentCtx, attempt, lastErr)
if isContextError(lastErr) {
if parentCtx.Err() != nil {
logTrace(parentCtx, "grpc_retry attempt: %d, parent context error: %v", attempt, parentCtx.Err())
// its the parent context deadline or cancellation.
return lastErr
} else if callOpts.perCallTimeout != 0 {
// We have set a perCallTimeout in the retry middleware, which would result in a context error if
// the deadline was exceeded, in which case try again.
logTrace(parentCtx, "grpc_retry attempt: %d, context error from retry call", attempt)
continue
}
}
if !isRetriable(lastErr, callOpts) {
return lastErr
}
}
return lastErr
}
}
func durationFromTrailer(md grpcMetadata.MD) (time.Duration, error) {
secStrings := md.Get("x-ratelimit-reset")
if len(secStrings) < 1 {
return 0, errors.New("x-ratelimit-reset metadata len < 1")
}
sec, err := strconv.Atoi(secStrings[0])
if err != nil {
return 0, err
}
return time.Second * time.Duration(sec), nil
}
func UnaryClientInterceptorRE(optFuncs ...CallOption) grpc.UnaryClientInterceptor {
intOpts := reuseOrNewWithCallOptions(defaultOptions, optFuncs)
return func(parentCtx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
grpcOpts, retryOpts := filterCallOptions(opts)
callOpts := reuseOrNewWithCallOptions(intOpts, retryOpts)
// short circuit for simplicity, and avoiding allocations.
if callOpts.max == 0 {
return invoker(parentCtx, method, req, reply, cc, grpcOpts...)
}
var lastErr error
var trailer grpcMetadata.MD
grpcOpts = append(grpcOpts, grpc.Trailer(&trailer))
for attempt := uint(0); attempt < callOpts.max; attempt++ {
if err := waitRetryBackoff(attempt, parentCtx, callOpts); err != nil {
return err
}
callCtx, cancel := perCallContext(parentCtx, callOpts, attempt)
defer cancel() // Clean up potential resources.
lastErr = invoker(callCtx, method, req, reply, cc, grpcOpts...)
if lastErr == nil {
return nil
}
switch {
case status.Code(lastErr) == codes.ResourceExhausted:
duration, err := durationFromTrailer(trailer)
// log.Printf("Resource Exhausted, sleep for %v...\n", duration.String())
// callOpts.onRetryCallback(parentCtx, attempt, lastErr)
callOpts.onRetryCallback(parentCtx, uint(duration.Seconds()), lastErr)
if err != nil {
return err
}
retryOpts = append(retryOpts, WithBackoff(BackoffExponential(duration)))
callOpts = reuseOrNewWithCallOptions(intOpts, retryOpts)
case isContextError(lastErr):
if parentCtx.Err() != nil {
logTrace(parentCtx, "grpc_retry attempt: %d, parent context error: %v", attempt, parentCtx.Err())
// its the parent context deadline or cancellation.
return lastErr
} else if callOpts.perCallTimeout != 0 {
// We have set a perCallTimeout in the retry middleware, which would result in a context error if
// the deadline was exceeded, in which case try again.
logTrace(parentCtx, "grpc_retry attempt: %d, context error from retry call", attempt)
continue
}
case !isRetriable(lastErr, callOpts):
return lastErr
}
}
return lastErr
}
}
// StreamClientInterceptor returns a new retrying stream client interceptor for server side streaming calls.
//
// The default configuration of the interceptor is to not retry *at all*. This behaviour can be
// changed through options (e.g. WithMax) on creation of the interceptor or on call (through grpc.CallOptions).
//
// Retry logic is available *only for ServerStreams*, i.e. 1:n streams, as the internal logic needs
// to buffer the messages sent by the client. If retry is enabled on any other streams (ClientStreams,
// BidiStreams), the retry interceptor will fail the call.
func StreamClientInterceptor(optFuncs ...CallOption) grpc.StreamClientInterceptor {
intOpts := reuseOrNewWithCallOptions(defaultOptions, optFuncs)
return func(parentCtx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
grpcOpts, retryOpts := filterCallOptions(opts)
callOpts := reuseOrNewWithCallOptions(intOpts, retryOpts)
// short circuit for simplicity, and avoiding allocations.
if callOpts.max == 0 {
return streamer(parentCtx, desc, cc, method, grpcOpts...)
}
//if desc.ClientStreams {
// return nil, status.Errorf(codes.Unimplemented, "grpc_retry: cannot retry on ClientStreams, set grpc_retry.Disable()")
//}
var lastErr error
for attempt := uint(0); attempt < callOpts.max; attempt++ {
if err := waitRetryBackoff(attempt, parentCtx, callOpts); err != nil {
return nil, err
}
callCtx, cancel := perCallContext(parentCtx, callOpts, 0)
defer cancel() // Clean up potential resources.
var newStreamer grpc.ClientStream
newStreamer, lastErr = streamer(callCtx, desc, cc, method, grpcOpts...)
if lastErr == nil {
retryingStreamer := &retryingStream{
ClientStream: newStreamer,
callOpts: callOpts,
parentCtx: parentCtx,
desc: desc,
streamerCall: func(ctx context.Context) (grpc.ClientStream, error) {
return streamer(ctx, desc, cc, method, grpcOpts...)
},
}
return retryingStreamer, nil
}
callOpts.onRetryCallback(parentCtx, attempt, lastErr)
if isContextError(lastErr) {
if parentCtx.Err() != nil {
logTrace(parentCtx, "grpc_retry attempt: %d, parent context error: %v", attempt, parentCtx.Err())
// its the parent context deadline or cancellation.
return nil, lastErr
} else if callOpts.perCallTimeout != 0 {
// We have set a perCallTimeout in the retry middleware, which would result in a context error if
// the deadline was exceeded, in which case try again.
logTrace(parentCtx, "grpc_retry attempt: %d, context error from retry call", attempt)
continue
}
}
if !isRetriable(lastErr, callOpts) {
return nil, lastErr
}
}
return nil, lastErr
}
}
// type retryingStream is the implementation of grpc.ClientStream that acts as a
// proxy to the underlying call. If any of the RecvMsg() calls fail, it will try to reestablish
// a new ClientStream according to the retry policy.
type retryingStream struct {
grpc.ClientStream
bufferedSends []any // single message that the client can send
wasClosedSend bool // indicates that CloseSend was closed
parentCtx context.Context
callOpts *options
streamerCall func(ctx context.Context) (grpc.ClientStream, error)
mu sync.RWMutex
desc *grpc.StreamDesc
}
func (s *retryingStream) setStream(clientStream grpc.ClientStream) {
s.mu.Lock()
s.ClientStream = clientStream
s.mu.Unlock()
}
func (s *retryingStream) getStream() grpc.ClientStream {
s.mu.RLock()
defer s.mu.RUnlock()
return s.ClientStream
}
func (s *retryingStream) SendMsg(m any) error {
s.mu.Lock()
s.bufferedSends = append(s.bufferedSends, m)
s.mu.Unlock()
return s.getStream().SendMsg(m)
}
func (s *retryingStream) CloseSend() error {
s.mu.Lock()
s.wasClosedSend = true
s.mu.Unlock()
return s.getStream().CloseSend()
}
func (s *retryingStream) Header() (grpcMetadata.MD, error) {
return s.getStream().Header()
}
func (s *retryingStream) Trailer() grpcMetadata.MD {
return s.getStream().Trailer()
}
func (s *retryingStream) RecvMsg(m any) error {
attemptRetry, lastErr := s.receiveMsgAndIndicateRetry(m)
if !attemptRetry {
return lastErr // success or hard failure
}
// We start off from attempt 1, because zeroth was already made on normal SendMsg().
for attempt := uint(1); attempt < s.callOpts.max; attempt++ {
if err := waitRetryBackoff(attempt, s.parentCtx, s.callOpts); err != nil {
return err
}
s.callOpts.onRetryCallback(s.parentCtx, attempt, lastErr)
callCtx, _ := perCallContext(s.parentCtx, s.callOpts, attempt) //nolint
newStream, err := s.reestablishStreamAndResendBuffer(callCtx)
if err != nil {
// Retry dial and transport errors of establishing stream as grpc doesn't retry.
if isRetriable(err, s.callOpts) {
continue
}
return err
}
s.setStream(newStream)
attemptRetry, lastErr = s.receiveMsgAndIndicateRetry(m)
if !attemptRetry {
return lastErr
}
}
return lastErr
}
func (s *retryingStream) receiveMsgAndIndicateRetry(m any) (bool, error) {
err := s.getStream().RecvMsg(m)
if err == nil || err == io.EOF {
return false, err
}
if isContextError(err) {
if s.parentCtx.Err() != nil {
logTrace(s.parentCtx, "grpc_retry parent context error: %v", s.parentCtx.Err())
return false, err
} else if s.callOpts.perCallTimeout != 0 {
// We have set a perCallTimeout in the retry middleware, which would result in a context error if
// the deadline was exceeded, in which case try again.
logTrace(s.parentCtx, "grpc_retry context error from retry call")
return true, err
}
}
return isRetriable(err, s.callOpts), err
}
func (s *retryingStream) reestablishStreamAndResendBuffer(callCtx context.Context) (grpc.ClientStream, error) {
s.mu.RLock()
bufferedSends := s.bufferedSends
s.mu.RUnlock()
newStream, err := s.streamerCall(callCtx)
if err != nil {
logTrace(callCtx, "grpc_retry failed redialing new stream: %v", err)
return nil, err
}
for _, msg := range bufferedSends {
if err := newStream.SendMsg(msg); err != nil {
logTrace(callCtx, "grpc_retry failed resending message: %v", err)
return nil, err
}
}
if !s.desc.ClientStreams {
if err := newStream.CloseSend(); err != nil {
logTrace(callCtx, "grpc_retry failed CloseSend on new stream %v", err)
return nil, err
}
}
return newStream, nil
}
func waitRetryBackoff(attempt uint, parentCtx context.Context, callOpts *options) error {
var waitTime time.Duration = 0
if attempt > 0 {
waitTime = callOpts.backoffFunc(parentCtx, attempt)
}
if waitTime > 0 {
logTrace(parentCtx, "grpc_retry attempt: %d, backoff for %v", attempt, waitTime)
timer := time.NewTimer(waitTime)
select {
case <-parentCtx.Done():
timer.Stop()
return contextErrToGrpcErr(parentCtx.Err())
case <-timer.C:
}
}
return nil
}
func isRetriable(err error, callOpts *options) bool {
errCode := status.Code(err)
if isContextError(err) {
// context errors are not retriable based on user settings.
return false
}
for _, code := range callOpts.codes {
if code == errCode {
return true
}
}
return false
}
func isContextError(err error) bool {
code := status.Code(err)
return code == codes.DeadlineExceeded || code == codes.Canceled
}
func perCallContext(parentCtx context.Context, callOpts *options, attempt uint) (context.Context, context.CancelFunc) {
cancel := context.CancelFunc(func() {})
ctx := parentCtx
if callOpts.perCallTimeout != 0 {
ctx, cancel = context.WithTimeout(ctx, callOpts.perCallTimeout)
}
if attempt > 0 && callOpts.includeHeader {
mdClone := metadata.ExtractOutgoing(ctx).Clone().Set(AttemptMetadataKey, fmt.Sprintf("%d", attempt))
ctx = mdClone.ToOutgoing(ctx)
}
return ctx, cancel
}
func contextErrToGrpcErr(err error) error {
switch err {
case context.DeadlineExceeded:
return status.Errorf(codes.DeadlineExceeded, err.Error())
case context.Canceled:
return status.Errorf(codes.Canceled, err.Error())
default:
return status.Errorf(codes.Unknown, err.Error())
}
}
func logTrace(ctx context.Context, format string, a ...any) {
tr, ok := trace.FromContext(ctx)
if !ok {
return
}
tr.LazyPrintf(format, a...)
}