Merge pull request facebook#12 from pikiwidb/async/download

baixin01 · web-flow · commit 47ed041da2fa · 2024-07-02T20:11:56.000+08:00
feat: support async download sst files when applying manifest
diff --git a/cloud/aws/aws_retry.cc b/cloud/aws/aws_retry.cc
@@ -154,6 +154,7 @@ Status AwsCloudOptions::GetClientConfiguration(
   config->useVirtualAddressing = false;
 
   config->region = ToAwsString(region);
+  config->throughputTargetGbps = cloud_fs_options.throughput_target_gbps;
   return Status::OK();
 }
 #else
diff --git a/cloud/aws/aws_s3.cc b/cloud/aws/aws_s3.cc
@@ -201,6 +201,12 @@ class AwsS3ClientWrapper {
     return outcome;
   }
 
+  void GetCloudObjectAsync(
+      const Aws::S3Crt::Model::GetObjectRequest& request,
+      const Aws::S3Crt::GetObjectResponseReceivedHandler& handler) {
+    client_->GetObjectAsync(request, handler, nullptr);
+  }
+
   template <class... Args>
   std::shared_ptr<Aws::Transfer::TransferHandle> DownloadFile(Args... args) {
     CloudRequestCallbackGuard guard(cloud_request_callback_.get(),
@@ -434,6 +440,11 @@ class S3StorageProvider : public CloudStorageProviderImpl {
       IODebugContext* dbg) override;
   Status PrepareOptions(const ConfigOptions& options) override;
  protected:
+  IOStatus DoGetCloudObjectAsync(
+      const std::string& bucket_name, const std::string& object_path,
+      const std::string& local_path,
+      std::shared_ptr<std::promise<bool>> prom_ptr) override;
+
   IOStatus DoGetCloudObject(const std::string& bucket_name,
                             const std::string& object_path,
                             const std::string& destination,
@@ -729,7 +740,7 @@ IOStatus S3StorageProvider::ExistsCloudObject(const std::string& bucket_name,
 IOStatus S3StorageProvider::GetCloudObjectSize(const std::string& bucket_name,
                                                const std::string& object_path,
                                                uint64_t* filesize) {
-  HeadObjectResult result;                                             
+  HeadObjectResult result;
   result.size = filesize;
   return HeadObject(bucket_name, object_path, &result);
 }
@@ -900,6 +911,15 @@ class WritableFileStreamBuf : public std::streambuf {
     *fileCloseStatus_ = fileWriter_->Close();
   }
 
+  // Flushes any buffered data
+  // move sync() from protected region to public
+  // because Close() is called after async handler been executed
+  // sync() must be called in async handler, or the file is incomplete
+  int sync() override {
+    auto st = fileWriter_->Flush();
+    return st.ok() ? 0 : -1;
+  }
+
  protected:
   // Appends a block of data to the stream. Must always write n if possible
   std::streamsize xsputn(const char* s, std::streamsize n) override {
@@ -924,17 +944,26 @@ class WritableFileStreamBuf : public std::streambuf {
     return ch;
   }
 
-  // Flushes any buffered data
-  int sync() override {
-    auto st = fileWriter_->Flush();
-    return st.ok() ? 0 : -1;
-  }
-
  private:
   IOStatus *fileCloseStatus_;
   std::unique_ptr<WritableFileWriter> fileWriter_;
 };
 
+// shared_ptr version of IOStreamWithOwnedBuf
+// because Writablefile pointer must be reachable in async callback
+// std::iostream takes a raw pointer to std::streambuf. This subclass
+// takes a shared_ptr to the streambuf, tying the std::streambuf's
+// lifetime to the iostream's.
+template <class T>
+class IOStreamWithOwnedBufSPtr : public std::iostream {
+ public:
+  IOStreamWithOwnedBufSPtr(std::shared_ptr<T> s)
+      : std::iostream(s.get()), s_(s) {}
+
+ private:
+  std::shared_ptr<T> s_;
+};
+
 // std::iostream takes a raw pointer to std::streambuf. This subclass
 // takes a unique_ptr to the streambuf, tying the std::streambuf's
 // lifetime to the iostream's.
@@ -950,6 +979,78 @@ class IOStreamWithOwnedBuf : public std::iostream {
 
 }  // namespace
 
+IOStatus S3StorageProvider::DoGetCloudObjectAsync(
+    const std::string& bucket_name, const std::string& object_path,
+    const std::string& local_path,
+    std::shared_ptr<std::promise<bool>> prom_ptr) {
+  std::string tmp_destination =
+      local_path + ".tmp-" + std::to_string(rng_.Next());
+
+  std::shared_ptr<IOStatus> fileCloseStatus = std::make_shared<IOStatus>();
+
+  FileOptions foptions;
+  foptions.use_direct_writes =
+      cfs_->GetCloudFileSystemOptions().use_direct_io_for_cloud_download;
+  std::unique_ptr<FSWritableFile> file;
+  auto st = NewWritableFile(cfs_->GetBaseFileSystem().get(), tmp_destination,
+                            &file, foptions);
+  if (!st.ok()) {
+    Log(InfoLogLevel::ERROR_LEVEL, cfs_->GetLogger(),
+        "create writeablefile for async download failed, msg: %s",
+        st.ToString().c_str());
+    prom_ptr->set_value(false);
+    return st;
+  }
+  std::shared_ptr<WritableFileStreamBuf> file_stream_buf =
+      std::make_shared<WritableFileStreamBuf>(
+          fileCloseStatus.get(),
+          std::unique_ptr<WritableFileWriter>(new WritableFileWriter(
+              std::move(file), tmp_destination, foptions)));
+
+  auto ioStreamFactory = [file_stream_buf,
+                          fileCloseStatus]() -> Aws::IOStream* {
+    return Aws::New<IOStreamWithOwnedBufSPtr<WritableFileStreamBuf>>(
+        Aws::Utils::ARRAY_ALLOCATION_TAG, file_stream_buf);
+  };
+
+  Aws::S3Crt::Model::GetObjectRequest request;
+  request.SetBucket(ToAwsString(bucket_name));
+  request.SetKey(ToAwsString(object_path));
+  request.SetResponseStreamFactory(std::move(ioStreamFactory));
+
+  auto handler = Aws::S3Crt::GetObjectResponseReceivedHandler{
+      [this, tmp_destination, local_path, prom_ptr, file_stream_buf](
+          const Aws::S3Crt::S3CrtClient*,
+          const Aws::S3Crt::Model::GetObjectRequest&,
+          Aws::S3Crt::Model::GetObjectOutcome outcome,
+          const std::shared_ptr<const Aws::Client::AsyncCallerContext>&) {
+        file_stream_buf->sync();
+        const auto& local_fs = cfs_->GetBaseFileSystem();
+        const IOOptions io_opts;
+        IODebugContext* dbg = nullptr;
+        auto remote_size = outcome.GetResult().GetContentLength();
+        uint64_t local_size{0};
+        auto s =
+            local_fs->GetFileSize(tmp_destination, io_opts, &local_size, dbg);
+        if (!outcome.IsSuccess() || !s.ok() ||
+            local_size != uint64_t(remote_size)) {
+          local_fs->DeleteFile(tmp_destination, io_opts, dbg);
+          Log(InfoLogLevel::ERROR_LEVEL, cfs_->GetLogger(),
+              "async download error, outcome: %d, local_size: %lu, "
+              "remote_size: %lu",
+              outcome.IsSuccess(), local_size, uint64_t(remote_size));
+          prom_ptr->set_value(false);
+          return;
+        }
+        local_fs->RenameFile(tmp_destination, local_path, io_opts, dbg);
+        cfs_->FileCacheInsert(local_path, local_size);
+        prom_ptr->set_value(true);
+      }};
+
+  s3client_->GetCloudObjectAsync(request, handler);
+  return IOStatus::OK();
+}
+
 IOStatus S3StorageProvider::DoGetCloudObject(const std::string& bucket_name,
                                              const std::string& object_path,
                                              const std::string& destination,
diff --git a/cloud/cloud_file_system_impl.cc b/cloud/cloud_file_system_impl.cc
@@ -78,6 +78,13 @@ IOStatus CloudFileSystemImpl::GetCloudObject(const std::string& fname) {
   return st;
 }
 
+IOStatus CloudFileSystemImpl::GetCloudObjectAsync(
+    const std::string& fname, std::shared_ptr<std::promise<bool>> prom_ptr) {
+  auto st = GetStorageProvider()->GetCloudObjectAsync(
+      GetDestBucketName(), destname(fname), fname, prom_ptr);
+  return st;
+}
+
 IOStatus CloudFileSystemImpl::GetCloudObjectSize(const std::string& fname,
                                                  uint64_t* remote_size) {
   auto st = IOStatus::NotFound();
@@ -134,6 +141,29 @@ IOStatus CloudFileSystemImpl::ListCloudObjects(
   return st;
 }
 
+IOStatus CloudFileSystemImpl::DownloadAsync(
+    const std::string& logical_fname,
+    std::shared_ptr<std::promise<bool>> prom_ptr) {
+  IOOptions io_opts;
+  auto fname = RemapFilename(logical_fname);
+  auto file_type = GetFileType(fname);
+  IOStatus st;
+  assert(file_type == RocksDBFileType::kSstFile);
+  Log(InfoLogLevel::INFO_LEVEL, info_log_, "async download fname: %s",
+      logical_fname.c_str(), fname.c_str());
+  if (cloud_fs_options.hasSstFileCache()) {
+    st = base_fs_->FileExists(fname, io_opts, nullptr);
+    if (st.ok()) {
+      FileCacheAccess(fname);
+      prom_ptr->set_value(true);
+      return st;
+    }
+  }
+
+  GetCloudObjectAsync(fname, prom_ptr);
+  return IOStatus::OK();
+}
+
 IOStatus CloudFileSystemImpl::NewCloudReadableFile(
     const std::string& fname, const FileOptions& options,
     std::unique_ptr<CloudStorageReadableFile>* result, IODebugContext* dbg) {
diff --git a/cloud/cloud_file_system_impl.h b/cloud/cloud_file_system_impl.h
@@ -59,6 +59,9 @@ class CloudFileSystemImpl : public CloudFileSystem {
                            std::unique_ptr<FSWritableFile>* result,
                            IODebugContext* dbg) override;
 
+  IOStatus DownloadAsync(const std::string& fname,
+                         std::shared_ptr<std::promise<bool>> prom_ptr) override;
+
   IOStatus ReopenWritableFile(const std::string& fname,
                               const FileOptions& options,
                               std::unique_ptr<FSWritableFile>* result,
@@ -155,7 +158,7 @@ class CloudFileSystemImpl : public CloudFileSystem {
 
   // Find all live files based on cloud_manifest_ and local MANIFEST FILE
   // If local MANIFEST file doesn't exist, it will pull from cloud
-  // 
+  //
   // REQUIRES: cloud_manifest_ is loaded
   // REQUIRES: cloud_manifest_ is not updated when calling this function
   IOStatus FindAllLiveFiles(const std::string& local_dbname,
@@ -381,6 +384,9 @@ class CloudFileSystemImpl : public CloudFileSystem {
  private:
   bool WaitPendingObjects() override;
 
+  IOStatus GetCloudObjectAsync(const std::string& fname,
+                               std::shared_ptr<std::promise<bool>> prom_ptr);
+
   // Delete all local files that are invisible
   IOStatus DeleteLocalInvisibleFiles(
       const std::string& dbname,
diff --git a/cloud/cloud_storage_provider.cc b/cloud/cloud_storage_provider.cc
@@ -3,6 +3,8 @@
 
 #include "rocksdb/cloud/cloud_storage_provider.h"
 
+#include <aws/s3-crt/S3CrtServiceClientModel.h>
+
 #include <cinttypes>
 #include <mutex>
 #include <set>
@@ -316,6 +318,14 @@ IOStatus CloudStorageProviderImpl::NewCloudReadableFile(
                                 options, result, dbg);
 }
 
+IOStatus CloudStorageProviderImpl::GetCloudObjectAsync(
+    const std::string& bucket_name, const std::string& object_path,
+    const std::string& local_path,
+    std::shared_ptr<std::promise<bool>> prom_ptr) {
+  const auto& local_fs = cfs_->GetBaseFileSystem();
+  return DoGetCloudObjectAsync(bucket_name, object_path, local_path, prom_ptr);
+}
+
 IOStatus CloudStorageProviderImpl::GetCloudObject(
     const std::string& bucket_name, const std::string& object_path,
     const std::string& local_destination) {
diff --git a/cloud/cloud_storage_provider_impl.h b/cloud/cloud_storage_provider_impl.h
@@ -128,6 +128,12 @@ class CloudStorageProviderImpl : public CloudStorageProvider {
   IOStatus PutCloudObject(const std::string& local_file,
                           const std::string& bucket_name,
                           const std::string& object_path) override;
+
+  IOStatus GetCloudObjectAsync(
+      const std::string& bucket_name, const std::string& object_path,
+      const std::string& local_path,
+      std::shared_ptr<std::promise<bool>> prom_ptr) override;
+
   IOStatus NewCloudReadableFile(
       const std::string& bucket, const std::string& fname,
       const FileOptions& options,
@@ -143,6 +149,11 @@ class CloudStorageProviderImpl : public CloudStorageProvider {
       std::unique_ptr<CloudStorageReadableFile>* result,
       IODebugContext* dbg) = 0;
 
+  virtual IOStatus DoGetCloudObjectAsync(
+      const std::string& bucket_name, const std::string& object_path,
+      const std::string& local_path,
+      std::shared_ptr<std::promise<bool>> prom_ptr) = 0;
+
   // Downloads object from the cloud into a local directory
   virtual IOStatus DoGetCloudObject(const std::string& bucket_name,
                                     const std::string& object_path,
diff --git a/db/version_builder.cc b/db/version_builder.cc
@@ -13,6 +13,7 @@
 #include <atomic>
 #include <cinttypes>
 #include <functional>
+#include <future>
 #include <map>
 #include <memory>
 #include <set>
@@ -233,7 +234,7 @@ class VersionBuilder::Rep {
   };
 
   const FileOptions& file_options_;
-  const ImmutableCFOptions* const ioptions_;
+  const ImmutableOptions* const ioptions_;
   TableCache* table_cache_;
   VersionStorageInfo* base_vstorage_;
   VersionSet* version_set_;
@@ -261,7 +262,7 @@ class VersionBuilder::Rep {
   std::shared_ptr<CacheReservationManager> file_metadata_cache_res_mgr_;
 
  public:
-  Rep(const FileOptions& file_options, const ImmutableCFOptions* ioptions,
+  Rep(const FileOptions& file_options, const ImmutableOptions* ioptions,
       TableCache* table_cache, VersionStorageInfo* base_vstorage,
       VersionSet* version_set,
       std::shared_ptr<CacheReservationManager> file_metadata_cache_res_mgr)
@@ -594,7 +595,7 @@ class VersionBuilder::Rep {
     // Note: we use C++11 for now but in C++14, this could be done in a more
     // elegant way using generalized lambda capture.
     VersionSet* const vs = version_set_;
-    const ImmutableCFOptions* const ioptions = ioptions_;
+    const ImmutableOptions* const ioptions = ioptions_;
 
     auto deleter = [vs, ioptions](SharedBlobFileMetaData* shared_meta) {
       if (vs) {
@@ -1245,12 +1246,20 @@ class VersionBuilder::Rep {
     // <file metadata, level>
     std::vector<std::pair<FileMetaData*, int>> files_meta;
     std::vector<Status> statuses;
+    std::vector<std::future<bool>> pending_downloads;
     for (int level = 0; level < num_levels_; level++) {
       for (auto& file_meta_pair : levels_[level].added_files) {
         auto* file_meta = file_meta_pair.second;
         // If the file has been opened before, just skip it.
         if (!file_meta->table_reader_handle) {
+          std::shared_ptr<std::promise<bool>> prom_ptr =
+              std::make_shared<std::promise<bool>>();
           files_meta.emplace_back(file_meta, level);
+          pending_downloads.emplace_back(prom_ptr->get_future());
+          std::string fname =
+              TableFileName(ioptions_->cf_paths, file_meta->fd.GetNumber(),
+                            file_meta->fd.GetPathId());
+          ioptions_->fs->DownloadAsync(fname, prom_ptr);
           statuses.emplace_back(Status::OK());
         }
         if (files_meta.size() >= max_load) {
@@ -1262,6 +1271,13 @@ class VersionBuilder::Rep {
       }
     }
 
+    // wait until all async download finished
+    // here we skip validating future returned value
+    // because it will redownload if local files were missing
+    for (auto& fut : pending_downloads) {
+      fut.get();
+    }
+
     std::atomic<size_t> next_file_meta_idx(0);
     std::function<void()> load_handlers_func([&]() {
       while (true) {
@@ -1309,7 +1325,7 @@ class VersionBuilder::Rep {
 };
 
 VersionBuilder::VersionBuilder(
-    const FileOptions& file_options, const ImmutableCFOptions* ioptions,
+    const FileOptions& file_options, const ImmutableOptions* ioptions,
     TableCache* table_cache, VersionStorageInfo* base_vstorage,
     VersionSet* version_set,
     std::shared_ptr<CacheReservationManager> file_metadata_cache_res_mgr)
diff --git a/db/version_builder.h b/db/version_builder.h
@@ -16,7 +16,7 @@
 
 namespace ROCKSDB_NAMESPACE {
 
-struct ImmutableCFOptions;
+struct ImmutableOptions;
 class TableCache;
 class VersionStorageInfo;
 class VersionEdit;
@@ -33,7 +33,7 @@ class CacheReservationManager;
 class VersionBuilder {
  public:
   VersionBuilder(const FileOptions& file_options,
-                 const ImmutableCFOptions* ioptions, TableCache* table_cache,
+                 const ImmutableOptions* ioptions, TableCache* table_cache,
                  VersionStorageInfo* base_vstorage, VersionSet* version_set,
                  std::shared_ptr<CacheReservationManager>
                      file_metadata_cache_res_mgr = nullptr);
diff --git a/include/rocksdb/cloud/cloud_storage_provider.h b/include/rocksdb/cloud/cloud_storage_provider.h
diff --git a/include/rocksdb/file_system.h b/include/rocksdb/file_system.h

Original file line number	Diff line number	Diff line change
`@@ -154,6 +154,7 @@ Status AwsCloudOptions::GetClientConfiguration(`
`154`	`154`	`config->useVirtualAddressing = false;`
`155`	`155`
`156`	`156`	`config->region = ToAwsString(region);`
	`157`	`+ config->throughputTargetGbps = cloud_fs_options.throughput_target_gbps;`
`157`	`158`	`return Status::OK();`
`158`	`159`	`}`
`159`	`160`	`#else`