feat(io): refresh vended storage credentials before they expire

mkdocs/docs/file-io.md

@@ -63,6 +63,7 @@ each file location's scheme.
 | `s3.access-key-id` | `admin` | Static access key ID; must be set together with the secret key |
 | `s3.secret-access-key` | `password` | Static secret access key |
 | `s3.session-token` | `AQoDYXdzEJr...` | Session token, for temporary credentials. Ignored unless both static keys are set |
+| `s3.session-token-expires-at-ms` | `1767225600000` | When the session token expires, in epoch milliseconds. Vended credentials are refreshed from the catalog five minutes before |
 | `client.region` | `us-east-1` | Region to sign requests for |
 | `s3.endpoint` | `https://127.0.0.1:9000` | Endpoint to use instead of the AWS one. When absent, the `AWS_ENDPOINT_URL_S3` / `AWS_ENDPOINT_URL` environment variables are consulted |
 | `s3.path-style-access` | `true` | Address buckets as a path (`endpoint/bucket`) instead of a virtual host (`bucket.endpoint`). Only takes effect together with a custom endpoint |

src/iceberg/arrow/s3/arrow_s3_file_io.cc

@@ -18,6 +18,8 @@
  */
 
 #include <algorithm>
+#include <chrono>
+#include <cstdint>
 #include <cstdlib>
 #include <memory>
 #include <mutex>
@@ -178,6 +180,71 @@ std::string CanonicalizeS3Scheme(std::string_view location) {
   return std::string(location);
 }
 
