frostfs-sdk-go/pool/tree/client.go
Pavel Pogodaev 6ce73790ea
All checks were successful
DCO / DCO (pull_request) Successful in 38s
Tests and linters / Tests (pull_request) Successful in 1m13s
Tests and linters / Lint (pull_request) Successful in 2m36s
[#276] Merge repo with frostfs-api-go
Signed-off-by: Pavel Pogodaev <p.pogodaev@yadro.com>
2024-10-22 14:05:12 +00:00

140 lines
3.2 KiB
Go

package tree
import (
"context"
"crypto/tls"
"errors"
"fmt"
"sync"
apiClient "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/api/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
}
// ErrUnhealthyEndpoint is returned when client in the pool considered unavailable.
var ErrUnhealthyEndpoint = errors.New("unhealthy endpoint")
// 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 = createClient(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 = createClient(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 createClient(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.NewClient(host, opts...)
if err != nil {
return nil, nil, fmt.Errorf("grpc create 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("%w: '%s'", ErrUnhealthyEndpoint, 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()
}