fix: fixes

Signed-off-by: Vladislav Oleshko <[email protected]>

Author: dranikpg

src/server/rdb_load.cc

@@ -2426,9 +2426,9 @@ void RdbLoader::FlushShardAsync(ShardId sid) {
 
     // Block, if tiered storage is active, but can't keep up
     while (EngineShard::tlocal()->ShouldThrottleForTiering()) {
-      this->blocked_shards_.fetch_add(memory_order_relaxed);  // stop adding items to shard queue
+      this->blocked_shards_.fetch_add(1, memory_order_relaxed);  // stop adding items to shard queue
       ThisFiber::SleepFor(100us);
-      this->blocked_shards_.fetch_sub(memory_order_relaxed);
+      this->blocked_shards_.fetch_sub(1, memory_order_relaxed);
     }
   };
 
@@ -2479,8 +2479,8 @@ void RdbLoader::LoadItemsBuffer(DbIndex db_ind, const ItemsBuf& ib) {
       LOG(WARNING) << "RDB has duplicated key '" << item->key << "' in DB " << db_ind;
     }
 
-    if (auto* ts = es->tiered_storage(); ts && ts->ShouldStash(res.it->second))
-      ts->Stash(db_cntx.db_index, item->key, &res.it->second);
+    if (auto* ts = es->tiered_storage(); ts)
+      ts->TryStash(db_cntx.db_index, item->key, &res.it->second);
   }
 
   for (auto* item : ib) {

src/server/string_family.cc

@@ -664,9 +664,8 @@ void SetCmd::PostEdit(const SetParams& params, std::string_view key, std::string
   EngineShard* shard = op_args_.shard;
 
   // Currently we always try to offload, but Stash may ignore it, if disk I/O is overloaded.
-  if (auto* ts = shard->tiered_storage(); ts && ts->ShouldStash(*pv)) {
-    ts->Stash(op_args_.db_cntx.db_index, key, pv);
-  }
+  if (auto* ts = shard->tiered_storage(); ts)
+    ts->TryStash(op_args_.db_cntx.db_index, key, pv);
 
   if (manual_journal_ && op_args_.shard->journal()) {
     RecordJournal(params, key, value);

src/server/tiered_storage.cc

@@ -271,14 +271,15 @@ template util::fb2::Future<size_t> TieredStorage::Modify(DbIndex dbid, std::stri
                                                          const PrimeValue& value,
                                                          std::function<size_t(std::string*)> modf);
 
-void TieredStorage::Stash(DbIndex dbid, string_view key, PrimeValue* value) {
-  CHECK(!value->IsExternal() && !value->HasIoPending());
+bool TieredStorage::TryStash(DbIndex dbid, string_view key, PrimeValue* value) {
+  if (!ShouldStash(*value))
+    return false;
 
   // TODO: When we are low on memory we should introduce a back-pressure, to avoid OOMs
   // with a lot of underutilized disk space.
   if (op_manager_->GetStats().pending_stash_cnt >= write_depth_limit_) {
     ++stats_.stash_overflow_cnt;
-    return;
+    return false;
   }
 
   string buf;
@@ -298,7 +299,10 @@ void TieredStorage::Stash(DbIndex dbid, string_view key, PrimeValue* value) {
   if (ec) {
     LOG(ERROR) << "Stash failed immediately" << ec.message();
     visit([this](auto id) { op_manager_->ClearIoPending(id); }, id);
+    return false;
   }
+
+  return true;
 }
 
 void TieredStorage::Delete(DbIndex dbid, PrimeValue* value) {
@@ -318,10 +322,6 @@ void TieredStorage::CancelStash(DbIndex dbid, std::string_view key, PrimeValue*
   value->SetIoPending(false);
 }
 
-bool TieredStorage::ShouldStash(const PrimeValue& pv) const {
-  return !pv.IsExternal() && pv.ObjType() == OBJ_STRING && pv.Size() >= kMinValueSize;
-}
-
 float TieredStorage::WriteDepthUsage() const {
   return 1.0f * op_manager_->GetStats().pending_stash_cnt / write_depth_limit_;
 }
@@ -364,30 +364,26 @@ void TieredStorage::RunOffloading(DbIndex dbid) {
   if (SliceSnapshot::IsSnaphotInProgress())
     return;
 
-  PrimeTable& table = op_manager_->db_slice_->GetDBTable(dbid)->prime;
-  int stash_limit = write_depth_limit_ - op_manager_->GetStats().pending_stash_cnt;
-  if (stash_limit <= 0)
-    return;
-
-  std::string tmp;
-  auto cb = [this, dbid, &tmp, &stash_limit](PrimeIterator it) {
-    if (it->second.HasIoPending() || it->second.IsExternal())
-      return;
-
-    if (ShouldStash(it->second)) {
-      Stash(dbid, it->first.GetSlice(&tmp), &it->second);
-      stash_limit--;
-    }
+  auto cb = [this, dbid, tmp = std::string{}](PrimeIterator it) mutable {
+    TryStash(dbid, it->first.GetSlice(&tmp), &it->second);
   };
 
+  PrimeTable& table = op_manager_->db_slice_->GetDBTable(dbid)->prime;
   PrimeTable::Cursor start_cursor{};
 
   // Loop while we haven't traversed all entries or reached our stash io device limit.
   // Keep number of iterations below resonable limit to keep datastore always responsive
   size_t iterations = 0;
   do {
+    if (op_manager_->GetStats().pending_stash_cnt >= write_depth_limit_)
+      break;
     offloading_cursor_ = table.TraverseBySegmentOrder(offloading_cursor_, cb);
-  } while (offloading_cursor_ != start_cursor && stash_limit > 0 && iterations++ < 500);
+  } while (offloading_cursor_ != start_cursor && iterations++ < 500);
+}
+
+bool TieredStorage::ShouldStash(const PrimeValue& pv) const {
+  return !pv.IsExternal() && !pv.HasIoPending() && pv.ObjType() == OBJ_STRING &&
+         pv.Size() >= kMinValueSize;
 }
 
 }  // namespace dfly

src/server/tiered_storage.h

@@ -56,17 +56,15 @@ class TieredStorage {
                               std::function<T(std::string*)> modf);
 
   // Stash value. Sets IO_PENDING flag and unsets it on error or when finished
-  void Stash(DbIndex dbid, std::string_view key, PrimeValue* value);
+  // Returns true if item was scheduled for stashing.
+  bool TryStash(DbIndex dbid, std::string_view key, PrimeValue* value);
 
   // Delete value, must be offloaded (external type)
   void Delete(DbIndex dbid, PrimeValue* value);
 
   // Cancel pending stash for value, must have IO_PENDING flag set
   void CancelStash(DbIndex dbid, std::string_view key, PrimeValue* value);
 
-  // Returns if a value should be stashed
-  bool ShouldStash(const PrimeValue& pv) const;
-
   // Percentage (0-1) of currently used storage_write_depth for ongoing stashes
   float WriteDepthUsage() const;
 
@@ -75,6 +73,10 @@ class TieredStorage {
   // Run offloading loop until i/o device is loaded or all entries were traversed
   void RunOffloading(DbIndex dbid);
 
+ private:
+  // Returns if a value should be stashed
+  bool ShouldStash(const PrimeValue& pv) const;
+
  private:
   PrimeTable::Cursor offloading_cursor_{};  // where RunOffloading left off
 
@@ -128,7 +130,7 @@ class TieredStorage {
     return {};
   }
 
-  void Stash(DbIndex dbid, std::string_view key, PrimeValue* value) {
+  void TryStash(DbIndex dbid, std::string_view key, PrimeValue* value) {
   }
 
   void Delete(PrimeValue* value) {
@@ -137,7 +139,7 @@ class TieredStorage {
   void CancelStash(DbIndex dbid, std::string_view key, PrimeValue* value) {
   }
 
-  bool ShouldStash(const PrimeValue& pv) {
+  bool ShouldStash(const PrimeValue& pv) const {
     return false;
   }
 

tests/dragonfly/snapshot_test.py

@@ -469,7 +469,7 @@ async def test_tiered_entries(async_client: aioredis.Redis):
         "dbfilename": "tiered-entries",
         "tiered_prefix": "tiering-test-backing",
         "tiered_offload_threshold": "0.5",  # ask to keep below 0.5 * 1G
-        "tiered_storage_write_depth": 5000,
+        "tiered_storage_write_depth": 50,
     }
 )
 async def test_tiered_entries_throttle(async_client: aioredis.Redis):
@@ -495,6 +495,7 @@ async def test_tiered_entries_throttle(async_client: aioredis.Redis):
 
     while not load_task.done():
         info = await async_client.info("ALL")
+        # print(info["used_memory_human"], info["used_memory_rss_human"])
         assert info["used_memory"] < 600e6  # less than 600mb,
         await asyncio.sleep(0.05)