1 files changed, 18 insertions, 15 deletions

diff --git a/include/framework/reconstruction/BackgroundTieringPolicy.h b/include/framework/reconstruction/BackgroundTieringPolicy.h
index 4c266fd..5a82695 100644
--- a/include/framework/reconstruction/BackgroundTieringPolicy.h
+++ b/include/framework/reconstruction/BackgroundTieringPolicy.h
@@ -24,14 +24,14 @@ public:
   BackgroundTieringPolicy(size_t scale_factor, size_t buffer_size)
       : m_scale_factor(scale_factor), m_buffer_size(buffer_size) {}
 
-  ReconstructionVector get_reconstruction_tasks(
-      const Version<ShardType, QueryType> *version) const override {
-    ReconstructionVector reconstructions;
+  std::vector<ReconstructionVector> get_reconstruction_tasks(
+      const Version<ShardType, QueryType> *version, LockManager &lock_mngr) const override {
+    std::vector<ReconstructionVector> reconstructions;
 
     auto levels = version->get_structure()->get_level_vector();
 
-    if (levels[0]->get_shard_count() < m_scale_factor) {
-      return reconstructions;
+    if (levels[0]->get_shard_count() == 0) {
+      return {};
     }
 
     level_index target_level = find_reconstruction_target(levels);
@@ -44,16 +44,19 @@ public:
     }
 
     for (level_index i = target_level; i > source_level; i--) {
-      size_t target_reccnt =
-          (i < (level_index)levels.size()) ? levels[i]->get_record_count() : 0;
-      size_t total_reccnt = levels[i - 1]->get_record_count() + target_reccnt;
-
-      std::vector<ShardID> shards;
-      for (ssize_t j=0; j<(ssize_t)levels[i-1]->get_shard_count(); j++) {
-        shards.push_back({i-1, j});
+      if (lock_mngr.take_lock(i-1)) {
+        ReconstructionVector recon;
+        size_t target_reccnt =
+            (i < (level_index)levels.size()) ? levels[i]->get_record_count() : 0;
+        size_t total_reccnt = levels[i - 1]->get_record_count() + target_reccnt;
+        std::vector<ShardID> shards;
+        for (ssize_t j=0; j<(ssize_t)levels[i-1]->get_shard_count(); j++) {
+          shards.push_back({i-1, j});
+        }
+
+        recon.add_reconstruction(shards, i, total_reccnt, ReconstructionType::Compact);
+        reconstructions.push_back(recon);
       }
-
-      reconstructions.add_reconstruction(shards, i, total_reccnt, ReconstructionType::Compact);
     }
     
     return reconstructions;
@@ -70,7 +73,7 @@ private:
   level_index find_reconstruction_target(LevelVector &levels) const {
     level_index target_level = invalid_level_idx;
 
-    for (level_index i = 0; i < (level_index)levels.size(); i++) {
+    for (level_index i = 1; i < (level_index)levels.size(); i++) {
       if (levels[i]->get_shard_count() + 1 <= capacity()) {
         target_level = i;
         break;