chore: disk backpressure class utility

Signed-off-by: Kostas Kyrimis <[email protected]>

Author: kostasrim

src/facade/CMakeLists.txt

@@ -3,7 +3,8 @@ cxx_link(dfly_parser_lib base strings_lib)
 
 add_library(dfly_facade conn_context.cc dragonfly_listener.cc dragonfly_connection.cc facade.cc
             memcache_parser.cc reply_builder.cc op_status.cc service_interface.cc
-            reply_capture.cc cmd_arg_parser.cc tls_helpers.cc socket_utils.cc)
+            reply_capture.cc cmd_arg_parser.cc tls_helpers.cc socket_utils.cc
+            disk_connection_backpressure.cc)
 
 if (DF_USE_SSL)
   set(TLS_LIB tls_lib)

src/facade/disk_connection_backpressure.cc

@@ -0,0 +1,132 @@
+// Copyright 2025, DragonflyDB authors.  All rights reserved.
+//
+// See LICENSE for licensing terms.
+//
+
+#include "facade/disk_connection_backpressure.h"
+
+#include <absl/strings/str_cat.h>
+
+#include <string>
+
+#include "base/flags.h"
+#include "base/logging.h"
+#include "facade/facade_types.h"
+#include "io/io.h"
+#include "util/fibers/uring_file.h"
+
+using facade::operator""_MB;
+
+ABSL_FLAG(std::string, disk_backpressure_folder, "/tmp/",
+          "Folder to store "
+          "disk backed connection backpressure");
+
+ABSL_FLAG(size_t, disk_backpressure_file_max_bytes, 50_MB,
+          "Maximum size of the backing file. When max size is reached, connection will "
+          "stop offloading backpressure to disk and block on client read.");
+
+ABSL_FLAG(size_t, disk_backpressure_load_size, 30,
+          "How many items to load in dispatch queue from the disk backed file.");
+
+namespace facade {
+
+DiskBackedBackpressureQueue::DiskBackedBackpressureQueue(uint32_t conn_id)
+    : max_backing_size_(absl::GetFlag(FLAGS_disk_backpressure_file_max_bytes)),
+      max_queue_load_size_(absl::GetFlag(FLAGS_disk_backpressure_load_size)),
+      id_(conn_id) {
+}
+
+std::error_code DiskBackedBackpressureQueue::Init() {
+  std::string backing_name = absl::StrCat(absl::GetFlag(FLAGS_disk_backpressure_folder), id_);
+  {
+    // Kernel transparently handles buffering via the page cache.
+    auto res = util::fb2::OpenWrite(backing_name, {} /* overwrite mode + non direct io */);
+    if (!res) {
+      return res.error();
+    }
+    writer_.reset(*res);
+  }
+
+  auto res = util::fb2::OpenRead(backing_name);
+  if (!res) {
+    return res.error();
+  }
+  reader_.reset(*res);
+
+  VLOG(3) << "Created backing for connection " << this << " " << backing_name;
+
+  return {};
+}
+
+DiskBackedBackpressureQueue::~DiskBackedBackpressureQueue() {
+  auto ec = writer_->Close();
+  LOG_IF(WARNING, ec) << ec.message();
+  ec = reader_->Close();
+  LOG_IF(WARNING, ec) << ec.message();
+}
+
+// Check if backing file is empty, i.e. backing file has 0 bytes.
+bool DiskBackedBackpressureQueue::Empty() const {
+  return total_backing_bytes_ == 0;
+}
+
+bool DiskBackedBackpressureQueue::HasEnoughBackingSpaceFor(size_t bytes) const {
+  return (bytes + total_backing_bytes_) < max_backing_size_;
+}
+
+size_t DiskBackedBackpressureQueue::TotalInMemoryBytes() const {
+  return offsets_.size() * sizeof(ItemOffset);
+}
+
+void DiskBackedBackpressureQueue::OffloadToBacking(std::string_view blob) {
+  ItemOffset item;
+  item.offset = next_offset_;
+  item.total_bytes = blob.size();
+
+  // TODO we should truncate as the file grows. That way we never end up with large files
+  // on disk.
+  auto res = writer_->Write(blob);
+  if (res) {
+    VLOG(2) << "Failed to offload connection " << this << " backpressure with offset "
+            << item.offset << " of size " << item.total_bytes << " to backing with error: " << res;
+    return;
+  }
+
+  total_backing_bytes_ += blob.size();
+  offsets_.push_back(item);
+  next_offset_ += item.total_bytes;
+
+  VLOG(2) << "Offload connection " << this << " backpressure of " << item.total_bytes
+          << " bytes to disk at offset: " << item.offset;
+  VLOG(3) << "Command offloaded: " << blob;
+}
+
+void DiskBackedBackpressureQueue::LoadFromDiskToQueue(std::function<void(io::MutableBytes)> f) {
+  std::string buffer;
+  size_t up_to = max_queue_load_size_;
+
+  while (!offsets_.empty() && up_to--) {
+    ItemOffset item = offsets_.front();
+
+    buffer.resize(item.total_bytes);
+
+    io::MutableBytes bytes{reinterpret_cast<uint8_t*>(buffer.data()), item.total_bytes};
+    auto result = reader_->Read(item.offset, bytes);
+    if (!result) {
+      LOG(ERROR) << "Could not load item at offset " << item.offset << " of size "
+                 << item.total_bytes << " from disk with error: " << result.error().value() << " "
+                 << result.error().message();
+      return;
+    }
+
+    VLOG(2) << "Loaded item with offset " << item.offset << " of size " << item.total_bytes
+            << " for connection " << this;
+
+    f(bytes);
+
+    offsets_.pop_front();
+    total_backing_bytes_ -= item.total_bytes;
+  }
+}
+
+}  // namespace facade

