fix: data race in save_stages_controller (#2647)

* fix: data race in save_stages_controller
dragonflydb · Feb 23, 2024 · a5af87e · a5af87e
1 parent 7f93a8f
commit a5af87e
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 6 deletions.
diff --git a/src/server/detail/save_stages_controller.cc b/src/server/detail/save_stages_controller.cc
@@ -149,7 +149,7 @@ SaveStagesController::SaveStagesController(SaveStagesInputs&& inputs)
 SaveStagesController::~SaveStagesController() {
 }
 
-SaveInfo SaveStagesController::Save() {
+std::optional<SaveInfo> SaveStagesController::InitResourcesAndStart() {
   if (auto err = BuildFullPath(); err) {
     shared_err_ = err;
     return GetSaveInfo();
@@ -162,8 +162,14 @@ SaveInfo SaveStagesController::Save() {
   else
     SaveRdb();
 
+  return {};
+}
+
+void SaveStagesController::WaitAllSnapshots() {
   RunStage(&SaveStagesController::SaveCb);
+}
 
+SaveInfo SaveStagesController::Finalize() {
   RunStage(&SaveStagesController::CloseCb);
 
   FinalizeFileMovement();

diff --git a/src/server/detail/save_stages_controller.h b/src/server/detail/save_stages_controller.h
@@ -71,10 +71,21 @@ class RdbSnapshot {
 
 struct SaveStagesController : public SaveStagesInputs {
   SaveStagesController(SaveStagesInputs&& input);
+  // Objects of this class are used concurrently. Call this function
+  // in a mutually exlusive context to avoid data races.
+  // Also call this function before any call to `WaitAllSnapshots`
+  // Returns empty optional on success and SaveInfo on failure
+  std::optional<SaveInfo> InitResourcesAndStart();
 
   ~SaveStagesController();
 
-  SaveInfo Save();
+  // Safe to call and no locks required
+  void WaitAllSnapshots();
+
+  // Same semantics as InitResourcesAndStart. Must be used in a mutually exclusive
+  // context. Call this function after you `WaitAllSnapshots`to finalize the chore.
+  // Performs cleanup of the object internally.
+  SaveInfo Finalize();
   size_t GetSaveBuffersSize();
   uint32_t GetCurrentSaveDuration();
 

diff --git a/src/server/server_family.cc b/src/server/server_family.cc
@@ -1310,19 +1310,29 @@ GenericError ServerFamily::DoSave(bool new_version, string_view basename, Transa
       return GenericError{make_error_code(errc::operation_in_progress),
                           "SAVING - can not save database"};
     }
+
     save_controller_ = make_unique<SaveStagesController>(detail::SaveStagesInputs{
         new_version, basename, trans, &service_, fq_threadpool_.get(), snapshot_storage_});
+
+    auto res = save_controller_->InitResourcesAndStart();
+
+    if (res) {
+      DCHECK_EQ(res->error, true);
+      last_save_info_.SetLastSaveError(*res);
+      save_controller_.reset();
+      return res->error;
+    }
   }
 
-  detail::SaveInfo save_info = save_controller_->Save();
+  save_controller_->WaitAllSnapshots();
+  detail::SaveInfo save_info;
 
   {
     std::lock_guard lk(save_mu_);
+    save_info = save_controller_->Finalize();
 
     if (save_info.error) {
-      last_save_info_.last_error = save_info.error;
-      last_save_info_.last_error_time = save_info.save_time;
-      last_save_info_.failed_duration_sec = save_info.duration_sec;
+      last_save_info_.SetLastSaveError(save_info);
     } else {
       last_save_info_.save_time = save_info.save_time;
       last_save_info_.success_duration_sec = save_info.duration_sec;

diff --git a/src/server/server_family.h b/src/server/server_family.h
@@ -104,6 +104,12 @@ struct Metrics {
 
 struct LastSaveInfo {
   // last success save info
+  void SetLastSaveError(const detail::SaveInfo& save_info) {
+    last_error = save_info.error;
+    last_error_time = save_info.save_time;
+    failed_duration_sec = save_info.duration_sec;
+  }
+
   time_t save_time = 0;  // epoch time in seconds.
   uint32_t success_duration_sec = 0;
   std::string file_name;                                      //