chore(tiering): Fixes

Signed-off-by: Vladislav Oleshko <[email protected]>
dragonflydb · Jul 5, 2024 · 1b1146b · 1b1146b
1 parent 8240c7f
commit 1b1146b
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 4 deletions.
diff --git a/src/server/tiered_storage.cc b/src/server/tiered_storage.cc
@@ -361,9 +361,17 @@ TieredStats TieredStorage::GetStats() const {
 }
 
 void TieredStorage::RunOffloading(DbIndex dbid) {
+  const size_t kMaxIterations = 500;
+
   if (SliceSnapshot::IsSnaphotInProgress())
     return;
 
+  // Don't run offloading if there's only very little space left
+  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 +386,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
@@ -11,6 +11,7 @@
 #include "io/io_buf.h"
 #include "server/error.h"
 #include "server/tiering/common.h"
+#include "server/tiering/external_alloc.h"
 #include "util/fibers/uring_proactor.h"
 
 using namespace ::dfly::tiering::literals;
@@ -165,7 +166,10 @@ std::error_code DiskStorage::Stash(io::Bytes bytes, StashCb cb) {
     backing_file_->WriteFixedAsync(buf.bytes, offset, *buf.buf_idx, std::move(io_cb));
   else
     backing_file_->WriteAsync(buf.bytes, offset, std::move(io_cb));
-  if (alloc_.allocated_bytes() > (size_ * 0.85) && !grow_pending_) {
+
+  // Grow in advance if needed and possible
+  if (alloc_.allocated_bytes() > (size_ * 0.85) &&
+      size_ + ExternalAllocator::kExtAlignment < static_cast<size_t>(max_size_) && !grow_pending_) {
     auto ec = Grow(265_MB);
     LOG_IF(ERROR, ec) << "Could not call grow :" << ec.message();
     return ec;
@@ -174,7 +178,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)>;