chore: run memory decommit after snapshot load/save

src/server/memory_cmd.cc

@@ -144,10 +144,8 @@ void MemoryCmd::Run(CmdArgList args) {
   }
 
   if (sub_cmd == "DECOMMIT") {
-    shard_set->pool()->AwaitBrief([](unsigned, auto* pb) {
-      ServerState::tlocal()->DecommitMemory(ServerState::kDataHeap | ServerState::kBackingHeap |
-                                            ServerState::kGlibcmalloc);
-    });
+    shard_set->pool()->AwaitBrief(
+        [](unsigned, auto* pb) { ServerState::tlocal()->DecommitMemory(ServerState::kAllMemory); });
     return cntx_->SendSimpleString("OK");
   }
 

src/server/rdb_load.cc

@@ -1991,6 +1991,12 @@ RdbLoader::~RdbLoader() {
       break;
     delete item;
   }
+
+  // Decommit local memory.
+  // We create an RdbLoader for each thread, so each one will Decommit for itself after
+  // full sync ends (since we explicitly reset the RdbLoader).
+  auto* tlocal = ServerState::tlocal();
+  tlocal->DecommitMemory(ServerState::kAllMemory);
 }
 
 error_code RdbLoader::Load(io::Source* src) {

src/server/rdb_save.cc

@@ -1406,6 +1406,10 @@ RdbSaver::RdbSaver(::io::Sink* sink, SaveMode save_mode, bool align_writes) {
 }
 
 RdbSaver::~RdbSaver() {
+  // Decommit local memory.
+  // We create an RdbSaver for each thread, so each one will Decommit for itself.
+  auto* tlocal = ServerState::tlocal();
+  tlocal->DecommitMemory(ServerState::kAllMemory);
 }
 
 void RdbSaver::StartSnapshotInShard(bool stream_journal, const Cancellation* cll,

src/server/rdb_test.cc

@@ -41,6 +41,15 @@ class RdbTest : public BaseFamilyTest {
   void SetUp();
 
   io::FileSource GetSource(string name);
+
+  std::error_code LoadRdb(const string& filename) {
+    return pp_->at(0)->Await([&] {
+      io::FileSource fs = GetSource(filename);
+
+      RdbLoader loader(service_.get());
+      return loader.Load(&fs);
+    });
+  }
 };
 
 void RdbTest::SetUp() {
@@ -84,19 +93,12 @@ TEST_F(RdbTest, Crc) {
 }
 
 TEST_F(RdbTest, LoadEmpty) {
-  io::FileSource fs = GetSource("empty.rdb");
-  RdbLoader loader(NULL);
-  auto ec = loader.Load(&fs);
+  auto ec = LoadRdb("empty.rdb");
   CHECK(!ec);
 }
 
 TEST_F(RdbTest, LoadSmall6) {
-  io::FileSource fs = GetSource("redis6_small.rdb");
-  RdbLoader loader{service_.get()};
-
-  // must run in proactor thread in order to avoid polluting the serverstate
-  // in the main, testing thread.
-  auto ec = pp_->at(0)->Await([&] { return loader.Load(&fs); });
+  auto ec = LoadRdb("redis6_small.rdb");
 
   ASSERT_FALSE(ec) << ec.message();
 
@@ -128,12 +130,7 @@ TEST_F(RdbTest, LoadSmall6) {
 }
 
 TEST_F(RdbTest, Stream) {
-  io::FileSource fs = GetSource("redis6_stream.rdb");
-  RdbLoader loader{service_.get()};
-
-  // must run in proactor thread in order to avoid polluting the serverstate
-  // in the main, testing thread.
-  auto ec = pp_->at(0)->Await([&] { return loader.Load(&fs); });
+  auto ec = LoadRdb("redis6_stream.rdb");
 
   ASSERT_FALSE(ec) << ec.message();
 
@@ -447,12 +444,7 @@ TEST_F(RdbTest, JsonTest) {
 class HllRdbTest : public RdbTest, public testing::WithParamInterface<string> {};
 
 TEST_P(HllRdbTest, Hll) {
-  io::FileSource fs = GetSource("hll.rdb");
-  RdbLoader loader{service_.get()};
-
-  // must run in proactor thread in order to avoid polluting the serverstate
-  // in the main, testing thread.
-  auto ec = pp_->at(0)->Await([&] { return loader.Load(&fs); });
+  auto ec = LoadRdb("hll.rdb");
 
   ASSERT_FALSE(ec) << ec.message();
 
@@ -478,12 +470,7 @@ TEST_F(RdbTest, LoadSmall7) {
   // 2. A hashtable called my-hset encoded as RDB_TYPE_HASH_LISTPACK
   // 3. A set called my-set encoded as RDB_TYPE_SET_LISTPACK
   // 4. A zset called my-zset encoded as RDB_TYPE_ZSET_LISTPACK
-  io::FileSource fs = GetSource("redis7_small.rdb");
-  RdbLoader loader{service_.get()};
-
-  // must run in proactor thread in order to avoid polluting the serverstate
-  // in the main, testing thread.
-  auto ec = pp_->at(0)->Await([&] { return loader.Load(&fs); });
+  auto ec = LoadRdb("redis7_small.rdb");
 
   ASSERT_FALSE(ec) << ec.message();
 
@@ -520,12 +507,7 @@ TEST_F(RdbTest, RedisJson) {
   // JSON.SET json-obj $
   // '{"company":"DragonflyDB","product":"Dragonfly","website":"https://dragondlydb.io","years-active":[2021,2022,2023,2024,"and
   // more!"]}'
-  io::FileSource fs = GetSource("redis_json.rdb");
-  RdbLoader loader{service_.get()};
-
-  // must run in proactor thread in order to avoid polluting the serverstate
-  // in the main, testing thread.
-  auto ec = pp_->at(0)->Await([&] { return loader.Load(&fs); });
+  auto ec = LoadRdb("redis_json.rdb");
 
   ASSERT_FALSE(ec) << ec.message();
 

src/server/replica.cc

@@ -157,6 +157,9 @@ void Replica::Stop() {
   // so we can freely release resources (connections).
   sync_fb_.JoinIfNeeded();
   acks_fb_.JoinIfNeeded();
+  for (auto& flow : shard_flows_) {
+    flow.reset();
+  }
 }
 
 void Replica::Pause(bool pause) {

src/server/server_state.h

@@ -282,7 +282,12 @@ class ServerState {  // public struct - to allow initialization.
   // Decommits 3 possible heaps according to the flags.
   // For decommit_glibcmalloc the heap is global for the process, for others it's specific only
   // for this thread.
-  enum { kDataHeap = 1, kBackingHeap = 2, kGlibcmalloc = 4 };
+  enum {
+    kDataHeap = 1,
+    kBackingHeap = 2,
+    kGlibcmalloc = 4,
+    kAllMemory = kDataHeap | kBackingHeap | kGlibcmalloc
+  };
   void DecommitMemory(uint8_t flags);
 
   // Exec descriptor frequency count for this thread.