distribution/registry/storage/driver/frostfs/wrapper/pool_wrapper.go

package wrapper

import (
	"context"
	"errors"
	"fmt"
	"io"
	"strings"

	cid "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/container/id"
	treepool "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/pool/tree"
	grpcService "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/pool/tree/service"
	"github.com/distribution/distribution/v3/registry/storage/driver/frostfs/tree"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

type GetNodeByPathResponseInfoWrapper struct {
	response *grpcService.GetNodeByPathResponse_Info
}

func (n GetNodeByPathResponseInfoWrapper) GetNodeID() uint64 {
	return n.response.GetNodeId()
}

func (n GetNodeByPathResponseInfoWrapper) GetParentID() uint64 {
	return n.response.GetParentId()
}

func (n GetNodeByPathResponseInfoWrapper) GetTimestamp() uint64 {
	return n.response.GetTimestamp()
}

func (n GetNodeByPathResponseInfoWrapper) GetMeta() []tree.Meta {
	res := make([]tree.Meta, len(n.response.Meta))
	for i, value := range n.response.Meta {
		res[i] = value
	}
	return res
}

type GetSubTreeResponseBodyWrapper struct {
	response *grpcService.GetSubTreeResponse_Body
}

func (n GetSubTreeResponseBodyWrapper) GetNodeID() uint64 {
	return n.response.GetNodeId()
}

func (n GetSubTreeResponseBodyWrapper) GetParentID() uint64 {
	return n.response.GetParentId()
}

func (n GetSubTreeResponseBodyWrapper) GetTimestamp() uint64 {
	return n.response.GetTimestamp()
}

func (n GetSubTreeResponseBodyWrapper) GetMeta() []tree.Meta {
	res := make([]tree.Meta, len(n.response.Meta))
	for i, value := range n.response.Meta {
		res[i] = value
	}
	return res
}

type PoolWrapper struct {
	p *treepool.Pool
}

func NewPoolWrapper(p *treepool.Pool) *PoolWrapper {
	return &PoolWrapper{p: p}
}

func (w *PoolWrapper) GetNodes(ctx context.Context, prm *tree.GetNodesParams) ([]tree.NodeResponse, error) {
	poolPrm := treepool.GetNodesParams{
		CID:           prm.ContainerID,
		TreeID:        prm.TreeID,
		Path:          prm.Path,
		Meta:          prm.Meta,
		PathAttribute: tree.FileNameKey,
		LatestOnly:    prm.LatestOnly,
		AllAttrs:      prm.AllAttrs,
	}

	nodes, err := w.p.GetNodes(ctx, poolPrm)
	if err != nil {
		return nil, handleError(err)
	}

	res := make([]tree.NodeResponse, len(nodes))
	for i, info := range nodes {
		res[i] = GetNodeByPathResponseInfoWrapper{info}
	}

	return res, nil
}

func (w *PoolWrapper) GetSubTree(ctx context.Context, containerID cid.ID, treeID string, rootID uint64, depth uint32) ([]tree.NodeResponse, error) {
	poolPrm := treepool.GetSubTreeParams{
		CID:    containerID,
		TreeID: treeID,
		RootID: rootID,
		Depth:  depth,
	}

	subTreeReader, err := w.p.GetSubTree(ctx, poolPrm)
	if err != nil {
		return nil, handleError(err)
	}

	var subtree []tree.NodeResponse

	node, err := subTreeReader.Next()
	for err == nil {
		subtree = append(subtree, GetSubTreeResponseBodyWrapper{node})
		node, err = subTreeReader.Next()
	}
	if err != nil && err != io.EOF {
		return nil, handleError(err)
	}

	return subtree, nil
}

func (w *PoolWrapper) AddNodeByPath(ctx context.Context, containerID cid.ID, treeID string, path []string, meta map[string]string) (uint64, error) {
	nodeID, err := w.p.AddNodeByPath(ctx, treepool.AddNodeByPathParams{
		CID:           containerID,
		TreeID:        treeID,
		Path:          path,
		Meta:          meta,
		PathAttribute: tree.FileNameKey,
	})
	return nodeID, handleError(err)
}

func (w *PoolWrapper) RemoveNode(ctx context.Context, containerID cid.ID, treeID string, nodeID uint64) error {
	return handleError(w.p.RemoveNode(ctx, treepool.RemoveNodeParams{
		CID:    containerID,
		TreeID: treeID,
		NodeID: nodeID,
	}))
}

func handleError(err error) error {
	if err == nil {
		return nil
	}
	if errors.Is(err, treepool.ErrNodeNotFound) {
		return fmt.Errorf("%w: %s", tree.ErrNodeNotFound, err.Error())
	}
	if errors.Is(err, treepool.ErrNodeAccessDenied) {
		return fmt.Errorf("%w: %s", tree.ErrNodeAccessDenied, err.Error())
	}
	if isTimeoutError(err) {
		return fmt.Errorf("%w: %s", tree.ErrGatewayTimeout, err.Error())
	}

	return err
}

func unwrapErr(err error) error {
	unwrappedErr := errors.Unwrap(err)
	for unwrappedErr != nil {
		err = unwrappedErr
		unwrappedErr = errors.Unwrap(err)
	}

	return err
}

func isTimeoutError(err error) bool {
	if strings.Contains(err.Error(), "timeout") ||
		errors.Is(err, context.DeadlineExceeded) {
		return true
	}

	return status.Code(unwrapErr(err)) == codes.DeadlineExceeded
}