feat(tiering): Two phase stash (#5984)

Author: dranikpg

src/server/tiered_storage.cc

@@ -380,14 +380,28 @@ std::optional<util::fb2::Future<bool>> TieredStorage::TryStash(DbIndex dbid, str
 
   tiering::OpManager::EntryId id;
   error_code ec;
+
+  // TODO(vlad): Replace with encoders for different types
+  auto stash_string = [&](std::string_view str) {
+    if (auto prepared = op_manager_->PrepareStash(str.size()); prepared) {
+      auto [offset, buf] = *prepared;
+      memcpy(buf.bytes.data(), str.data(), str.size());
+      tiering::DiskSegment segment{offset, str.size()};
+      op_manager_->Stash(id, segment, buf);
+    } else {
+      ec = prepared.error();
+    }
+  };
+
   if (OccupiesWholePages(value->Size())) {  // large enough for own page
     id = KeyRef(dbid, key);
-    ec = op_manager_->Stash(id, raw_string.view());
+    stash_string(raw_string.view());
   } else if (auto bin = bins_->Stash(dbid, key, raw_string.view()); bin) {
     id = bin->first;
-    ec = op_manager_->Stash(id, bin->second);
+    // TODO(vlad): Write bin to prepared buffer instead of allocating one
+    stash_string(bin->second);
   } else {
-    return {};  // Silently added to bin
+    return {};  // silently added to bin
   }
 
   if (ec) {

src/server/tiering/disk_storage.cc

@@ -144,33 +144,40 @@ void DiskStorage::MarkAsFree(DiskSegment segment) {
   alloc_.Free(segment.offset, segment.length);
 }
 
-std::error_code DiskStorage::Stash(io::Bytes bytes, StashCb cb) {
-  DCHECK_GT(bytes.length(), 0u);
+io::Result<std::pair<size_t, UringBuf>> DiskStorage::PrepareStash(size_t length) {
+  using namespace nonstd;
 
-  size_t len = bytes.size();
-  int64_t offset = alloc_.Malloc(len);
+  int64_t offset = alloc_.Malloc(length);
+  if (offset >= 0)
+    return std::make_pair(offset, PrepareBuf(length));
 
   // If we don't have enough space, request grow and return to avoid blocking
   if (offset < 0) {
     auto ec = RequestGrow(-offset);
-    return ec ? ec : make_error_code(errc::operation_would_block);
+    return make_unexpected(ec ? ec : make_error_code(errc::operation_would_block));
   }
 
-  UringBuf buf = PrepareBuf(len);
-  memcpy(buf.bytes.data(), bytes.data(), bytes.length());
+  offset = alloc_.Malloc(length);
+  if (offset < 0)  // we can't fit it even after resizing
+    return make_unexpected(make_error_code(errc::file_too_large));
+
+  return std::make_pair(offset, PrepareBuf(length));
+}
 
-  auto io_cb = [this, cb, offset, buf, len](int io_res) {
+void DiskStorage::Stash(DiskSegment segment, UringBuf buf, StashCb cb) {
+  auto io_cb = [this, cb, buf, segment](int io_res) {
     if (io_res < 0) {
-      MarkAsFree({size_t(offset), len});
-      cb(nonstd::make_unexpected(error_code{-io_res, std::system_category()}));
+      MarkAsFree(segment);
+      cb(error_code{-io_res, std::system_category()});
     } else {
-      cb(DiskSegment{size_t(offset), len});
+      cb({});
     }
     ReturnBuf(buf);
     pending_ops_--;
   };
 
   pending_ops_++;
+  size_t offset = segment.offset;
   if (buf.buf_idx)
     backing_file_->WriteFixedAsync(buf.bytes, offset, *buf.buf_idx, std::move(io_cb));
   else
@@ -183,8 +190,6 @@ std::error_code DiskStorage::Stash(io::Bytes bytes, StashCb cb) {
     auto ec = RequestGrow(256_MB);
     LOG_IF(ERROR, ec && ec != errc::file_too_large) << "Could not call grow :" << ec.message();
   }
-
-  return {};  // Must succeed after the operation was scheduled to run cleanup
 }
 
 DiskStorage::Stats DiskStorage::GetStats() const {

src/server/tiering/disk_storage.h

@@ -29,7 +29,7 @@ class DiskStorage {
   };
 
   using ReadCb = std::function<void(io::Result<std::string_view>)>;
-  using StashCb = std::function<void(io::Result<DiskSegment>)>;
+  using StashCb = std::function<void(std::error_code)>;
 
   explicit DiskStorage(size_t max_size);
 
@@ -42,11 +42,13 @@ class DiskStorage {
   // Mark segment as free, performed immediately
   void MarkAsFree(DiskSegment segment);
 
-  // Request bytes to be stored, cb will be called with assigned segment on completion. Can block to
-  // grow backing file. Returns error code if operation failed  immediately (most likely it failed
-  // to grow the backing file) or passes an empty segment if the final write operation failed.
-  // Bytes are copied and can be dropped before cb is resolved
-  std::error_code Stash(io::Bytes bytes, StashCb cb);
+  // Allocate segment of at least given length and prepare buffer. Might block to grow backing file.
+  // Return error if not enough space is available or growing failed.
+  // Every successful preparation must end in a Stash(), otherwise resources are leaked.
+  io::Result<std::pair<size_t /* offset */, util::fb2::UringBuf>> PrepareStash(size_t length);
+
+  // Write prepared buffer to given segment and resolve completion callback when write is done.
+  void Stash(DiskSegment segment, util::fb2::UringBuf buf, StashCb cb);
 
   Stats GetStats() const;
 

src/server/tiering/disk_storage_test.cc

@@ -39,11 +39,14 @@ struct DiskStorageTest : public PoolTestBase {
 
   void Stash(size_t index, string value) {
     pending_ops_++;
-    auto buf = make_shared<string>(value);
-    storage_->Stash(io::Buffer(*buf), [this, index, buf](io::Result<DiskSegment> segment) {
-      if (segment.has_value()) {
-        EXPECT_GT(segment->length, 0u);
-      }
+
+    auto prepared = storage_->PrepareStash(value.length());
+    EXPECT_TRUE(prepared.has_value());
+    auto [offset, buf] = *prepared;
+    memcpy(buf.bytes.data(), value.data(), value.size());
+
+    DiskSegment segment{offset, value.size()};
+    storage_->Stash({offset, value.size()}, buf, [this, index, segment](std::error_code ec) {
       segments_[index] = segment;
       pending_ops_--;
     });

src/server/tiering/op_manager.cc

@@ -87,20 +87,17 @@ void OpManager::DeleteOffloaded(DiskSegment segment) {
   }
 }
 
-std::error_code OpManager::Stash(EntryId id_ref, std::string_view value) {
+void OpManager::Stash(EntryId id_ref, tiering::DiskSegment segment, util::fb2::UringBuf buf) {
   auto id = ToOwned(id_ref);
   unsigned version = pending_stash_ver_[id] = ++pending_stash_counter_;
 
-  io::Bytes buf_view = io::Buffer(value);
-  auto io_cb = [this, version, id = std::move(id)](io::Result<DiskSegment> segment) {
-    ProcessStashed(Borrowed(id), version, segment);
+  auto io_cb = [this, version, id = std::move(id), segment](std::error_code ec) {
+    ProcessStashed(Borrowed(id), version,
+                   ec ? nonstd::make_unexpected(ec) : io::Result<DiskSegment>(segment));
   };
 
   // May block due to blocking call to Grow.
-  auto ec = storage_.Stash(buf_view, std::move(io_cb));
-  if (ec)
-    pending_stash_ver_.erase(ToOwned(id_ref));
-  return ec;
+  storage_.Stash(segment, buf, std::move(io_cb));
 }
 
 OpManager::ReadOp& OpManager::PrepareRead(DiskSegment aligned_segment) {

src/server/tiering/op_manager.h

@@ -63,8 +63,12 @@ class OpManager {
   // Delete offloaded entry located at the segment.
   void DeleteOffloaded(DiskSegment segment);
 
+  auto PrepareStash(size_t length) {
+    return storage_.PrepareStash(length);
+  }
+
   // Stash value to be offloaded. It is opaque to OpManager.
-  std::error_code Stash(EntryId id, std::string_view value);
+  void Stash(EntryId id, tiering::DiskSegment segment, util::fb2::UringBuf buf);
 
   Stats GetStats() const;
 

src/server/tiering/op_manager_test.cc

@@ -82,6 +82,18 @@ struct OpManagerTest : PoolTestBase, OpManager {
     return true;
   }
 
+  std::error_code Stash(EntryId id, std::string_view value) {
+    auto prepared = storage_.PrepareStash(value.size());
+    if (!prepared.has_value())
+      return prepared.error();
+
+    auto [offset, buf] = *prepared;
+    memcpy(buf.bytes.data(), value.data(), value.size());
+    DiskSegment segment{offset, value.size()};
+    OpManager::Stash(id, segment, buf);
+    return {};
+  }
+
   absl::flat_hash_map<EntryId, std::string> fetched_;
   absl::flat_hash_map<EntryId, DiskSegment> stashed_;
 };