Add checks to terminate early when backup is stopped (#14129)

Summary:
I want to reduce the time from when we call `StopBackup` to `CreateNewBackup` returning `BackupStopped`. We already check for the `stop_backup_` inside `CopyOrCreateFile` and `ReadFileAndComputeChecksum`, but we should add a check at the top of these methods to abort immediately. This could help save some latency from the file system metadata operations, like creating the sequential file and writable file.

We also want to update the API documentation for `StopBackup` which currently does not indicate that once it is called, all subsequent requests to create backups will fail.

In a follow up PR, we should also add coverage of `StopBackup` to the crash tests.

Pull Request resolved: https://github.com/facebook/rocksdb/pull/14129

Test Plan:
We were missing unit test coverage for `StopBackup`. I added test cases which cancel backups at different points in time.

Once this change is rolled out to production, we can monitor the DB close latencies, which depend on first cancelling ongoing backups

Reviewed By: pdillinger

Differential Revision: D87356536

Pulled By: archang19

fbshipit-source-id: 687094a41f096f6a156be65b2cce0b5054fb26f2

include/rocksdb/utilities/backup_engine.h

@@ -621,7 +621,14 @@ class BackupEngineAppendOnlyBase {
   // The backup will stop ASAP and the call to CreateNewBackup will
   // return Status::Incomplete(). It will not clean up after itself, but
   // the state will remain consistent. The state will be cleaned up the
-  // next time you call CreateNewBackup or GarbageCollect.
+  // next time you call CreateNewBackup or GarbageCollect for the same backup
+  // directory on a new BackupEngine object.
+  //
+  // NOTE: This is a one-way operation. Once StopBackup() is called on a
+  // BackupEngine instance, all subsequent backup requests (CreateNewBackup,
+  // CreateNewBackupWithMetadata) will fail with Status::Incomplete().
+  // To create new backups after calling StopBackup(), you must open a new
+  // BackupEngine instance.
   virtual void StopBackup() = 0;
 
   // Will delete any files left over from incomplete creation or deletion of

utilities/backup/backup_engine.cc

@@ -615,6 +615,10 @@ class BackupEngineImpl {
                                       std::string* checksum_hex,
                                       const Temperature src_temperature) const;
 
+  // Helper method to check if backup should be stopped. Can be overridden
+  // via sync points for testing.
+  bool ShouldStopBackup() const;
+
   // Obtain db_id and db_session_id from the table properties of file_path
   Status GetFileDbIdentities(Env* src_env, const EnvOptions& src_env_options,
                              const std::string& file_path,
@@ -2353,6 +2357,10 @@ IOStatus BackupEngineImpl::CopyOrCreateFile(
     Temperature dst_temperature, uint64_t* bytes_toward_next_callback,
     uint64_t* size, std::string* checksum_hex) {
   assert(src.empty() != contents.empty());
+  if (ShouldStopBackup()) {
+    return status_to_io_status(Status::Incomplete("Backup stopped"));
+  }
+
   IOStatus io_s;
   std::unique_ptr<FSWritableFile> dst_file;
   std::unique_ptr<FSSequentialFile> src_file;
@@ -2413,7 +2421,7 @@ IOStatus BackupEngineImpl::CopyOrCreateFile(
   Slice data;
   const IOOptions opts;
   do {
-    if (stop_backup_.load(std::memory_order_acquire)) {
+    if (ShouldStopBackup()) {
       return status_to_io_status(Status::Incomplete("Backup stopped"));
     }
     if (!src.empty()) {
@@ -2749,6 +2757,12 @@ IOStatus BackupEngineImpl::AddBackupFileWorkItem(
   return IOStatus::OK();
 }
 
+bool BackupEngineImpl::ShouldStopBackup() const {
+  bool should_stop = stop_backup_.load(std::memory_order_acquire);
+  TEST_SYNC_POINT_CALLBACK("BackupEngineImpl::ShouldStopBackup", &should_stop);
+  return should_stop;
+}
+
 IOStatus BackupEngineImpl::ReadFileAndComputeChecksum(
     const std::string& src, const std::shared_ptr<FileSystem>& src_fs,
     const EnvOptions& src_env_options, uint64_t size_limit,
@@ -2756,6 +2770,9 @@ IOStatus BackupEngineImpl::ReadFileAndComputeChecksum(
   if (checksum_hex == nullptr) {
     return status_to_io_status(Status::Aborted("Checksum pointer is null"));
   }
+  if (ShouldStopBackup()) {
+    return status_to_io_status(Status::Incomplete("Backup stopped"));
+  }
   uint32_t checksum_value = 0;
   if (size_limit == 0) {
     size_limit = std::numeric_limits<uint64_t>::max();
@@ -2783,7 +2800,7 @@ IOStatus BackupEngineImpl::ReadFileAndComputeChecksum(
   Slice data;
 
   do {
-    if (stop_backup_.load(std::memory_order_acquire)) {
+    if (ShouldStopBackup()) {
       return status_to_io_status(Status::Incomplete("Backup stopped"));
     }
     size_t buffer_to_read =

utilities/backup/backup_engine_test.cc

@@ -43,6 +43,7 @@
 #include "test_util/sync_point.h"
 #include "test_util/testharness.h"
 #include "test_util/testutil.h"
+#include "util/atomic.h"
 #include "util/cast_util.h"
 #include "util/mutexlock.h"
 #include "util/random.h"
@@ -4790,6 +4791,79 @@ TEST_F(BackupEngineTest, IOBufferSize) {
   ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->ClearAllCallBacks();
 }
 
+// Test stopping backup at different points in the backup lifecycle
+// Uses randomized stop points with geometric distribution to better catch
+// edge cases across multiple iterations.
+TEST_F(BackupEngineTest, StopBackupAtDifferentStages) {
+  const int keys_iteration = 5000;
+  const int num_iterations = 10;
+
+  // Enable multi-threaded backup
+  engine_options_->max_background_operations = 7;
+
+  // Generate DB once and reuse across iterations
+  OpenDBAndBackupEngine(true);
+  FillDB(db_.get(), 0, keys_iteration);
+
+  Random rnd(301);
+
+  for (int iteration = 0; iteration < num_iterations; iteration++) {
+    // Generate stop threshold using skewed distribution
+    // Smaller numbers are more likely, which is more interesting for testing
+    // Range: [0, 2^7-1] = [0, 127] with exponential bias towards 0
+    int stop_after_calls = rnd.Skewed(7);
+
+    RelaxedAtomic<int> call_count{0};
+    ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->SetCallBack(
+        "BackupEngineImpl::ShouldStopBackup", [&](void* arg) {
+          call_count.FetchAddRelaxed(1);
+          if (call_count.LoadRelaxed() > stop_after_calls) {
+            bool* should_stop = static_cast<bool*>(arg);
+            *should_stop = true;
+          }
+        });
+    ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->EnableProcessing();
+
+    // Create backup - it may complete successfully or be stopped
+    IOStatus s = backup_engine_->CreateNewBackup(db_.get());
+
+    // Verify that ShouldStopBackup was called
+    ASSERT_GT(call_count.LoadRelaxed(), 0);
+
+    if (s.IsIncomplete()) {
+      // Backup was stopped - verify it's the expected error
+      ASSERT_TRUE(s.ToString().find("Backup stopped") != std::string::npos)
+          << "Unexpected incomplete status for threshold " << stop_after_calls
+          << ": " << s.ToString();
+      ASSERT_GT(call_count.LoadRelaxed(), stop_after_calls)
+          << "Expected call_count > stop_after_calls";
+
+      // Verify that no valid backup was created
+      std::vector<BackupInfo> backup_info;
+      backup_engine_->GetBackupInfo(&backup_info);
+      ASSERT_EQ(0, backup_info.size());
+    } else {
+      // Backup completed successfully before reaching the stop threshold
+      ASSERT_OK(s) << "Unexpected error for threshold " << stop_after_calls;
+      ASSERT_LE(call_count.LoadRelaxed(), stop_after_calls)
+          << "Backup completed but call_count exceeded threshold";
+
+      // Verify a backup was created
+      std::vector<BackupInfo> backup_info;
+      backup_engine_->GetBackupInfo(&backup_info);
+      ASSERT_EQ(1, backup_info.size());
+
+      // Clean up the successful backup for next iteration
+      ASSERT_OK(backup_engine_->DeleteBackup(backup_info[0].backup_id));
+    }
+
+    ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->DisableProcessing();
+    ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->ClearAllCallBacks();
+  }
+
+  CloseDBAndBackupEngine();
+}
+
 }  // namespace
 
 }  // namespace ROCKSDB_NAMESPACE