frostfs-s3-gw/pkg/service/tree/tree_client_in_memory.go
Denis Kirillov 377fa127b5 [#370] Fix fetching attributes from tree
Port TrueCloudLab/frostfs-s3-gw#374

Signed-off-by: Denis Kirillov <d.kirillov@yadro.com>
2024-04-19 17:33:19 +03:00

482 lines
9.8 KiB
Go

package tree
import (
"context"
"fmt"
"io"
"sort"
"time"
"git.frostfs.info/TrueCloudLab/frostfs-s3-gw/api/data"
"golang.org/x/exp/slices"
)
type nodeMeta struct {
key string
value []byte
}
func (m nodeMeta) GetKey() string {
return m.key
}
func (m nodeMeta) GetValue() []byte {
return m.value
}
type nodeResponse struct {
meta []nodeMeta
nodeID uint64
parentID uint64
timestamp uint64
}
func (n nodeResponse) GetNodeID() uint64 {
return n.nodeID
}
func (n nodeResponse) GetParentID() uint64 {
return n.parentID
}
func (n nodeResponse) GetTimestamp() uint64 {
return n.timestamp
}
func (n nodeResponse) GetMeta() []Meta {
res := make([]Meta, len(n.meta))
for i, value := range n.meta {
res[i] = value
}
return res
}
func (n nodeResponse) getValue(key string) string {
for _, value := range n.meta {
if value.key == key {
return string(value.value)
}
}
return ""
}
type ServiceClientMemory struct {
containers map[string]containerInfo
}
type containerInfo struct {
bkt *data.BucketInfo
trees map[string]memoryTree
}
type memoryTree struct {
idCounter uint64
treeData *treeNodeMemory
}
type treeNodeMemory struct {
data nodeResponse
parent *treeNodeMemory
children []*treeNodeMemory
}
func (t *treeNodeMemory) getNode(nodeID uint64) *treeNodeMemory {
if t.data.nodeID == nodeID {
return t
}
for _, child := range t.children {
if node := child.getNode(nodeID); node != nil {
return node
}
}
return nil
}
func (t *memoryTree) getNodesByPath(path []string) []nodeResponse {
if len(path) == 0 {
return nil
}
var res []nodeResponse
for _, child := range t.treeData.children {
res = child.listNodesByPath(res, path)
}
return res
}
func (t *treeNodeMemory) listNodesByPath(res []nodeResponse, path []string) []nodeResponse {
if len(path) == 0 || t.data.getValue(FileNameKey) != path[0] {
return res
}
if len(path) == 1 {
return append(res, t.data)
}
for _, ch := range t.children {
res = ch.listNodesByPath(res, path[1:])
}
return res
}
func (t *memoryTree) createPathIfNotExist(parent *treeNodeMemory, path []string) *treeNodeMemory {
if len(path) == 0 {
return parent
}
var node *treeNodeMemory
for _, child := range parent.children {
if len(child.data.meta) == 1 && child.data.getValue(FileNameKey) == path[0] {
node = child
break
}
}
if node == nil {
node = &treeNodeMemory{
data: nodeResponse{
meta: []nodeMeta{{key: FileNameKey, value: []byte(path[0])}},
nodeID: t.idCounter,
parentID: parent.data.nodeID,
timestamp: uint64(time.Now().UnixMicro()),
},
parent: parent,
}
t.idCounter++
parent.children = append(parent.children, node)
}
return t.createPathIfNotExist(node, path[1:])
}
func (t *treeNodeMemory) removeChild(nodeID uint64) {
ind := -1
for i, ch := range t.children {
if ch.data.nodeID == nodeID {
ind = i
break
}
}
if ind != -1 {
t.children = append(t.children[:ind], t.children[ind+1:]...)
}
}
func (t *treeNodeMemory) listNodes(res []NodeResponse, depth uint32) []NodeResponse {
res = append(res, t.data)
if depth == 0 {
return res
}
for _, ch := range t.children {
res = ch.listNodes(res, depth-1)
}
return res
}
func NewTreeServiceClientMemory() (*ServiceClientMemory, error) {
return &ServiceClientMemory{
containers: make(map[string]containerInfo),
}, nil
}
type nodeResponseWrapper struct {
nodeResponse
allAttr bool
attrs []string
}
func (n nodeResponseWrapper) GetMeta() []Meta {
res := make([]Meta, 0, len(n.meta))
for _, value := range n.meta {
if n.allAttr || slices.Contains(n.attrs, value.key) {
res = append(res, value)
}
}
return res
}
func (c *ServiceClientMemory) GetNodes(_ context.Context, p *GetNodesParams) ([]NodeResponse, error) {
cnr, ok := c.containers[p.BktInfo.CID.EncodeToString()]
if !ok {
return nil, nil
}
tr, ok := cnr.trees[p.TreeID]
if !ok {
return nil, nil
}
res := tr.getNodesByPath(p.Path)
sort.Slice(res, func(i, j int) bool {
return res[i].timestamp < res[j].timestamp
})
if p.LatestOnly && len(res) != 0 {
res = res[len(res)-1:]
}
res2 := make([]NodeResponse, len(res))
for i, n := range res {
res2[i] = nodeResponseWrapper{
nodeResponse: n,
allAttr: p.AllAttrs,
attrs: p.Meta,
}
}
return res2, nil
}
func (c *ServiceClientMemory) GetSubTree(_ context.Context, bktInfo *data.BucketInfo, treeID string, rootID uint64, depth uint32) ([]NodeResponse, error) {
cnr, ok := c.containers[bktInfo.CID.EncodeToString()]
if !ok {
return nil, nil
}
tr, ok := cnr.trees[treeID]
if !ok {
return nil, ErrNodeNotFound
}
node := tr.treeData.getNode(rootID)
if node == nil {
return nil, ErrNodeNotFound
}
// we depth-1 in case of uint32 and 0 as mark to get all subtree leads to overflow and depth is getting quite big to walk all tree levels
return node.listNodes(nil, depth-1), nil
}
type SubTreeStreamMemoryImpl struct {
res []NodeResponse
offset int
err error
}
func (s *SubTreeStreamMemoryImpl) Next() (NodeResponse, error) {
if s.err != nil {
return nil, s.err
}
if s.offset > len(s.res)-1 {
return nil, io.EOF
}
s.offset++
return s.res[s.offset-1], nil
}
func (c *ServiceClientMemory) GetSubTreeStream(_ context.Context, bktInfo *data.BucketInfo, treeID string, rootID uint64, depth uint32) (SubTreeStream, error) {
cnr, ok := c.containers[bktInfo.CID.EncodeToString()]
if !ok {
return &SubTreeStreamMemoryImpl{err: ErrNodeNotFound}, nil
}
tr, ok := cnr.trees[treeID]
if !ok {
return nil, ErrNodeNotFound
}
node := tr.treeData.getNode(rootID)
if node == nil {
return nil, ErrNodeNotFound
}
sortNode(tr.treeData)
return &SubTreeStreamMemoryImpl{
res: node.listNodes(nil, depth-1),
offset: 0,
}, nil
}
func newContainerInfo(bktInfo *data.BucketInfo, treeID string) containerInfo {
return containerInfo{
bkt: bktInfo,
trees: map[string]memoryTree{
treeID: {
idCounter: 1,
treeData: &treeNodeMemory{
data: nodeResponse{
timestamp: uint64(time.Now().UnixMicro()),
},
},
},
},
}
}
func newMemoryTree() memoryTree {
return memoryTree{
idCounter: 1,
treeData: &treeNodeMemory{
data: nodeResponse{
timestamp: uint64(time.Now().UnixMicro()),
},
},
}
}
func (c *ServiceClientMemory) AddNode(ctx context.Context, bktInfo *data.BucketInfo, treeID string, parent uint64, meta map[string]string) (uint64, error) {
return c.AddNodeBase(ctx, bktInfo, treeID, parent, meta, true)
}
func (c *ServiceClientMemory) AddNodeBase(_ context.Context, bktInfo *data.BucketInfo, treeID string, parent uint64, meta map[string]string, needSort bool) (uint64, error) {
cnr, ok := c.containers[bktInfo.CID.EncodeToString()]
if !ok {
cnr = newContainerInfo(bktInfo, treeID)
c.containers[bktInfo.CID.EncodeToString()] = cnr
}
tr, ok := cnr.trees[treeID]
if !ok {
tr = newMemoryTree()
cnr.trees[treeID] = tr
}
parentNode := tr.treeData.getNode(parent)
if parentNode == nil {
return 0, ErrNodeNotFound
}
newID := tr.idCounter
tr.idCounter++
tn := &treeNodeMemory{
data: nodeResponse{
meta: metaToNodeMeta(meta),
nodeID: newID,
parentID: parent,
timestamp: uint64(time.Now().UnixMicro()),
},
parent: parentNode,
}
parentNode.children = append(parentNode.children, tn)
if needSort {
sortNodes(parentNode.children)
}
cnr.trees[treeID] = tr
return newID, nil
}
func (c *ServiceClientMemory) AddNodeByPath(_ context.Context, bktInfo *data.BucketInfo, treeID string, path []string, meta map[string]string) (uint64, error) {
cnr, ok := c.containers[bktInfo.CID.EncodeToString()]
if !ok {
cnr = newContainerInfo(bktInfo, treeID)
c.containers[bktInfo.CID.EncodeToString()] = cnr
}
tr, ok := cnr.trees[treeID]
if !ok {
tr = newMemoryTree()
cnr.trees[treeID] = tr
}
parentNode := tr.createPathIfNotExist(tr.treeData, path)
if parentNode == nil {
return 0, fmt.Errorf("create path '%s'", path)
}
newID := tr.idCounter
tr.idCounter++
tn := &treeNodeMemory{
data: nodeResponse{
meta: metaToNodeMeta(meta),
nodeID: newID,
parentID: parentNode.data.nodeID,
timestamp: uint64(time.Now().UnixMicro()),
},
parent: parentNode,
}
parentNode.children = append(parentNode.children, tn)
cnr.trees[treeID] = tr
return newID, nil
}
func (c *ServiceClientMemory) MoveNode(_ context.Context, bktInfo *data.BucketInfo, treeID string, nodeID, parentID uint64, meta map[string]string) error {
cnr, ok := c.containers[bktInfo.CID.EncodeToString()]
if !ok {
return ErrNodeNotFound
}
tr, ok := cnr.trees[treeID]
if !ok {
return ErrNodeNotFound
}
node := tr.treeData.getNode(nodeID)
if node == nil {
return ErrNodeNotFound
}
newParent := tr.treeData.getNode(parentID)
if newParent == nil {
return ErrNodeNotFound
}
node.data.meta = metaToNodeMeta(meta)
node.data.parentID = parentID
newParent.children = append(newParent.children, node)
node.parent.removeChild(nodeID)
return nil
}
func sortNode(node *treeNodeMemory) {
if node == nil {
return
}
sortNodes(node.children)
for _, child := range node.children {
sortNode(child)
}
}
func sortNodes(list []*treeNodeMemory) {
sort.Slice(list, func(i, j int) bool {
return list[i].data.getValue(FileNameKey) < list[j].data.getValue(FileNameKey)
})
}
func (c *ServiceClientMemory) RemoveNode(_ context.Context, bktInfo *data.BucketInfo, treeID string, nodeID uint64) error {
cnr, ok := c.containers[bktInfo.CID.EncodeToString()]
if !ok {
return ErrNodeNotFound
}
tr, ok := cnr.trees[treeID]
if !ok {
return ErrNodeNotFound
}
node := tr.treeData.getNode(nodeID)
if node == nil {
return ErrNodeNotFound
}
node.parent.removeChild(nodeID)
return nil
}
func metaToNodeMeta(m map[string]string) []nodeMeta {
result := make([]nodeMeta, 0, len(m))
for key, value := range m {
result = append(result, nodeMeta{key: key, value: []byte(value)})
}
return result
}