src/facade/disk_connection_backpressure.h

@@ -0,0 +1,65 @@
+// Copyright 2025, DragonflyDB authors.  All rights reserved.
+// See LICENSE for licensing terms.
+//
+
+#pragma once
+
+#include <util/fibers/uring_file.h>
+
+#include <deque>
+#include <functional>
+#include <memory>
+#include <string_view>
+#include <system_error>
+
+#include "io/io.h"
+
+namespace facade {
+
+class DiskBackedBackpressureQueue {
+ public:
+  explicit DiskBackedBackpressureQueue(uint32_t conn_id);
+  ~DiskBackedBackpressureQueue();
+
+  std::error_code Init();
+
+  // Check if we can offload bytes to backing file.
+  bool HasEnoughBackingSpaceFor(size_t bytes) const;
+
+  // Total size of internal buffers/structures.
+  size_t TotalInMemoryBytes() const;
+
+  void OffloadToBacking(std::string_view blob);
+
+  // For each item loaded from disk it calls f(item) to consume it.
+  // Reads up to max_queue_load_size_ items on each call
+  void LoadFromDiskToQueue(std::function<void(io::MutableBytes)> f);
+
+  // Check if backing file is empty, i.e. backing file has 0 bytes.
+  bool Empty() const;
+
+ private:
+  // File Reader/Writer
+  std::unique_ptr<io::WriteFile> writer_;
+  std::unique_ptr<io::ReadonlyFile> reader_;
+
+  // In memory backed file map
+  struct ItemOffset {
+    size_t offset = 0;
+    size_t total_bytes = 0;
+  };
+
+  std::deque<ItemOffset> offsets_;
+
+  size_t total_backing_bytes_ = 0;
+  size_t next_offset_ = 0;
+
+  // Read only constants
+  const size_t max_backing_size_ = 0;
+  const size_t max_queue_load_size_ = 0;
+
+  // unique id for the file backed
+  const size_t id_ = 0;
+};
+
+}  // namespace facade

src/facade/dragonfly_connection.cc

@@ -1,4 +1,4 @@
-// Copyright 2022, DragonflyDB authors.  All rights reserved.
+// Copyright 2025, DragonflyDB authors.  All rights reserved.
 //
 // See LICENSE for licensing terms.
 //
@@ -13,6 +13,7 @@
 #include <numeric>
 #include <variant>
 
+#include "absl/flags/internal/flag.h"
 #include "base/cycle_clock.h"
 #include "base/flag_utils.h"
 #include "base/flags.h"
@@ -22,6 +23,7 @@
 #include "base/stl_util.h"
 #include "core/heap_size.h"
 #include "facade/conn_context.h"
+#include "facade/disk_connection_backpressure.h"
 #include "facade/dragonfly_listener.h"
 #include "facade/memcache_parser.h"
 #include "facade/redis_parser.h"
@@ -112,6 +114,9 @@ ABSL_FLAG(uint32_t, pipeline_wait_batch_usec, 0,
           "If non-zero, waits for this time for more I/O "
           " events to come for the connection in case there is only one command in the pipeline. ");
 
+ABSL_FLAG(size_t, disk_backpressure_offload_watermark, 0,
+          "Offload backpressure to disk when dispatch queue size crosses the watermark.");
+
 using namespace util;
 using namespace std;
 using absl::GetFlag;
@@ -676,6 +681,16 @@ Connection::Connection(Protocol protocol, util::HttpListenerBase* http_listener,
 #endif
 
   UpdateLibNameVerMap(lib_name_, lib_ver_, +1);
+
+  backpressure_to_disk_watermark_ = absl::GetFlag(FLAGS_disk_backpressure_offload_watermark);
+  if (backpressure_to_disk_watermark_ > 0) {
+    backing_queue_ = std::make_unique<DiskBackedBackpressureQueue>(id_);
+    auto ec = backing_queue_->Init();
+    if (ec) {
+      LOG(ERROR) << "Error while initializing backpressure file " << ec.message();
+      backing_queue_.reset();
+    }
+  }
 }
 
 Connection::~Connection() {

src/facade/dragonfly_connection.h

@@ -1,4 +1,4 @@
-// Copyright 2022, DragonflyDB authors.  All rights reserved.
+// Copyright 2025, DragonflyDB authors.  All rights reserved.
 // See LICENSE for licensing terms.
 //
 
@@ -45,6 +45,7 @@ class ConnectionContext;
 class RedisParser;
 class ServiceInterface;
 class SinkReplyBuilder;
+class DiskBackedBackpressureQueue;
 
 // Connection represents an active connection for a client.
 //
@@ -503,6 +504,9 @@ class Connection : public util::Connection {
     };
   };
 
+  std::unique_ptr<DiskBackedBackpressureQueue> backing_queue_;
+  size_t backpressure_to_disk_watermark_ = 0;
+
   bool request_shutdown_ = false;
 };