From 3a91383f24a9d40104f0ad141e89a7b643b03809 Mon Sep 17 00:00:00 2001
From: Leonard Lyubich <ctulhurider@gmail.com>
Date: Mon, 3 Oct 2022 13:34:22 +0400
Subject: [PATCH] [#419] rpc/client: Use provided context for client dial

In previous implementation `Client` passed `context.Background()` to
`grpc.DialContext` function. This didn't allow to abort dial stage by
the given context.

Base dial context on the one provided with `WithContext` option. Fall
back to using `context.Background` if context is not specified.

Signed-off-by: Leonard Lyubich <ctulhurider@gmail.com>
---
 rpc/client/call_options.go | 14 ++++++++------
 rpc/client/connect.go      | 14 +++++++++-----
 rpc/client/init.go         |  4 ++--
 3 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/rpc/client/call_options.go b/rpc/client/call_options.go
index e443a89f..bf99c929 100644
--- a/rpc/client/call_options.go
+++ b/rpc/client/call_options.go
@@ -2,26 +2,28 @@ package client
 
 import (
 	"context"
-
-	"github.com/nspcc-dev/neofs-api-go/v2/rpc/grpc"
 )
 
 // CallOption is a messaging session option within Protobuf RPC.
 type CallOption func(*callParameters)
 
 type callParameters struct {
-	callOpts []grpc.CallOption
+	ctx context.Context
 }
 
 func defaultCallParameters() *callParameters {
 	return &callParameters{
-		callOpts: make([]grpc.CallOption, 0, 1),
+		ctx: context.Background(),
 	}
 }
 
-// WithContext return options to specify call context.
+// WithContext returns option to specify call context. If provided, all network
+// communications will be based on this context. Otherwise, context.Background()
+// is used.
+//
+// Context SHOULD NOT be nil.
 func WithContext(ctx context.Context) CallOption {
 	return func(prm *callParameters) {
-		prm.callOpts = append(prm.callOpts, grpc.WithContext(ctx))
+		prm.ctx = ctx
 	}
 }
diff --git a/rpc/client/connect.go b/rpc/client/connect.go
index a5403e33..cdbbbea8 100644
--- a/rpc/client/connect.go
+++ b/rpc/client/connect.go
@@ -13,9 +13,10 @@ import (
 	"google.golang.org/grpc/credentials/insecure"
 )
 
-func (c *Client) createGRPCClient() (err error) {
+func (c *Client) createGRPCClient(ctx context.Context) (err error) {
 	c.gRPCClientOnce.Do(func() {
-		if err = c.openGRPCConn(); err != nil {
+		if err = c.openGRPCConn(ctx); err != nil {
+			err = fmt.Errorf("open gRPC connection: %w", err)
 			return
 		}
 
@@ -30,7 +31,7 @@ func (c *Client) createGRPCClient() (err error) {
 
 var errInvalidEndpoint = errors.New("invalid endpoint options")
 
-func (c *Client) openGRPCConn() error {
+func (c *Client) openGRPCConn(ctx context.Context) error {
 	if c.conn != nil {
 		return nil
 	}
@@ -47,15 +48,18 @@ func (c *Client) openGRPCConn() error {
 		creds = insecure.NewCredentials()
 	}
 
-	dialCtx, cancel := context.WithTimeout(context.Background(), c.dialTimeout)
+	dialCtx, cancel := context.WithTimeout(ctx, c.dialTimeout)
 	var err error
+
 	c.conn, err = grpcstd.DialContext(dialCtx, c.addr,
 		grpcstd.WithTransportCredentials(creds),
 		grpcstd.WithBlock(),
 	)
+
 	cancel()
+
 	if err != nil {
-		return fmt.Errorf("open gRPC client connection: %w", err)
+		return fmt.Errorf("gRPC dial: %w", err)
 	}
 
 	return nil
diff --git a/rpc/client/init.go b/rpc/client/init.go
index 27c61998..bad7ed0a 100644
--- a/rpc/client/init.go
+++ b/rpc/client/init.go
@@ -68,11 +68,11 @@ func (g rwGRPC) WriteMessage(m message.Message) error {
 }
 
 func (c *Client) initGRPC(info common.CallMethodInfo, prm *callParameters) (MessageReadWriter, error) {
-	if err := c.createGRPCClient(); err != nil {
+	if err := c.createGRPCClient(prm.ctx); err != nil {
 		return nil, err
 	}
 
-	rw, err := c.gRPCClient.Init(info, prm.callOpts...)
+	rw, err := c.gRPCClient.Init(info, grpc.WithContext(prm.ctx))
 	if err != nil {
 		return nil, err
 	}