fixes

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

Author: dranikpg

src/server/string_family.cc

@@ -936,8 +936,9 @@ 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 we are beyound the offloading threshold, TryStash might return a backpressure future.
   if (auto* ts = shard->tiered_storage(); ts) {
-    auto bp = ts->TryStash(op_args_.db_cntx.db_index, key, pv);
+    auto bp = ts->TryStash(op_args_.db_cntx.db_index, key, pv, true);
     if (bp && params.backpressure)
       *params.backpressure = std::move(*bp);
   }
@@ -1068,8 +1069,10 @@ void StringFamily::Set(CmdArgList args, const CommandContext& cmnd_cntx) {
     return builder->SendError(kWrongTypeErr);
   }
 
-  if (backpressure)
-    std::move(backpressure)->GetFor(100ms);
+  // If backpressure was provided, wait with reasonable limit (to avoid client deadlocking).
+  if (backpressure) {
+    std::move(backpressure)->GetFor(10ms);
+  }
 
   if (sparams.flags & SetCmd::SET_GET) {
     return GetReplies{cmnd_cntx.rb}.Send(std::move(prev));

src/server/tiered_storage.cc

@@ -172,11 +172,8 @@ class TieredStorage::ShardOpManager : public tiering::OpManager {
   }
 
   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);