+// Lead time before expiry, matching Java's VendedCredentialsProvider.
+constexpr auto kRefreshLeadTime = std::chrono::minutes(5);
+
+// After a failed refresh, how long to keep the current credentials before
+// asking again, so an unreachable catalog is not queried per file operation.
+constexpr auto kRefreshRetryBackoff = std::chrono::seconds(30);
+
+// Floor on a backoff shortened to land on the expiry.
+constexpr auto kMinRefreshRetryBackoff = std::chrono::seconds(1);
+
+// How long an operation with expired credentials waits for a refresh already
+// under way. Bounded: the catalog request behind it has no deadline of its own.
+constexpr auto kExpiredCredentialWait = std::chrono::seconds(10);
+
+// When the earliest of `credentials` stops being valid, or nullopt if none of
+// them does. No session token means static keys, which never expire; a token
+// with no usable expiry is reported as already expired so it gets replaced
+// rather than used until it fails, as Java does.
+std::optional<std::chrono::system_clock::time_point> EarliestExpiry(
+    const std::vector<StorageCredential>& credentials) {
+  std::optional<std::chrono::system_clock::time_point> earliest;
+  const auto note = [&earliest](std::chrono::system_clock::time_point expires_at) {
+    if (!earliest.has_value() || expires_at < *earliest) {
+      earliest = expires_at;
+    }
+  };
+  for (const auto& credential : credentials) {
+    if (!IsS3CredentialPrefix(credential.prefix) ||
+        FindProperty(credential.config, S3Properties::kSessionToken) == nullptr) {
+      continue;
+    }
+    const auto* value =
+        FindProperty(credential.config, S3Properties::kSessionTokenExpiresAtMs);
+    if (value == nullptr) {
+      ICEBERG_LOG_WARN("Credential \"{}\" has a session token but no \"{}\"",
+                       credential.prefix, S3Properties::kSessionTokenExpiresAtMs);
+      note(std::chrono::system_clock::now());
+      continue;
+    }
+    auto millis = StringUtils::ParseNumber<int64_t>(*value);
+    if (!millis.has_value()) {
+      ICEBERG_LOG_WARN(
+          "Credential \"{}\" has a session token but an unparseable \"{}\" value \"{}\"",
+          credential.prefix, S3Properties::kSessionTokenExpiresAtMs, *value);
+      note(std::chrono::system_clock::now());
+      continue;
+    }
+    // Beyond what the clock can hold, converting would overflow it.
+    constexpr auto kMaxMillis = std::chrono::duration_cast<std::chrono::milliseconds>(
+                                    std::chrono::system_clock::duration::max())
+                                    .count();
+    constexpr auto kMinMillis = std::chrono::duration_cast<std::chrono::milliseconds>(
+                                    std::chrono::system_clock::duration::min())
+                                    .count();
+    if (*millis > kMaxMillis || *millis < kMinMillis) {
+      ICEBERG_LOG_WARN("Credential \"{}\" has an out-of-range \"{}\" value \"{}\"",
+                       credential.prefix, S3Properties::kSessionTokenExpiresAtMs, *value);
+      note(std::chrono::system_clock::now());
+      continue;
+    }
+    note(std::chrono::system_clock::time_point(std::chrono::milliseconds(*millis)));
+  }
+  return earliest;
+}
+
 class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials {
  public:
   ArrowS3FileIO(std::shared_ptr<::arrow::fs::FileSystem> arrow_fs,
@@ -204,6 +271,13 @@ class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials {
     return storage_credentials_;
   }
 
+  void SetCredentialRefresher(StorageCredentialRefresher refresher) override {
+    std::unique_lock lock(mutex_);
+    refresher_ = std::move(refresher);
+    // A refresh in flight was started for the refresher just replaced.
+    ++credential_generation_;
+  }
+
   SupportsStorageCredentials* AsSupportsStorageCredentials() override { return this; }
 
  private:
@@ -235,12 +309,46 @@ class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials {
   void InstallCredentials(std::vector<StorageCredential>& storage_credentials,
                           DelegatesByPrefix& delegates);
 
+  /// \brief Whether the installed credentials are close enough to expiring to
+  /// be replaced, and no backoff is in effect.
+  ///
+  /// Callers must hold `mutex_`, at least shared.
+  bool RefreshDue() const;
+
+  /// \brief Whether the installed credentials have already stopped being valid.
+  ///
+  /// Callers must hold `mutex_`, at least shared.
+  bool Expired() const;
+
+  /// \brief When the next refresh attempt becomes allowed after a failure.
+  ///
+  /// Never past the point the credentials stop being valid.
+  ///
+  /// Callers must hold `mutex_`, at least shared.
+  std::chrono::steady_clock::time_point BackoffUntil() const;
+
+  /// \brief Replace the credentials once they are close to expiring.
+  ///
+  /// Called before each handle is created; a handle keeps the delegate it was
+  /// built from, so I/O on an open one is not re-checked. A failure keeps the
+  /// current credentials rather than failing the read.
+  void MaybeRefreshCredentials();
+
   std::shared_ptr<ArrowFileSystemFileIO> default_file_io_;
   std::unordered_map<std::string, std::string> default_properties_;
   // Guards everything below; shared because reads happen per file operation.
   mutable std::shared_mutex mutex_;
   std::vector<StorageCredential> storage_credentials_;
   DelegatesByPrefix file_io_by_prefix_;
+  StorageCredentialRefresher refresher_;
+  std::optional<std::chrono::system_clock::time_point> expires_at_;
+  std::chrono::steady_clock::time_point retry_refresh_at_;
+  // Bumped whenever the credentials or the refresher change, so a refresh that
+  // fetched before one of those happened can tell its result is already stale.
+  uint64_t credential_generation_ = 0;
+  // Held across a refresh so concurrent operations skip it. Timed, so waiting
+  // on it is bounded.
+  std::timed_mutex refresh_mutex_;
 };
 
 Status ArrowS3FileIO::SetStorageCredentials(
@@ -260,7 +368,6 @@ Result<ArrowS3FileIO::DelegatesByPrefix> ArrowS3FileIO::BuildDelegates(
     const std::vector<StorageCredential>& storage_credentials) const {
   DelegatesByPrefix delegates;
   delegates.reserve(storage_credentials.size());
-  // TODO(gangwu): Refresh vended credentials via credentials.uri before tokens expire.
   for (const auto& credential : storage_credentials) {
     ICEBERG_RETURN_UNEXPECTED(credential.Validate());
     // A server may vend credentials for several storage systems at once;
@@ -291,9 +398,132 @@ Result<ArrowS3FileIO::DelegatesByPrefix> ArrowS3FileIO::BuildDelegates(
 void ArrowS3FileIO::InstallCredentials(
     std::vector<StorageCredential>& storage_credentials, DelegatesByPrefix& delegates) {
   file_io_by_prefix_.swap(delegates);
+  expires_at_ = EarliestExpiry(storage_credentials);
+  retry_refresh_at_ = {};
+  ++credential_generation_;
   storage_credentials_.swap(storage_credentials);
 }
 
+bool ArrowS3FileIO::RefreshDue() const {
+  return expires_at_.has_value() &&
+         std::chrono::system_clock::now() + kRefreshLeadTime >= *expires_at_ &&
+         std::chrono::steady_clock::now() >= retry_refresh_at_;
+}
+
+bool ArrowS3FileIO::Expired() const {
+  return expires_at_.has_value() && std::chrono::system_clock::now() >= *expires_at_;
+}
+
+std::chrono::steady_clock::time_point ArrowS3FileIO::BackoffUntil() const {
+  auto delay =
+      std::chrono::duration_cast<std::chrono::milliseconds>(kRefreshRetryBackoff);
+  if (expires_at_.has_value()) {
+    const auto remaining = std::chrono::duration_cast<std::chrono::milliseconds>(
+        *expires_at_ - std::chrono::system_clock::now());
+    // Worth retrying before they run out; once they have, faster retries only
+    // hammer a catalog that is already failing.
+    if (remaining > std::chrono::milliseconds::zero()) {
+      delay = std::clamp(
+          remaining,
+          std::chrono::duration_cast<std::chrono::milliseconds>(kMinRefreshRetryBackoff),
+          delay);
+    }
+  }
+  return std::chrono::steady_clock::now() + delay;
+}
+
+void ArrowS3FileIO::MaybeRefreshCredentials() {
+  {
+    // Cheap pre-check, so the common case costs one shared lock and no more.
+    std::shared_lock lock(mutex_);
+    if (!refresher_ || !RefreshDue()) {
+      return;
+    }
+  }
+
+  std::unique_lock refresh_lock(refresh_mutex_, std::defer_lock);
+  if (!refresh_lock.try_lock()) {
+    // Another operation is already fetching; normally just use what we have.
+    {
+      std::shared_lock lock(mutex_);
+      if (!Expired()) {
+        return;
+      }
+    }
+    // Expired credentials leave nothing to proceed with, so wait instead.
+    if (!refresh_lock.try_lock_for(kExpiredCredentialWait)) {
+      return;
+    }
+  }
+  // Read together: pairing this refresher with a generation bumped by another
+  // one installed in between would make its result look current.
+  StorageCredentialRefresher refresher;
+  uint64_t generation = 0;
+  {
+    // Whoever held the lock may also have just finished, leaving nothing to do.
+    std::shared_lock lock(mutex_);
+    if (!refresher_ || !RefreshDue()) {
+      return;
+    }
+    refresher = refresher_;
+    generation = credential_generation_;
+  }
+
+  // Outside `mutex_`: both are slow and must not block readers.
+  Status status;
+  DelegatesByPrefix delegates;
+  auto refreshed = refresher();
+  if (refreshed.has_value()) {
+    auto built = BuildDelegates(*refreshed);
+    if (!built.has_value()) {
+      status = std::unexpected(built.error());
+    } else if (built->empty()) {
+      // Installing this would drop working credentials for whatever ambient
+      // identity the AWS chain finds. Java refuses an empty list too.
+      status = NotFound("Refreshed credentials contain no S3-compatible prefix");
+    } else {
+      delegates = std::move(built).value();
+    }
+  } else {
+    status = std::unexpected(refreshed.error());
+  }
+
+  std::unique_lock lock(mutex_);
+  // Credentials installed meanwhile supersede this refresh: what it fetched is
+  // by now the older set.
+  const bool superseded = credential_generation_ != generation;
+  if (!status.has_value()) {
+    // Reported either way, so a failing catalog stays visible.
+    if (superseded) {
+      ICEBERG_LOG_WARN(
+          "Failed to refresh vended storage credentials ({}); they have since been "
+          "replaced",
+          status.error().message);
+      return;
+    }
+    retry_refresh_at_ = BackoffUntil();
+    ICEBERG_LOG_WARN(
+        "Failed to refresh vended storage credentials ({}); keeping the current "
+        "ones and retrying in {}ms",
+        status.error().message,
+        std::chrono::duration_cast<std::chrono::milliseconds>(
+            retry_refresh_at_ - std::chrono::steady_clock::now())
+            .count());
+    return;
+  }
+  if (superseded) {
+    return;
+  }
+
+  // `delegates`/`*refreshed` take the retired generation; both are declared
+  // before the lock, so it destructs only after the lock releases.
+  InstallCredentials(*refreshed, delegates);
+  if (RefreshDue()) {
+    // Tokens shorter-lived than the lead time come back due again at once.
+    retry_refresh_at_ = BackoffUntil();
+  }
+}
+
 std::shared_ptr<ArrowFileSystemFileIO> ArrowS3FileIO::MatchDelegate(
     const std::shared_ptr<ArrowFileSystemFileIO>& fallback,
     const DelegatesByPrefix& by_prefix, std::string_view location) {
@@ -314,6 +544,8 @@ std::shared_ptr<ArrowFileSystemFileIO> ArrowS3FileIO::MatchDelegate(
 
 std::shared_ptr<ArrowFileSystemFileIO> ArrowS3FileIO::FileIOForPath(
     std::string_view location) {
+  MaybeRefreshCredentials();
+
   std::shared_lock lock(mutex_);
   return MatchDelegate(default_file_io_, file_io_by_prefix_, location);
 }
@@ -338,6 +570,8 @@ Status ArrowS3FileIO::DeleteFile(const std::string& file_location) {
 }
 
 Status ArrowS3FileIO::DeleteFiles(const std::vector<std::string>& file_locations) {
+  MaybeRefreshCredentials();
+
   // One snapshot so the whole batch matches the same delegate generation; only
   // ever a handful of delegates, so a linear scan beats hashing.
   std::shared_ptr<ArrowFileSystemFileIO> fallback;

src/iceberg/arrow/s3/s3_properties.h

@@ -42,6 +42,9 @@ struct S3Properties {
   static constexpr std::string_view kSecretAccessKey = "s3.secret-access-key";
   /// AWS session token (for temporary credentials)
   static constexpr std::string_view kSessionToken = "s3.session-token";
+  /// Epoch milliseconds at which a vended session token stops being valid
+  static constexpr std::string_view kSessionTokenExpiresAtMs =
+      "s3.session-token-expires-at-ms";
   /// AWS region, standard Iceberg client property.
   static constexpr std::string_view kClientRegion = "client.region";
   /// Custom endpoint override (for S3-compatible object stores)

src/iceberg/catalog/rest/json_serde.cc

@@ -152,6 +152,25 @@ Result<StorageCredential> StorageCredentialFromJson(const nlohmann::json& json)
   return credential;
 }
 
+/// \brief Reads the optional `storage-credentials` array shared by the
+/// LoadTable and LoadCredentials responses.
+Result<std::vector<StorageCredential>> StorageCredentialsFromJson(
+    const nlohmann::json& json) {
+  std::vector<StorageCredential> credentials;
+  auto it = json.find(kStorageCredentials);
+  if (it == json.end() || it->is_null()) {
+    return credentials;
+  }
+  if (!it->is_array()) {
+    return JsonParseError("Cannot parse storage credentials from non-array");
+  }
+  for (const auto& entry : *it) {
+    ICEBERG_ASSIGN_OR_RAISE(auto credential, StorageCredentialFromJson(entry));
+    credentials.push_back(std::move(credential));
+  }
+  return credentials;
+}
+
 template <typename Value>
 Result<std::map<int32_t, Value>> KeyValueMapFromJson(const nlohmann::json& json,
                                                      std::string_view key) {
@@ -738,19 +757,24 @@ Result<LoadTableResult> LoadTableResultFromJson(const nlohmann::json& json) {
   ICEBERG_ASSIGN_OR_RAISE(result.metadata, TableMetadataFromJson(metadata_json));
   ICEBERG_ASSIGN_OR_RAISE(result.config,
                           GetJsonValueOrDefault<decltype(result.config)>(json, kConfig));
-  if (auto it = json.find(kStorageCredentials); it != json.end() && !it->is_null()) {
-    if (!it->is_array()) {
-      return JsonParseError("Cannot parse storage credentials from non-array");
-    }
-    for (const auto& entry : *it) {
-      ICEBERG_ASSIGN_OR_RAISE(auto cred, StorageCredentialFromJson(entry));
-      result.storage_credentials.push_back(std::move(cred));
-    }
-  }
+  ICEBERG_ASSIGN_OR_RAISE(result.storage_credentials, StorageCredentialsFromJson(json));
   ICEBERG_RETURN_UNEXPECTED(result.Validate());
   return result;
 }
 
+Result<LoadCredentialsResponse> LoadCredentialsResponseFromJson(
+    const nlohmann::json& json) {
+  // Required here, unlike in LoadTable: reading a malformed response as "no
+  // credentials" would look like a refresh that succeeded and dropped them.
+  if (auto it = json.find(kStorageCredentials); it == json.end() || it->is_null()) {
+    return JsonParseError("Missing '{}'", kStorageCredentials);
+  }
+  LoadCredentialsResponse response;
+  ICEBERG_ASSIGN_OR_RAISE(response.storage_credentials, StorageCredentialsFromJson(json));
+  ICEBERG_RETURN_UNEXPECTED(response.Validate());
+  return response;
+}
+
 nlohmann::json ToJson(const ListNamespacesResponse& response) {
   nlohmann::json json;
   SetOptionalStringField(json, kNextPageToken, response.next_page_token);

src/iceberg/catalog/rest/json_serde_internal.h

@@ -71,6 +71,10 @@ template <>
 ICEBERG_REST_EXPORT Result<LoadTableResult> FromJson(const nlohmann::json& json);
 ICEBERG_REST_EXPORT Result<nlohmann::json> ToJson(const LoadTableResult& model);
 
+// Response-only model: a client never serializes it, so no ToJson.
+ICEBERG_REST_EXPORT Result<LoadCredentialsResponse> LoadCredentialsResponseFromJson(
+    const nlohmann::json& json);
+
 ICEBERG_REST_EXPORT Result<CreateTableRequest> CreateTableRequestFromJson(
     const nlohmann::json& json);
 template <>