This is an automated email from the ASF dual-hosted git repository.
wgtmac pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg-cpp.git
The following commit(s) were added to refs/heads/main by this push:
new 50454271 refactor(auth): refresh OAuth2 tokens on a caller-provided
executor (#1033)
50454271 is described below
commit 504542718687f097c9b3e0bd65b804c85bbd6e92
Author: Gang Wu <[email protected]>
AuthorDate: Fri Oct 23 15:44:04 2026 +0800
refactor(auth): refresh OAuth2 tokens on a caller-provided executor (#1033)
Replace TokenRefreshScheduler with an internal ScheduledExecutor using
one process-wide timer. Pass a borrowed auth executor through
RestCatalog and AuthManager.
Use RetryRunner for HTTP retries within each refresh task, matching
Java's retry settings. Stop refreshing when submission fails.
Cancel pending refreshes and wait for active refreshes when closing
the session.
---
src/iceberg/CMakeLists.txt | 1 +
src/iceberg/catalog/rest/CMakeLists.txt | 1 -
src/iceberg/catalog/rest/auth/auth_manager.cc | 33 ++-
src/iceberg/catalog/rest/auth/auth_manager.h | 4 +
src/iceberg/catalog/rest/auth/auth_session.cc | 18 +-
src/iceberg/catalog/rest/auth/auth_session.h | 23 +-
.../catalog/rest/auth/auth_session_internal.h | 174 +++++++++-----
.../rest/auth/sigv4_auth_manager_internal.h | 2 +
src/iceberg/catalog/rest/auth/sigv4_manager.cc | 4 +
.../catalog/rest/auth/token_refresh_scheduler.cc | 125 ----------
.../catalog/rest/auth/token_refresh_scheduler.h | 112 ---------
src/iceberg/catalog/rest/rest_catalog.cc | 5 +-
src/iceberg/catalog/rest/rest_catalog.h | 13 +-
src/iceberg/cpp20_compatibility_internal.h | 1 -
src/iceberg/test/CMakeLists.txt | 4 +
src/iceberg/test/arrow_s3_file_io_test.cc | 49 +---
src/iceberg/test/auth_manager_test.cc | 265 ++++++++++++++-------
src/iceberg/test/executor.h | 36 +++
src/iceberg/test/rest_catalog_unit_test.cc | 73 +++++-
src/iceberg/test/rest_metrics_reporter_test.cc | 49 ++--
src/iceberg/test/scheduled_executor_test.cc | 188 +++++++++++++++
src/iceberg/util/executor.h | 2 +
src/iceberg/util/retry_util.h | 5 +-
src/iceberg/util/scheduled_executor.cc | 213 +++++++++++++++++
src/iceberg/util/scheduled_executor_internal.h | 73 ++++++
25 files changed, 962 insertions(+), 511 deletions(-)
diff --git a/src/iceberg/CMakeLists.txt b/src/iceberg/CMakeLists.txt
index c751cd7a..5af4fdae 100644
--- a/src/iceberg/CMakeLists.txt
+++ b/src/iceberg/CMakeLists.txt
@@ -155,6 +155,7 @@ set(ICEBERG_SOURCES
util/murmurhash3_internal.cc
util/property_util.cc
util/retry_util.cc
+ util/scheduled_executor.cc
util/snapshot_util.cc
util/string_util.cc
util/struct_like_set.cc
diff --git a/src/iceberg/catalog/rest/CMakeLists.txt
b/src/iceberg/catalog/rest/CMakeLists.txt
index 0736290b..a8146d46 100644
--- a/src/iceberg/catalog/rest/CMakeLists.txt
+++ b/src/iceberg/catalog/rest/CMakeLists.txt
@@ -24,7 +24,6 @@ set(ICEBERG_REST_SOURCES
auth/auth_session.cc
auth/oauth2_util.cc
auth/sigv4_manager.cc
- auth/token_refresh_scheduler.cc
catalog_properties.cc
endpoint.cc
error_handlers.cc
diff --git a/src/iceberg/catalog/rest/auth/auth_manager.cc
b/src/iceberg/catalog/rest/auth/auth_manager.cc
index 5facfa3e..df5d7304 100644
--- a/src/iceberg/catalog/rest/auth/auth_manager.cc
+++ b/src/iceberg/catalog/rest/auth/auth_manager.cc
@@ -140,6 +140,10 @@ Result<std::unique_ptr<AuthManager>> MakeBasicAuthManager(
/// \brief OAuth2 authentication manager.
class OAuth2Manager : public AuthManager {
public:
+ void SetAuthRefreshExecutor(Executor* executor) override {
+ refresh_executor_ = executor;
+ }
+
Result<std::shared_ptr<AuthSession>> InitSession(
std::shared_ptr<HttpClient> init_client,
const std::unordered_map<std::string, std::string>& properties) override
{
@@ -178,10 +182,7 @@ class OAuth2Manager : public AuthManager {
refresh_client_ = std::move(shared_client);
// Reuse the token response and start time from the init phase.
if (auth_response_.has_value()) {
- return internal::MakeOAuth2Session(
- *auth_response_, config.oauth2_server_uri(), config.client_id(),
- config.client_secret(), config.scope(), config.keep_refreshed(),
- config.optional_oauth_params(), refresh_client_, start_time_);
+ return MakeSession(*auth_response_, config, config.keep_refreshed(),
start_time_);
}
// TODO(lishuxu): Honor token-refresh-enabled for catalog bearer tokens,
matching
@@ -192,10 +193,7 @@ class OAuth2Manager : public AuthManager {
.token_type = "bearer",
.issued_token_type = AuthProperties::kAccessTokenType,
};
- return AuthSession::MakeOAuth2(token_response,
config.oauth2_server_uri(),
- config.client_id(),
config.client_secret(),
- config.scope(), /*keep_refreshed=*/false,
- config.optional_oauth_params(),
refresh_client_);
+ return MakeSession(token_response, config, /*keep_refreshed=*/false);
}
// Fetch a new token using client_credentials grant.
@@ -206,10 +204,7 @@ class OAuth2Manager : public AuthManager {
ICEBERG_ASSIGN_OR_RAISE(
token_response,
OAuth2Util::FetchToken(*refresh_client_, *base_session, config));
- return AuthSession::MakeOAuth2(token_response,
config.oauth2_server_uri(),
- config.client_id(),
config.client_secret(),
- config.scope(), config.keep_refreshed(),
- config.optional_oauth_params(),
refresh_client_);
+ return MakeSession(token_response, config, config.keep_refreshed());
}
return MakeSession(AccessTokenResponse(""), config,
/*keep_refreshed=*/false);
@@ -256,13 +251,16 @@ class OAuth2Manager : public AuthManager {
Result<std::shared_ptr<AuthSession>> MakeSession(
const OAuthTokenResponse& token_response, const AuthProperties& config,
- bool keep_refreshed) const {
+ bool keep_refreshed,
+ std::optional<std::chrono::steady_clock::time_point>
token_request_started_at =
+ std::nullopt) const {
ICEBERG_PRECHECK(refresh_client_ != nullptr,
"OAuth2 catalog session must be initialized before child
sessions");
- return AuthSession::MakeOAuth2(token_response, config.oauth2_server_uri(),
- config.client_id(), config.client_secret(),
- config.scope(), keep_refreshed,
- config.optional_oauth_params(),
refresh_client_);
+ return internal::MakeOAuth2Session(token_response,
config.oauth2_server_uri(),
+ config.client_id(),
config.client_secret(),
+ config.scope(), keep_refreshed,
+ config.optional_oauth_params(),
refresh_client_,
+ token_request_started_at,
refresh_executor_);
}
Result<std::shared_ptr<AuthSession>> MaybeCreateChildSession(
@@ -319,6 +317,7 @@ class OAuth2Manager : public AuthManager {
std::optional<OAuthTokenResponse> auth_response_;
std::optional<std::chrono::steady_clock::time_point> start_time_;
std::shared_ptr<HttpClient> refresh_client_;
+ Executor* refresh_executor_ = nullptr;
};
Result<std::unique_ptr<AuthManager>> MakeOAuth2Manager(
diff --git a/src/iceberg/catalog/rest/auth/auth_manager.h
b/src/iceberg/catalog/rest/auth/auth_manager.h
index 6ba25c28..d390d0be 100644
--- a/src/iceberg/catalog/rest/auth/auth_manager.h
+++ b/src/iceberg/catalog/rest/auth/auth_manager.h
@@ -97,6 +97,10 @@ class ICEBERG_REST_EXPORT AuthManager {
///
/// \return Status of the close operation.
virtual Status Close() { return {}; }
+
+ /// \brief Set the borrowed refresh executor before InitSession.
+ /// It must outlive this manager and its sessions. Wrappers must forward
this call.
+ virtual void SetAuthRefreshExecutor(Executor* /*executor*/) {}
};
} // namespace iceberg::rest::auth
diff --git a/src/iceberg/catalog/rest/auth/auth_session.cc
b/src/iceberg/catalog/rest/auth/auth_session.cc
index 60aa0a3b..bd74b5bb 100644
--- a/src/iceberg/catalog/rest/auth/auth_session.cc
+++ b/src/iceberg/catalog/rest/auth/auth_session.cc
@@ -19,18 +19,13 @@
#include "iceberg/catalog/rest/auth/auth_session.h"
-#include <algorithm>
#include <chrono>
-#include <condition_variable>
#include <memory>
-#include <mutex>
-#include <shared_mutex>
#include <utility>
#include "iceberg/catalog/rest/auth/auth_properties.h"
#include "iceberg/catalog/rest/auth/auth_session_internal.h"
#include "iceberg/catalog/rest/auth/oauth2_util.h"
-#include "iceberg/catalog/rest/auth/token_refresh_scheduler.h"
#include "iceberg/catalog/rest/http_client.h"
#include "iceberg/util/macros.h"
@@ -67,10 +62,10 @@ Result<std::shared_ptr<AuthSession>>
AuthSession::MakeOAuth2(
const std::string& client_id, const std::string& client_secret,
const std::string& scope, bool keep_refreshed,
const std::unordered_map<std::string, std::string>& optional_oauth_params,
- std::shared_ptr<HttpClient> client) {
+ std::shared_ptr<HttpClient> client, Executor* auth_executor) {
return internal::MakeOAuth2Session(
initial_token, token_endpoint, client_id, client_secret, scope,
keep_refreshed,
- optional_oauth_params, std::move(client), std::nullopt);
+ optional_oauth_params, std::move(client), std::nullopt, auth_executor);
}
Result<std::shared_ptr<AuthSession>> internal::MakeOAuth2Session(
@@ -79,14 +74,19 @@ Result<std::shared_ptr<AuthSession>>
internal::MakeOAuth2Session(
const std::string& scope, bool keep_refreshed,
const std::unordered_map<std::string, std::string>& optional_oauth_params,
std::shared_ptr<HttpClient> client,
- std::optional<std::chrono::steady_clock::time_point>
token_request_started_at) {
+ std::optional<std::chrono::steady_clock::time_point>
token_request_started_at,
+ Executor* refresh_executor) {
+ ICEBERG_PRECHECK(!keep_refreshed || refresh_executor != nullptr,
+ "OAuth2 token refresh is enabled but no auth executor was
provided; "
+ "pass an executor or set {}=false",
+ AuthProperties::kKeepRefreshed.key());
internal::OAuth2Session::Config config{
.token_endpoint = token_endpoint,
.client_id = client_id,
.client_secret = client_secret,
.scope = scope,
.optional_oauth_params = optional_oauth_params,
- .keep_refreshed = keep_refreshed,
+ .refresh_executor = keep_refreshed ? refresh_executor : nullptr,
};
ICEBERG_ASSIGN_OR_RAISE(auto session, internal::OAuth2Session::Make(
initial_token, std::move(config),
diff --git a/src/iceberg/catalog/rest/auth/auth_session.h
b/src/iceberg/catalog/rest/auth/auth_session.h
index bdc77ebc..7e330706 100644
--- a/src/iceberg/catalog/rest/auth/auth_session.h
+++ b/src/iceberg/catalog/rest/auth/auth_session.h
@@ -28,6 +28,7 @@
#include "iceberg/catalog/rest/iceberg_rest_export.h"
#include "iceberg/catalog/rest/type_fwd.h"
#include "iceberg/result.h"
+#include "iceberg/type_fwd.h"
/// \file iceberg/catalog/rest/auth/auth_session.h
/// \brief Authentication session interface for REST catalog.
@@ -68,13 +69,8 @@ class ICEBERG_REST_EXPORT AuthSession {
/// \brief Return OAuth2 metadata when this is an OAuth2 session.
virtual std::optional<OAuth2SessionInfo> OAuth2Info() const { return
std::nullopt; }
- /// \brief Close the session and release any resources.
- ///
- /// This method is called when the session is no longer needed. For stateful
- /// sessions (e.g., OAuth2 with token refresh), this should stop any
background
- /// threads and release resources.
- ///
- /// \return Status indicating success or failure of closing the session.
+ /// \brief Stop background refresh and release session resources.
+ /// \return Status of the close operation.
virtual Status Close() { return {}; }
/// \brief Create a session with static headers.
@@ -90,10 +86,8 @@ class ICEBERG_REST_EXPORT AuthSession {
/// \brief Create an OAuth2 session with automatic token refresh.
///
- /// This factory method creates a session that holds an access token and,
- /// when keep_refreshed is enabled, schedules background refresh based on
- /// token expiration. Authenticate() uses the latest cached Authorization
- /// header and does not perform a synchronous token refresh.
+ /// Refreshes expiring tokens in the background when keep_refreshed is true.
+ /// Authenticate() uses the cached token without refreshing it.
///
/// \param initial_token The initial token response from FetchToken().
/// \param token_endpoint Full URL of the OAuth2 token endpoint for refresh.
@@ -104,13 +98,16 @@ class ICEBERG_REST_EXPORT AuthSession {
/// \param optional_oauth_params Optional OAuth params (audience, resource)
for refresh.
/// \param client HTTP client for making refresh requests. The session
retains
/// ownership of the client.
- /// \return A new session that manages token lifecycle automatically.
+ /// \param auth_executor Borrowed refresh executor; required if
keep_refreshed is true.
+ /// Ignored otherwise. Must enqueue asynchronously and outlive the session
and its
+ /// tasks.
+ /// \return The session, or InvalidArgument if refresh is enabled without an
executor.
static Result<std::shared_ptr<AuthSession>> MakeOAuth2(
const OAuthTokenResponse& initial_token, const std::string&
token_endpoint,
const std::string& client_id, const std::string& client_secret,
const std::string& scope, bool keep_refreshed,
const std::unordered_map<std::string, std::string>&
optional_oauth_params,
- std::shared_ptr<HttpClient> client);
+ std::shared_ptr<HttpClient> client, Executor* auth_executor = nullptr);
};
} // namespace iceberg::rest::auth
diff --git a/src/iceberg/catalog/rest/auth/auth_session_internal.h
b/src/iceberg/catalog/rest/auth/auth_session_internal.h
index b8db3836..fa2f7ede 100644
--- a/src/iceberg/catalog/rest/auth/auth_session_internal.h
+++ b/src/iceberg/catalog/rest/auth/auth_session_internal.h
@@ -20,9 +20,9 @@
#pragma once
#include <algorithm>
-#include <atomic>
#include <chrono>
#include <condition_variable>
+#include <exception>
#include <memory>
#include <mutex>
#include <optional>
@@ -34,9 +34,11 @@
#include "iceberg/catalog/rest/auth/auth_properties.h"
#include "iceberg/catalog/rest/auth/auth_session.h"
#include "iceberg/catalog/rest/auth/oauth2_util.h"
-#include "iceberg/catalog/rest/auth/token_refresh_scheduler.h"
#include "iceberg/catalog/rest/http_client.h"
+#include "iceberg/logging/log_macros.h"
#include "iceberg/util/macros.h"
+#include "iceberg/util/retry_util.h"
+#include "iceberg/util/scheduled_executor_internal.h"
namespace iceberg::rest::auth::internal {
@@ -67,7 +69,8 @@ class OAuth2Session final : public AuthSession,
std::string client_secret;
std::string scope;
std::unordered_map<std::string, std::string> optional_oauth_params;
- bool keep_refreshed;
+ /// Borrowed executor for background refresh. Null disables refresh.
+ Executor* refresh_executor = nullptr;
};
static Result<std::shared_ptr<OAuth2Session>> Make(
@@ -111,17 +114,23 @@ class OAuth2Session final : public AuthSession,
std::shared_ptr<HttpClient> client)
: config_(std::move(config)),
refresh_properties_(std::move(refresh_properties)),
- client_(std::move(client)) {}
+ client_(std::move(client)) {
+ if (config_.refresh_executor != nullptr) {
+ scheduler_.emplace(*config_.refresh_executor);
+ }
+ }
Status CloseImpl() {
- bool expected = false;
- if (!closed_.compare_exchange_strong(expected, true)) {
- return {};
+ // Stop queued refreshes before shutting down the scheduler.
+ {
+ std::lock_guard lock(refresh_mutex_);
+ closed_ = true;
+ }
+ if (scheduler_) {
+ scheduler_->Shutdown();
}
- TokenRefreshScheduler::Instance().Cancel(scheduled_task_id_.exchange(0));
std::unique_lock lock(refresh_mutex_);
refresh_cv_.wait(lock, [this] { return active_refresh_count_ == 0; });
- TokenRefreshScheduler::Instance().Cancel(scheduled_task_id_.exchange(0));
return {};
}
@@ -143,10 +152,16 @@ class OAuth2Session final : public AuthSession,
public:
explicit RefreshAttemptGuard(OAuth2Session& session) : session_(session) {
std::lock_guard lock(session_.refresh_mutex_);
- ++session_.active_refresh_count_;
+ active_ = !session_.closed_;
+ if (active_) {
+ ++session_.active_refresh_count_;
+ }
}
~RefreshAttemptGuard() {
+ if (!active_) {
+ return;
+ }
bool notify = false;
{
std::lock_guard lock(session_.refresh_mutex_);
@@ -157,8 +172,11 @@ class OAuth2Session final : public AuthSession,
}
}
+ bool IsActive() const { return active_; }
+
private:
OAuth2Session& session_;
+ bool active_ = false;
};
void UpdateTokenState(const OAuthTokenResponse& token_response,
@@ -182,73 +200,103 @@ class OAuth2Session final : public AuthSession,
const OAuthTokenResponse& token_response,
std::optional<std::chrono::steady_clock::time_point>
token_request_started_at) {
UpdateTokenState(token_response, token_request_started_at);
- if (config_.keep_refreshed &&
- expires_at_ != std::chrono::steady_clock::time_point{}) {
- ScheduleRefresh();
+ if (!scheduler_) {
+ return;
+ }
+ // Refresh an already-expired initial token immediately.
+ if (auto delay = CalculateRefreshDelay()) {
+ ScheduleRefresh(*delay);
}
}
- void DoRefresh() {
- DoRefreshAttempt(0, std::chrono::milliseconds(200),
std::chrono::steady_clock::now());
- }
-
- void DoRefreshAttempt(int attempt, std::chrono::milliseconds backoff,
- std::chrono::steady_clock::time_point
refresh_started_at) {
- static constexpr int kMaxRetries = 5;
- static constexpr auto kMaxBackoff = std::chrono::milliseconds(10'000);
+ // Complete one refresh, including HTTP retries, on the executor's worker.
+ std::optional<std::chrono::milliseconds> Refresh() {
RefreshAttemptGuard guard(*this);
- if (closed_.load()) return;
+ if (!guard.IsActive()) {
+ return std::nullopt;
+ }
+
+ auto result = RetryRunner<retry::StopRetryOn<>>({.num_retries =
5}).Run([this] {
+ return RefreshToken();
+ });
+ if (!result.has_value()) {
+ ICEBERG_LOG_WARN("Failed to refresh OAuth2 token: {}",
result.error().message);
+ return std::nullopt;
+ }
+ return *result;
+ }
- auto empty_session = AuthSession::MakeDefault({});
- // TODO(lishuxu): Honor token-exchange-enabled and refresh via token
exchange,
- // matching Java.
- auto result = OAuth2Util::FetchToken(*client_, *empty_session,
refresh_properties_);
- if (result.has_value()) {
- auto& response = result.value();
+ Result<std::optional<std::chrono::milliseconds>> RefreshToken() {
+ {
+ std::lock_guard lock(refresh_mutex_);
+ if (closed_) {
+ return std::nullopt;
+ }
+ }
+ try {
+ auto request_started_at = std::chrono::steady_clock::now();
+ auto empty_session = AuthSession::MakeDefault({});
+ // TODO(lishuxu): Honor token-exchange-enabled when refreshing tokens.
+ ICEBERG_ASSIGN_OR_RAISE(
+ auto response,
+ OAuth2Util::FetchToken(*client_, *empty_session,
refresh_properties_));
{
std::unique_lock lock(mutex_);
- UpdateTokenState(response, refresh_started_at);
+ UpdateTokenState(response, request_started_at);
}
- ScheduleRefresh();
- return;
- }
-
- if (attempt + 1 < kMaxRetries && !closed_.load()) {
- auto next_backoff =
-
std::min(std::chrono::duration_cast<std::chrono::milliseconds>(backoff * 2),
- kMaxBackoff);
- std::weak_ptr<OAuth2Session> weak_self = shared_from_this();
- auto retry_id = TokenRefreshScheduler::Instance().Schedule(
- backoff, [weak_self = std::move(weak_self), next_attempt = attempt +
1,
- next_backoff, refresh_started_at] {
- if (auto self = weak_self.lock()) {
- self->DoRefreshAttempt(next_attempt, next_backoff,
refresh_started_at);
- }
- });
- scheduled_task_id_.store(retry_id);
+ auto delay = CalculateRefreshDelay();
+ if (delay.has_value() && *delay <= std::chrono::milliseconds::zero()) {
+ return TokenExpired("Refreshed OAuth2 token is already expired");
+ }
+ return delay;
+ } catch (const std::exception& e) {
+ return UnknownError("OAuth2 token refresh threw: {}", e.what());
+ } catch (...) {
+ return UnknownError("OAuth2 token refresh threw a non-standard
exception");
}
}
- void ScheduleRefresh() {
- if (!config_.keep_refreshed || closed_.load()) return;
- auto delay = CalculateRefreshDelay();
- if (delay < std::chrono::milliseconds::zero()) return;
+ static constexpr int kMaxSubmissionRetries = 5;
- std::weak_ptr<OAuth2Session> weak_self = shared_from_this();
- auto new_id = TokenRefreshScheduler::Instance().Schedule(
- delay, [weak_self = std::move(weak_self)] {
- if (auto self = weak_self.lock()) self->DoRefresh();
+ void ScheduleRefresh(std::chrono::milliseconds delay, int submission_retries
= 0) {
+ auto weak_self = weak_from_this();
+ ExecutorTask refresh([weak_self] {
+ if (auto self = weak_self.lock()) {
+ if (auto next_delay = self->Refresh()) {
+ self->ScheduleRefresh(*next_delay);
+ }
+ }
+ });
+ ::iceberg::internal::TaskSubmissionCallback on_submission(
+ [weak_self, submission_retries](Status status) {
+ if (status.has_value()) {
+ return;
+ }
+ if (submission_retries >= kMaxSubmissionRetries) {
+ ICEBERG_LOG_WARN(
+ "Stopping OAuth2 token refresh after repeated submission
failures: {}",
+ status.error().message);
+ return;
+ }
+ if (auto self = weak_self.lock()) {
+ // Requeue on the timer; HTTP retries start only after Submit
succeeds.
+ auto retry_delay = std::chrono::milliseconds(200) * (1 <<
submission_retries);
+ self->ScheduleRefresh(retry_delay, submission_retries + 1);
+ }
});
- scheduled_task_id_.store(new_id);
+ scheduler_->Schedule(delay, std::move(refresh), std::move(on_submission));
}
- std::chrono::milliseconds CalculateRefreshDelay() const {
+ /// Returns nullopt if expiration is unknown.
+ std::optional<std::chrono::milliseconds> CalculateRefreshDelay() const {
std::shared_lock lock(mutex_);
- auto now = std::chrono::steady_clock::now();
if (expires_at_ == std::chrono::steady_clock::time_point{}) {
- return std::chrono::milliseconds(-1);
+ return std::nullopt;
+ }
+ auto now = std::chrono::steady_clock::now();
+ if (expires_at_ <= now) {
+ return std::chrono::milliseconds::zero();
}
- if (expires_at_ <= now) return std::chrono::milliseconds::zero();
auto expires_in =
std::chrono::duration_cast<std::chrono::milliseconds>(expires_at_ -
now);
auto refresh_window = std::min(expires_in / 10,
std::chrono::milliseconds(300'000));
@@ -264,19 +312,21 @@ class OAuth2Session final : public AuthSession,
Config config_;
AuthProperties refresh_properties_;
std::shared_ptr<HttpClient> client_;
- std::atomic<uint64_t> scheduled_task_id_{0};
- std::atomic<bool> closed_{false};
+ bool closed_ = false;
std::mutex refresh_mutex_;
std::condition_variable refresh_cv_;
int active_refresh_count_ = 0;
+ std::optional<::iceberg::internal::ScheduledExecutor> scheduler_;
};
+/// Returns InvalidArgument if keep_refreshed is true without an executor.
Result<std::shared_ptr<AuthSession>> MakeOAuth2Session(
const OAuthTokenResponse& initial_token, const std::string& token_endpoint,
const std::string& client_id, const std::string& client_secret,
const std::string& scope, bool keep_refreshed,
const std::unordered_map<std::string, std::string>& optional_oauth_params,
std::shared_ptr<HttpClient> client,
- std::optional<std::chrono::steady_clock::time_point>
token_request_started_at);
+ std::optional<std::chrono::steady_clock::time_point>
token_request_started_at,
+ Executor* refresh_executor);
} // namespace iceberg::rest::auth::internal
diff --git a/src/iceberg/catalog/rest/auth/sigv4_auth_manager_internal.h
b/src/iceberg/catalog/rest/auth/sigv4_auth_manager_internal.h
index 7bb7a3aa..feb4da5e 100644
--- a/src/iceberg/catalog/rest/auth/sigv4_auth_manager_internal.h
+++ b/src/iceberg/catalog/rest/auth/sigv4_auth_manager_internal.h
@@ -117,6 +117,8 @@ class ICEBERG_REST_EXPORT SigV4AuthManager : public
AuthManager {
explicit SigV4AuthManager(std::unique_ptr<AuthManager> delegate);
~SigV4AuthManager() override;
+ void SetAuthRefreshExecutor(Executor* executor) override;
+
Result<std::shared_ptr<AuthSession>> InitSession(
std::shared_ptr<HttpClient> init_client,
const std::unordered_map<std::string, std::string>& properties) override;
diff --git a/src/iceberg/catalog/rest/auth/sigv4_manager.cc
b/src/iceberg/catalog/rest/auth/sigv4_manager.cc
index 6e0dce86..65689974 100644
--- a/src/iceberg/catalog/rest/auth/sigv4_manager.cc
+++ b/src/iceberg/catalog/rest/auth/sigv4_manager.cc
@@ -379,6 +379,10 @@
SigV4AuthManager::SigV4AuthManager(std::unique_ptr<AuthManager> delegate)
SigV4AuthManager::~SigV4AuthManager() = default;
+void SigV4AuthManager::SetAuthRefreshExecutor(Executor* executor) {
+ delegate_->SetAuthRefreshExecutor(executor);
+}
+
Result<std::shared_ptr<AuthSession>> SigV4AuthManager::InitSession(
std::shared_ptr<HttpClient> init_client,
const std::unordered_map<std::string, std::string>& properties) {
diff --git a/src/iceberg/catalog/rest/auth/token_refresh_scheduler.cc
b/src/iceberg/catalog/rest/auth/token_refresh_scheduler.cc
deleted file mode 100644
index d66ed7d9..00000000
--- a/src/iceberg/catalog/rest/auth/token_refresh_scheduler.cc
+++ /dev/null
@@ -1,125 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-#include "iceberg/catalog/rest/auth/token_refresh_scheduler.h"
-
-#include <algorithm>
-#include <ranges>
-
-namespace iceberg::rest::auth {
-
-TokenRefreshScheduler& TokenRefreshScheduler::Instance() {
- // Intentionally leaked to avoid destruction-order races at process exit.
- static auto* instance = new TokenRefreshScheduler();
- return *instance;
-}
-
-TokenRefreshScheduler::TokenRefreshScheduler() : worker_([this] { Run(); }) {}
-
-TokenRefreshScheduler::~TokenRefreshScheduler() { Shutdown(); }
-
-uint64_t TokenRefreshScheduler::Schedule(std::chrono::milliseconds delay,
- std::function<void()> callback) {
- std::lock_guard lock(mutex_);
- if (shutdown_) {
- return 0;
- }
- uint64_t id = next_id_++;
- tasks_.push_back(Task{.id = id,
- .fire_at = std::chrono::steady_clock::now() + delay,
- .callback = std::move(callback)});
- cv_.notify_one();
- return id;
-}
-
-void TokenRefreshScheduler::Cancel(uint64_t handle) {
- if (handle == 0) return;
- std::lock_guard lock(mutex_);
- std::erase_if(tasks_, [handle](const Task& t) { return t.id == handle; });
-}
-
-void TokenRefreshScheduler::Shutdown() {
- {
- std::lock_guard lock(mutex_);
- if (shutdown_) return;
- shutdown_ = true;
- tasks_.clear();
- }
- cv_.notify_one();
- if (worker_.joinable()) {
- worker_.join();
- }
-}
-
-void TokenRefreshScheduler::Run() {
- while (true) {
- std::function<void()> callback;
-
- {
- std::unique_lock lock(mutex_);
-
- if (tasks_.empty() && !shutdown_) {
- // Wait until a task is added or shutdown is requested
- cv_.wait(lock, [this] { return !tasks_.empty() || shutdown_; });
- }
-
- if (shutdown_) break;
- if (tasks_.empty()) continue;
-
- // Find the task with the earliest fire_at
- auto earliest_it = std::ranges::min_element(
- tasks_, [](const Task& a, const Task& b) { return a.fire_at <
b.fire_at; });
-
- auto fire_at = earliest_it->fire_at;
- auto target_id = earliest_it->id;
-
- // Wait until fire_at or until woken (new task, cancel, or shutdown).
- // Note: The predicate does O(n) scan on each spurious wakeup. This is
- // acceptable for the expected task count (< 10). If task count grows
- // significantly, consider replacing vector with a priority queue.
- cv_.wait_until(lock, fire_at, [&] {
- // Wake up if: shutdown, task list changed, or time is up
- if (shutdown_) return true;
- if (tasks_.empty()) return true;
- // Check if the earliest task has changed (new task added or cancelled)
- auto new_earliest = std::ranges::min_element(
- tasks_, [](const Task& a, const Task& b) { return a.fire_at <
b.fire_at; });
- return new_earliest->id != target_id;
- });
-
- if (shutdown_) break;
-
- // If we were woken because the earliest task changed, loop again
- auto now = std::chrono::steady_clock::now();
- auto due_it =
- std::ranges::find_if(tasks_, [now](const Task& t) { return t.fire_at
<= now; });
- if (due_it == tasks_.end()) continue;
-
- callback = std::move(due_it->callback);
- tasks_.erase(due_it);
- }
-
- // Execute callback outside the lock
- if (callback) {
- callback();
- }
- }
-}
-
-} // namespace iceberg::rest::auth
diff --git a/src/iceberg/catalog/rest/auth/token_refresh_scheduler.h
b/src/iceberg/catalog/rest/auth/token_refresh_scheduler.h
deleted file mode 100644
index 5ef20ed9..00000000
--- a/src/iceberg/catalog/rest/auth/token_refresh_scheduler.h
+++ /dev/null
@@ -1,112 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-#pragma once
-
-#include <chrono>
-#include <condition_variable>
-#include <cstdint>
-#include <functional>
-#include <mutex>
-#include <thread>
-#include <vector>
-
-#include "iceberg/catalog/rest/iceberg_rest_export.h"
-
-/// \file iceberg/catalog/rest/auth/token_refresh_scheduler.h
-/// \brief Global scheduler for OAuth2 token refresh tasks.
-
-namespace iceberg::rest::auth {
-
-/// \brief A process-global scheduler for delayed token refresh tasks.
-///
-/// Uses a single background thread that sleeps until the next task is due.
-/// All OAuth2Session instances share this scheduler. Tasks are lightweight
-/// (a single HTTP POST to refresh a token), so one thread is sufficient.
-///
-/// Thread safety: All public methods are thread-safe.
-///
-/// TODO(lishuxu): Migrate to the shared thread pool abstraction once available
-/// (see
https://github.com/apache/iceberg-cpp/pull/646#discussion_r3304315308).
-class ICEBERG_REST_EXPORT TokenRefreshScheduler {
- public:
- /// \brief Get the global singleton instance.
- ///
- /// Lazily created on first access and intentionally leaked: the worker
- /// thread is reclaimed by the OS at process exit. Tests needing
- /// deterministic shutdown should use a local instance.
- static TokenRefreshScheduler& Instance();
-
- /// \brief Schedule a callback to run after a delay.
- ///
- /// \param delay Time to wait before executing the callback.
- /// \param callback Function to execute when the delay expires.
- /// \return A unique handle that can be used to cancel the task.
- uint64_t Schedule(std::chrono::milliseconds delay, std::function<void()>
callback);
-
- /// \brief Cancel a previously scheduled task.
- ///
- /// If the task has already fired or does not exist, this is a no-op.
- ///
- /// \param handle The handle returned by Schedule().
- void Cancel(uint64_t handle);
-
- /// \brief Shutdown the scheduler, cancelling all pending tasks.
- ///
- /// After shutdown, Schedule() calls are no-ops (return 0).
- /// This is called automatically on destruction.
- ///
- /// WARNING: Do not call this on the global Instance() unless you intend to
- /// permanently stop all token refresh for the entire process. This is mainly
- /// useful for testing with locally-constructed scheduler instances.
- void Shutdown();
-
- ~TokenRefreshScheduler();
-
- // Non-copyable, non-movable
- TokenRefreshScheduler(const TokenRefreshScheduler&) = delete;
- TokenRefreshScheduler& operator=(const TokenRefreshScheduler&) = delete;
- TokenRefreshScheduler(TokenRefreshScheduler&&) = delete;
- TokenRefreshScheduler& operator=(TokenRefreshScheduler&&) = delete;
-
- /// \brief Construct a scheduler (prefer Instance() for production use).
- ///
- /// This constructor is public to allow testing with isolated instances.
- /// In production code, use Instance() to get the global singleton.
- TokenRefreshScheduler();
-
- private:
- /// \brief Worker loop that processes tasks.
- void Run();
-
- struct Task {
- uint64_t id;
- std::chrono::steady_clock::time_point fire_at;
- std::function<void()> callback;
- };
-
- std::mutex mutex_;
- std::condition_variable cv_;
- std::vector<Task> tasks_;
- uint64_t next_id_ = 1; // 0 is reserved as "invalid handle"
- bool shutdown_ = false;
- std::thread worker_;
-};
-
-} // namespace iceberg::rest::auth
diff --git a/src/iceberg/catalog/rest/rest_catalog.cc
b/src/iceberg/catalog/rest/rest_catalog.cc
index a945e25d..c8d88d06 100644
--- a/src/iceberg/catalog/rest/rest_catalog.cc
+++ b/src/iceberg/catalog/rest/rest_catalog.cc
@@ -447,8 +447,8 @@ RestCatalog::~RestCatalog() {
}
Result<std::shared_ptr<RestCatalog>> RestCatalog::Make(
- const RestCatalogProperties& config, Executor* metrics_executor,
- ClientFactory client_factory) {
+ const RestCatalogProperties& config, ClientFactory client_factory,
+ Executor* auth_executor, Executor* metrics_executor) {
if (!client_factory) {
client_factory = [](const RestCatalogProperties& properties)
-> Result<std::shared_ptr<HttpClient>> {
@@ -467,6 +467,7 @@ Result<std::shared_ptr<RestCatalog>> RestCatalog::Make(
std::string catalog_name = config.Get(RestCatalogProperties::kName);
ICEBERG_ASSIGN_OR_RAISE(auto auth_manager,
auth::AuthManagers::Load(catalog_name,
config.configs()));
+ auth_manager->SetAuthRefreshExecutor(auth_executor);
ICEBERG_ASSIGN_OR_RAISE(
auto paths,
ResourcePaths::Make(std::string(TrimTrailingSlash(uri)),
diff --git a/src/iceberg/catalog/rest/rest_catalog.h
b/src/iceberg/catalog/rest/rest_catalog.h
index 1e713a50..d648fa76 100644
--- a/src/iceberg/catalog/rest/rest_catalog.h
+++ b/src/iceberg/catalog/rest/rest_catalog.h
@@ -58,14 +58,19 @@ class ICEBERG_REST_EXPORT RestCatalog final
std::function<Result<std::shared_ptr<HttpClient>>(const
RestCatalogProperties&)>;
/// \brief Create a RestCatalog instance.
+ /// \param client_factory Optional HTTP client factory, called with the
initial config
+ /// and then the merged server config. When empty, uses the default HTTP
client.
+ /// \param auth_executor Borrowed OAuth2 refresh executor. With OAuth2
credentials,
+ /// token-refresh-enabled=true (default) requires it; null returns
InvalidArgument.
+ /// Must enqueue asynchronously and outlive the catalog, its auth sessions
and their
+ /// tasks.
/// \param metrics_executor Optional non-owning executor for asynchronous
REST metrics.
/// It must outlive the catalog and all tables created by it. When null,
REST metrics
/// are reported synchronously.
- /// \param client_factory Optional HTTP client factory, called with the
initial config
- /// and then the merged server config. When empty, uses the default HTTP
client.
static Result<std::shared_ptr<RestCatalog>> Make(const
RestCatalogProperties& config,
- Executor* metrics_executor
= nullptr,
- ClientFactory
client_factory = {});
+ ClientFactory
client_factory = {},
+ Executor* auth_executor =
nullptr,
+ Executor* metrics_executor
= nullptr);
std::string_view name() const override;
diff --git a/src/iceberg/cpp20_compatibility_internal.h
b/src/iceberg/cpp20_compatibility_internal.h
index bd9a9475..9f5d2451 100644
--- a/src/iceberg/cpp20_compatibility_internal.h
+++ b/src/iceberg/cpp20_compatibility_internal.h
@@ -235,7 +235,6 @@
# include "iceberg/catalog/rest/auth/auth_properties.h"
# include "iceberg/catalog/rest/auth/auth_session.h"
# include "iceberg/catalog/rest/auth/oauth2_util.h"
-# include "iceberg/catalog/rest/auth/token_refresh_scheduler.h"
# include "iceberg/catalog/rest/catalog_properties.h"
# include "iceberg/catalog/rest/constant.h"
# include "iceberg/catalog/rest/endpoint.h"
diff --git a/src/iceberg/test/CMakeLists.txt b/src/iceberg/test/CMakeLists.txt
index 8622a1a6..37e20e33 100644
--- a/src/iceberg/test/CMakeLists.txt
+++ b/src/iceberg/test/CMakeLists.txt
@@ -124,6 +124,9 @@ add_iceberg_test(json_serde_test
metadata_serde_test.cc
schema_json_test.cc)
+add_iceberg_test(scheduled_executor_test SOURCES scheduled_executor_test.cc)
+set_tests_properties(scheduled_executor_test PROPERTIES TIMEOUT 60)
+
add_iceberg_test(util_test
USE_DATA
SOURCES
@@ -342,6 +345,7 @@ if(ICEBERG_BUILD_REST)
rest_metrics_reporter_test.cc
rest_table_scan_test.cc
rest_util_test.cc)
+ set_tests_properties(rest_catalog_test PROPERTIES TIMEOUT 120)
if(ICEBERG_SIGV4)
add_rest_iceberg_test(sigv4_auth_test SOURCES sigv4_auth_test.cc)
diff --git a/src/iceberg/test/arrow_s3_file_io_test.cc
b/src/iceberg/test/arrow_s3_file_io_test.cc
index 9a5eace0..d0d61386 100644
--- a/src/iceberg/test/arrow_s3_file_io_test.cc
+++ b/src/iceberg/test/arrow_s3_file_io_test.cc
@@ -202,18 +202,7 @@ Result<std::string> ReadAll(InputFile& file, size_t
length) {
return content;
}
-class InlineDeleteExecutor : public Executor {
- public:
- Status Submit(ExecutorTask task) override {
- ++submissions;
- std::move(task)();
- return {};
- }
-
- int submissions = 0;
-};
-
-enum class SubmitFailure { kStatus, kException, kAcceptedException };
+enum class SubmitFailure { kStatus, kException };
class FailingDeleteExecutor : public Executor {
public:
@@ -232,18 +221,11 @@ class FailingDeleteExecutor : public Executor {
Status Submit(ExecutorTask task) override {
++submissions_;
- if (submissions_ == 1 || failure_ == SubmitFailure::kAcceptedException) {
- const bool first = submissions_ == 1;
- const bool wait = !first || failure_ !=
SubmitFailure::kAcceptedException;
- threads_.emplace_back([this, first, wait, task = std::move(task)]()
mutable {
- if (wait) {
- release_.wait();
- }
+ if (submissions_ == 1) {
+ threads_.emplace_back([this, task = std::move(task)]() mutable {
+ release_.wait();
std::move(task)();
++completed;
- if (first) {
- first_completed.set_value();
- }
});
}
if (submissions_ == 2) {
@@ -257,7 +239,6 @@ class FailingDeleteExecutor : public Executor {
}
std::promise<void> submission_failed;
- std::promise<void> first_completed;
std::atomic<int> completed = 0;
private:
@@ -270,7 +251,7 @@ class FailingDeleteExecutor : public Executor {
} // namespace
TEST(DeleteExecutorStateTest, InitializationIsImmutable) {
- InlineDeleteExecutor executor;
+ test::InlineExecutor executor;
DeleteExecutorState state;
EXPECT_THAT(state.Initialize(&executor), IsOk());
EXPECT_THAT(state.Initialize(&executor), IsOk());
@@ -283,8 +264,8 @@ TEST(DeleteExecutorStateTest, InitializationIsImmutable) {
}
TEST(DeleteExecutorStateTest, ConcurrentInitializationChoosesOneExecutor) {
- InlineDeleteExecutor first;
- InlineDeleteExecutor second;
+ test::InlineExecutor first;
+ test::InlineExecutor second;
DeleteExecutorState state;
std::barrier start(2);
std::atomic<int> initialized = 0;
@@ -389,13 +370,11 @@ TEST(BulkDeleteFilesTest,
CountsThrownDeletesAndAttemptsRemainingFiles) {
}
TEST(BulkDeleteFilesTest, DrainsAcceptedTasksWhenSubmissionFails) {
- for (auto failure : {SubmitFailure::kStatus, SubmitFailure::kException,
- SubmitFailure::kAcceptedException}) {
+ for (auto failure : {SubmitFailure::kStatus, SubmitFailure::kException}) {
SCOPED_TRACE(static_cast<int>(failure));
std::promise<void> release;
FailingDeleteExecutor executor(failure, release.get_future().share());
auto submission_failed = executor.submission_failed.get_future();
- auto first_completed = executor.first_completed.get_future();
const std::vector<std::string> paths = {"first", "second", "third"};
std::atomic<int> attempted = 0;
auto run = std::async(std::launch::async, [&] {
@@ -410,18 +389,12 @@ TEST(BulkDeleteFilesTest,
DrainsAcceptedTasksWhenSubmissionFails) {
});
});
const auto submitted = submission_failed.wait_for(std::chrono::seconds(5));
- auto first = std::future_status::ready;
- if (failure == SubmitFailure::kAcceptedException) {
- // Leave only the last task blocked.
- first = first_completed.wait_for(std::chrono::seconds(5));
- }
const auto returned = run.wait_for(std::chrono::milliseconds(100));
// Unblock workers before assertions.
release.set_value();
auto status = run.get();
executor.Wait();
EXPECT_EQ(submitted, std::future_status::ready);
- EXPECT_EQ(first, std::future_status::ready);
EXPECT_EQ(returned, std::future_status::timeout);
EXPECT_THAT(status, HasErrorMessage("executor rejected deletion"));
EXPECT_THAT(status, HasErrorMessage("failed to delete 1 of 3 files"));
@@ -429,7 +402,7 @@ TEST(BulkDeleteFilesTest,
DrainsAcceptedTasksWhenSubmissionFails) {
IsError(failure == SubmitFailure::kStatus ?
ErrorKind::kServiceUnavailable
:
ErrorKind::kIOError));
EXPECT_EQ(attempted, 3);
- EXPECT_EQ(executor.completed, failure == SubmitFailure::kAcceptedException
? 2 : 1);
+ EXPECT_EQ(executor.completed, 1);
}
}
@@ -448,7 +421,7 @@ TEST(BulkDeleteFilesTest,
DoesNotDeleteOnTheCallerWhenSubmissionFails) {
}
TEST(BulkDeleteFilesTest, EmptyDeleteDoesNotSubmitTasks) {
- InlineDeleteExecutor executor;
+ test::InlineExecutor executor;
EXPECT_THAT(BulkDeleteFiles({}, &executor, 3,
[](const std::string&) -> Status {
ADD_FAILURE()
@@ -456,7 +429,7 @@ TEST(BulkDeleteFilesTest, EmptyDeleteDoesNotSubmitTasks) {
return {};
}),
IsOk());
- EXPECT_EQ(executor.submissions, 0);
+ EXPECT_EQ(executor.submit_count(), 0);
}
TEST_F(ArrowS3FileIOTest, Create) {
diff --git a/src/iceberg/test/auth_manager_test.cc
b/src/iceberg/test/auth_manager_test.cc
index fd9605fc..f6f1400b 100644
--- a/src/iceberg/test/auth_manager_test.cc
+++ b/src/iceberg/test/auth_manager_test.cc
@@ -19,13 +19,14 @@
#include "iceberg/catalog/rest/auth/auth_manager.h"
+#include <array>
#include <atomic>
#include <chrono>
-#include <condition_variable>
+#include <functional>
+#include <future>
#include <memory>
-#include <mutex>
+#include <stdexcept>
#include <string>
-#include <thread>
#include <unordered_map>
#include <vector>
@@ -38,7 +39,6 @@
#include "iceberg/catalog/rest/auth/auth_session.h"
#include "iceberg/catalog/rest/auth/auth_session_internal.h"
#include "iceberg/catalog/rest/auth/oauth2_util.h"
-#include "iceberg/catalog/rest/auth/token_refresh_scheduler.h"
#include "iceberg/catalog/rest/catalog_properties.h"
#include "iceberg/catalog/rest/error_handlers.h"
#include "iceberg/catalog/rest/http_client.h"
@@ -46,7 +46,9 @@
#include "iceberg/catalog/session_context.h"
#include "iceberg/json_serde_internal.h"
#include "iceberg/table_identifier.h"
+#include "iceberg/test/executor.h"
#include "iceberg/test/matchers.h"
+#include "iceberg/test/retry.h"
#include "iceberg/util/base64.h"
namespace iceberg::rest::auth {
@@ -692,102 +694,195 @@ TEST_F(AuthManagerTest,
ExpiresAtMillisInvalidTokensReturnNullopt) {
}
}
-// ---- TokenRefreshScheduler tests ----
-
-// Verifies that a scheduled task fires after the specified delay
-TEST(TokenRefreshSchedulerTest, ScheduleFiresAfterDelay) {
- TokenRefreshScheduler scheduler;
- std::mutex mutex;
- std::condition_variable cv;
- bool fired = false;
+// ---- OAuth2Session tests ----
- scheduler.Schedule(std::chrono::milliseconds(50), [&] {
- {
- std::lock_guard lock(mutex);
- fired = true;
- }
- cv.notify_one();
- });
+namespace {
- {
- std::unique_lock lock(mutex);
- EXPECT_TRUE(cv.wait_for(lock, std::chrono::seconds(5), [&] { return fired;
}));
+class RefreshHttpClient : public HttpClient {
+ public:
+ Result<HttpResponse> PostForm(const std::string&,
+ const std::unordered_map<std::string,
std::string>&,
+ const std::unordered_map<std::string,
std::string>&,
+ const ErrorHandler&, AuthSession&) override {
+ ++requests;
+ return response();
}
+ std::function<Result<HttpResponse>()> response = [] {
+ return IOError("refresh failed");
+ };
+ std::atomic<int> requests{0};
+};
- scheduler.Shutdown();
-}
-
-// Verifies that cancelling a task prevents it from executing
-TEST(TokenRefreshSchedulerTest, CancelPreventsExecution) {
- TokenRefreshScheduler scheduler;
- std::atomic<bool> fired{false};
-
- auto handle =
- scheduler.Schedule(std::chrono::milliseconds(100), [&] {
fired.store(true); });
-
- // Cancel before it fires
- scheduler.Cancel(handle);
-
- // Wait past the scheduled time
- std::this_thread::sleep_for(std::chrono::milliseconds(200));
- EXPECT_FALSE(fired.load());
+using test::MockExecutor;
- scheduler.Shutdown();
+// Already expired, so the session schedules its first refresh immediately.
+OAuthTokenResponse ExpiredInitialToken() {
+ return {.access_token = "initial", .token_type = "bearer", .expires_in_secs
= 0};
}
-// Verifies that shutdown with pending tasks does not crash
-TEST(TokenRefreshSchedulerTest, ShutdownWithPendingTasks) {
- TokenRefreshScheduler scheduler;
- std::atomic<bool> fired{false};
-
- scheduler.Schedule(std::chrono::milliseconds(5000), [&] { fired.store(true);
});
-
- // Shutdown immediately — should not crash and task should not fire
- scheduler.Shutdown();
-
- std::this_thread::sleep_for(std::chrono::milliseconds(50));
- EXPECT_FALSE(fired.load());
+Result<std::shared_ptr<AuthSession>> RefreshingSession(
+ const std::shared_ptr<HttpClient>& client, Executor& executor) {
+ return AuthSession::MakeOAuth2(ExpiredInitialToken(),
"https://auth.example/token";,
+ "client", "secret", "catalog", true, {},
client,
+ &executor);
}
-// Verifies that Schedule after shutdown returns invalid handle (0)
-TEST(TokenRefreshSchedulerTest, ScheduleAfterShutdownIsNoop) {
- TokenRefreshScheduler scheduler;
- scheduler.Shutdown();
+} // namespace
- auto handle = scheduler.Schedule(std::chrono::milliseconds(10), [] {});
- EXPECT_EQ(0u, handle);
+TEST(OAuth2SessionTest, RefreshesAndReschedulesThroughPublicFactory) {
+ MockExecutor executor;
+ std::array<std::promise<ExecutorTask>, 2> queued;
+ int submitted = 0;
+ EXPECT_CALL(executor, Submit).Times(2).WillRepeatedly([&](ExecutorTask task)
-> Status {
+ queued.at(submitted++).set_value(std::move(task));
+ return {};
+ });
+ auto client = std::make_shared<RefreshHttpClient>();
+ client->response = [&]() -> Result<HttpResponse> {
+ return HttpResponse(
+ 200, client->requests == 1
+ ?
R"({"access_token":"first","token_type":"bearer","expires_in":1})"
+ : R"({"access_token":"second","token_type":"bearer"})");
+ };
+ ICEBERG_UNWRAP_OR_FAIL(auto session, RefreshingSession(client, executor));
+ for (auto& slot : queued) {
+ auto task = slot.get_future();
+ ASSERT_EQ(task.wait_for(std::chrono::seconds(5)),
std::future_status::ready);
+ std::move(task.get())();
+ }
+ ASSERT_THAT(session->Close(), IsOk());
+ ICEBERG_UNWRAP_OR_FAIL(auto request, session->Authenticate({}));
+ EXPECT_EQ(request.headers.at("Authorization"), "Bearer second");
+ EXPECT_EQ(client->requests, 2);
+}
+
+TEST(OAuth2SessionTest, AlreadyExpiredRefreshedTokenBacksOff) {
+ MockExecutor executor;
+ std::promise<ExecutorTask> queued;
+ auto task = queued.get_future();
+ EXPECT_CALL(executor, Submit).WillOnce([&](ExecutorTask work) -> Status {
+ queued.set_value(std::move(work));
+ return {};
+ });
+ auto client = std::make_shared<RefreshHttpClient>();
+ client->response = [&]() -> Result<HttpResponse> {
+ return HttpResponse(
+ 200, client->requests == 1
+ ?
R"({"access_token":"stale","token_type":"bearer","expires_in":0})"
+ : R"({"access_token":"fresh","token_type":"bearer"})");
+ };
+ ICEBERG_UNWRAP_OR_FAIL(auto session, RefreshingSession(client, executor));
+ ASSERT_EQ(task.wait_for(std::chrono::seconds(5)), std::future_status::ready);
+ test::FakeRetryEnvironment fake_retry;
+ ScopedRetryTestHooks hooks(fake_retry.hooks());
+ std::move(task.get())();
+ EXPECT_THAT(fake_retry.sleep_durations(),
+ ::testing::ElementsAre(std::chrono::milliseconds(100)));
+ EXPECT_EQ(client->requests, 2);
+ ICEBERG_UNWRAP_OR_FAIL(auto request, session->Authenticate({}));
+ EXPECT_EQ(request.headers.at("Authorization"), "Bearer fresh");
+ ASSERT_THAT(session->Close(), IsOk());
}
-// Verifies that cancelling an invalid handle does not crash
-TEST(TokenRefreshSchedulerTest, CancelInvalidHandleIsNoop) {
- TokenRefreshScheduler scheduler;
- // Should not crash
- scheduler.Cancel(0);
- scheduler.Cancel(999);
- scheduler.Shutdown();
+TEST(OAuth2SessionTest, ThrowingRefreshIsRetried) {
+ MockExecutor executor;
+ std::promise<ExecutorTask> queued;
+ auto task = queued.get_future();
+ EXPECT_CALL(executor, Submit).WillOnce([&](ExecutorTask work) -> Status {
+ queued.set_value(std::move(work));
+ return {};
+ });
+ auto client = std::make_shared<RefreshHttpClient>();
+ client->response = [&]() -> Result<HttpResponse> {
+ if (client->requests == 1) {
+ throw std::runtime_error("HTTP library failure");
+ }
+ return HttpResponse(200,
R"({"access_token":"recovered","token_type":"bearer"})");
+ };
+ ICEBERG_UNWRAP_OR_FAIL(auto session, RefreshingSession(client, executor));
+ ASSERT_EQ(task.wait_for(std::chrono::seconds(5)), std::future_status::ready);
+ test::FakeRetryEnvironment fake_retry;
+ ScopedRetryTestHooks hooks(fake_retry.hooks());
+ std::move(task.get())();
+ ICEBERG_UNWRAP_OR_FAIL(auto request, session->Authenticate({}));
+ EXPECT_EQ(request.headers.at("Authorization"), "Bearer recovered");
+ EXPECT_EQ(client->requests, 2);
+ ASSERT_THAT(session->Close(), IsOk());
}
-// ---- OAuth2Session tests ----
-
-TEST(OAuth2SessionTest, InitialTokenIsUsed) {
- auto client = std::make_shared<HttpClient>();
- OAuthTokenResponse token_response;
- token_response.access_token = "initial-token-123";
- token_response.token_type = "bearer";
- token_response.expires_in_secs = 3600;
-
- // Create session (refresh will fail since there's no real server, but
- // initial token should work)
- auto session_result =
- AuthSession::MakeOAuth2(token_response, "http://localhost/oauth/tokens";,
- "client_id", "client_secret", "catalog", true,
{}, client);
- ASSERT_THAT(session_result, IsOk());
- auto session = session_result.value();
-
- auto auth_result = session->Authenticate({});
- ASSERT_THAT(auth_result, IsOk());
- EXPECT_EQ(auth_result.value().headers.at("Authorization"), "Bearer
initial-token-123");
+TEST(OAuth2SessionTest, CloseSkipsAlreadyQueuedRefresh) {
+ MockExecutor executor;
+ std::promise<ExecutorTask> queued;
+ auto task = queued.get_future();
+ EXPECT_CALL(executor, Submit).WillOnce([&](ExecutorTask work) -> Status {
+ queued.set_value(std::move(work));
+ return {};
+ });
+ auto client = std::make_shared<RefreshHttpClient>();
+ ICEBERG_UNWRAP_OR_FAIL(auto session, RefreshingSession(client, executor));
+ ASSERT_EQ(task.wait_for(std::chrono::seconds(5)), std::future_status::ready);
+ ASSERT_THAT(session->Close(), IsOk());
+ std::move(task.get())();
+ EXPECT_EQ(client->requests, 0);
+}
+TEST(OAuth2SessionTest, ConcurrentCloseCallsWaitForCurrentHttpRequest) {
+ MockExecutor executor;
+ std::promise<ExecutorTask> queued;
+ auto task = queued.get_future();
+ EXPECT_CALL(executor, Submit).WillOnce([&](ExecutorTask work) -> Status {
+ queued.set_value(std::move(work));
+ return {};
+ });
+ auto client = std::make_shared<RefreshHttpClient>();
+ std::promise<void> entered;
+ std::promise<void> release;
+ auto entered_future = entered.get_future();
+ auto release_future = release.get_future().share();
+ client->response = [&]() -> Result<HttpResponse> {
+ entered.set_value();
+ release_future.wait();
+ // Closing the session must stop retries of this failed request.
+ return IOError("refresh failed during Close");
+ };
+ ICEBERG_UNWRAP_OR_FAIL(auto session, RefreshingSession(client, executor));
+ ASSERT_EQ(task.wait_for(std::chrono::seconds(5)), std::future_status::ready);
+ auto refresh = std::async(std::launch::async, [&] { std::move(task.get())();
});
+ auto entered_status = entered_future.wait_for(std::chrono::seconds(5));
+ if (entered_status != std::future_status::ready) {
+ release.set_value();
+ }
+ ASSERT_EQ(entered_status, std::future_status::ready);
+ auto first = std::async(std::launch::async, [&] { return session->Close();
});
+ auto second = std::async(std::launch::async, [&] { return session->Close();
});
+ EXPECT_EQ(first.wait_for(std::chrono::milliseconds(30)),
std::future_status::timeout);
+ EXPECT_EQ(second.wait_for(std::chrono::milliseconds(30)),
std::future_status::timeout);
+ release.set_value();
+ refresh.get();
+ ASSERT_THAT(first.get(), IsOk());
+ ASSERT_THAT(second.get(), IsOk());
+ EXPECT_EQ(client->requests, 1);
+}
+
+TEST(OAuth2SessionTest, StopsAfterHttpRetryBudget) {
+ MockExecutor executor;
+ std::promise<ExecutorTask> queued;
+ auto task = queued.get_future();
+ EXPECT_CALL(executor, Submit).WillOnce([&](ExecutorTask work) -> Status {
+ queued.set_value(std::move(work));
+ return {};
+ });
+ auto client = std::make_shared<RefreshHttpClient>();
+ ICEBERG_UNWRAP_OR_FAIL(auto session, RefreshingSession(client, executor));
+ ASSERT_EQ(task.wait_for(std::chrono::seconds(5)), std::future_status::ready);
+ test::FakeRetryEnvironment fake_retry;
+ ScopedRetryTestHooks hooks(fake_retry.hooks());
+ std::move(task.get())();
+ EXPECT_EQ(client->requests, 6);
+ EXPECT_THAT(fake_retry.observed_base_delays_ms(),
+ ::testing::ElementsAre(100, 200, 400, 800, 1600));
+ ICEBERG_UNWRAP_OR_FAIL(auto request, session->Authenticate({}));
+ EXPECT_EQ(request.headers.at("Authorization"), "Bearer initial");
ASSERT_THAT(session->Close(), IsOk());
}
diff --git a/src/iceberg/test/executor.h b/src/iceberg/test/executor.h
index a3a801b9..236c2d0e 100644
--- a/src/iceberg/test/executor.h
+++ b/src/iceberg/test/executor.h
@@ -20,15 +20,40 @@
#pragma once
#include <atomic>
+#include <chrono>
+#include <future>
#include <thread>
#include <utility>
#include <vector>
+#include <gmock/gmock.h>
+
#include "iceberg/result.h"
#include "iceberg/util/executor.h"
+#include "iceberg/util/scheduled_executor_internal.h"
namespace iceberg::test {
+class MockExecutor : public Executor {
+ public:
+ MOCK_METHOD(Status, Submit, (ExecutorTask task), (override));
+};
+
+/// Runs each task on the submitting thread.
+class InlineExecutor final : public Executor {
+ public:
+ Status Submit(ExecutorTask task) override {
+ submit_count_.fetch_add(1, std::memory_order_relaxed);
+ std::move(task)();
+ return {};
+ }
+
+ int submit_count() const { return
submit_count_.load(std::memory_order_relaxed); }
+
+ private:
+ std::atomic<int> submit_count_{0};
+};
+
class ThreadExecutor final : public Executor {
public:
explicit ThreadExecutor(Status submit_status = {})
@@ -59,4 +84,15 @@ class ThreadExecutor final : public Executor {
std::vector<std::thread> threads_;
};
+/// Waits until the process-wide ScheduledExecutor timer has dispatched every
task due
+/// within `delay` from now. Returns false on timeout.
+inline bool WaitForSharedTimer(std::chrono::milliseconds delay = {}) {
+ InlineExecutor executor;
+ std::promise<void> ran;
+ auto done = ran.get_future();
+ ::iceberg::internal::ScheduledExecutor timer(executor);
+ timer.Schedule(delay, ExecutorTask([&] { ran.set_value(); }));
+ return done.wait_for(std::chrono::seconds(5)) == std::future_status::ready;
+}
+
} // namespace iceberg::test
diff --git a/src/iceberg/test/rest_catalog_unit_test.cc
b/src/iceberg/test/rest_catalog_unit_test.cc
index 86865a66..75bea714 100644
--- a/src/iceberg/test/rest_catalog_unit_test.cc
+++ b/src/iceberg/test/rest_catalog_unit_test.cc
@@ -29,6 +29,7 @@
#include <nlohmann/json.hpp>
#include "iceberg/catalog.h"
+#include "iceberg/catalog/rest/auth/auth_session.h"
#include "iceberg/catalog/rest/catalog_properties.h"
#include "iceberg/catalog/rest/endpoint.h"
#include "iceberg/catalog/rest/error_handlers.h"
@@ -47,6 +48,7 @@
#include "iceberg/table_identifier.h"
#include "iceberg/table_requirement.h"
#include "iceberg/table_update.h"
+#include "iceberg/test/executor.h"
#include "iceberg/test/matchers.h"
namespace iceberg::rest {
@@ -76,6 +78,13 @@ class MockHttpClient : public HttpClient {
const ErrorHandler& error_handler, auth::AuthSession& session),
(override));
+ MOCK_METHOD(Result<HttpResponse>, PostForm,
+ (const std::string& path,
+ (const std::unordered_map<std::string, std::string>&)form_data,
+ (const std::unordered_map<std::string, std::string>&)headers,
+ const ErrorHandler& error_handler, auth::AuthSession& session),
+ (override));
+
MOCK_METHOD(Result<HttpResponse>, Delete,
(const std::string& path,
(const std::unordered_map<std::string, std::string>&)params,
@@ -207,7 +216,7 @@ class RestCatalogLoadTableTest : public ::testing::Test {
EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/v1/config"), _, _, _,
_))
.WillOnce(Return(HttpResponse(200, ToJson(server_config).dump())));
auto config = RestCatalogProperties::FromMap(std::move(props));
- return RestCatalog::Make(config, /*metrics_executor=*/nullptr,
+ return RestCatalog::Make(config,
[client = mock_client_](const auto&) { return
client; });
}
@@ -300,8 +309,7 @@ TEST_F(RestCatalogLoadTableTest,
InitializationPropagatesClientFactoryErrors) {
}
int calls = 0;
auto result = RestCatalog::Make(
- config, /*metrics_executor=*/nullptr,
- [&](const RestCatalogProperties&) ->
Result<std::shared_ptr<HttpClient>> {
+ config, [&](const RestCatalogProperties&) ->
Result<std::shared_ptr<HttpClient>> {
if (++calls == fail_on_call) {
return IOError("HTTP client setup failed");
}
@@ -314,6 +322,54 @@ TEST_F(RestCatalogLoadTableTest,
InitializationPropagatesClientFactoryErrors) {
}
}
+TEST_F(RestCatalogLoadTableTest, OAuthRefreshRunsOnAuthExecutorWhenEnabled) {
+ for (bool refresh : {false, true}) {
+ SCOPED_TRACE(::testing::Message() << "token-refresh-enabled=" << refresh);
+ test::MockExecutor executor;
+ auto& submits = EXPECT_CALL(executor, Submit).Times(refresh ? 1 : 0);
+ if (refresh) {
+ submits.WillOnce([](ExecutorTask task) -> Status {
+ std::move(task)();
+ return {};
+ });
+ }
+ auto config = RestCatalogProperties::FromMap(
+ {{"uri", "http://test-server"},
+ {"io-impl", "rest-catalog-unit-test-noop"},
+ {"credential", "client:secret"},
+ {"token-refresh-enabled", refresh ? "true" : "false"}});
+ auto& tokens =
+ EXPECT_CALL(*mock_client_, PostForm(_, _, _, _, _)).Times(refresh ? 2
: 1);
+ tokens.WillOnce(Return(HttpResponse(
+ 200,
R"({"access_token":"initial","token_type":"bearer","expires_in":0})")));
+ if (refresh) {
+ tokens.WillOnce(Return(
+ HttpResponse(200,
R"({"access_token":"updated","token_type":"bearer"})")));
+ }
+ EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/v1/config"), _, _, _,
_))
+ .WillOnce(Return(HttpResponse(200, ToJson(CatalogConfig{}).dump())));
+ ICEBERG_UNWRAP_OR_FAIL(
+ auto root,
+ RestCatalog::Make(
+ config, [client = mock_client_](const auto&) { return client; },
&executor));
+
+ // Let any immediately due catalog refresh dispatch first.
+ ASSERT_TRUE(test::WaitForSharedTimer());
+
+ EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/namespaces"), _, _,
_, _))
+ .WillOnce([&](const auto&, const auto&, const auto&, const auto&,
+ auth::AuthSession& session) -> Result<HttpResponse> {
+ ICEBERG_ASSIGN_OR_RAISE(auto request, session.Authenticate({}));
+ EXPECT_EQ(request.headers.at("Authorization"),
+ refresh ? "Bearer updated" : "Bearer initial");
+ return HttpResponse(200, R"({"namespaces":[]})");
+ });
+ ICEBERG_UNWRAP_OR_FAIL(auto catalog, root->AsCatalog());
+ ASSERT_THAT(catalog->ListNamespaces({}), IsOk());
+
ASSERT_TRUE(::testing::Mock::VerifyAndClearExpectations(mock_client_.get()));
+ }
+}
+
TEST_F(RestCatalogLoadTableTest, InitializationUsesMergedServerConfig) {
CatalogConfig server_config{
.defaults = {{"header.X-Default", "default"}, {"header.X-Shared",
"default"}},
@@ -344,12 +400,11 @@ TEST_F(RestCatalogLoadTableTest,
InitializationUsesMergedServerConfig) {
{"scan-planning-mode", "client"},
{"header.X-Shared", "client"}});
std::vector<RestCatalogProperties> client_configs;
- ICEBERG_UNWRAP_OR_FAIL(auto root,
- RestCatalog::Make(config,
/*metrics_executor=*/nullptr,
- [&](const RestCatalogProperties&
properties) {
-
client_configs.push_back(properties);
- return mock_client_;
- }));
+ ICEBERG_UNWRAP_OR_FAIL(
+ auto root, RestCatalog::Make(config, [&](const RestCatalogProperties&
properties) {
+ client_configs.push_back(properties);
+ return mock_client_;
+ }));
ASSERT_EQ(client_configs.size(), 2);
EXPECT_EQ(client_configs[0].ExtractHeaders().at("X-Shared"), "client");
EXPECT_FALSE(client_configs[0].ExtractHeaders().contains("X-Default"));
diff --git a/src/iceberg/test/rest_metrics_reporter_test.cc
b/src/iceberg/test/rest_metrics_reporter_test.cc
index 29bc3ab0..47d06f34 100644
--- a/src/iceberg/test/rest_metrics_reporter_test.cc
+++ b/src/iceberg/test/rest_metrics_reporter_test.cc
@@ -33,6 +33,7 @@
#include "iceberg/metrics/metrics_reporter.h"
#include "iceberg/metrics/scan_report.h"
#include "iceberg/result.h"
+#include "iceberg/test/executor.h"
#include "iceberg/test/matchers.h"
#include "iceberg/util/executor.h"
@@ -40,31 +41,13 @@ namespace iceberg::rest {
namespace {
-class ManualExecutor final : public Executor {
- public:
- Status Submit(ExecutorTask task) override {
- task_ = std::move(task);
+// A Submit action that keeps the task so the test decides when it runs.
+auto CaptureTask(std::optional<ExecutorTask>& slot) {
+ return [&slot](ExecutorTask task) -> Status {
+ slot.emplace(std::move(task));
return {};
- }
-
- bool HasTask() const { return task_.has_value(); }
-
- void Run() {
- auto task = std::move(*task_);
- task_.reset();
- std::move(task)();
- }
-
- private:
- std::optional<ExecutorTask> task_;
-};
-
-class RejectingExecutor final : public Executor {
- public:
- Status Submit(ExecutorTask /*task*/) override {
- return IOError("executor rejected task");
- }
-};
+ };
+}
} // namespace
@@ -121,7 +104,9 @@ class RestMetricsReporterPayloadTest
public ::testing::WithParamInterface<ReportTestCase> {};
TEST_F(RestMetricsReporterTest, SuppressesHttpErrors) {
- ManualExecutor executor;
+ test::MockExecutor executor;
+ std::optional<ExecutorTask> task;
+ EXPECT_CALL(executor, Submit).WillOnce(CaptureTask(task));
RestMetricsReporter reporter(
"http://localhost:0/v1/ns/tables/tbl/metrics";, session_, &executor,
[](const std::string&, const std::string&, auth::AuthSession&) {
@@ -129,13 +114,15 @@ TEST_F(RestMetricsReporterTest, SuppressesHttpErrors) {
});
MetricsReport report = MakeScanReport();
EXPECT_THAT(reporter.Report(report), IsOk());
- ASSERT_TRUE(executor.HasTask());
- EXPECT_NO_THROW(executor.Run());
+ ASSERT_TRUE(task.has_value());
+ EXPECT_NO_THROW(std::move(*task)());
}
TEST_P(RestMetricsReporterPayloadTest,
PostsSerializedPayloadToConfiguredEndpoint) {
const auto& test_case = GetParam();
- ManualExecutor executor;
+ test::MockExecutor executor;
+ std::optional<ExecutorTask> task;
+ EXPECT_CALL(executor, Submit).WillOnce(CaptureTask(task));
const std::string endpoint = "http://mock-host/v1/ns/tables/tbl/metrics";;
std::string captured_path;
@@ -149,8 +136,8 @@ TEST_P(RestMetricsReporterPayloadTest,
PostsSerializedPayloadToConfiguredEndpoin
EXPECT_THAT(reporter.Report(test_case.report), IsOk());
EXPECT_TRUE(captured_path.empty());
- ASSERT_TRUE(executor.HasTask());
- executor.Run();
+ ASSERT_TRUE(task.has_value());
+ std::move (*task)();
EXPECT_EQ(captured_path, endpoint);
auto json = nlohmann::json::parse(captured_body);
@@ -172,7 +159,7 @@ TEST_F(RestMetricsReporterTest,
PostsSynchronouslyWithoutExecutor) {
}
TEST_F(RestMetricsReporterTest, SuppressesExecutorRejection) {
- RejectingExecutor executor;
+ test::ThreadExecutor executor(IOError("executor rejected task"));
RestMetricsReporter reporter(
"http://mock-host/v1/ns/tables/tbl/metrics";, session_, &executor,
[](const std::string&, const std::string&, auth::AuthSession&) {});
diff --git a/src/iceberg/test/scheduled_executor_test.cc
b/src/iceberg/test/scheduled_executor_test.cc
new file mode 100644
index 00000000..64d9f7e8
--- /dev/null
+++ b/src/iceberg/test/scheduled_executor_test.cc
@@ -0,0 +1,188 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include <atomic>
+#include <chrono>
+#include <future>
+#include <memory>
+#include <stdexcept>
+#include <thread>
+
+#include <gtest/gtest.h>
+
+#include "iceberg/test/executor.h"
+#include "iceberg/test/matchers.h"
+#include "iceberg/util/scheduled_executor_internal.h"
+
+namespace iceberg::internal {
+
+// Inline execution lets the tests observe the timer's submitting thread.
+using test::InlineExecutor;
+
+TEST(ScheduledExecutorTest, SharesTimerAndHonorsDeadlinesAcrossExecutors) {
+ InlineExecutor first_executor;
+ InlineExecutor second_executor;
+ std::promise<void> entered;
+ std::promise<void> release;
+ auto entered_done = entered.get_future();
+ auto release_done = release.get_future().share();
+ std::promise<std::thread::id> earlier;
+ std::promise<std::thread::id> later;
+ auto earlier_done = earlier.get_future();
+ auto later_done = later.get_future();
+ std::atomic<int> order{0};
+ ScheduledExecutor first(first_executor);
+ ScheduledExecutor second(second_executor);
+ // Install both deadlines while the shared timer is still in the first
callback.
+ first.Schedule(std::chrono::milliseconds::zero(), ExecutorTask([&] {
+ entered.set_value();
+ release_done.wait();
+ }));
+ auto entered_status = entered_done.wait_for(std::chrono::seconds(5));
+ if (entered_status != std::future_status::ready) {
+ release.set_value();
+ }
+ ASSERT_EQ(entered_status, std::future_status::ready);
+ const auto now = std::chrono::steady_clock::now();
+ first.Schedule(std::chrono::milliseconds(150), ExecutorTask([&] {
+ EXPECT_GE(std::chrono::steady_clock::now(),
+ now + std::chrono::milliseconds(150));
+ EXPECT_EQ(order++, 1);
+ later.set_value(std::this_thread::get_id());
+ }));
+ second.Schedule(std::chrono::milliseconds(10), ExecutorTask([&] {
+ EXPECT_GE(std::chrono::steady_clock::now(),
+ now + std::chrono::milliseconds(10));
+ EXPECT_EQ(order++, 0);
+ earlier.set_value(std::this_thread::get_id());
+ }));
+ release.set_value();
+ ASSERT_EQ(earlier_done.wait_for(std::chrono::seconds(5)),
std::future_status::ready);
+ ASSERT_EQ(later_done.wait_for(std::chrono::seconds(5)),
std::future_status::ready);
+ auto timer_thread = earlier_done.get();
+ EXPECT_NE(timer_thread, std::this_thread::get_id());
+ EXPECT_EQ(timer_thread, later_done.get());
+ EXPECT_EQ(first_executor.submit_count(), 2);
+ EXPECT_EQ(second_executor.submit_count(), 1);
+}
+
+TEST(ScheduledExecutorTest, ExceptionsDoNotEscapeTimerThread) {
+ // Simulates a user executor that throws, e.g. when it cannot start a thread.
+ class ThrowingExecutor : public Executor {
+ public:
+ Status Submit(ExecutorTask) override { throw std::runtime_error("no
threads"); }
+ } executor;
+ std::promise<Status> submitted;
+ auto done = submitted.get_future();
+ ScheduledExecutor scheduler(executor);
+ scheduler.Schedule(std::chrono::milliseconds::zero(),
+ ExecutorTask([] { ADD_FAILURE(); }),
+ TaskSubmissionCallback([&](Status status) {
+ submitted.set_value(std::move(status));
+ throw std::runtime_error("callback failed");
+ }));
+ ASSERT_EQ(done.wait_for(std::chrono::seconds(5)), std::future_status::ready);
+ auto status = done.get();
+ EXPECT_THAT(status, IsError(ErrorKind::kUnknownError));
+ EXPECT_THAT(status, HasErrorMessage("no threads"));
+ // The shared timer survived both exceptions.
+ EXPECT_TRUE(test::WaitForSharedTimer());
+}
+
+TEST(ScheduledExecutorTest, ShutdownCancelsOnlyItsOwnTasks) {
+ test::ThreadExecutor first_executor(ServiceUnavailable("rejected"));
+ InlineExecutor second_executor;
+ std::promise<void> dispatched;
+ auto done = dispatched.get_future();
+ std::atomic<int> notifications{0};
+ ScheduledExecutor first(first_executor);
+ ScheduledExecutor second(second_executor);
+ auto capture = std::make_shared<int>(42);
+ std::weak_ptr<int> weak_capture = capture;
+ first.Schedule(
+ std::chrono::milliseconds(3'600'000),
+ ExecutorTask([capture = std::move(capture)] { ADD_FAILURE() << *capture;
}),
+ TaskSubmissionCallback([&](Status) { ++notifications; }));
+ first.Shutdown();
+ first.Shutdown();
+ EXPECT_TRUE(weak_capture.expired());
+ // Scheduling after Shutdown is refused and drops the task without a
callback.
+ auto late_capture = std::make_shared<int>(7);
+ std::weak_ptr<int> weak_late_capture = late_capture;
+ first.Schedule(
+ std::chrono::milliseconds::zero(),
+ ExecutorTask([late_capture = std::move(late_capture)] { ADD_FAILURE();
}),
+ TaskSubmissionCallback([&](Status) { ++notifications; }));
+ EXPECT_TRUE(weak_late_capture.expired());
+
+ second.Schedule(std::chrono::milliseconds::zero(),
+ ExecutorTask([&] { dispatched.set_value(); }));
+ ASSERT_EQ(done.wait_for(std::chrono::seconds(5)), std::future_status::ready);
+ EXPECT_EQ(first_executor.submit_count(), 0);
+ EXPECT_EQ(notifications, 0);
+ EXPECT_EQ(second_executor.submit_count(), 1);
+}
+
+TEST(ScheduledExecutorTest, SubmissionFailureCanDestroySchedulerOnTimerThread)
{
+ test::ThreadExecutor executor(ServiceUnavailable("rejected"));
+ struct Context {
+ std::unique_ptr<ScheduledExecutor> scheduler;
+ std::promise<void> rejected;
+ };
+ auto context = std::make_shared<Context>();
+ context->scheduler = std::make_unique<ScheduledExecutor>(executor);
+ auto rejected = context->rejected.get_future();
+ context->scheduler->Schedule(std::chrono::milliseconds::zero(),
+ ExecutorTask([] { ADD_FAILURE(); }),
+ TaskSubmissionCallback([context](Status status)
{
+ EXPECT_FALSE(status.has_value());
+ context->scheduler.reset();
+ context->rejected.set_value();
+ }));
+ ASSERT_EQ(rejected.wait_for(std::chrono::seconds(5)),
std::future_status::ready);
+ EXPECT_EQ(context->scheduler, nullptr);
+ EXPECT_EQ(executor.submit_count(), 1);
+}
+
+TEST(ScheduledExecutorTest, ShutdownWaitsForActiveSubmit) {
+ InlineExecutor executor;
+ std::promise<void> release;
+ std::promise<void> entered;
+ auto release_done = release.get_future().share();
+ auto entered_done = entered.get_future();
+ ScheduledExecutor scheduler(executor);
+ scheduler.Schedule(std::chrono::milliseconds::zero(), ExecutorTask([&] {
+ entered.set_value();
+ release_done.wait();
+ }));
+ auto entered_status = entered_done.wait_for(std::chrono::seconds(5));
+ if (entered_status != std::future_status::ready) {
+ release.set_value();
+ }
+ ASSERT_EQ(entered_status, std::future_status::ready);
+
+ auto shutdown = std::async(std::launch::async, [&] { scheduler.Shutdown();
});
+ EXPECT_EQ(shutdown.wait_for(std::chrono::milliseconds(50)),
+ std::future_status::timeout);
+ release.set_value();
+ ASSERT_EQ(shutdown.wait_for(std::chrono::seconds(5)),
std::future_status::ready);
+ shutdown.get();
+}
+
+} // namespace iceberg::internal
diff --git a/src/iceberg/util/executor.h b/src/iceberg/util/executor.h
index 625c833f..4b81ca44 100644
--- a/src/iceberg/util/executor.h
+++ b/src/iceberg/util/executor.h
@@ -48,6 +48,8 @@ class ICEBERG_EXPORT Executor {
virtual ~Executor() = default;
/// \brief Schedule a task for execution.
+ /// \note Success means the task was accepted, not that it has finished.
+ /// On an error or exception, the task must not run, including after Submit
returns.
virtual Status Submit(ExecutorTask task) = 0;
};
diff --git a/src/iceberg/util/retry_util.h b/src/iceberg/util/retry_util.h
index 8d2ed2a1..afd35004 100644
--- a/src/iceberg/util/retry_util.h
+++ b/src/iceberg/util/retry_util.h
@@ -94,8 +94,8 @@ template <RetryPolicyMode Mode, ErrorKind... Kinds>
struct RetryPolicy {
static_assert(Mode != RetryPolicyMode::kNoRetry || sizeof...(Kinds) == 0,
"NoRetry must not include error kinds");
- static_assert(Mode == RetryPolicyMode::kNoRetry || sizeof...(Kinds) > 0,
- "RetryPolicy must include at least one error kind");
+ static_assert(Mode != RetryPolicyMode::kOnlyRetryOn || sizeof...(Kinds) > 0,
+ "OnlyRetryOn must include at least one error kind");
static constexpr RetryPolicyMode kMode = Mode;
static constexpr bool kEnabled = Mode != RetryPolicyMode::kNoRetry;
@@ -116,6 +116,7 @@ using NoRetry = RetryPolicy<RetryPolicyMode::kNoRetry>;
template <ErrorKind... Kinds>
using OnlyRetryOn = RetryPolicy<RetryPolicyMode::kOnlyRetryOn, Kinds...>;
+/// An empty stop list retries all errors.
template <ErrorKind... Kinds>
using StopRetryOn = RetryPolicy<RetryPolicyMode::kStopRetryOn, Kinds...>;
diff --git a/src/iceberg/util/scheduled_executor.cc
b/src/iceberg/util/scheduled_executor.cc
new file mode 100644
index 00000000..0c143920
--- /dev/null
+++ b/src/iceberg/util/scheduled_executor.cc
@@ -0,0 +1,213 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include <condition_variable>
+#include <exception>
+#include <iterator>
+#include <map>
+#include <mutex>
+#include <thread>
+#include <utility>
+
+#include "iceberg/util/scheduled_executor_internal.h"
+
+namespace iceberg::internal {
+
+// Kept alive so static destructors can still schedule tasks.
+class SharedTimer {
+ public:
+ static SharedTimer& Instance() {
+ static auto* timer = new SharedTimer();
+ return *timer;
+ }
+
+ void Schedule(std::weak_ptr<ScheduledExecutorState> owner,
+ std::chrono::steady_clock::time_point deadline, ExecutorTask
task,
+ std::optional<TaskSubmissionCallback> on_submission) {
+ std::lock_guard lock(mutex_);
+ // Equal deadlines keep insertion order.
+ auto it = tasks_.emplace(deadline, Task{.owner = std::move(owner),
+ .task = std::move(task),
+ .on_submission =
std::move(on_submission)});
+ // Wake the timer only if its next deadline changed.
+ if (it == tasks_.begin()) {
+ cv_.notify_one();
+ }
+ }
+
+ void CancelAll(const std::weak_ptr<ScheduledExecutorState>& owner) {
+ Queue cancelled; // Release task captures outside the queue lock.
+ {
+ std::lock_guard lock(mutex_);
+ for (auto it = tasks_.begin(); it != tasks_.end();) {
+ auto next = std::next(it);
+ if (SameOwner(it->second.owner, owner)) {
+ cancelled.insert(tasks_.extract(it));
+ }
+ it = next;
+ }
+ }
+ if (!cancelled.empty()) {
+ cv_.notify_one();
+ }
+ }
+
+ bool IsTimerThread() const {
+ return timer_thread_.get_id() == std::this_thread::get_id();
+ }
+
+ private:
+ SharedTimer() : timer_thread_([this] { Run(); }) {}
+
+ struct Task {
+ std::weak_ptr<ScheduledExecutorState> owner;
+ ExecutorTask task;
+ std::optional<TaskSubmissionCallback> on_submission;
+ };
+
+ using Queue = std::multimap<std::chrono::steady_clock::time_point, Task>;
+
+ static bool SameOwner(const std::weak_ptr<ScheduledExecutorState>& a,
+ const std::weak_ptr<ScheduledExecutorState>& b) {
+ return !a.owner_before(b) && !b.owner_before(a);
+ }
+
+ void Run();
+
+ std::mutex mutex_;
+ std::condition_variable cv_;
+ Queue tasks_;
+ std::thread timer_thread_;
+};
+
+class ScheduledExecutorState
+ : public std::enable_shared_from_this<ScheduledExecutorState> {
+ public:
+ explicit ScheduledExecutorState(Executor& executor)
+ : executor_(executor), timer_(SharedTimer::Instance()) {}
+
+ ScheduledExecutorState(const ScheduledExecutorState&) = delete;
+ ScheduledExecutorState& operator=(const ScheduledExecutorState&) = delete;
+
+ void Schedule(std::chrono::steady_clock::time_point deadline, ExecutorTask
task,
+ std::optional<TaskSubmissionCallback> on_submission) {
+ std::lock_guard lock(mutex_);
+ if (shutdown_) {
+ return;
+ }
+ timer_.Schedule(weak_from_this(), deadline, std::move(task),
+ std::move(on_submission));
+ }
+
+ void Shutdown() {
+ {
+ std::lock_guard lock(mutex_);
+ shutdown_ = true;
+ }
+ timer_.CancelAll(weak_from_this());
+ if (!timer_.IsTimerThread()) {
+ std::unique_lock lock(mutex_);
+ cv_.wait(lock, [this] { return active_dispatches_ == 0; });
+ }
+ }
+
+ private:
+ friend class SharedTimer;
+
+ void Dispatch(ExecutorTask task, std::optional<TaskSubmissionCallback>
on_submission) {
+ {
+ std::lock_guard lock(mutex_);
+ // Shutdown may race with popping a task.
+ if (shutdown_) {
+ return;
+ }
+ ++active_dispatches_;
+ }
+ struct DispatchGuard {
+ ScheduledExecutorState& state;
+ ~DispatchGuard() {
+ std::lock_guard lock(state.mutex_);
+ if (--state.active_dispatches_ == 0) {
+ state.cv_.notify_all();
+ }
+ }
+ } guard{*this};
+
+ Status status;
+ try {
+ status = executor_.Submit(std::move(task));
+ } catch (const std::exception& e) {
+ status = UnknownError("Executor::Submit threw: {}", e.what());
+ } catch (...) {
+ status = UnknownError("Executor::Submit threw a non-standard exception");
+ }
+ if (!on_submission.has_value()) {
+ return;
+ }
+ try {
+ std::move (*on_submission)(std::move(status));
+ } catch (...) {
+ // Keep the shared timer running.
+ }
+ }
+
+ Executor& executor_;
+ SharedTimer& timer_;
+ std::mutex mutex_;
+ std::condition_variable cv_;
+ size_t active_dispatches_ = 0;
+ bool shutdown_ = false;
+};
+
+void SharedTimer::Run() {
+ while (true) {
+ Queue::node_type node;
+ {
+ std::unique_lock lock(mutex_);
+ cv_.wait(lock, [this] { return !tasks_.empty(); });
+
+ auto deadline = tasks_.begin()->first;
+ if (cv_.wait_until(lock, deadline, [this, deadline] {
+ return tasks_.empty() || tasks_.begin()->first != deadline;
+ })) {
+ continue;
+ }
+ node = tasks_.extract(tasks_.begin());
+ }
+ auto& task = node.mapped();
+ if (auto owner = task.owner.lock()) {
+ owner->Dispatch(std::move(task.task), std::move(task.on_submission));
+ }
+ }
+}
+
+ScheduledExecutor::ScheduledExecutor(Executor& executor)
+ : state_(std::make_shared<ScheduledExecutorState>(executor)) {}
+
+ScheduledExecutor::~ScheduledExecutor() { Shutdown(); }
+
+void ScheduledExecutor::Schedule(std::chrono::milliseconds delay, ExecutorTask
task,
+ std::optional<TaskSubmissionCallback>
on_submission) {
+ state_->Schedule(std::chrono::steady_clock::now() + delay, std::move(task),
+ std::move(on_submission));
+}
+
+void ScheduledExecutor::Shutdown() { state_->Shutdown(); }
+
+} // namespace iceberg::internal
diff --git a/src/iceberg/util/scheduled_executor_internal.h
b/src/iceberg/util/scheduled_executor_internal.h
new file mode 100644
index 00000000..15521ef3
--- /dev/null
+++ b/src/iceberg/util/scheduled_executor_internal.h
@@ -0,0 +1,73 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#pragma once
+
+#include <chrono>
+#include <memory>
+#include <optional>
+
+#include "iceberg/iceberg_export.h"
+#include "iceberg/result.h"
+#include "iceberg/util/executor.h"
+
+namespace iceberg::internal {
+
+class ScheduledExecutorState;
+
+/// Called once after Executor::Submit returns or throws, to report whether
the task
+/// was submitted. This does not report the task's result or wait for it to
finish.
+///
+/// Submit exceptions are reported as an error Status. Under Executor's
contract,
+/// an error means the task will not run, so the caller can safely retry.
+/// Runs on the shared timer thread and must not block. Cancelled tasks skip
the callback.
+using TaskSubmissionCallback = FnOnce<void(Status)>;
+
+/// Adds delayed scheduling to an Executor. One timer thread shared by all
instances
+/// waits for deadlines, then submits each due task to that instance's
executor.
+///
+/// The executor is not owned and must outlive this object and its submitted
tasks.
+/// Submit must queue work promptly: running tasks inline would block the
timer and
+/// delay scheduled work for every instance. Destruction calls Shutdown().
+class ICEBERG_EXPORT ScheduledExecutor {
+ public:
+ explicit ScheduledExecutor(Executor& executor);
+ ~ScheduledExecutor();
+
+ ScheduledExecutor(const ScheduledExecutor&) = delete;
+ ScheduledExecutor& operator=(const ScheduledExecutor&) = delete;
+
+ /// Queues `task` for submission once `delay` has passed. The executor
decides when
+ /// it actually runs. Equal deadlines are submitted in scheduling order.
+ /// After Shutdown, drops the task without calling on_submission.
+ void Schedule(std::chrono::milliseconds delay, ExecutorTask task,
+ std::optional<TaskSubmissionCallback> on_submission =
std::nullopt);
+
+ /// Stops accepting tasks and cancels this instance's tasks still waiting on
the timer.
+ /// Already submitted tasks are not cancelled or waited for.
+ ///
+ /// Waits for active Submit calls and their callbacks to finish, unless
called on the
+ /// timer thread itself. Safe to call more than once; other instances are
unaffected.
+ void Shutdown();
+
+ private:
+ std::shared_ptr<ScheduledExecutorState> state_;
+};
+
+} // namespace iceberg::internal
