frostfs-sdk-go/pool/tree/client.go

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 {
	mu      sync.RWMutex
	address string
	opts    []grpc.DialOption
	conn    *grpc.ClientConn
	service grpcService.TreeServiceClient
	healthy bool
}

// newTreeClient creates new tree client with auto dial.
func newTreeClient(addr string, opts ...grpc.DialOption) *treeClient {
	return &treeClient{
		address: addr,
		opts:    opts,
	}
}

func (c *treeClient) dial(ctx context.Context) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.conn != nil {
		return fmt.Errorf("couldn't dial '%s': connection already established", c.address)
	}

	var err error
	if c.conn, c.service, err = dialClient(ctx, c.address, c.opts...); err != nil {
		return err
	}

	if _, err = c.service.Healthcheck(ctx, &grpcService.HealthcheckRequest{}); err != nil {
		return fmt.Errorf("healthcheck tree service: %w", err)
	}

	c.healthy = true

	return nil
}

func (c *treeClient) redialIfNecessary(ctx context.Context) (healthHasChanged bool, err error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.conn == nil {
		if c.conn, c.service, err = dialClient(ctx, c.address, c.opts...); err != nil {
			return false, err
		}
	}

	wasHealthy := c.healthy
	if _, err = c.service.Healthcheck(ctx, &grpcService.HealthcheckRequest{}); err != nil {
		c.healthy = false
		return wasHealthy, fmt.Errorf("healthcheck tree service: %w", err)
	}

	c.healthy = true

	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()

	if c.conn == nil || !c.healthy {
		return nil, fmt.Errorf("unhealthy endpoint: '%s'", c.address)
	}

	return c.service, nil
}

func (c *treeClient) endpoint() string {
	return c.address
}

func (c *treeClient) isHealthy() bool {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.healthy
}

func (c *treeClient) setHealthy(val bool) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.healthy = val
}

func (c *treeClient) close() error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.conn == nil {
		return nil
	}

	return c.conn.Close()
}