From ba3db7fed5a2b20784576e71e5c8ac3da1f7b999 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@morphbits.ru>
Date: Thu, 6 Oct 2022 19:06:19 +0300
Subject: [PATCH] [#1996] engine: Always select proper shard for a tree

Currently there is a possibility for modifying operations to fail
because of I/O errors and a new tree to be created on another shard.
This commit adds existence check for modifying operations.
Read operations remain as they are, not to slow things.
`TreeDrop` is an exception, because this is a tree removal and trying
multiple shards is not an unwanted behaviour.

Signed-off-by: Evgenii Stratonikov <evgeniy@morphbits.ru>
---
 CHANGELOG.md                                  |  1 +
 pkg/local_object_storage/engine/tree.go       | 98 ++++++++++++-------
 pkg/local_object_storage/pilorama/boltdb.go   | 13 +++
 pkg/local_object_storage/pilorama/forest.go   |  7 ++
 .../pilorama/forest_test.go                   | 36 +++++++
 .../pilorama/interface.go                     |  3 +
 pkg/local_object_storage/shard/tree.go        |  8 ++
 7 files changed, 128 insertions(+), 38 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f2cd7c3d4..32b1c0fd0 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,7 @@ Changelog for NeoFS Node
 - Open FSTree in sync mode by default (#1992)
 - `neofs-cli container nodes`'s output (#1991)
 - Do not panic with bad inputs for `GET_RANGE` (#2007)
+- Correctly select the shard for applying tree service operations (#1996)
 
 ### Removed
 ### Updated
diff --git a/pkg/local_object_storage/engine/tree.go b/pkg/local_object_storage/engine/tree.go
index e0fc50f57..6271bca16 100644
--- a/pkg/local_object_storage/engine/tree.go
+++ b/pkg/local_object_storage/engine/tree.go
@@ -13,62 +13,60 @@ var _ pilorama.Forest = (*StorageEngine)(nil)
 
 // TreeMove implements the pilorama.Forest interface.
 func (e *StorageEngine) TreeMove(d pilorama.CIDDescriptor, treeID string, m *pilorama.Move) (*pilorama.LogMove, error) {
-	var err error
-	var lm *pilorama.LogMove
-	for _, sh := range e.sortShardsByWeight(d.CID) {
-		lm, err = sh.TreeMove(d, treeID, m)
-		if err != nil {
-			if errors.Is(err, shard.ErrReadOnlyMode) || err == shard.ErrPiloramaDisabled {
-				return nil, err
-			}
-			e.reportShardError(sh, "can't perform `TreeMove`", err,
+	index, lst, err := e.getTreeShard(d.CID, treeID)
+	if err != nil && !errors.Is(err, pilorama.ErrTreeNotFound) {
+		return nil, err
+	}
+
+	lm, err := lst[index].TreeMove(d, treeID, m)
+	if err != nil {
+		if !errors.Is(err, shard.ErrReadOnlyMode) && err != shard.ErrPiloramaDisabled {
+			e.reportShardError(lst[index], "can't perform `TreeMove`", err,
 				zap.Stringer("cid", d.CID),
 				zap.String("tree", treeID))
-			continue
 		}
-		return lm, nil
+
+		return nil, err
 	}
-	return nil, err
+	return lm, nil
 }
 
 // TreeAddByPath implements the pilorama.Forest interface.
 func (e *StorageEngine) TreeAddByPath(d pilorama.CIDDescriptor, treeID string, attr string, path []string, m []pilorama.KeyValue) ([]pilorama.LogMove, error) {
-	var err error
-	var lm []pilorama.LogMove
-	for _, sh := range e.sortShardsByWeight(d.CID) {
-		lm, err = sh.TreeAddByPath(d, treeID, attr, path, m)
-		if err != nil {
-			if errors.Is(err, shard.ErrReadOnlyMode) || err == shard.ErrPiloramaDisabled {
-				return nil, err
-			}
-			e.reportShardError(sh, "can't perform `TreeAddByPath`", err,
+	index, lst, err := e.getTreeShard(d.CID, treeID)
+	if err != nil && !errors.Is(err, pilorama.ErrTreeNotFound) {
+		return nil, err
+	}
+
+	lm, err := lst[index].TreeAddByPath(d, treeID, attr, path, m)
+	if err != nil {
+		if !errors.Is(err, shard.ErrReadOnlyMode) && err != shard.ErrPiloramaDisabled {
+			e.reportShardError(lst[index], "can't perform `TreeAddByPath`", err,
 				zap.Stringer("cid", d.CID),
 				zap.String("tree", treeID))
-			continue
 		}
-		return lm, nil
+		return nil, err
 	}
-	return nil, err
+	return lm, nil
 }
 
 // TreeApply implements the pilorama.Forest interface.
 func (e *StorageEngine) TreeApply(d pilorama.CIDDescriptor, treeID string, m *pilorama.Move) error {
-	var err error
-	for _, sh := range e.sortShardsByWeight(d.CID) {
-		err = sh.TreeApply(d, treeID, m)
-		if err != nil {
-			if errors.Is(err, shard.ErrReadOnlyMode) || err == shard.ErrPiloramaDisabled {
-				return err
-			}
-			e.reportShardError(sh, "can't perform `TreeApply`", err,
-				zap.Stringer("cid", d.CID),
-				zap.String("tree", treeID))
-			continue
-		}
-		return nil
+	index, lst, err := e.getTreeShard(d.CID, treeID)
+	if err != nil && !errors.Is(err, pilorama.ErrTreeNotFound) {
+		return err
 	}
 
-	return err
+	err = lst[index].TreeApply(d, treeID, m)
+	if err != nil {
+		if !errors.Is(err, shard.ErrReadOnlyMode) && err != shard.ErrPiloramaDisabled {
+			e.reportShardError(lst[index], "can't perform `TreeApply`", err,
+				zap.Stringer("cid", d.CID),
+				zap.String("tree", treeID))
+		}
+		return err
+	}
+	return nil
 }
 
 // TreeGetByPath implements the pilorama.Forest interface.
@@ -205,3 +203,27 @@ func (e *StorageEngine) TreeList(cid cidSDK.ID) ([]string, error) {
 
 	return resIDs, nil
 }
+
+// TreeExists implements the pilorama.Forest interface.
+func (e *StorageEngine) TreeExists(cid cidSDK.ID, treeID string) (bool, error) {
+	_, _, err := e.getTreeShard(cid, treeID)
+	if errors.Is(err, pilorama.ErrTreeNotFound) {
+		return false, nil
+	}
+	return err == nil, err
+}
+
+func (e *StorageEngine) getTreeShard(cid cidSDK.ID, treeID string) (int, []hashedShard, error) {
+	lst := e.sortShardsByWeight(cid)
+	for i, sh := range lst {
+		exists, err := sh.TreeExists(cid, treeID)
+		if err != nil {
+			return 0, nil, err
+		}
+		if exists {
+			return i, lst, err
+		}
+	}
+
+	return 0, lst, pilorama.ErrTreeNotFound
+}
diff --git a/pkg/local_object_storage/pilorama/boltdb.go b/pkg/local_object_storage/pilorama/boltdb.go
index 9266f7fe4..ddc6353ff 100644
--- a/pkg/local_object_storage/pilorama/boltdb.go
+++ b/pkg/local_object_storage/pilorama/boltdb.go
@@ -154,6 +154,19 @@ func (t *boltForest) TreeMove(d CIDDescriptor, treeID string, m *Move) (*LogMove
 	})
 }
 
+// TreeExists implements the Forest interface.
+func (t *boltForest) TreeExists(cid cidSDK.ID, treeID string) (bool, error) {
+	var exists bool
+
+	err := t.db.View(func(tx *bbolt.Tx) error {
+		treeRoot := tx.Bucket(bucketName(cid, treeID))
+		exists = treeRoot != nil
+		return nil
+	})
+
+	return exists, err
+}
+
 // TreeAddByPath implements the Forest interface.
 func (t *boltForest) TreeAddByPath(d CIDDescriptor, treeID string, attr string, path []string, meta []KeyValue) ([]LogMove, error) {
 	if !d.checkValid() {
diff --git a/pkg/local_object_storage/pilorama/forest.go b/pkg/local_object_storage/pilorama/forest.go
index e4d4f5f24..ba765bd35 100644
--- a/pkg/local_object_storage/pilorama/forest.go
+++ b/pkg/local_object_storage/pilorama/forest.go
@@ -209,3 +209,10 @@ func (f *memoryForest) TreeList(cid cidSDK.ID) ([]string, error) {
 
 	return res, nil
 }
+
+// TreeExists implements the pilorama.Forest interface.
+func (f *memoryForest) TreeExists(cid cidSDK.ID, treeID string) (bool, error) {
+	fullID := cid.EncodeToString() + "/" + treeID
+	_, ok := f.treeMap[fullID]
+	return ok, nil
+}
diff --git a/pkg/local_object_storage/pilorama/forest_test.go b/pkg/local_object_storage/pilorama/forest_test.go
index 1da666ad3..9dd9ecae7 100644
--- a/pkg/local_object_storage/pilorama/forest_test.go
+++ b/pkg/local_object_storage/pilorama/forest_test.go
@@ -478,6 +478,42 @@ func testForestTreeGetOpLog(t *testing.T, constructor func(t testing.TB) Forest)
 	})
 }
 
+func TestForest_TreeExists(t *testing.T) {
+	for i := range providers {
+		t.Run(providers[i].name, func(t *testing.T) {
+			testForestTreeExists(t, providers[i].construct)
+		})
+	}
+}
+
+func testForestTreeExists(t *testing.T, constructor func(t testing.TB) Forest) {
+	s := constructor(t)
+
+	checkExists := func(t *testing.T, expected bool, cid cidSDK.ID, treeID string) {
+		actual, err := s.TreeExists(cid, treeID)
+		require.NoError(t, err)
+		require.Equal(t, expected, actual)
+	}
+
+	cid := cidtest.ID()
+	treeID := "version"
+	d := CIDDescriptor{cid, 0, 1}
+
+	t.Run("empty state, no panic", func(t *testing.T) {
+		checkExists(t, false, cid, treeID)
+	})
+
+	require.NoError(t, s.TreeApply(d, treeID, &Move{Parent: 0, Child: 1}))
+	checkExists(t, true, cid, treeID)
+	checkExists(t, false, cidtest.ID(), treeID) // different CID, same tree
+	checkExists(t, false, cid, "another tree")  // same CID, different tree
+
+	t.Run("can be removed", func(t *testing.T) {
+		require.NoError(t, s.TreeDrop(cid, treeID))
+		checkExists(t, false, cid, treeID)
+	})
+}
+
 func TestForest_ApplyRandom(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 087761f39..7426ca11d 100644
--- a/pkg/local_object_storage/pilorama/interface.go
+++ b/pkg/local_object_storage/pilorama/interface.go
@@ -39,6 +39,9 @@ type Forest interface {
 	// TreeList returns all the tree IDs that have been added to the
 	// passed container ID. Nil slice should be returned if no tree found.
 	TreeList(cid cidSDK.ID) ([]string, error)
+	// TreeExists checks if a tree exists locally.
+	// If the tree is not found, false and a nil error should be returned.
+	TreeExists(cid cidSDK.ID, treeID string) (bool, error)
 }
 
 type ForestStorage interface {
diff --git a/pkg/local_object_storage/shard/tree.go b/pkg/local_object_storage/shard/tree.go
index 25300314c..dc4030254 100644
--- a/pkg/local_object_storage/shard/tree.go
+++ b/pkg/local_object_storage/shard/tree.go
@@ -91,3 +91,11 @@ func (s *Shard) TreeList(cid cidSDK.ID) ([]string, error) {
 	}
 	return s.pilorama.TreeList(cid)
 }
+
+// TreeExists implements the pilorama.Forest interface.
+func (s *Shard) TreeExists(cid cidSDK.ID, treeID string) (bool, error) {
+	if s.pilorama == nil {
+		return false, ErrPiloramaDisabled
+	}
+	return s.pilorama.TreeExists(cid, treeID)
+}