Fix concurrent S3 credential refresh and preloaded split stalls

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

@@ -324,7 +324,8 @@ class S3FileSystem::Impl {
   // Return a default AWSCredentialsProvider.
   std::shared_ptr<Aws::Auth::AWSCredentialsProvider>
   getDefaultCredentialsProvider() const {
-    return std::make_shared<Aws::Auth::DefaultAWSCredentialsProviderChain>();
+    return makeSynchronizedCachingCredentialsProvider(
+        std::make_shared<Aws::Auth::DefaultAWSCredentialsProviderChain>());
   }
 
   // Configure and return an AWSCredentialsProvider with S3 IAM Role.

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

@@ -24,7 +24,68 @@
 
 #include "velox/connectors/hive/storage_adapters/s3fs/S3Util.h"
 
+#include <mutex>
+
 namespace facebook::velox::filesystems {
+namespace {
+
+class SynchronizedCachingCredentialsProvider final
+    : public Aws::Auth::AWSCredentialsProvider {
+ public:
+  explicit SynchronizedCachingCredentialsProvider(
+      std::shared_ptr<Aws::Auth::AWSCredentialsProvider> source)
+      : source_(std::move(source)) {
+    refresh();
+  }
+
+  Aws::Auth::AWSCredentials GetAWSCredentials() override {
+    {
+      std::lock_guard<std::mutex> lock(mutex_);
+      if (isUsable(credentials_)) {
+        return credentials_;
+      }
+    }
+    refresh();
+    std::lock_guard<std::mutex> lock(mutex_);
+    return credentials_;
+  }
+
+ private:
+  static bool isUsable(const Aws::Auth::AWSCredentials& credentials) {
+    return !credentials.IsEmpty() &&
+        !credentials.ExpiresSoon(5 * 60 * 1000);
+  }
+
+  void refresh() {
+    std::lock_guard<std::mutex> refreshLock(refreshMutex_);
+    {
+      std::lock_guard<std::mutex> lock(mutex_);
+      if (isUsable(credentials_)) {
+        return;
+      }
+    }
+    auto refreshed = source_->GetAWSCredentials();
+    std::lock_guard<std::mutex> lock(mutex_);
+    if (!refreshed.IsEmpty()) {
+      credentials_ = std::move(refreshed);
+    }
+  }
+
+  const std::shared_ptr<Aws::Auth::AWSCredentialsProvider> source_;
+  std::mutex mutex_;
+  std::mutex refreshMutex_;
+  Aws::Auth::AWSCredentials credentials_;
+};
+
+} // namespace
+
+std::shared_ptr<Aws::Auth::AWSCredentialsProvider>
+makeSynchronizedCachingCredentialsProvider(
+    std::shared_ptr<Aws::Auth::AWSCredentialsProvider> source) {
+  VELOX_CHECK_NOT_NULL(source);
+  return std::make_shared<SynchronizedCachingCredentialsProvider>(
+      std::move(source));
+}
 
 std::string getErrorStringFromS3Error(
     const Aws::Client::AWSError<Aws::S3::S3Errors>& error) {

velox/connectors/hive/storage_adapters/s3fs/S3Util.h

@@ -21,6 +21,7 @@
 
 #pragma once
 
+#include <aws/core/auth/AWSCredentialsProvider.h>
 #include <aws/s3/S3Errors.h>
 #include <aws/s3/model/HeadObjectResult.h>
 #include <fmt/format.h>
@@ -32,6 +33,17 @@
 
 namespace facebook::velox::filesystems {
 
+/// Wraps a refreshable AWS credentials provider with synchronized caching.
+///
+/// A refreshable provider chain can invoke the underlying identity provider
+/// once per concurrent S3 signer when an empty or expiring credential is
+/// observed. The wrapper serializes refresh and serves the same credential
+/// snapshot until five minutes before expiration. The underlying provider
+/// remains responsible for normal credential rotation.
+std::shared_ptr<Aws::Auth::AWSCredentialsProvider>
+makeSynchronizedCachingCredentialsProvider(
+    std::shared_ptr<Aws::Auth::AWSCredentialsProvider> source);
+
 namespace {
 static std::string_view kSep{"/"};
 // AWS S3 EMRFS, Hadoop block storage filesystem on-top of Amazon S3 buckets.

velox/connectors/hive/storage_adapters/s3fs/tests/S3UtilTest.cpp

@@ -18,8 +18,44 @@
 
 #include "gtest/gtest.h"
 
+#include <atomic>
+#include <chrono>
+#include <thread>
+#include <vector>
+
 namespace facebook::velox::filesystems {
 
+TEST(S3UtilTest, synchronizedCachingCredentialsProvider) {
+  class CountingProvider final : public Aws::Auth::AWSCredentialsProvider {
+   public:
+    Aws::Auth::AWSCredentials GetAWSCredentials() override {
+      const auto call = ++calls;
+      if (call == 1) {
+        return {};
+      }
+      std::this_thread::sleep_for(std::chrono::milliseconds(10));
+      return {"access", "secret", "session"};
+    }
+
+    std::atomic<size_t> calls{0};
+  };
+
+  auto source = std::make_shared<CountingProvider>();
+  auto cached = makeSynchronizedCachingCredentialsProvider(source);
+  std::vector<std::thread> threads;
+  for (size_t index = 0; index < 128; ++index) {
+    threads.emplace_back([cached] {
+      EXPECT_EQ(cached->GetAWSCredentials().GetAWSAccessKeyId(), "access");
+    });
+  }
+  for (auto& thread : threads) {
+    thread.join();
+  }
+  // The eager construction attempt is empty. All concurrent callers share
+  // the single subsequent refresh.
+  EXPECT_EQ(source->calls.load(), 2);
+}
+
 // TODO: Each prefix should be implemented as its own filesystem.
 TEST(S3UtilTest, isS3File) {
   EXPECT_FALSE(isS3File("ss3://"));

velox/exec/TableScan.cpp

@@ -508,10 +508,15 @@ void TableScan::checkPreload() {
         [ioExecutor,
          this](const std::shared_ptr<connector::ConnectorSplit>& split) {
           preload(split);
-          ioExecutor->add([connectorSplit = split]() mutable {
-            connectorSplit->dataSource->prepare();
-            connectorSplit.reset();
-          });
+          ioExecutor->add(
+              [connectorSplit = split,
+               task = operatorCtx_->task(),
+               splitGroupId = driverCtx_->splitGroupId,
+               planNodeId = planNodeId()]() mutable {
+                connectorSplit->dataSource->prepare();
+                connectorSplit.reset();
+                task->splitPreloadFinished(splitGroupId, planNodeId);
+              });
         };
   }
 }

velox/exec/Task.cpp

@@ -282,8 +282,11 @@ class QueueSplitsStore : public SplitsStore {
       Split& split,
       ContinueFuture& future) override {
     if (!splits_.empty()) {
-      split = getSplit(maxPreloadSplits, preload);
-      return true;
+      if (getSplit(maxPreloadSplits, preload, split)) {
+        return true;
+      }
+      future = makeFuture();
+      return false;
     }
     if (tryGetBarrier(driverId, split)) {
       return true;
@@ -2307,6 +2310,29 @@ BlockingReason Task::getSplitOrFuture(
                     : BlockingReason::kWaitForSplit;
 }
 
+void Task::splitPreloadFinished(
+    uint32_t splitGroupId,
+    const core::PlanNodeId& planNodeId) {
+  std::vector<ContinuePromise> promises;
+  {
+    std::lock_guard<std::timed_mutex> l(mutex_);
+    const auto stateIt = splitsStates_.find(planNodeId);
+    if (stateIt == splitsStates_.end()) {
+      return;
+    }
+    const auto storeIt =
+        stateIt->second.groupSplitsStores.find(splitGroupId);
+    if (storeIt == stateIt->second.groupSplitsStores.end() ||
+        storeIt->second == nullptr) {
+      return;
+    }
+    promises = storeIt->second->splitPreloadFinished();
+  }
+  for (auto& promise : promises) {
+    promise.setValue();
+  }
+}
+
 bool Task::testingHasDriverWaitForSplit() const {
   std::lock_guard<std::timed_mutex> l(mutex_);
   for (const auto& splitState : splitsStates_) {

velox/exec/Task.h

@@ -527,6 +527,12 @@ class Task : public std::enable_shared_from_this<Task> {
       exec::Split& split,
       ContinueFuture& future);
 
+  /// Notifies scan drivers that an asynchronously preloaded split has
+  /// completed and ready-first split selection should be retried.
+  void splitPreloadFinished(
+      uint32_t splitGroupId,
+      const core::PlanNodeId& planNodeId);
+
   /// Returns the scaled scan controller for a given table scan node if the
   /// query has configured.
   std::shared_ptr<ScaledScanController> getScaledScanControllerLocked(

velox/exec/TaskStructs.cpp

@@ -49,9 +49,10 @@ ContinueFuture SplitsStore::makeFuture() {
   return std::move(future);
 }
 
-Split SplitsStore::getSplit(
+bool SplitsStore::getSplit(
     int maxPreloadSplits,
-    const ConnectorSplitPreloadFunc& preload) {
+    const ConnectorSplitPreloadFunc& preload,
+    Split& split) {
   int readySplitIndex = -1;
   if (maxPreloadSplits > 0) {
     for (int i = 0, end = std::min<size_t>(maxPreloadSplits, splits_.size());
@@ -72,12 +73,19 @@ Split SplitsStore::getSplit(
         preloadingSplits_->erase(connectorSplit);
       }
     }
+    // Do not bind a scan driver to an arbitrary in-flight preload. The
+    // completion path wakes a waiter, which retries and takes whichever split
+    // is ready first. This keeps I/O and compute pipelined without
+    // head-of-line blocking on the queue front.
+    if (readySplitIndex == -1) {
+      return false;
+    }
   }
   if (readySplitIndex == -1) {
     readySplitIndex = 0;
   }
   VELOX_CHECK(!splits_.empty());
-  auto split = std::move(splits_[readySplitIndex]);
+  split = std::move(splits_[readySplitIndex]);
   splits_.erase(splits_.begin() + readySplitIndex);
   --taskStats_->numQueuedSplits;
   ++taskStats_->numRunningSplits;
@@ -93,7 +101,7 @@ Split SplitsStore::getSplit(
   if (taskStats_->firstSplitStartTimeMs == 0) {
     taskStats_->firstSplitStartTimeMs = taskStats_->lastSplitStartTimeMs;
   }
-  return split;
+  return true;
 }
 
 bool SplitsStore::tryGetBarrier(

velox/exec/TaskStructs.h

@@ -136,6 +136,16 @@ class SplitsStore {
     return std::move(promises_);
   }
 
+  /// Wakes drivers waiting for a preloaded split to become ready.
+  ///
+  /// Readiness is level-triggered: by the time this completion arrives there
+  /// may already be multiple ready splits, and completions that happened
+  /// before waiters registered do not leave an edge to wake them later.
+  /// Wake all current waiters so each rechecks the ready queue.
+  std::vector<ContinuePromise> splitPreloadFinished() {
+    return std::move(promises_);
+  }
+
   void setTaskStats(TaskStats& taskStats) {
     taskStats_ = &taskStats;
   }
@@ -147,9 +157,10 @@ class SplitsStore {
   }
 
  protected:
-  Split getSplit(
+  bool getSplit(
       int maxPreloadSplits,
-      const ConnectorSplitPreloadFunc& preload);
+      const ConnectorSplitPreloadFunc& preload,
+      Split& split);
 
   ContinueFuture makeFuture();