chore(tiering): Fixes

dragonflydb · Jul 5, 2024 · 0e9026c · 0e9026c
1 parent 8240c7f
commit 0e9026c
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 3 deletions.
diff --git a/src/server/tiered_storage.cc b/src/server/tiered_storage.cc
@@ -361,9 +361,16 @@ TieredStats TieredStorage::GetStats() const {
 }
 
 void TieredStorage::RunOffloading(DbIndex dbid) {
+  const size_t kMaxIterations = 500;
+
   if (SliceSnapshot::IsSnaphotInProgress())
     return;
 
+  auto disk_stats = op_manager_->GetStats().disk_stats;
+  if (disk_stats.allocated_bytes + kMaxIterations / 2 * tiering::kPageSize >
+      disk_stats.max_file_size)
+    return;
+
   auto cb = [this, dbid, tmp = std::string{}](PrimeIterator it) mutable {
     TryStash(dbid, it->first.GetSlice(&tmp), &it->second);
   };
@@ -378,12 +385,14 @@ void TieredStorage::RunOffloading(DbIndex dbid) {
     if (op_manager_->GetStats().pending_stash_cnt >= write_depth_limit_)
       break;
     offloading_cursor_ = table.TraverseBySegmentOrder(offloading_cursor_, cb);
-  } while (offloading_cursor_ != start_cursor && iterations++ < 500);
+  } while (offloading_cursor_ != start_cursor && iterations++ < kMaxIterations);
 }
 
 bool TieredStorage::ShouldStash(const PrimeValue& pv) const {
+  auto disk_stats = op_manager_->GetStats().disk_stats;
   return !pv.IsExternal() && !pv.HasIoPending() && pv.ObjType() == OBJ_STRING &&
-         pv.Size() >= kMinValueSize;
+         pv.Size() >= kMinValueSize &&
+         disk_stats.allocated_bytes + tiering::kPageSize + pv.Size() < disk_stats.max_file_size;
 }
 
 }  // namespace dfly
diff --git a/src/server/tiering/disk_storage.cc b/src/server/tiering/disk_storage.cc
@@ -174,7 +174,8 @@ std::error_code DiskStorage::Stash(io::Bytes bytes, StashCb cb) {
 }
 
 DiskStorage::Stats DiskStorage::GetStats() const {
-  return {alloc_.allocated_bytes(), alloc_.capacity(), heap_buf_alloc_cnt_, reg_buf_alloc_cnt_};
+  return {alloc_.allocated_bytes(), alloc_.capacity(), heap_buf_alloc_cnt_, reg_buf_alloc_cnt_,
+          static_cast<size_t>(max_size_)};
 }
 
 std::error_code DiskStorage::Grow(off_t grow_size) {

diff --git a/src/server/tiering/disk_storage.h b/src/server/tiering/disk_storage.h
@@ -22,6 +22,7 @@ class DiskStorage {
     size_t capacity_bytes = 0;
     uint64_t heap_buf_alloc_count = 0;
     uint64_t registered_buf_alloc_count = 0;
+    size_t max_file_size = 0;
   };
 
   using ReadCb = std::function<void(std::string_view, std::error_code)>;