From 44ebc24dcae1a2f2fd0100a673887f2461e50d5d Mon Sep 17 00:00:00 2001
From: Zhichao Cao <zhichao@fb.com>
Date: Sat, 26 Dec 2020 22:05:42 -0800
Subject: [PATCH] Add rate_limiter to GenerateOneFileChecksum (#7811)

Summary:
In GenerateOneFileChecksum(), RocksDB reads the file and computes its checksum. A rate limiter can be passed to the constructor of RandomAccessFileReader so that read I/O can be rate limited.

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

Test Plan: make check

Reviewed By: cheng-chang

Differential Revision: D25699896

Pulled By: zhichao-cao

fbshipit-source-id: e2688bc1126c543979a3bcf91dda784bd7b74164
---
 db/db_impl/db_impl.cc                 | 2 +-
 db/external_sst_file_ingestion_job.cc | 5 +++--
 file/file_util.cc                     | 5 +++--
 file/file_util.h                      | 2 +-
 4 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/db/db_impl/db_impl.cc b/db/db_impl/db_impl.cc
index 3e417545d..4c204e468 100644
--- a/db/db_impl/db_impl.cc
+++ b/db/db_impl/db_impl.cc
@@ -4872,7 +4872,7 @@ Status DBImpl::VerifySstFileChecksum(const FileMetaData& fmeta,
       fs_.get(), fname, immutable_db_options_.file_checksum_gen_factory.get(),
       fmeta.file_checksum_func_name, &file_checksum, &func_name,
       read_options.readahead_size, immutable_db_options_.allow_mmap_reads,
-      io_tracer_);
+      io_tracer_, immutable_db_options_.rate_limiter.get());
   if (s.ok()) {
     assert(fmeta.file_checksum_func_name == func_name);
     if (file_checksum != fmeta.file_checksum) {
diff --git a/db/external_sst_file_ingestion_job.cc b/db/external_sst_file_ingestion_job.cc
index a08081b87..9bd064c80 100644
--- a/db/external_sst_file_ingestion_job.cc
+++ b/db/external_sst_file_ingestion_job.cc
@@ -201,7 +201,8 @@ Status ExternalSstFileIngestionJob::Prepare(
             requested_checksum_func_name, &generated_checksum,
             &generated_checksum_func_name,
             ingestion_options_.verify_checksums_readahead_size,
-            db_options_.allow_mmap_reads, io_tracer_);
+            db_options_.allow_mmap_reads, io_tracer_,
+            db_options_.rate_limiter.get());
         if (!io_s.ok()) {
           status = io_s;
           ROCKS_LOG_WARN(db_options_.info_log,
@@ -856,7 +857,7 @@ IOStatus ExternalSstFileIngestionJob::GenerateChecksumForIngestedFile(
       db_options_.file_checksum_gen_factory.get(), requested_checksum_func_name,
       &file_checksum, &file_checksum_func_name,
       ingestion_options_.verify_checksums_readahead_size,
-      db_options_.allow_mmap_reads, io_tracer_);
+      db_options_.allow_mmap_reads, io_tracer_, db_options_.rate_limiter.get());
   if (!io_s.ok()) {
     return io_s;
   }
diff --git a/file/file_util.cc b/file/file_util.cc
index 0e6f4a514..451ef7728 100644
--- a/file/file_util.cc
+++ b/file/file_util.cc
@@ -132,7 +132,7 @@ IOStatus GenerateOneFileChecksum(
     const std::string& requested_checksum_func_name, std::string* file_checksum,
     std::string* file_checksum_func_name,
     size_t verify_checksums_readahead_size, bool allow_mmap_reads,
-    std::shared_ptr<IOTracer>& io_tracer) {
+    std::shared_ptr<IOTracer>& io_tracer, RateLimiter* rate_limiter) {
   if (checksum_factory == nullptr) {
     return IOStatus::InvalidArgument("Checksum factory is invalid");
   }
@@ -173,7 +173,8 @@ IOStatus GenerateOneFileChecksum(
       return io_s;
     }
     reader.reset(new RandomAccessFileReader(std::move(r_file), file_path,
-                                            nullptr /*Env*/, io_tracer));
+                                            nullptr /*Env*/, io_tracer, nullptr,
+                                            0, nullptr, rate_limiter));
   }
 
   // Found that 256 KB readahead size provides the best performance, based on
diff --git a/file/file_util.h b/file/file_util.h
index a9b0a9509..de92c5e21 100644
--- a/file/file_util.h
+++ b/file/file_util.h
@@ -39,7 +39,7 @@ extern IOStatus GenerateOneFileChecksum(
     const std::string& requested_checksum_func_name, std::string* file_checksum,
     std::string* file_checksum_func_name,
     size_t verify_checksums_readahead_size, bool allow_mmap_reads,
-    std::shared_ptr<IOTracer>& io_tracer);
+    std::shared_ptr<IOTracer>& io_tracer, RateLimiter* rate_limiter = nullptr);
 
 inline IOStatus PrepareIOFromReadOptions(const ReadOptions& ro, Env* env,
                                          IOOptions& opts) {