diff options
Diffstat (limited to 'weed/util/grpc_client_server.go')
| -rw-r--r-- | weed/util/grpc_client_server.go | 9 |
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) |