make TObjectDistributions.AddNode idempotent KIKIMR-19696

2 files changed, 34 insertions, 1 deletions

diff --git a/ydb/core/mind/hive/object_distribution.h b/ydb/core/mind/hive/object_distribution.h
index 7e5f9e3ee1..8f76faf56f 100644
--- a/ydb/core/mind/hive/object_distribution.h
+++ b/ydb/core/mind/hive/object_distribution.h
@@ -77,6 +77,12 @@ struct TObjectDistribution {
         Mean = newMean;
     }
 
+    void SetCount(const TNodeInfo& node, i64 value) {
+        auto it = Distribution.find(node.Id);
+        i64 oldValue = (it == Distribution.end()) ? 0 : it->second;
+        UpdateCount(node, value - oldValue);
+    }
+
     void RemoveNode(TNodeId node) {
         auto it = Distribution.find(node);
         if (it == Distribution.end()) {
@@ -202,7 +208,11 @@ struct TObjectDistributions {
             UpdateCount(obj, node, 0);
         }
         for (const auto& [obj, tablets] : node.TabletsOfObject) {
-            UpdateCount(obj, node, tablets.size());
+            ui64 cnt = tablets.size();
+            auto updateFunc = [&](TObjectDistribution& dist) {
+                dist.SetCount(node, cnt);
+            };
+            UpdateDistribution(obj, updateFunc);
         }
     }
 
diff --git a/ydb/core/mind/hive/object_distribution_ut.cpp b/ydb/core/mind/hive/object_distribution_ut.cpp
index e0b447eb15..4dd5b68bac 100644
--- a/ydb/core/mind/hive/object_distribution_ut.cpp
+++ b/ydb/core/mind/hive/object_distribution_ut.cpp
@@ -210,4 +210,27 @@ Y_UNIT_TEST_SUITE(ObjectDistribuiton) {
         UNIT_ASSERT_DOUBLES_EQUAL(distributions.at(tablets[2]->ObjectId)->GetImbalance(), getTrueImbalance(tablets[2]->ObjectId, 0, NUM_NODES), 1e-5);
 
     }
+
+    Y_UNIT_TEST(TestAddSameNode) {
+        TIntrusivePtr<TTabletStorageInfo> hiveStorage = new TTabletStorageInfo;
+        hiveStorage->TabletType = TTabletTypes::Hive;
+        THive hive(hiveStorage.Get(), TActorId());
+
+        std::unordered_map<TNodeId, TNodeInfo> nodes;
+        TObjectDistributions objectDistributions(nodes);
+
+        TNodeInfo node1(1, hive);
+        TNodeInfo node2(2, hive);
+        TTabletInfo tablet(NKikimr::NHive::TTabletInfo::ETabletRole::Leader, hive);
+        node1.TabletsOfObject[{1, 1}].insert(&tablet);
+        objectDistributions.AddNode(node1);
+        objectDistributions.AddNode(node2);
+
+        auto imbalance = objectDistributions.GetMaxImbalance();
+        for (int i = 0; i < 10; ++i) {
+            objectDistributions.AddNode(node1);
+        }
+        UNIT_ASSERT_VALUES_EQUAL(imbalance, objectDistributions.GetMaxImbalance());
+
+    }
 }