diff --git a/pkg/local_object_storage/engine/tree.go b/pkg/local_object_storage/engine/tree.go index 986bea15e1..04362b8c6c 100644 --- a/pkg/local_object_storage/engine/tree.go +++ b/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 +} diff --git a/pkg/local_object_storage/pilorama/boltdb.go b/pkg/local_object_storage/pilorama/boltdb.go index 6eb12a6c6e..808ab18e62 100644 --- a/pkg/local_object_storage/pilorama/boltdb.go +++ b/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 diff --git a/pkg/local_object_storage/pilorama/forest.go b/pkg/local_object_storage/pilorama/forest.go index 53d044866a..f6d79413d7 100644 --- a/pkg/local_object_storage/pilorama/forest.go +++ b/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 +} diff --git a/pkg/local_object_storage/pilorama/forest_test.go b/pkg/local_object_storage/pilorama/forest_test.go index f3a498fe50..d3503bf3b1 100644 --- a/pkg/local_object_storage/pilorama/forest_test.go +++ b/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) { diff --git a/pkg/local_object_storage/pilorama/interface.go b/pkg/local_object_storage/pilorama/interface.go index 005ce5ce5a..9ffa529a81 100644 --- a/pkg/local_object_storage/pilorama/interface.go +++ b/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 { diff --git a/pkg/local_object_storage/shard/tree.go b/pkg/local_object_storage/shard/tree.go index 2a4120a12d..13cd8fcf00 100644 --- a/pkg/local_object_storage/shard/tree.go +++ b/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) +} diff --git a/pkg/services/tree/service.go b/pkg/services/tree/service.go index 00de22b8ad..8cf6236cd6 100644 --- a/pkg/services/tree/service.go +++ b/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 } diff --git a/pkg/services/tree/service.pb.go b/pkg/services/tree/service.pb.go index bf1f26904e..5f0afdb9a2 100644 Binary files a/pkg/services/tree/service.pb.go and b/pkg/services/tree/service.pb.go differ diff --git a/pkg/services/tree/service.proto b/pkg/services/tree/service.proto index f18b835f7e..0da34f652d 100644 --- a/pkg/services/tree/service.proto +++ b/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; diff --git a/pkg/services/tree/service_grpc.pb.go b/pkg/services/tree/service_grpc.pb.go index 47aac0f788..dc6fe4fd19 100644 Binary files a/pkg/services/tree/service_grpc.pb.go and b/pkg/services/tree/service_grpc.pb.go differ