[#1326] services/tree: Implement GetSubTree RPC

Signed-off-by: Evgenii Stratonikov <evgeniy@nspcc.ru>

pkg/local_object_storage/engine/tree.go

@@ -76,3 +76,17 @@ func (e *StorageEngine) TreeGetMeta(cid cidSDK.ID, treeID string, nodeID piloram
 	}
 	return pilorama.Meta{}, err
 }
+
+// TreeGetChildren implements the pilorama.Forest interface.
+func (e *StorageEngine) TreeGetChildren(cid cidSDK.ID, treeID string, nodeID pilorama.Node) ([]uint64, error) {
+	var err error
+	var nodes []uint64
+	for _, sh := range e.sortShardsByWeight(cid) {
+		nodes, err = sh.TreeGetChildren(cid, treeID, nodeID)
+		if err != nil {
+			continue
+		}
+		return nodes, nil
+	}
+	return nil, err
+}

pkg/local_object_storage/pilorama/boltdb.go

@@ -416,6 +416,31 @@ func (t *boltForest) TreeGetMeta(cid cidSDK.ID, treeID string, nodeID Node) (Met
 	return m, err
 }
 
+// TreeGetChildren implements the Forest interface.
+func (t *boltForest) TreeGetChildren(cid cidSDK.ID, treeID string, nodeID Node) ([]uint64, error) {
+	key := make([]byte, 9)
+	key[0] = 'c'
+	binary.LittleEndian.PutUint64(key[1:], nodeID)
+
+	var children []uint64
+
+	err := t.db.View(func(tx *bbolt.Tx) error {
+		treeRoot := tx.Bucket(bucketName(cid, treeID))
+		if treeRoot == nil {
+			return ErrTreeNotFound
+		}
+
+		b := treeRoot.Bucket(dataBucket)
+		c := b.Cursor()
+		for k, _ := c.Seek(key); len(k) == 17 && binary.LittleEndian.Uint64(k[1:]) == nodeID; k, _ = c.Next() {
+			children = append(children, binary.LittleEndian.Uint64(k[9:]))
+		}
+		return nil
+	})
+
+	return children, err
+}
+
 func (t *boltForest) getPathPrefix(bTree *bbolt.Bucket, attr string, path []string) (int, Node, error) {
 	var key [9]byte
 

pkg/local_object_storage/pilorama/forest.go

@@ -112,3 +112,21 @@ func (f *memoryForest) TreeGetMeta(cid cidSDK.ID, treeID string, nodeID Node) (M
 
 	return s.getMeta(nodeID), nil
 }
+
+// TreeGetChildren implements the Forest interface.
+func (f *memoryForest) TreeGetChildren(cid cidSDK.ID, treeID string, nodeID Node) ([]uint64, error) {
+	fullID := cid.String() + "/" + treeID
+	s, ok := f.treeMap[fullID]
+	if !ok {
+		return nil, ErrTreeNotFound
+	}
+
+	children, ok := s.childMap[nodeID]
+	if !ok {
+		return nil, nil
+	}
+
+	res := make([]Node, len(children))
+	copy(res, children)
+	return res, nil
+}

pkg/local_object_storage/pilorama/forest_test.go

@@ -98,6 +98,62 @@ func testForestTreeMove(t *testing.T, s Forest) {
 	})
 }
 
