chore: fix

dragonflydb · May 25, 2024 · 8cc36d7 · 8cc36d7
1 parent a4b60dc
commit 8cc36d7
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 11 deletions.
diff --git a/src/server/tiered_storage_test.cc b/src/server/tiered_storage_test.cc
@@ -47,15 +47,6 @@ class TieredStorageTest : public BaseFamilyTest {
 
     BaseFamilyTest::SetUp();
   }
-
-  void TearDown() override {
-    TieredStats stats;
-    do {
-      util::ThisFiber::SleepFor(20ms);
-      stats = GetMetrics().tiered_stats;
-    } while (stats.pending_read_cnt + stats.pending_stash_cnt > 0);
-    BaseFamilyTest::TearDown();
-  }
 };
 
 // Perform simple series of SET, GETSET and GET

diff --git a/src/server/tiering/disk_storage.cc b/src/server/tiering/disk_storage.cc
@@ -41,7 +41,6 @@ UringBuf PrepareBuf(size_t size) {
   DCHECK_EQ(ProactorBase::me()->GetKind(), ProactorBase::IOURING);
   auto* up = static_cast<UringProactor*>(ProactorBase::me());
 
-  UringBuf buf;
   if (auto borrowed = up->RequestBuffer(size); borrowed)
     return *borrowed;
   else
@@ -77,6 +76,10 @@ std::error_code DiskStorage::Open(std::string_view path) {
 }
 
 void DiskStorage::Close() {
+  using namespace std::chrono_literals;
+  while (pending_ops_ > 0)
+    util::ThisFiber::SleepFor(10ms);
+
   io_mgr_.Shutdown();
 }
 
@@ -85,13 +88,16 @@ void DiskStorage::Read(DiskSegment segment, ReadCb cb) {
   DCHECK_EQ(segment.offset % kPageSize, 0u);
 
   UringBuf buf = PrepareBuf(segment.length);
-  auto io_cb = [cb = std::move(cb), buf, segment](int io_res) {
+  auto io_cb = [this, cb = std::move(cb), buf, segment](int io_res) {
     if (io_res < 0)
       cb("", std::error_code{-io_res, std::system_category()});
     else
       cb(std::string_view{reinterpret_cast<char*>(buf.bytes.data()), segment.length}, {});
     ReturnBuf(buf);
+    pending_ops_--;
   };
+
+  pending_ops_++;
   io_mgr_.ReadAsync(segment.offset, buf, std::move(io_cb));
 }
 
@@ -135,8 +141,10 @@ std::error_code DiskStorage::Stash(io::Bytes bytes, StashCb cb) {
       cb({size_t(offset), len}, {});
     }
     ReturnBuf(buf);
+    pending_ops_--;
   };
 
+  pending_ops_++;
   io_mgr_.WriteAsync(offset, buf, std::move(io_cb));
   return {};
 }

diff --git a/src/server/tiering/disk_storage.h b/src/server/tiering/disk_storage.h
@@ -45,6 +45,7 @@ class DiskStorage {
   Stats GetStats() const;
 
  private:
+  size_t pending_ops_ = 0;
   size_t max_size_;
   IoMgr io_mgr_;
   ExternalAllocator alloc_;