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