go-zero/zrpc/internal/clientinterceptors/timeoutinterceptor.go
2021-07-24 22:29:02 +08:00

24 lines
592 B
Go

package clientinterceptors
import (
"context"
"time"
"google.golang.org/grpc"
)
// TimeoutInterceptor is an interceptor that controls timeout.
func TimeoutInterceptor(timeout time.Duration) grpc.UnaryClientInterceptor {
return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn,
invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
if timeout <= 0 {
return invoker(ctx, method, req, reply, cc, opts...)
}
ctx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
return invoker(ctx, method, req, reply, cc, opts...)
}
}