fix(tiering): Basic stash backpressure

Author: dranikpg

src/server/string_family.cc

@@ -84,6 +84,7 @@ class SetCmd {
     uint32_t memcache_flags = 0;
     uint64_t expire_after_ms = 0;  // Relative value based on now. 0 means no expiration.
     optional<StringResult>* prev_val = nullptr;  // if set, previous value will be stored if found
+    optional<util::fb2::Future<bool>>* backpressure = nullptr;
 
     constexpr bool IsConditionalSet() const {
       return flags & SET_IF_NOTEXIST || flags & SET_IF_EXISTS;
@@ -935,8 +936,11 @@ 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->TryStash(op_args_.db_cntx.db_index, key, pv);
+  if (auto* ts = shard->tiered_storage(); ts) {
+    auto bp = ts->TryStash(op_args_.db_cntx.db_index, key, pv);
+    if (bp && params.backpressure)
+      *params.backpressure = std::move(*bp);
+  }
 
   if (explicit_journal_ && op_args_.shard->journal()) {
     RecordJournal(params, key, value);
@@ -1056,11 +1060,17 @@ void StringFamily::Set(CmdArgList args, const CommandContext& cmnd_cntx) {
   if (sparams.flags & SetCmd::SET_GET)
     sparams.prev_val = &prev;
 
+  optional<util::fb2::Future<bool>> backpressure;
+  sparams.backpressure = &backpressure;
+
   OpStatus result = SetGeneric(sparams, key, value, cmnd_cntx);
   if (result == OpStatus::WRONG_TYPE) {
     return builder->SendError(kWrongTypeErr);
   }
 
+  if (backpressure)
+    std::move(backpressure)->GetFor(100ms);
+
   if (sparams.flags & SetCmd::SET_GET) {
     return GetReplies{cmnd_cntx.rb}.Send(std::move(prev));
   }

src/server/tiered_storage.cc

@@ -85,6 +85,7 @@ class TieredStorage::ShardOpManager : public tiering::OpManager {
 
   // Clear IO pending flag for entry
   void ClearIoPending(OpManager::KeyRef key) {
+    FlagBackpressure(key, false);
     if (auto pv = Find(key); pv) {
       pv->SetStashPending(false);
       stats_.total_cancels++;
@@ -143,6 +144,7 @@ class TieredStorage::ShardOpManager : public tiering::OpManager {
   // Find entry by key in db_slice and store external segment in place of original value.
   // Update memory stats
   void SetExternal(OpManager::KeyRef key, tiering::DiskSegment segment) {
+    FlagBackpressure(key, true);
     if (auto* pv = Find(key); pv) {
       auto* stats = GetDbTableStats(key.first);
 
@@ -169,6 +171,14 @@ class TieredStorage::ShardOpManager : public tiering::OpManager {
       SetExternal({sub_dbid, sub_key}, sub_segment);
   }
 
+  void FlagBackpressure(OpManager::KeyRef id, bool result) {
+    const auto& [dbid, key] = id;
+    if (auto it = ts_->backpressure_.find({dbid, string{key}}); it != ts_->backpressure_.end()) {
+      it->second.Resolve(result);
+      ts_->backpressure_.erase(it);
+    }
+  }
+
   struct {
     uint64_t total_stashes = 0, total_cancels = 0, total_fetches = 0;
     uint64_t total_defrags = 0;
@@ -309,6 +319,8 @@ error_code TieredStorage::Open(string_view base_path) {
 }
 
 void TieredStorage::Close() {
+  for (auto& [_, f] : backpressure_)
+    f.Resolve(false);
   op_manager_->Close();
 }
 
@@ -350,9 +362,10 @@ template TieredStorage::TResult<size_t> TieredStorage::Modify(
     DbIndex dbid, std::string_view key, const PrimeValue& value,
     std::function<size_t(std::string*)> modf);
 
-bool TieredStorage::TryStash(DbIndex dbid, string_view key, PrimeValue* value) {
+std::optional<util::fb2::Future<bool>> TieredStorage::TryStash(DbIndex dbid, string_view key,
+                                                               PrimeValue* value) {
   if (!ShouldStash(*value))
-    return false;
+    return {};
 
   // This invariant should always hold because ShouldStash tests for IoPending flag.
   CHECK(!bins_->IsPending(dbid, key));
@@ -361,7 +374,7 @@ bool TieredStorage::TryStash(DbIndex dbid, string_view key, PrimeValue* value) {
   // with a lot of underutilized disk space.
   if (op_manager_->GetStats().pending_stash_cnt >= config_.write_depth_limit) {
     ++stats_.stash_overflow_cnt;
-    return false;
+    return {};
   }
 
   StringOrView raw_string = value->GetRawString();
@@ -375,15 +388,24 @@ bool TieredStorage::TryStash(DbIndex dbid, string_view key, PrimeValue* value) {
   } else if (auto bin = bins_->Stash(dbid, key, raw_string.view()); bin) {
     id = bin->first;
     ec = op_manager_->Stash(id, bin->second);
+  } else {
+    return {};  // Silently added to bin
   }
 
   if (ec) {
     LOG_IF(ERROR, ec != errc::file_too_large) << "Stash failed immediately" << ec.message();
     visit([this](auto id) { op_manager_->ClearIoPending(id); }, id);
-    return false;
+    return {};
+  }
+
+  // Throttle stashes over the offload boundary
+  if (ShouldOffload()) {
+    util::fb2::Future<bool> fut;
+    backpressure_[{dbid, string{key}}] = fut;
+    return fut;
   }
 
-  return true;
+  return {};
 }
 
 void TieredStorage::Delete(DbIndex dbid, PrimeValue* value) {
@@ -405,6 +427,11 @@ void TieredStorage::Delete(DbIndex dbid, PrimeValue* value) {
 
 void TieredStorage::CancelStash(DbIndex dbid, std::string_view key, PrimeValue* value) {
   DCHECK(value->HasStashPending());
+
+  // If any previous write was happening, it has been cancelled
+  if (auto node = backpressure_.extract({dbid, string{key}}); !node.empty())
+    std::move(node.mapped()).Resolve(false);
+
   if (OccupiesWholePages(value->Size())) {
     op_manager_->Delete(KeyRef(dbid, key));
   } else if (auto bin = bins_->Delete(dbid, key); bin) {

src/server/tiered_storage.h

@@ -60,8 +60,9 @@ class TieredStorage {
                     std::function<T(std::string*)> modf);
 
   // Stash value. Sets IO_PENDING flag and unsets it on error or when finished
-  // Returns true if item was scheduled for stashing.
-  bool TryStash(DbIndex dbid, std::string_view key, PrimeValue* value);
+  // Returns opional backpressure.
+  std::optional<util::fb2::Future<bool>> TryStash(DbIndex dbid, std::string_view key,
+                                                  PrimeValue* value);
 
   // Delete value, must be offloaded (external type)
   void Delete(DbIndex dbid, PrimeValue* value);
@@ -105,6 +106,8 @@ class TieredStorage {
 
   PrimeTable::Cursor offloading_cursor_{};  // where RunOffloading left off
 
+  absl::flat_hash_map<std::pair<DbIndex, std::string>, ::util::fb2::Future<bool>> backpressure_;
+
   std::unique_ptr<ShardOpManager> op_manager_;
   std::unique_ptr<tiering::SmallBins> bins_;