From 388d1ca1de0cd37a874147b81a4ee79dc68c1c81 Mon Sep 17 00:00:00 2001
From: Denis Kirillov <d.kirillov@yadro.com>
Date: Fri, 7 Jul 2023 12:57:53 +0300
Subject: [PATCH] [#107] pool/tree: Support grpc schemas

Signed-off-by: Denis Kirillov <d.kirillov@yadro.com>
---
 pool/tree/client.go | 41 ++++++++++++++++++++++++++++++++---------
 pool/tree/pool.go   |  2 +-
 2 files changed, 33 insertions(+), 10 deletions(-)

diff --git a/pool/tree/client.go b/pool/tree/client.go
index 60377dbf..326a8ccd 100644
--- a/pool/tree/client.go
+++ b/pool/tree/client.go
@@ -2,11 +2,15 @@ package tree
 
 import (
 	"context"
+	"crypto/tls"
 	"fmt"
 	"sync"
 
+	apiClient "git.frostfs.info/TrueCloudLab/frostfs-api-go/v2/rpc/client"
 	grpcService "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/pool/tree/service"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/credentials/insecure"
 )
 
 type treeClient struct {
@@ -35,12 +39,10 @@ func (c *treeClient) dial(ctx context.Context) error {
 	}
 
 	var err error
-	c.conn, err = grpc.DialContext(ctx, c.address, c.opts...)
-	if err != nil {
-		return fmt.Errorf("grpc dial node tree service: %w", err)
+	if c.conn, c.service, err = dialClient(ctx, c.address, c.opts...); err != nil {
+		return err
 	}
 
-	c.service = grpcService.NewTreeServiceClient(c.conn)
 	if _, err = c.service.Healthcheck(ctx, &grpcService.HealthcheckRequest{}); err != nil {
 		return fmt.Errorf("healthcheck tree service: %w", err)
 	}
@@ -55,12 +57,9 @@ func (c *treeClient) redialIfNecessary(ctx context.Context) (healthHasChanged bo
 	defer c.mu.Unlock()
 
 	if c.conn == nil {
-		c.conn, err = grpc.DialContext(ctx, c.address, c.opts...)
-		if err != nil {
-			return false, fmt.Errorf("grpc dial node tree service: %w", err)
+		if c.conn, c.service, err = dialClient(ctx, c.address, c.opts...); err != nil {
+			return false, err
 		}
-
-		c.service = grpcService.NewTreeServiceClient(c.conn)
 	}
 
 	wasHealthy := c.healthy
@@ -74,6 +73,30 @@ func (c *treeClient) redialIfNecessary(ctx context.Context) (healthHasChanged bo
 	return !wasHealthy, nil
 }
 
+func dialClient(ctx context.Context, addr string, clientOptions ...grpc.DialOption) (*grpc.ClientConn, grpcService.TreeServiceClient, error) {
+	host, tlsEnable, err := apiClient.ParseURI(addr)
+	if err != nil {
+		return nil, nil, fmt.Errorf("parse address: %w", err)
+	}
+
+	creds := insecure.NewCredentials()
+	if tlsEnable {
+		creds = credentials.NewTLS(&tls.Config{})
+	}
+
+	options := []grpc.DialOption{grpc.WithTransportCredentials(creds)}
+
+	// the order is matter, we want client to be able to overwrite options.
+	opts := append(options, clientOptions...)
+
+	conn, err := grpc.DialContext(ctx, host, opts...)
+	if err != nil {
+		return nil, nil, fmt.Errorf("grpc dial node tree service: %w", err)
+	}
+
+	return conn, grpcService.NewTreeServiceClient(conn), nil
+}
+
 func (c *treeClient) serviceClient() (grpcService.TreeServiceClient, error) {
 	c.mu.RLock()
 	defer c.mu.RUnlock()
diff --git a/pool/tree/pool.go b/pool/tree/pool.go
index ff155dcb..7fca21be 100644
--- a/pool/tree/pool.go
+++ b/pool/tree/pool.go
@@ -196,7 +196,7 @@ func (p *Pool) Dial(ctx context.Context) error {
 		for j, node := range nodes {
 			clients[j] = newTreeClient(node.Address(), p.dialOptions...)
 			if err := clients[j].dial(ctx); err != nil {
-				p.log(zap.WarnLevel, "failed to build client", zap.String("address", node.Address()), zap.Error(err))
+				p.log(zap.WarnLevel, "failed to dial tree client", zap.String("address", node.Address()), zap.Error(err))
 				continue
 			}