From 3945b7e4fa340b06e0de16b88d985070f80a5ce4 Mon Sep 17 00:00:00 2001
From: Roman Gershman <roman@dragonflydb.io>
Date: Thu, 26 Sep 2024 18:42:28 +0300
Subject: [PATCH] chore: tune TieredStorageTest.MemoryPressure (#3805)

* chore: tune TieredStorageTest.MemoryPressure

* chore: print more stats on failure
---
 src/server/tiered_storage.cc          |  5 +++--
 src/server/tiered_storage.h           |  2 +-
 src/server/tiered_storage_test.cc     |  5 +++--
 src/server/tiering/op_manager_test.cc | 12 +++++++++++-
 4 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/src/server/tiered_storage.cc b/src/server/tiered_storage.cc
index 76c1b1253..eb323aab9 100644
--- a/src/server/tiered_storage.cc
+++ b/src/server/tiered_storage.cc
@@ -312,7 +312,7 @@ TieredStorage::TieredStorage(size_t max_size, DbSlice* db_slice)
       bins_{make_unique<tiering::SmallBins>()} {
   write_depth_limit_ = absl::GetFlag(FLAGS_tiered_storage_write_depth);
   size_t mem_per_shard = max_memory_limit / shard_set->size();
-  SetMemoryLowLimit(absl::GetFlag(FLAGS_tiered_low_memory_factor) * mem_per_shard);
+  SetMemoryLowWatermark(absl::GetFlag(FLAGS_tiered_low_memory_factor) * mem_per_shard);
 }
 
 TieredStorage::~TieredStorage() {
@@ -329,7 +329,7 @@ void TieredStorage::Close() {
   op_manager_->Close();
 }
 
-void TieredStorage::SetMemoryLowLimit(size_t mem_limit) {
+void TieredStorage::SetMemoryLowWatermark(size_t mem_limit) {
   op_manager_->memory_low_limit_ = mem_limit;
   VLOG(1) << "Memory low limit is " << mem_limit;
 }
@@ -338,6 +338,7 @@ util::fb2::Future<string> TieredStorage::Read(DbIndex dbid, string_view key,
                                               const PrimeValue& value) {
   util::fb2::Future<std::string> fut;
   Read(dbid, key, value, [fut](const std::string& value) mutable { fut.Resolve(value); });
+
   return fut;
 }
 
diff --git a/src/server/tiered_storage.h b/src/server/tiered_storage.h
index 0fc2b530b..c0a84efb1 100644
--- a/src/server/tiered_storage.h
+++ b/src/server/tiered_storage.h
@@ -44,7 +44,7 @@ class TieredStorage {
   std::error_code Open(std::string_view path);
   void Close();
 
-  void SetMemoryLowLimit(size_t mem_limit);
+  void SetMemoryLowWatermark(size_t mem_limit);
 
   // Read offloaded value. It must be of external type
   util::fb2::Future<std::string> Read(DbIndex dbid, std::string_view key, const PrimeValue& value);
diff --git a/src/server/tiered_storage_test.cc b/src/server/tiered_storage_test.cc
index 50313933c..038d72e9c 100644
--- a/src/server/tiered_storage_test.cc
+++ b/src/server/tiered_storage_test.cc
@@ -301,7 +301,8 @@ TEST_F(TieredStorageTest, FlushPending) {
 
 TEST_F(TieredStorageTest, MemoryPressure) {
   max_memory_limit = 20_MB;
-  pp_->at(0)->AwaitBrief([] { EngineShard::tlocal()->tiered_storage()->SetMemoryLowLimit(2_MB); });
+  pp_->at(0)->AwaitBrief(
+      [] { EngineShard::tlocal()->tiered_storage()->SetMemoryLowWatermark(2_MB); });
 
   constexpr size_t kNum = 10000;
   for (size_t i = 0; i < kNum; i++) {
@@ -310,7 +311,7 @@ TEST_F(TieredStorageTest, MemoryPressure) {
       resp = Run({"INFO", "ALL"});
       ASSERT_FALSE(true) << i << "\nInfo ALL:\n" << resp.GetString();
     }
-    ThisFiber::SleepFor(300us);
+    ThisFiber::SleepFor(500us);
   }
 
   EXPECT_LT(used_mem_peak.load(), 20_MB);
diff --git a/src/server/tiering/op_manager_test.cc b/src/server/tiering/op_manager_test.cc
index c2aa5cf52..baffc62e5 100644
--- a/src/server/tiering/op_manager_test.cc
+++ b/src/server/tiering/op_manager_test.cc
@@ -20,6 +20,16 @@ namespace dfly::tiering {
 using namespace std;
 using namespace std::string_literals;
 
+ostream& operator<<(ostream& os, const OpManager::Stats& stats) {
+  return os << "pending_read_cnt: " << stats.pending_read_cnt
+            << ", pending_stash_cnt: " << stats.pending_stash_cnt
+            << ", alloc_bytes: " << stats.disk_stats.allocated_bytes
+            << ", capacity_bytes: " << stats.disk_stats.capacity_bytes
+            << ", heap_buf_allocs: " << stats.disk_stats.heap_buf_alloc_count
+            << ", registered_buf_allocs: " << stats.disk_stats.registered_buf_alloc_count
+            << ", max_file_size: " << stats.disk_stats.max_file_size;
+}
+
 struct OpManagerTest : PoolTestBase, OpManager {
   OpManagerTest() : OpManager(256_MB) {
   }
@@ -76,7 +86,7 @@ TEST_F(OpManagerTest, SimpleStashesWithReads) {
     while (stashed_.size() < 100)
       util::ThisFiber::SleepFor(1ms);
 
-    EXPECT_EQ(GetStats().disk_stats.allocated_bytes, 100 * kPageSize);
+    EXPECT_EQ(GetStats().disk_stats.allocated_bytes, 100 * kPageSize) << GetStats();
 
     for (unsigned i = 0; i < 100; i++) {
       EXPECT_GE(stashed_[i].offset, i > 0);