+func TestMemoryForest_TreeGetChildren(t *testing.T) {
+	for i := range providers {
+		t.Run(providers[i].name, func(t *testing.T) {
+			testForestTreeGetChildren(t, providers[i].construct(t))
+		})
+	}
+}
+
+func testForestTreeGetChildren(t *testing.T, s Forest) {
+	cid := cidtest.ID()
+	treeID := "version"
+
+	treeAdd := func(t *testing.T, child, parent Node) {
+		_, err := s.TreeMove(cid, treeID, &Move{
+			Parent: parent,
+			Child:  child,
+		})
+		require.NoError(t, err)
+	}
+
+	// 0
+	// |- 10
+	// |  |- 3
+	// |  |- 6
+	// |     |- 11
+	// |- 2
+	// |- 7
+	treeAdd(t, 10, 0)
+	treeAdd(t, 3, 10)
+	treeAdd(t, 6, 10)
+	treeAdd(t, 11, 6)
+	treeAdd(t, 2, 0)
+	treeAdd(t, 7, 0)
+
+	testGetChildren := func(t *testing.T, nodeID Node, expected []Node) {
+		actual, err := s.TreeGetChildren(cid, treeID, nodeID)
+		require.NoError(t, err)
+		require.ElementsMatch(t, expected, actual)
+	}
+
+	testGetChildren(t, 0, []uint64{10, 2, 7})
+	testGetChildren(t, 10, []uint64{3, 6})
+	testGetChildren(t, 3, nil)
+	testGetChildren(t, 6, []uint64{11})
+	testGetChildren(t, 11, nil)
+	testGetChildren(t, 2, nil)
+	testGetChildren(t, 7, nil)
+	t.Run("missing node", func(t *testing.T) {
+		testGetChildren(t, 42, nil)
+	})
+	t.Run("missing tree", func(t *testing.T) {
+		_, err := s.TreeGetChildren(cid, treeID+"123", 0)
+		require.ErrorIs(t, err, ErrTreeNotFound)
+	})
+}
+
 func TestForest_TreeAdd(t *testing.T) {
 	for i := range providers {
 		t.Run(providers[i].name, func(t *testing.T) {

pkg/local_object_storage/pilorama/interface.go

@@ -20,6 +20,8 @@ type Forest interface {
 	TreeGetByPath(cid cidSDK.ID, treeID string, attr string, path []string, latest bool) ([]Node, error)
 	// TreeGetMeta returns meta information of the node with the specified ID.
 	TreeGetMeta(cid cidSDK.ID, treeID string, nodeID Node) (Meta, error)
+	// TreeGetChildren returns children of the node with the specified ID. The order is arbitrary.
+	TreeGetChildren(cid cidSDK.ID, treeID string, nodeID Node) ([]uint64, error)
 }
 
 type ForestStorage interface {

pkg/local_object_storage/shard/tree.go

@@ -31,3 +31,8 @@ func (s *Shard) TreeGetByPath(cid cidSDK.ID, treeID string, attr string, path []
 func (s *Shard) TreeGetMeta(cid cidSDK.ID, treeID string, nodeID pilorama.Node) (pilorama.Meta, error) {
 	return s.pilorama.TreeGetMeta(cid, treeID, nodeID)
 }
+
+// TreeGetChildren implements the pilorama.Forest interface.
+func (s *Shard) TreeGetChildren(cid cidSDK.ID, treeID string, nodeID pilorama.Node) ([]uint64, error) {
+	return s.pilorama.TreeGetChildren(cid, treeID, nodeID)
+}

pkg/services/tree/service.go

@@ -21,6 +21,9 @@ type Service struct {
 	closeCh     chan struct{}
 }
 
+// MaxGetSubTreeDepth represents maximum allowed traversal depth in GetSubTree RPC.
+const MaxGetSubTreeDepth = 10
+
 var _ TreeServiceServer = (*Service)(nil)
 
 // New creates new tree service instance.
@@ -66,7 +69,7 @@ func (s *Service) Add(_ context.Context, req *AddRequest) (*AddResponse, error)
 	log, err := s.forest.TreeMove(cid, b.GetTreeId(), &pilorama.Move{
 		Parent: b.GetParentId(),
 		Child:  pilorama.RootID,
-		Meta:   pilorama.Meta{Items: constructMeta(b.GetMeta())},
+		Meta:   pilorama.Meta{Items: protoToMeta(b.GetMeta())},
 	})
 	if err != nil {
 		return nil, err
@@ -93,7 +96,7 @@ func (s *Service) AddByPath(_ context.Context, req *AddByPathRequest) (*AddByPat
 		return nil, err
 	}
 
-	meta := constructMeta(b.GetMeta())
+	meta := protoToMeta(b.GetMeta())
 
 	attr := b.GetPathAttribute()
 	if len(attr) == 0 {
@@ -173,7 +176,7 @@ func (s *Service) Move(_ context.Context, req *MoveRequest) (*MoveResponse, erro
 	log, err := s.forest.TreeMove(cid, b.GetTreeId(), &pilorama.Move{
 		Parent: b.GetParentId(),
 		Child:  b.GetNodeId(),
-		Meta:   pilorama.Meta{Items: constructMeta(b.GetMeta())},
+		Meta:   pilorama.Meta{Items: protoToMeta(b.GetMeta())},
 	})
 	if err != nil {
 		return nil, err
@@ -211,23 +214,20 @@ func (s *Service) GetNodeByPath(_ context.Context, req *GetNodeByPathRequest) (*
 		var x GetNodeByPathResponse_Info
 		x.NodeId = node
 		x.Timestamp = m.Time
-		for _, kv := range m.Items {
-			needAttr := b.AllAttributes
-			if !needAttr {
+		if b.AllAttributes {
+			x.Meta = metaToProto(m.Items)
+		} else {
+			for _, kv := range m.Items {
 				for _, attr := range b.GetAttributes() {
 					if kv.Key == attr {
-						needAttr = true
+						x.Meta = append(x.Meta, &KeyValue{
+							Key:   kv.Key,
+							Value: kv.Value,
+						})
 						break
 					}
 				}
 			}
-
-			if needAttr {
-				x.Meta = append(x.Meta, &KeyValue{
-					Key:   kv.Key,
-					Value: kv.Value,
-				})
-			}
 		}
 		info = append(info, &x)
 	}
@@ -239,8 +239,56 @@ func (s *Service) GetNodeByPath(_ context.Context, req *GetNodeByPathRequest) (*
 	}, nil
 }
 
-func (s *Service) GetSubTree(_ context.Context, req *GetSubTreeRequest) (*GetSubTreeResponse, error) {
-	return nil, errors.New("GetSubTree is unimplemented")
+type nodeDepthPair struct {
+	nodes []uint64
+	depth uint32
+}
+
+func (s *Service) GetSubTree(req *GetSubTreeRequest, srv TreeService_GetSubTreeServer) error {
+	b := req.GetBody()
+	if b.GetDepth() > MaxGetSubTreeDepth {
+		return fmt.Errorf("too big depth: max=%d, got=%d", MaxGetSubTreeDepth, b.GetDepth())
+	}
+
+	var cid cidSDK.ID
+	if err := cid.Decode(b.GetContainerId()); err != nil {
+		return err
+	}
+
+	queue := []nodeDepthPair{{[]uint64{b.GetRootId()}, 0}}
+
+	for len(queue) != 0 {
+		for _, nodeID := range queue[0].nodes {
+			m, err := s.forest.TreeGetMeta(cid, b.GetTreeId(), nodeID)
+			if err != nil {
+				return err
+			}
+			err = srv.Send(&GetSubTreeResponse{
+				Body: &GetSubTreeResponse_Body{
+					NodeId:    b.GetRootId(),
+					ParentId:  b.GetRootId(),
+					Timestamp: m.Time,
+					Meta:      metaToProto(m.Items),
+				},
+			})
+			if err != nil {
+				return err
+			}
+		}
+
+		if queue[0].depth < b.GetDepth() {
+			for _, nodeID := range queue[0].nodes {
+				children, err := s.forest.TreeGetChildren(cid, b.GetTreeId(), nodeID)
+				if err != nil {
+					return err
+				}
+				queue = append(queue, nodeDepthPair{children, queue[0].depth + 1})
+			}
+		}
+
+		queue = queue[1:]
+	}
+	return nil
 }
 
 // Apply locally applies operation from the remote node to the tree.
@@ -284,11 +332,24 @@ loop:
 	})
 }
 
-func constructMeta(arr []*KeyValue) []pilorama.KeyValue {
+func protoToMeta(arr []*KeyValue) []pilorama.KeyValue {
 	meta := make([]pilorama.KeyValue, len(arr))
 	for i, kv := range arr {
-		meta[i].Key = kv.Key
-		meta[i].Value = kv.Value
+		if kv != nil {
+			meta[i].Key = kv.Key
+			meta[i].Value = kv.Value
+		}
+	}
+	return meta
+}
+
+func metaToProto(arr []pilorama.KeyValue) []*KeyValue {
+	meta := make([]*KeyValue, len(arr))
+	for i, kv := range arr {
+		meta[i] = &KeyValue{
+			Key:   kv.Key,
+			Value: kv.Value,
+		}
 	}
 	return meta
 }

pkg/services/tree/service.proto

@@ -21,7 +21,7 @@ service TreeService {
   // GetNodeByPath returns list of IDs corresponding to a specific filepath.
   rpc GetNodeByPath (GetNodeByPathRequest) returns (GetNodeByPathResponse);
   // GetSubTree returns tree corresponding to a specific node.
-  rpc GetSubTree (GetSubTreeRequest) returns (GetSubTreeResponse);
+  rpc GetSubTree (GetSubTreeRequest) returns (stream GetSubTreeResponse);
 
   /* Synchronization API */
 
@@ -153,7 +153,10 @@ message GetSubTreeRequest {
   message Body {
     bytes container_id = 1;
     string tree_id = 2;
-    repeated uint64 nodes = 3;
+    uint64 root_id = 3;
+    // Optional depth of the traversal. Zero means return only root.
+    // Maximum depth is 10.
+    uint32 depth = 4;
   }
 
   Body body = 1;
@@ -161,12 +164,11 @@ message GetSubTreeRequest {
 }
 
 message GetSubTreeResponse {
-  message Info {
-    uint64 node_id = 1;
-    repeated KeyValue meta = 2;
-  }
   message Body {
-    repeated Info info = 1;
+    uint64 node_id = 1;
+    uint64 parent_id = 2;
+    uint64 timestamp = 3;
+    repeated KeyValue meta = 4;
   }
 
   Body body = 1;