Add io_buffer_size to BackupEngineOptions (facebook#13236)

Summary:
The RocksDB backup engine code currently derives the IO buffer size based on the following criteria:

1. If specified, use the rate limiter burst size
2. Otherwise, use the default size (5 MiB)

We want to be able to explicitly choose the IO size based on the storage backend. We want the new criteria to be:

1. If specified, use the size in `BackupEngineOptions`
2. If specified, use the rate limiter burst size
3. Otherwise, use the default size (5 MiB)

This PR adds a new option called `io_buffer_size` to `BackupEngineOptions` and updates the logic used to set the buffer size.

Pull Request resolved: facebook#13236

Test Plan:
I added a separate unit test and verified that we can either use the `io_buffer_size`, rate limiter burst size, or the default size.

I decided to use a `TEST_SYNC_POINT_CALLBACK`. I considered the alternative of updating the `Read` implementation of `DummySequentialFile` / `CheckIOOptsSequentialFile` to check the value of `n`. However, that would have considerably complicated the whole test code, and we also do not need to be checking for this in every single test case. I think the `TEST_SYNC_POINT_CALLBACK` turned out to be quite elegant.

Reviewed By: sushilpa

Differential Revision: D67765000

Pulled By: archang19

fbshipit-source-id: 2122fab7379335de44ba4423af47aa0563635688

include/rocksdb/utilities/backup_engine.h

@@ -77,6 +77,15 @@ struct BackupEngineOptions {
   // Default: true
   bool backup_log_files;
 
+  // Size of the buffer in bytes used for reading files.
+  // Enables optimally configuring the IO size based on the storage backend.
+  // If specified, takes precendence over the rate limiter burst size (if
+  // specified) as well as kDefaultCopyFileBufferSize.
+  // If 0, the rate limiter burst size (if specified) or
+  // kDefaultCopyFileBufferSize will be used.
+  // Default: 0
+  uint64_t io_buffer_size;
+
   // Max bytes that can be transferred in a second during backup.
   // If 0, go as fast as you can
   // This limit only applies to writes. To also limit reads,
@@ -228,8 +237,9 @@ struct BackupEngineOptions {
       const std::string& _backup_dir, Env* _backup_env = nullptr,
       bool _share_table_files = true, Logger* _info_log = nullptr,
       bool _sync = true, bool _destroy_old_data = false,
-      bool _backup_log_files = true, uint64_t _backup_rate_limit = 0,
-      uint64_t _restore_rate_limit = 0, int _max_background_operations = 1,
+      bool _backup_log_files = true, uint64_t _io_buffer_size = 0,
+      uint64_t _backup_rate_limit = 0, uint64_t _restore_rate_limit = 0,
+      int _max_background_operations = 1,
       uint64_t _callback_trigger_interval_size = 4 * 1024 * 1024,
       int _max_valid_backups_to_open = INT_MAX,
       ShareFilesNaming _share_files_with_checksum_naming =
@@ -241,6 +251,7 @@ struct BackupEngineOptions {
         sync(_sync),
         destroy_old_data(_destroy_old_data),
         backup_log_files(_backup_log_files),
+        io_buffer_size(_io_buffer_size),
         backup_rate_limit(_backup_rate_limit),
         restore_rate_limit(_restore_rate_limit),
         share_files_with_checksum(true),

unreleased_history/public_api_changes/backup_engine_options_io_buffer_size.md

@@ -0,0 +1 @@
+Add `io_buffer_size` to BackupEngineOptions to enable optimal configuration of IO size

utilities/backup/backup_engine.cc

@@ -581,6 +581,8 @@ class BackupEngineImpl {
                             uint64_t* bytes_toward_next_callback,
                             uint64_t* size, std::string* checksum_hex);
 
+  uint64_t CalculateIOBufferSize(RateLimiter* rate_limiter) const;
+
   IOStatus ReadFileAndComputeChecksum(const std::string& src,
                                       const std::shared_ptr<FileSystem>& src_fs,
                                       const EnvOptions& src_env_options,
@@ -2212,9 +2214,9 @@ IOStatus BackupEngineImpl::CopyOrCreateFile(
     return io_s;
   }
 
-  size_t buf_size =
-      rate_limiter ? static_cast<size_t>(rate_limiter->GetSingleBurstBytes())
-                   : kDefaultCopyFileBufferSize;
+  size_t buf_size = CalculateIOBufferSize(rate_limiter);
+  TEST_SYNC_POINT_CALLBACK(
+      "BackupEngineImpl::CopyOrCreateFile:CalculateIOBufferSize", &buf_size);
 
   // TODO: pass in Histograms if the destination file is sst or blob
   std::unique_ptr<WritableFileWriter> dest_writer(
@@ -2307,6 +2309,16 @@ IOStatus BackupEngineImpl::CopyOrCreateFile(
   return io_s;
 }
 
+uint64_t BackupEngineImpl::CalculateIOBufferSize(
+    RateLimiter* rate_limiter) const {
+  if (options_.io_buffer_size > 0) {
+    return options_.io_buffer_size;
+  }
+  return rate_limiter != nullptr
+             ? static_cast<size_t>(rate_limiter->GetSingleBurstBytes())
+             : kDefaultCopyFileBufferSize;
+}
+
 // fname will always start with "/"
 IOStatus BackupEngineImpl::AddBackupFileWorkItem(
     std::unordered_set<std::string>& live_dst_paths,
@@ -2588,7 +2600,7 @@ IOStatus BackupEngineImpl::ReadFileAndComputeChecksum(
     return io_s;
   }
 
-  size_t buf_size = kDefaultCopyFileBufferSize;
+  size_t buf_size = CalculateIOBufferSize(rate_limiter);
   std::unique_ptr<char[]> buf(new char[buf_size]);
   Slice data;
 

utilities/backup/backup_engine_test.cc

@@ -4494,6 +4494,60 @@ TEST_F(BackupEngineTest, ExcludeFiles) {
   delete alt_backup_engine;
 }
 
+TEST_F(BackupEngineTest, IOBufferSize) {
+  size_t expected_buffer_size = 5 * 1024 * 1024;
+  std::atomic<bool> io_buffer_size_calculated{false};
+  SyncPoint::GetInstance()->SetCallBack(
+      "BackupEngineImpl::CopyOrCreateFile:CalculateIOBufferSize",
+      [&](void* data) {
+        if (data != nullptr) {
+          EXPECT_EQ(expected_buffer_size, *static_cast<uint64_t*>(data));
+        }
+        io_buffer_size_calculated = true;
+      });
+  ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->EnableProcessing();
+
+  const int keys_iteration = 5000;
+  Random rnd(6);
+  // With no overrides, fall back to the default buffer size of 5 MB
+  OpenDBAndBackupEngine(true /* destroy_old_data */);
+  FillDB(db_.get(), 0, keys_iteration);
+  ASSERT_OK(backup_engine_->CreateNewBackup(db_.get(), false));
+  ASSERT_TRUE(io_buffer_size_calculated);
+  CloseDBAndBackupEngine();
+
+  // Override the default buffer size to 64 MB through BackupEngineOptions
+  expected_buffer_size = 64 * 1024 * 1024;
+  engine_options_->io_buffer_size = expected_buffer_size;
+  io_buffer_size_calculated = false;
+  OpenDBAndBackupEngine(true /* destroy_old_data */);
+  FillDB(db_.get(), 0, keys_iteration);
+  ASSERT_OK(backup_engine_->CreateNewBackup(db_.get(), false));
+  ASSERT_TRUE(io_buffer_size_calculated);
+  CloseDBAndBackupEngine();
+
+  // Without io_buffer_size specified, the rate limiter burst bytes value will
+  // be used (16 MB in this example)
+  engine_options_->io_buffer_size = 0;
+  size_t single_burst_bytes = 16 * 1024 * 1024;
+  expected_buffer_size = single_burst_bytes;
+  std::shared_ptr<RateLimiter> backup_rate_limiter(NewGenericRateLimiter(
+      5 * 1024 * 1024 /* rate_bytes_per_sec */,
+      100 * 1000 /* refill_period_us */, 10 /* fairness */,
+      RateLimiter::Mode::kWritesOnly /* mode */, false /* auto_tuned */,
+      single_burst_bytes /* single_burst_bytes */));
+  engine_options_->backup_rate_limiter = backup_rate_limiter;
+  io_buffer_size_calculated = false;
+  OpenDBAndBackupEngine(true /* destroy_old_data */);
+  FillDB(db_.get(), 0, keys_iteration);
+  ASSERT_OK(backup_engine_->CreateNewBackup(db_.get(), false));
+  ASSERT_TRUE(io_buffer_size_calculated);
+  CloseDBAndBackupEngine();
+
+  ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->DisableProcessing();
+  ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->ClearAllCallBacks();
+}
+
 }  // namespace
 
 }  // namespace ROCKSDB_NAMESPACE