group tablets by object in drain KIKIMR-18695

2 files changed, 14 insertions, 4 deletions

diff --git a/ydb/core/mind/hive/drain.cpp b/ydb/core/mind/hive/drain.cpp
index 180cb2dd06..c608074406 100644
--- a/ydb/core/mind/hive/drain.cpp
+++ b/ydb/core/mind/hive/drain.cpp
@@ -158,10 +158,13 @@ public:
         TNodeInfo* nodeInfo = Hive->FindNode(NodeId);
         if (nodeInfo != nullptr) {
             {
-                const auto& tablets = nodeInfo->Tablets[TTabletInfo::EVolatileState::TABLET_VOLATILE_STATE_RUNNING];
-                Tablets.reserve(tablets.size());
-                for (TTabletInfo* tabletInfo : tablets) {
-                    Tablets.push_back(tabletInfo->GetFullTabletId());
+                Tablets.reserve(nodeInfo->GetTabletsRunning());
+                for (const auto& [object, tablets] : nodeInfo->TabletsOfObject) {
+                    for (TTabletInfo* tabletInfo : tablets) {
+                        if (tabletInfo->GetVolatileState() == TTabletInfo::EVolatileState::TABLET_VOLATILE_STATE_RUNNING) {
+                            Tablets.push_back(tabletInfo->GetFullTabletId());
+                        }
+                    }
                 }
             }
             NextKick = Tablets.begin();
diff --git a/ydb/core/mind/hive/node_info.h b/ydb/core/mind/hive/node_info.h
index c6633df8c0..15d7054a68 100644
--- a/ydb/core/mind/hive/node_info.h
+++ b/ydb/core/mind/hive/node_info.h
@@ -114,6 +114,13 @@ public:
         return totalSize;
     }
 
+    ui32 GetTabletsRunning() const {
+        auto it = Tablets.find(TTabletInfo::EVolatileState::TABLET_VOLATILE_STATE_RUNNING);
+        if (it != Tablets.end())
+            return it->second.size();
+        return 0;
+    }
+
     ui32 GetTabletNeighboursCount(const TTabletInfo& tablet) const {
         auto it = TabletsOfObject.find(tablet.GetObjectId());
         if (it != TabletsOfObject.end()) {