[#3] Add tests for tree service

Signed-off-by: Roman Loginov <r.loginov@yadro.com>
2024-03-27 16:42:46 +03:00 · 2024-03-27 16:42:46 +03:00 · 7488fe2d43
commit 7488fe2d43
parent ca2da5e23d
3 changed files with 593 additions and 12 deletions
--- a/go.mod
+++ b/go.mod
@ -43,16 +43,23 @@ require (
 )

 require (
+	cloud.google.com/go v0.110.7 // indirect
+	cloud.google.com/go/compute v1.23.0 // indirect
+	cloud.google.com/go/compute/metadata v0.2.3 // indirect
+	cloud.google.com/go/iam v1.1.1 // indirect
 	dario.cat/mergo v1.0.0 // indirect
 	git.frostfs.info/TrueCloudLab/frostfs-api-go/v2 v2.16.1-0.20240327095603-491a47e7fe24 // indirect
 	git.frostfs.info/TrueCloudLab/frostfs-contract v0.0.0-20230307110621-19a8ef2d02fb // indirect
 	git.frostfs.info/TrueCloudLab/frostfs-crypto v0.6.0 // indirect
 	git.frostfs.info/TrueCloudLab/hrw v1.2.1 // indirect
 	git.frostfs.info/TrueCloudLab/rfc6979 v0.4.0 // indirect
+	github.com/Azure/azure-sdk-for-go/sdk/internal v1.3.0 // indirect
 	github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1 // indirect
+	github.com/AzureAD/microsoft-authentication-library-for-go v1.0.0 // indirect
 	github.com/Microsoft/go-winio v0.6.1 // indirect
 	github.com/Microsoft/hcsshim v0.11.4 // indirect
 	github.com/antlr4-go/antlr/v4 v4.13.0 // indirect
+	github.com/benbjohnson/clock v1.1.0 // indirect
 	github.com/containerd/containerd v1.7.12 // indirect
 	github.com/containerd/log v0.1.0 // indirect
 	github.com/cpuguy83/dockercfg v0.3.1 // indirect
@ -63,6 +70,12 @@ require (
 	github.com/docker/go-units v0.5.0 // indirect
 	github.com/go-ole/go-ole v1.2.6 // indirect
 	github.com/gogo/protobuf v1.3.2 // indirect
+	github.com/golang-jwt/jwt/v4 v4.5.0 // indirect
+	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
+	github.com/golang/protobuf v1.5.3 // indirect
+	github.com/google/go-cmp v0.6.0 // indirect
+	github.com/googleapis/enterprise-certificate-proxy v0.2.3 // indirect
+	github.com/googleapis/gax-go/v2 v2.11.0 // indirect
 	github.com/gorilla/websocket v1.5.0 // indirect
 	github.com/hashicorp/golang-lru v0.6.0 // indirect
 	github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 // indirect
@ -88,19 +101,13 @@ require (
 	github.com/yusufpapurcu/wmi v1.2.3 // indirect
 	go.uber.org/atomic v1.10.0 // indirect
 	go.uber.org/multierr v1.11.0 // indirect
-	go.uber.org/zap v1.24.0 // indirect
+	go.uber.org/zap v1.24.0
 	golang.org/x/exp v0.0.0-20230515195305-f3d0a9c9a5cc // indirect
 	golang.org/x/mod v0.16.0 // indirect
 	golang.org/x/tools v0.13.0 // indirect
 )

 require (
-	cloud.google.com/go v0.110.7 // indirect
-	cloud.google.com/go/compute v1.23.0 // indirect
-	cloud.google.com/go/compute/metadata v0.2.3 // indirect
-	cloud.google.com/go/iam v1.1.1 // indirect
-	github.com/Azure/azure-sdk-for-go/sdk/internal v1.3.0 // indirect
-	github.com/AzureAD/microsoft-authentication-library-for-go v1.0.0 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/cenkalti/backoff/v4 v4.2.1 // indirect
 	github.com/cespare/xxhash/v2 v2.2.0 // indirect
@ -109,12 +116,7 @@ require (
 	github.com/felixge/httpsnoop v1.0.4 // indirect
 	github.com/go-logr/logr v1.3.0 // indirect
 	github.com/go-logr/stdr v1.2.2 // indirect
-	github.com/golang-jwt/jwt/v4 v4.5.0 // indirect
-	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
-	github.com/golang/protobuf v1.5.3 // indirect
 	github.com/google/s2a-go v0.1.4 // indirect
-	github.com/googleapis/enterprise-certificate-proxy v0.2.3 // indirect
-	github.com/googleapis/gax-go/v2 v2.11.0 // indirect
 	github.com/grpc-ecosystem/grpc-gateway/v2 v2.16.0 // indirect
 	github.com/hashicorp/golang-lru/v2 v2.0.5 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
--- a/registry/storage/driver/frostfs/tree/tree_client_in_memory.go
+++ b/registry/storage/driver/frostfs/tree/tree_client_in_memory.go
@ -0,0 +1,347 @@
+package tree
+
+import (
+	"context"
+	"fmt"
+	"sort"
+	"time"
+
+	cid "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/container/id"
+)
+
+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 {
+	containerID cid.ID
+	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
+}
+
+func (c *ServiceClientMemory) GetNodes(_ context.Context, p *GetNodesParams) ([]NodeResponse, error) {
+	cnr, ok := c.containers[p.ContainerID.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] = n
+	}
+
+	return res2, nil
+}
+
+func (c *ServiceClientMemory) GetSubTree(_ context.Context, containerID cid.ID, treeID string, rootID uint64, depth uint32) ([]NodeResponse, error) {
+	cnr, ok := c.containers[containerID.EncodeToString()]
+	if !ok {
+		return nil, nil
+	}
+
+	tr, ok := cnr.trees[treeID]
+	if !ok {
+		return nil, ErrNodeNotFound
+	}
+
+	sortNode(tr.treeData)
+
+	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
+}
+
+func newContainerInfo(containerID cid.ID, treeID string) containerInfo {
+	return containerInfo{
+		containerID: containerID,
+		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) AddNodeByPath(_ context.Context, containerID cid.ID, treeID string, path []string, meta map[string]string) (uint64, error) {
+	cnr, ok := c.containers[containerID.EncodeToString()]
+	if !ok {
+		cnr = newContainerInfo(containerID, treeID)
+		c.containers[containerID.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 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, containerID cid.ID, treeID string, nodeID uint64) error {
+	cnr, ok := c.containers[containerID.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
+}
--- a/registry/storage/driver/frostfs/tree/tree_test.go
+++ b/registry/storage/driver/frostfs/tree/tree_test.go
@ -0,0 +1,232 @@
+package tree
+
+import (
+	"context"
+	"testing"
+
+	cidtest "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/container/id/test"
+	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object"
+	oidtest "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object/id/test"
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/zap/zaptest"
+)
+
+func TestGetObjectByPath(t *testing.T) {
+	ctx := context.Background()
+
+	memCli, err := NewTreeServiceClientMemory()
+	require.NoError(t, err)
+	treeService := NewTree(memCli, zaptest.NewLogger(t))
+
+	cidTest := cidtest.ID()
+	oidTest1 := oidtest.ID()
+	oidTest2 := oidtest.ID()
+	testSize1 := uint64(10)
+	testSize2 := uint64(20)
+
+	nodeID1, err := treeService.AddObject(ctx, cidTest, "/a/b", oidTest1, testSize1)
+	require.NoError(t, err)
+
+	nodeID2, err := treeService.AddObject(ctx, cidTest, "/a/b/c/d", oidTest2, testSize2)
+	require.NoError(t, err)
+
+	node1, err := treeService.GetObjectByPath(ctx, cidTest, "/a/b")
+	require.NoError(t, err)
+	require.Equal(t, nodeID1, node1.ID)
+	require.Equal(t, oidTest1, node1.ObjID)
+	require.Equal(t, testSize1, node1.PayloadSize)
+
+	node2, err := treeService.GetObjectByPath(ctx, cidTest, "/a/b/c/d")
+	require.NoError(t, err)
+	require.Equal(t, nodeID2, node2.ID)
+	require.Equal(t, oidTest2, node2.ObjID)
+	require.Equal(t, testSize2, node2.PayloadSize)
+
+	_, err = treeService.GetObjectByPath(ctx, cidTest, "/g")
+	require.ErrorIs(t, err, ErrNodeNotFound)
+
+	_, err = treeService.GetObjectByPath(ctx, cidTest, "/a/b/c")
+	require.ErrorIs(t, err, ErrNodeNotFound)
+}
+
+func TestGetObjectByPathDir(t *testing.T) {
+	ctx := context.Background()
+
+	memCli, err := NewTreeServiceClientMemory()
+	require.NoError(t, err)
+	treeService := NewTree(memCli, zaptest.NewLogger(t))
+
+	cidTest := cidtest.ID()
+	oidTest1 := oidtest.ID()
+	oidTest2 := oidtest.ID()
+	testSize1 := uint64(10)
+	testSize2 := uint64(20)
+
+	nodeID1, err := treeService.AddObject(ctx, cidTest, "/a/b", oidTest1, testSize1)
+	require.NoError(t, err)
+
+	nodeID2, err := treeService.AddObject(ctx, cidTest, "/a/b/c/d", oidTest2, testSize2)
+	require.NoError(t, err)
+
+	node1, err := treeService.GetObjectByPath(ctx, cidTest, "/a/b")
+	require.NoError(t, err)
+	require.Equal(t, nodeID1, node1.ID)
+	require.Equal(t, oidTest1, node1.ObjID)
+	require.Equal(t, testSize1, node1.PayloadSize)
+
+	node2, err := treeService.GetObjectByPath(ctx, cidTest, "/a/b/c/d")
+	require.NoError(t, err)
+	require.Equal(t, nodeID2, node2.ID)
+	require.Equal(t, oidTest2, node2.ObjID)
+	require.Equal(t, testSize2, node2.PayloadSize)
+
+	_, err = treeService.GetObjectByPathDir(ctx, cidTest, "/g")
+	require.ErrorIs(t, err, ErrNodeNotFound)
+
+	_, err = treeService.GetObjectByPathDir(ctx, cidTest, "/a/b/c")
+	require.ErrorIs(t, err, ErrOnlyDirFound)
+
+}
+
+func TestGetListObjectByPrefix(t *testing.T) {
+	ctx := context.Background()
+
+	memCli, err := NewTreeServiceClientMemory()
+	require.NoError(t, err)
+	treeService := NewTree(memCli, zaptest.NewLogger(t))
+
+	cidTest := cidtest.ID()
+	oidTest1 := oidtest.ID()
+	oidTest2 := oidtest.ID()
+	oidTest3 := oidtest.ID()
+	oidTest4 := oidtest.ID()
+	testSize1 := uint64(10)
+	testSize2 := uint64(20)
+	testSize3 := uint64(30)
+	testSize4 := uint64(40)
+
+	_, err = treeService.AddObject(ctx, cidTest, "/a/b", oidTest1, testSize1)
+	require.NoError(t, err)
+
+	nodeID2, err := treeService.AddObject(ctx, cidTest, "/a/b/c", oidTest2, testSize2)
+	require.NoError(t, err)
+
+	nodeID3, err := treeService.AddObject(ctx, cidTest, "/a/b/c/d", oidTest3, testSize3)
+	require.NoError(t, err)
+
+	nodeID4, err := treeService.AddObject(ctx, cidTest, "/a/b/c/d/e", oidTest4, testSize4)
+	require.NoError(t, err)
+
+	nodes, err := treeService.GetListObjectByPrefix(ctx, cidTest, "/a/b")
+	require.NoError(t, err)
+	require.Equal(t, nodeID2, nodes[0].ID)
+	require.Equal(t, nodeID3, nodes[1].ID)
+	require.Equal(t, nodeID4, nodes[2].ID)
+	require.Equal(t, testSize2, nodes[0].PayloadSize)
+	require.Equal(t, testSize3, nodes[1].PayloadSize)
+	require.Equal(t, testSize4, nodes[2].PayloadSize)
+
+	nodes, err = treeService.GetListObjectByPrefix(ctx, cidTest, "/g/s")
+	require.NoError(t, err)
+	require.Equal(t, 0, len(nodes))
+}
+
+func TestGetListOIDBySplitID(t *testing.T) {
+	ctx := context.Background()
+
+	memCli, err := NewTreeServiceClientMemory()
+	require.NoError(t, err)
+	treeService := NewTree(memCli, zaptest.NewLogger(t))
+
+	cidTest := cidtest.ID()
+	oidTest1 := oidtest.ID()
+	oidTest2 := oidtest.ID()
+	oidTest3 := oidtest.ID()
+
+	splitID := object.NewSplitID()
+	splitID.SetUUID(uuid.New())
+
+	_, err = treeService.AddPHYObject(ctx, cidTest, "/a/b", oidTest1, splitID)
+	require.NoError(t, err)
+
+	_, err = treeService.AddPHYObject(ctx, cidTest, "/a/b", oidTest2, splitID)
+	require.NoError(t, err)
+
+	_, err = treeService.AddPHYObject(ctx, cidTest, "/a/b", oidTest3, splitID)
+	require.NoError(t, err)
+
+	ids, err := treeService.GetListOIDBySplitID(ctx, cidTest, "/a/b", splitID)
+	require.NoError(t, err)
+	require.Equal(t, oidTest1, ids[0])
+	require.Equal(t, oidTest2, ids[1])
+	require.Equal(t, oidTest3, ids[2])
+
+	ids, err = treeService.GetListOIDBySplitID(ctx, cidTest, "/c/d", splitID)
+	require.NoError(t, err)
+	require.Equal(t, 0, len(ids))
+}
+
+func TestDeleteObject(t *testing.T) {
+	ctx := context.Background()
+
+	memCli, err := NewTreeServiceClientMemory()
+	require.NoError(t, err)
+	treeService := NewTree(memCli, zaptest.NewLogger(t))
+
+	cidTest := cidtest.ID()
+	oidTest1 := oidtest.ID()
+	testSize := uint64(10)
+
+	nodeID1, err := treeService.AddObject(ctx, cidTest, "/a/b", oidTest1, testSize)
+	require.NoError(t, err)
+
+	err = treeService.DeleteObject(ctx, cidTest, nodeID1)
+	require.NoError(t, err)
+
+	_, err = treeService.GetObjectByPath(ctx, cidTest, "/a/b")
+	require.ErrorIs(t, err, ErrNodeNotFound)
+
+	err = treeService.DeleteObject(ctx, cidTest, nodeID1+1)
+	require.ErrorIs(t, err, ErrNodeNotFound)
+}
+
+func TestDeleteObjectsBySplitID(t *testing.T) {
+	ctx := context.Background()
+
+	memCli, err := NewTreeServiceClientMemory()
+	require.NoError(t, err)
+	treeService := NewTree(memCli, zaptest.NewLogger(t))
+
+	cidTest := cidtest.ID()
+	oidTest1 := oidtest.ID()
+	oidTest2 := oidtest.ID()
+	oidTest3 := oidtest.ID()
+
+	splitID := object.NewSplitID()
+	splitID.SetUUID(uuid.New())
+
+	_, err = treeService.AddPHYObject(ctx, cidTest, "/a/b", oidTest1, splitID)
+	require.NoError(t, err)
+
+	_, err = treeService.AddPHYObject(ctx, cidTest, "/a/b", oidTest2, splitID)
+	require.NoError(t, err)
+
+	_, err = treeService.AddPHYObject(ctx, cidTest, "/a/b", oidTest3, splitID)
+	require.NoError(t, err)
+
+	ids, err := treeService.GetListOIDBySplitID(ctx, cidTest, "/a/b", splitID)
+	require.NoError(t, err)
+	require.Equal(t, 3, len(ids))
+
+	err = treeService.DeleteObjectsBySplitID(ctx, cidTest, "/a/b", splitID)
+	require.NoError(t, err)
+
+	ids, err = treeService.GetListOIDBySplitID(ctx, cidTest, "/a/b", splitID)
+	require.NoError(t, err)
+	require.Equal(t, 0, len(ids))
+
+	ids, err = treeService.GetListOIDBySplitID(ctx, cidTest, "/c/d", splitID)
+	require.NoError(t, err)
+	require.Equal(t, 0, len(ids))
+}