git.frostfs.info/TrueCloudLab/frostfs-sdk-go@v0.0.0-20241022124111-5361f0ecebd3/pool/tree/client.go (about)

     1  package tree
     2  
     3  import (
     4  	"context"
     5  	"crypto/tls"
     6  	"errors"
     7  	"fmt"
     8  	"sync"
     9  
    10  	apiClient "git.frostfs.info/TrueCloudLab/frostfs-api-go/v2/rpc/client"
    11  	grpcService "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/pool/tree/service"
    12  	"google.golang.org/grpc"
    13  	"google.golang.org/grpc/credentials"
    14  	"google.golang.org/grpc/credentials/insecure"
    15  )
    16  
    17  type treeClient struct {
    18  	mu      sync.RWMutex
    19  	address string
    20  	opts    []grpc.DialOption
    21  	conn    *grpc.ClientConn
    22  	service grpcService.TreeServiceClient
    23  	healthy bool
    24  }
    25  
    26  // ErrUnhealthyEndpoint is returned when client in the pool considered unavailable.
    27  var ErrUnhealthyEndpoint = errors.New("unhealthy endpoint")
    28  
    29  // newTreeClient creates new tree client with auto dial.
    30  func newTreeClient(addr string, opts ...grpc.DialOption) *treeClient {
    31  	return &treeClient{
    32  		address: addr,
    33  		opts:    opts,
    34  	}
    35  }
    36  
    37  func (c *treeClient) dial(ctx context.Context) error {
    38  	c.mu.Lock()
    39  	defer c.mu.Unlock()
    40  
    41  	if c.conn != nil {
    42  		return fmt.Errorf("couldn't dial '%s': connection already established", c.address)
    43  	}
    44  
    45  	var err error
    46  	if c.conn, c.service, err = createClient(c.address, c.opts...); err != nil {
    47  		return err
    48  	}
    49  
    50  	if _, err = c.service.Healthcheck(ctx, &grpcService.HealthcheckRequest{}); err != nil {
    51  		return fmt.Errorf("healthcheck tree service: %w", err)
    52  	}
    53  
    54  	c.healthy = true
    55  
    56  	return nil
    57  }
    58  
    59  func (c *treeClient) redialIfNecessary(ctx context.Context) (healthHasChanged bool, err error) {
    60  	c.mu.Lock()
    61  	defer c.mu.Unlock()
    62  
    63  	if c.conn == nil {
    64  		if c.conn, c.service, err = createClient(c.address, c.opts...); err != nil {
    65  			return false, err
    66  		}
    67  	}
    68  
    69  	wasHealthy := c.healthy
    70  	if _, err = c.service.Healthcheck(ctx, &grpcService.HealthcheckRequest{}); err != nil {
    71  		c.healthy = false
    72  		return wasHealthy, fmt.Errorf("healthcheck tree service: %w", err)
    73  	}
    74  
    75  	c.healthy = true
    76  
    77  	return !wasHealthy, nil
    78  }
    79  
    80  func createClient(addr string, clientOptions ...grpc.DialOption) (*grpc.ClientConn, grpcService.TreeServiceClient, error) {
    81  	host, tlsEnable, err := apiClient.ParseURI(addr)
    82  	if err != nil {
    83  		return nil, nil, fmt.Errorf("parse address: %w", err)
    84  	}
    85  
    86  	creds := insecure.NewCredentials()
    87  	if tlsEnable {
    88  		creds = credentials.NewTLS(&tls.Config{})
    89  	}
    90  
    91  	options := []grpc.DialOption{grpc.WithTransportCredentials(creds)}
    92  
    93  	// the order is matter, we want client to be able to overwrite options.
    94  	opts := append(options, clientOptions...)
    95  
    96  	conn, err := grpc.NewClient(host, opts...)
    97  	if err != nil {
    98  		return nil, nil, fmt.Errorf("grpc create node tree service: %w", err)
    99  	}
   100  
   101  	return conn, grpcService.NewTreeServiceClient(conn), nil
   102  }
   103  
   104  func (c *treeClient) serviceClient() (grpcService.TreeServiceClient, error) {
   105  	c.mu.RLock()
   106  	defer c.mu.RUnlock()
   107  
   108  	if c.conn == nil || !c.healthy {
   109  		return nil, fmt.Errorf("%w: '%s'", ErrUnhealthyEndpoint, c.address)
   110  	}
   111  
   112  	return c.service, nil
   113  }
   114  
   115  func (c *treeClient) endpoint() string {
   116  	return c.address
   117  }
   118  
   119  func (c *treeClient) isHealthy() bool {
   120  	c.mu.RLock()
   121  	defer c.mu.RUnlock()
   122  	return c.healthy
   123  }
   124  
   125  func (c *treeClient) setHealthy(val bool) {
   126  	c.mu.Lock()
   127  	defer c.mu.Unlock()
   128  	c.healthy = val
   129  }
   130  
   131  func (c *treeClient) close() error {
   132  	c.mu.Lock()
   133  	defer c.mu.Unlock()
   134  
   135  	if c.conn == nil {
   136  		return nil
   137  	}
   138  
   139  	return c.conn.Close()
   140  }