/
timeoutinterceptor.go
47 lines (38 loc) · 1.15 KB
/
timeoutinterceptor.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
package clientinterceptors
import (
"context"
"time"
"google.golang.org/grpc"
)
// TimeoutCallOption is a call option that controls timeout.
type TimeoutCallOption struct {
grpc.EmptyCallOption
timeout time.Duration
}
// TimeoutInterceptor is an interceptor that controls timeout.
func TimeoutInterceptor(timeout time.Duration) grpc.UnaryClientInterceptor {
return func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn,
invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
t := getTimeoutFromCallOptions(opts, timeout)
if t <= 0 {
return invoker(ctx, method, req, reply, cc, opts...)
}
ctx, cancel := context.WithTimeout(ctx, t)
defer cancel()
return invoker(ctx, method, req, reply, cc, opts...)
}
}
// WithCallTimeout returns a call option that controls method call timeout.
func WithCallTimeout(timeout time.Duration) grpc.CallOption {
return TimeoutCallOption{
timeout: timeout,
}
}
func getTimeoutFromCallOptions(opts []grpc.CallOption, defaultTimeout time.Duration) time.Duration {
for _, opt := range opts {
if o, ok := opt.(TimeoutCallOption); ok {
return o.timeout
}
}
return defaultTimeout
}