add context.Context

1 files changed, 5 insertions, 4 deletions

diff --git a/weed/util/grpc_client_server.go b/weed/util/grpc_client_server.go
index b989a35d1..361d245b8 100644
--- a/weed/util/grpc_client_server.go
+++ b/weed/util/grpc_client_server.go
@@ -1,6 +1,7 @@
 package util
 
 import (
+	"context"
 	"fmt"
 	"strconv"
 	"strings"
@@ -33,7 +34,7 @@ func NewGrpcServer(opts ...grpc.ServerOption) *grpc.Server {
 	return grpc.NewServer(options...)
 }
 
-func GrpcDial(address string, opts ...grpc.DialOption) (*grpc.ClientConn, error) {
+func GrpcDial(ctx context.Context, address string, opts ...grpc.DialOption) (*grpc.ClientConn, error) {
 	// opts = append(opts, grpc.WithBlock())
 	// opts = append(opts, grpc.WithTimeout(time.Duration(5*time.Second)))
 	var options []grpc.DialOption
@@ -48,10 +49,10 @@ func GrpcDial(address string, opts ...grpc.DialOption) (*grpc.ClientConn, error)
 			options = append(options, opt)
 		}
 	}
-	return grpc.Dial(address, options...)
+	return grpc.DialContext(ctx, address, options...)
 }
 
-func WithCachedGrpcClient(fn func(*grpc.ClientConn) error, address string, opts ...grpc.DialOption) error {
+func WithCachedGrpcClient(ctx context.Context, fn func(*grpc.ClientConn) error, address string, opts ...grpc.DialOption) error {
 
 	grpcClientsLock.Lock()
 
@@ -61,7 +62,7 @@ func WithCachedGrpcClient(fn func(*grpc.ClientConn) error, address string, opts
 		return fn(existingConnection)
 	}
 
-	grpcConnection, err := GrpcDial(address, opts...)
+	grpcConnection, err := GrpcDial(ctx, address, opts...)
 	if err != nil {
 		grpcClientsLock.Unlock()
 		return fmt.Errorf("fail to dial %s: %v", address, err)