[Perf] Replaced the S3 connector implementation to use the AWS transfer manager

CMakeLists.txt

@@ -164,7 +164,7 @@ if(VELOX_ENABLE_S3)
   if(AWSSDK_ROOT_DIR)
     set(CMAKE_PREFIX_PATH ${AWSSDK_ROOT_DIR})
   endif()
-  find_package(AWSSDK REQUIRED COMPONENTS s3;identity-management)
+  find_package(AWSSDK REQUIRED COMPONENTS s3;identity-management;transfer)
   add_definitions(-DVELOX_ENABLE_S3)
 endif()
 

scripts/setup-adapters.sh

@@ -25,10 +25,11 @@ DEPENDENCY_DIR=${DEPENDENCY_DIR:-$(pwd)}
 
 function install_aws-sdk-cpp {
   local AWS_REPO_NAME="aws/aws-sdk-cpp"
-  local AWS_SDK_VERSION="1.9.96"
+  local AWS_SDK_VERSION="1.11.37"
 
   github_checkout $AWS_REPO_NAME $AWS_SDK_VERSION --depth 1 --recurse-submodules
-  cmake_install -DCMAKE_BUILD_TYPE=Debug -DBUILD_SHARED_LIBS:BOOL=OFF -DMINIMIZE_SIZE:BOOL=ON -DENABLE_TESTING:BOOL=OFF -DBUILD_ONLY:STRING="s3;identity-management" -DCMAKE_INSTALL_PREFIX="${DEPENDENCY_DIR}/install"
+  cmake_install -DCMAKE_BUILD_TYPE=Debug -DBUILD_SHARED_LIBS:BOOL=OFF -DMINIMIZE_SIZE:BOOL=ON -DENABLE_TESTING:BOOL=OFF -DBUILD_ONLY:STRING="s3;identity-management;transfer" -DCMAKE_INSTALL_PREFIX="${DEPENDENCY_DIR}/install"
+  mkdir -p /root/adapter-deps/install/bin
 }
 
 function install_libhdfs3 {

velox/connectors/hive/storage_adapters/s3fs/S3FileSystem.cpp

@@ -21,48 +21,46 @@
 
 #include <fmt/format.h>
 #include <glog/logging.h>
+#include <chrono>
+#include <fstream>
 #include <memory>
 #include <stdexcept>
 
 #include <aws/core/Aws.h>
 #include <aws/core/auth/AWSCredentialsProviderChain.h>
 #include <aws/core/http/HttpResponse.h>
+#include <aws/core/utils/StringUtils.h>
 #include <aws/core/utils/logging/ConsoleLogSystem.h>
+#include <aws/core/utils/memory/AWSMemory.h>
+#include <aws/core/utils/memory/stl/AWSStreamFwd.h>
 #include <aws/core/utils/stream/PreallocatedStreamBuf.h>
+#include <aws/core/utils/threading/Executor.h>
 #include <aws/identity-management/auth/STSAssumeRoleCredentialsProvider.h>
-#include <aws/s3/S3Client.h>
-#include <aws/s3/model/GetObjectRequest.h>
 #include <aws/s3/model/HeadObjectRequest.h>
+#include <aws/transfer/TransferHandle.h>
+#include <aws/transfer/TransferManager.h>
 
 namespace facebook::velox {
 namespace {
-// Reference: https://issues.apache.org/jira/browse/ARROW-8692
-// https://github.com/apache/arrow/blob/master/cpp/src/arrow/filesystem/s3fs.cc#L843
-// A non-copying iostream. See
-// https://stackoverflow.com/questions/35322033/aws-c-sdk-uploadpart-times-out
-// https://stackoverflow.com/questions/13059091/creating-an-input-stream-from-constant-memory
-class StringViewStream : Aws::Utils::Stream::PreallocatedStreamBuf,
-                         public std::iostream {
- public:
-  StringViewStream(const void* data, int64_t nbytes)
-      : Aws::Utils::Stream::PreallocatedStreamBuf(
-            reinterpret_cast<unsigned char*>(const_cast<void*>(data)),
-            static_cast<size_t>(nbytes)),
-        std::iostream(this) {}
-};
 
 // By default, the AWS SDK reads object data into an auto-growing StringStream.
 // To avoid copies, read directly into a pre-allocated buffer instead.
 // See https://github.com/aws/aws-sdk-cpp/issues/64 for an alternative but
 // functionally similar recipe.
-Aws::IOStreamFactory AwsWriteableStreamFactory(void* data, int64_t nbytes) {
-  return [=]() { return Aws::New<StringViewStream>("", data, nbytes); };
-}
+class UnderlyingStreamWrapper : public Aws::IOStream {
+ public:
+  explicit UnderlyingStreamWrapper(std::streambuf* buf) : Aws::IOStream(buf) {}
+  ~UnderlyingStreamWrapper() override = default;
+};
 
 class S3ReadFile final : public ReadFile {
  public:
-  S3ReadFile(const std::string& path, Aws::S3::S3Client* client)
-      : client_(client) {
+  S3ReadFile(
+      const std::string& path,
+      std::shared_ptr<Aws::S3::S3Client> client,
+      const std::shared_ptr<Aws::Utils::Threading::PooledThreadExecutor>&
+          executor)
+      : client_(client), executor_(executor) {
     bucketAndKeyFromS3Path(path, bucket_, key_);
   }
 
@@ -83,6 +81,9 @@ class S3ReadFile final : public ReadFile {
         outcome, "Failed to get metadata for S3 object", bucket_, key_);
     length_ = outcome.GetResult().GetContentLength();
     VELOX_CHECK_GE(length_, 0);
+    Aws::Transfer::TransferManagerConfiguration transferConfig(executor_.get());
+    transferConfig.s3Client = client_;
+    transferManager_ = Aws::Transfer::TransferManager::Create(transferConfig);
   }
 
   std::string_view pread(uint64_t offset, uint64_t length, void* buffer)
@@ -151,30 +152,56 @@ class S3ReadFile final : public ReadFile {
   // bytes.
   void preadInternal(uint64_t offset, uint64_t length, char* position) const {
     // Read the desired range of bytes.
-    Aws::S3::Model::GetObjectRequest request;
-    Aws::S3::Model::GetObjectResult result;
-
-    request.SetBucket(awsString(bucket_));
-    request.SetKey(awsString(key_));
-    std::stringstream ss;
-    ss << "bytes=" << offset << "-" << offset + length - 1;
-    request.SetRange(awsString(ss.str()));
-    request.SetResponseStreamFactory(
-        AwsWriteableStreamFactory(position, length));
-    auto outcome = client_->GetObject(request);
-    VELOX_CHECK_AWS_OUTCOME(outcome, "Failed to get S3 object", bucket_, key_);
+    Aws::Utils::Stream::PreallocatedStreamBuf streamBuffer(
+        (unsigned char*)position, length);
+    auto downloadHandle =
+        transferManager_->DownloadFile(bucket_, key_, offset, length, [&]() {
+          return Aws::New<UnderlyingStreamWrapper>("TestTag", &streamBuffer);
+        });
+    int count = 1;
+    using namespace std::chrono;
+    uint64_t startTime =
+        duration_cast<milliseconds>(system_clock::now().time_since_epoch())
+            .count();
+    uint64_t now = startTime;
+    uint64_t logAt = startTime + 5000; // wait 5 seconds; then warn and backoff
+    auto status = downloadHandle->GetStatus();
+    while (status == Aws::Transfer::TransferStatus::NOT_STARTED ||
+           status == Aws::Transfer::TransferStatus::IN_PROGRESS) {
+      if (now > logAt) {
+        VLOG(1) << "S3 Request is taking a long time: "
+                << (now - startTime) / 1000 << "s!\n";
+        count++;
+        logAt += count * 5000; // backoff for next warning
+      }
+      status = downloadHandle->GetStatus();
+      now = duration_cast<milliseconds>(system_clock::now().time_since_epoch())
+                .count();
+    }
+    if (downloadHandle->GetStatus() != Aws::Transfer::TransferStatus::COMPLETED)
+      VELOX_FAIL(
+          "{} from location: {}:{}",
+          "Failed to get S3 object using the transfer manager",
+          bucket_,
+          key_);
   }
 
-  Aws::S3::S3Client* client_;
+  std::shared_ptr<Aws::S3::S3Client> client_;
+  std::shared_ptr<Aws::Transfer::TransferManager> transferManager_;
   std::string bucket_;
   std::string key_;
   int64_t length_ = -1;
+  std::shared_ptr<Aws::Utils::Threading::PooledThreadExecutor> executor_;
 };
+
 } // namespace
 
 namespace filesystems {
 
 class S3Config {
+ public: // Constants
+  static const char* kS3TransferManagerMaxThreadsConfig;
+
  public:
   S3Config(const Config* config) : config_(config) {}
 
@@ -222,6 +249,10 @@ class S3Config {
         "hive.s3.iam-role-session-name", std::string("velox-session"));
   }
 
+  int getTranferManagerMaxThreads() const {
+    return config_->get<int>(kS3TransferManagerMaxThreadsConfig, 25);
+  }
+
   Aws::Utils::Logging::LogLevel getLogLevel() const {
     auto level = config_->get("hive.s3.log-level", std::string("FATAL"));
     // Convert to upper case.
@@ -251,6 +282,9 @@ class S3Config {
   const Config* FOLLY_NONNULL config_;
 };
 
+const char* S3Config::kS3TransferManagerMaxThreadsConfig =
+    "hive.s3.transfer-manager-max-threads";
+
 class S3FileSystem::Impl {
  public:
   Impl(const Config* config) : s3Config_(config) {
@@ -267,6 +301,9 @@ class S3FileSystem::Impl {
       // log a message.
       awsOptions.httpOptions.installSigPipeHandler = true;
       Aws::InitAPI(awsOptions);
+      S3FileSystem::Impl::executor_ =
+          std::make_shared<Aws::Utils::Threading::PooledThreadExecutor>(
+              s3Config_.getTranferManagerMaxThreads());
     }
   }
 
@@ -275,6 +312,7 @@ class S3FileSystem::Impl {
     if (newCount == 0) {
       Aws::SDKOptions awsOptions;
       awsOptions.loggingOptions.logLevel = s3Config_.getLogLevel();
+      client_.reset();
       Aws::ShutdownAPI(awsOptions);
     }
   }
@@ -366,8 +404,17 @@ class S3FileSystem::Impl {
   // Make it clear that the S3FileSystem instance owns the S3Client.
   // Once the S3FileSystem is destroyed, the S3Client fails to work
   // due to the Aws::ShutdownAPI invocation in the destructor.
-  Aws::S3::S3Client* s3Client() const {
-    return client_.get();
+  const std::shared_ptr<Aws::S3::S3Client>& s3Client() const {
+    return client_;
+  }
+
+  const S3Config& s3Config() const {
+    return s3Config_;
+  }
+
+  const std::shared_ptr<Aws::Utils::Threading::PooledThreadExecutor>&
+  getExecutor() const {
+    return S3FileSystem::Impl::executor_;
   }
 
   std::string getLogLevelName() const {
@@ -378,9 +425,12 @@ class S3FileSystem::Impl {
   const S3Config s3Config_;
   std::shared_ptr<Aws::S3::S3Client> client_;
   static std::atomic<size_t> initCounter_;
+  static std::shared_ptr<Aws::Utils::Threading::PooledThreadExecutor> executor_;
 };
 
 std::atomic<size_t> S3FileSystem::Impl::initCounter_(0);
+std::shared_ptr<Aws::Utils::Threading::PooledThreadExecutor>
+    S3FileSystem::Impl::executor_(nullptr);
 folly::once_flag S3FSInstantiationFlag;
 
 S3FileSystem::S3FileSystem(std::shared_ptr<const Config> config)
@@ -398,7 +448,8 @@ std::string S3FileSystem::getLogLevelName() const {
 
 std::unique_ptr<ReadFile> S3FileSystem::openFileForRead(std::string_view path) {
   const std::string file = s3Path(path);
-  auto s3file = std::make_unique<S3ReadFile>(file, impl_->s3Client());
+  auto s3file = std::make_unique<S3ReadFile>(
+      file, impl_->s3Client(), impl_->getExecutor());
   s3file->initialize();
   return s3file;
 }