-    }
+    if (auto node = ts_->stash_backpressure_.extract(id); !node.empty())
+      node.mapped().Resolve(result);
   }
 
   struct {
@@ -319,7 +316,7 @@ error_code TieredStorage::Open(string_view base_path) {
 }
 
 void TieredStorage::Close() {
-  for (auto& [_, f] : backpressure_)
+  for (auto& [_, f] : stash_backpressure_)
     f.Resolve(false);
   op_manager_->Close();
 }
@@ -363,7 +360,7 @@ template TieredStorage::TResult<size_t> TieredStorage::Modify(
     std::function<size_t(std::string*)> modf);
 
 std::optional<util::fb2::Future<bool>> TieredStorage::TryStash(DbIndex dbid, string_view key,
-                                                               PrimeValue* value) {
+                                                               PrimeValue* value, bool provide_bp) {
   if (!ShouldStash(*value))
     return {};
 
@@ -398,12 +395,9 @@ std::optional<util::fb2::Future<bool>> TieredStorage::TryStash(DbIndex dbid, str
     return {};
   }
 
-  // Throttle stashes over the offload boundary
-  if (ShouldOffload()) {
-    util::fb2::Future<bool> fut;
-    backpressure_[{dbid, string{key}}] = fut;
-    return fut;
-  }
+  // If we are in the active offloading phase, throttle stashes by providing backpressure future
+  if (provide_bp && ShouldOffload())
+    return stash_backpressure_[{dbid, string{key}}];
 
   return {};
 }
@@ -429,7 +423,7 @@ void TieredStorage::CancelStash(DbIndex dbid, std::string_view key, PrimeValue*
   DCHECK(value->HasStashPending());
 
   // If any previous write was happening, it has been cancelled
-  if (auto node = backpressure_.extract({dbid, string{key}}); !node.empty())
+  if (auto node = stash_backpressure_.extract(make_pair(dbid, key)); !node.empty())
     std::move(node.mapped()).Resolve(false);
 
   if (OccupiesWholePages(value->Size())) {

src/server/tiered_storage.h

@@ -11,6 +11,7 @@
 #include "server/common.h"
 #include "server/table.h"
 #include "server/tiering/common.h"
+#include "server/tiering/entry_map.h"
 #include "server/tx_base.h"
 #include "util/fibers/future.h"
 
@@ -59,10 +60,10 @@ class TieredStorage {
   TResult<T> Modify(DbIndex dbid, std::string_view key, const PrimeValue& value,
                     std::function<T(std::string*)> modf);
 
-  // Stash value. Sets IO_PENDING flag and unsets it on error or when finished
-  // Returns opional backpressure.
+  // Stash value. Sets IO_PENDING flag and unsets it on error or when finished.
+  // Returns optional future for backpressure if `provide_bp` is set and conditons are met.
   std::optional<util::fb2::Future<bool>> TryStash(DbIndex dbid, std::string_view key,
-                                                  PrimeValue* value);
+                                                  PrimeValue* value, bool provide_bp = false);
 
   // Delete value, must be offloaded (external type)
   void Delete(DbIndex dbid, PrimeValue* value);
@@ -106,7 +107,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_;
+  // Stash operations waiting for completion to throttle
+  tiering::EntryMap<::util::fb2::Future<bool>> stash_backpressure_;
 
   std::unique_ptr<ShardOpManager> op_manager_;
   std::unique_ptr<tiering::SmallBins> bins_;

src/server/tiering/entry_map.h

@@ -0,0 +1,37 @@
+// Copyright 2025, DragonflyDB authors.  All rights reserved.
+// See LICENSE for licensing terms.
+//
+
+#pragma once
+
+#include <absl/container/flat_hash_map.h>
+
+#include <string>
+
+#include "server/tx_base.h"
+
+namespace dfly::tiering {
+
+namespace detail {
+struct Hasher {
+  using is_transparent = void;
+  template <typename S> size_t operator()(const std::pair<DbIndex, S>& p) const {
+    return absl::Hash(p)();
+  }
+};
+
+struct Eq {
+  using is_transparent = void;
+  template <typename S1, typename S2>
+  bool operator()(const std::pair<DbIndex, S1>& l, const std::pair<DbIndex, S2>& r) const {
+    return l == r;
+  }
+};
+}  // namespace detail
+
+// Map of key (db index, string key) -> T with heterogeneous lookup
+template <typename T>
+using EntryMap =
+    absl::flat_hash_map<std::pair<DbIndex, std::string>, T, detail::Hasher, detail::Eq>;
+
+}  // namespace dfly::tiering

src/server/tiering/small_bins.cc

@@ -122,10 +122,7 @@ std::vector<std::pair<DbIndex, std::string>> SmallBins::ReportStashAborted(BinId
 }
 
 std::optional<SmallBins::BinId> SmallBins::Delete(DbIndex dbid, std::string_view key) {
-  std::pair<DbIndex, std::string> key_pair{dbid, key};
-  auto it = current_bin_.find(key_pair);
-
-  if (it != current_bin_.end()) {
+  if (auto it = current_bin_.find(make_pair(dbid, key)); it != current_bin_.end()) {
     size_t stashed_size = StashedValueSize(it->second);
     DCHECK_GE(current_bin_bytes_, stashed_size);
 
@@ -135,7 +132,7 @@ std::optional<SmallBins::BinId> SmallBins::Delete(DbIndex dbid, std::string_view
   }
 
   for (auto& [id, keys] : pending_bins_) {
-    if (keys.erase(key_pair))
+    if (keys.erase(make_pair(dbid, key)))
       return keys.empty() ? std::make_optional(id) : std::nullopt;
   }
   return std::nullopt;

src/server/tiering/small_bins.h

@@ -11,6 +11,7 @@
 #include <vector>
 
 #include "server/tiering/disk_storage.h"
+#include "server/tiering/entry_map.h"
 
 namespace dfly::tiering {
 
@@ -84,12 +85,10 @@ class SmallBins {
   BinId last_bin_id_ = 0;
 
   unsigned current_bin_bytes_ = 0;
-  absl::flat_hash_map<std::pair<DbIndex, std::string>, std::string> current_bin_;
+  tiering::EntryMap<std::string> current_bin_;
 
   // Pending stashes, their keys and value sizes
-  absl::flat_hash_map<unsigned /* id */,
-                      absl::flat_hash_map<std::pair<DbIndex, std::string> /* key*/, DiskSegment>>
-      pending_bins_;
+  absl::flat_hash_map<unsigned /* id */, tiering::EntryMap<DiskSegment>> pending_bins_;
 
   // Map of bins that were stashed and should be deleted when number of entries reaches 0
   absl::flat_hash_map<size_t /*offset*/, StashInfo> stashed_bins_;