Просмотр исходного кода

add option to close grpc client connection

Alexander Belanger 4 лет назад
Родитель
Сommit
096ecb0fae
1 измененных файлов с 7 добавлено и 0 удалено
  1. 7 0
      provisioner/client/client.go

+ 7 - 0
provisioner/client/client.go

@@ -23,6 +23,8 @@ type Client struct {
 	TokenID    uint
 	HTTPClient *http.Client
 	GRPCClient pb.ProvisionerClient
+
+	conn *grpc.ClientConn
 }
 
 func NewClient(baseURL, token string, tokenID uint) (*Client, error) {
@@ -48,6 +50,7 @@ func NewClient(baseURL, token string, tokenID uint) (*Client, error) {
 			Timeout: time.Minute,
 		},
 		GRPCClient: gClient,
+		conn:       conn,
 	}
 
 	return client, nil
@@ -70,6 +73,10 @@ func (c *Client) NewGRPCContext(workspaceID string) (context.Context, context.Ca
 	return context.WithCancel(ctx)
 }
 
+func (c *Client) CloseConnection() error {
+	return c.conn.Close()
+}
+
 func (c *Client) getRequest(relPath string, data interface{}, response interface{}) error {
 	vals := make(map[string][]string)
 	err := schema.NewEncoder().Encode(data, vals)