grpc: Add client-side timeouts

Author: monstermunchkin

grpc/client.go

@@ -19,6 +19,10 @@ import (
 	grpc_retry "github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/retry"
 )
 
+const (
+	timeout = time.Minute
+)
+
 // Deprecated: Use NewClient instead.
 func DialContext(_ context.Context, addr string) (*grpc.ClientConn, error) {
 	return NewClient(addr)
@@ -44,7 +48,11 @@ func NewClient(addr string) (*grpc.ClientConn, error) {
 			grpc_retry.StreamClientInterceptor(opts...),
 			func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
 				start := time.Now()
-				cs, err := streamer(prepareClientContext(ctx), desc, cc, method, opts...)
+
+				ctx, cancel := context.WithTimeout(prepareClientContext(ctx), timeout)
+				defer cancel()
+
+				cs, err := streamer(ctx, desc, cc, method, opts...)
 				log.Ctx(ctx).Debug().Str("method", method).
 					Dur("duration", time.Since(start)).
 					Str("type", "stream").
@@ -58,7 +66,11 @@ func NewClient(addr string) (*grpc.ClientConn, error) {
 			grpc_retry.UnaryClientInterceptor(opts...),
 			func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
 				start := time.Now()
-				err := invoker(prepareClientContext(ctx), method, req, reply, cc, opts...)
+
+				ctx, cancel := context.WithTimeout(prepareClientContext(ctx), timeout)
+				defer cancel()
+
+				err := invoker(ctx, method, req, reply, cc, opts...)
 				log.Ctx(ctx).Debug().Str("method", method).
 					Dur("duration", time.Since(start)).
 					Str("type", "unary").