From 7a63bead2a9edd582ac4ca2be6f7e015035294f8 Mon Sep 17 00:00:00 2001 From: Gang Wu Date: Tue, 20 Oct 2026 22:21:27 +0800 Subject: [PATCH] refactor: refresh OAuth2 tokens on a caller-provided executor 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 | 164 ++++++----- .../rest/auth/sigv4_auth_manager_internal.h | 2 + .../catalog/rest/auth/sigv4_manager.cc | 4 + .../rest/auth/token_refresh_scheduler.cc | 125 --------- .../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 ++++- .../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 ++++++++++++++ .../util/scheduled_executor_internal.h | 73 +++++ 25 files changed, 950 insertions(+), 513 deletions(-) delete mode 100644 src/iceberg/catalog/rest/auth/token_refresh_scheduler.cc delete mode 100644 src/iceberg/catalog/rest/auth/token_refresh_scheduler.h create mode 100644 src/iceberg/test/scheduled_executor_test.cc create mode 100644 src/iceberg/util/scheduled_executor.cc create mode 100644 src/iceberg/util/scheduled_executor_internal.h diff --git a/src/iceberg/CMakeLists.txt b/src/iceberg/CMakeLists.txt index 091acd3e78..11278ecbda 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 0736290bc5..a8146d468a 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 5facfa3eb5..df5d7304a8 100644 --- a/src/iceberg/catalog/rest/auth/auth_manager.cc +++ b/src/iceberg/catalog/rest/auth/auth_manager.cc @@ -140,6 +140,10 @@ Result> MakeBasicAuthManager( /// \brief OAuth2 authentication manager. class OAuth2Manager : public AuthManager { public: + void SetAuthRefreshExecutor(Executor* executor) override { + refresh_executor_ = executor; + } + Result> InitSession( std::shared_ptr init_client, const std::unordered_map& 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> MakeSession( const OAuthTokenResponse& token_response, const AuthProperties& config, - bool keep_refreshed) const { + bool keep_refreshed, + std::optional 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> MaybeCreateChildSession( @@ -319,6 +317,7 @@ class OAuth2Manager : public AuthManager { std::optional auth_response_; std::optional start_time_; std::shared_ptr refresh_client_; + Executor* refresh_executor_ = nullptr; }; Result> MakeOAuth2Manager( diff --git a/src/iceberg/catalog/rest/auth/auth_manager.h b/src/iceberg/catalog/rest/auth/auth_manager.h index 6ba25c2840..d390d0beb3 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 60aa0a3b48..bd74b5bb60 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 #include -#include #include -#include -#include #include #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> AuthSession::MakeOAuth2( const std::string& client_id, const std::string& client_secret, const std::string& scope, bool keep_refreshed, const std::unordered_map& optional_oauth_params, - std::shared_ptr client) { + std::shared_ptr 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> internal::MakeOAuth2Session( @@ -79,14 +74,19 @@ Result> internal::MakeOAuth2Session( const std::string& scope, bool keep_refreshed, const std::unordered_map& optional_oauth_params, std::shared_ptr client, - std::optional token_request_started_at) { + std::optional 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 bdc77ebc4c..7e330706c2 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 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> 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& optional_oauth_params, - std::shared_ptr client); + std::shared_ptr 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 b8db3836e0..537f2a9145 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 -#include #include #include +#include #include #include #include @@ -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 optional_oauth_params; - bool keep_refreshed; + /// Borrowed executor for background refresh. Null disables refresh. + Executor* refresh_executor = nullptr; }; static Result> Make( @@ -111,17 +114,23 @@ class OAuth2Session final : public AuthSession, std::shared_ptr 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,89 @@ class OAuth2Session final : public AuthSession, const OAuthTokenResponse& token_response, std::optional 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 Refresh() { RefreshAttemptGuard guard(*this); - if (closed_.load()) return; + if (!guard.IsActive()) { + return std::nullopt; + } + + auto result = RetryRunner>({.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> 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(backoff * 2), - kMaxBackoff); - std::weak_ptr 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; - - std::weak_ptr 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(); - }); - scheduled_task_id_.store(new_id); + void ScheduleRefresh(std::chrono::milliseconds delay) { + ExecutorTask refresh([weak_self = weak_from_this()] { + if (auto self = weak_self.lock()) { + if (auto next_delay = self->Refresh()) { + self->ScheduleRefresh(*next_delay); + } + } + }); + ::iceberg::internal::TaskSubmissionCallback on_submission([](Status status) { + if (!status.has_value()) { + ICEBERG_LOG_WARN("Failed to submit OAuth2 token refresh: {}", + status.error().message); + } + }); + scheduler_->Schedule(delay, std::move(refresh), std::move(on_submission)); } - std::chrono::milliseconds CalculateRefreshDelay() const { + /// Returns nullopt if expiration is unknown. + std::optional 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(expires_at_ - now); auto refresh_window = std::min(expires_in / 10, std::chrono::milliseconds(300'000)); @@ -264,19 +298,21 @@ class OAuth2Session final : public AuthSession, Config config_; AuthProperties refresh_properties_; std::shared_ptr client_; - std::atomic scheduled_task_id_{0}; - std::atomic 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> 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& optional_oauth_params, std::shared_ptr client, - std::optional token_request_started_at); + std::optional 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 7bb7a3aa3f..feb4da5ec6 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 delegate); ~SigV4AuthManager() override; + void SetAuthRefreshExecutor(Executor* executor) override; + Result> InitSession( std::shared_ptr init_client, const std::unordered_map& properties) override; diff --git a/src/iceberg/catalog/rest/auth/sigv4_manager.cc b/src/iceberg/catalog/rest/auth/sigv4_manager.cc index 6e0dce86b2..6568997498 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 delegate) SigV4AuthManager::~SigV4AuthManager() = default; +void SigV4AuthManager::SetAuthRefreshExecutor(Executor* executor) { + delegate_->SetAuthRefreshExecutor(executor); +} + Result> SigV4AuthManager::InitSession( std::shared_ptr init_client, const std::unordered_map& 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 d66ed7d906..0000000000 --- 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 -#include - -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 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 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 5ef20ed94a..0000000000 --- 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 -#include -#include -#include -#include -#include -#include - -#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 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 callback; - }; - - std::mutex mutex_; - std::condition_variable cv_; - std::vector 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 a945e25d4c..c8d88d06fb 100644 --- a/src/iceberg/catalog/rest/rest_catalog.cc +++ b/src/iceberg/catalog/rest/rest_catalog.cc @@ -447,8 +447,8 @@ RestCatalog::~RestCatalog() { } Result> 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> { @@ -467,6 +467,7 @@ Result> 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 1e713a5040..d648fa7628 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>(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> 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 359e5ddff8..64b3d932ad 100644 --- a/src/iceberg/cpp20_compatibility_internal.h +++ b/src/iceberg/cpp20_compatibility_internal.h @@ -231,7 +231,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 891ce374a5..cbea23f15b 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 @@ -337,6 +340,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 9a5eace0dd..d0d61386c7 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 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 submission_failed; - std::promise first_completed; std::atomic 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 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(failure)); std::promise 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 paths = {"first", "second", "third"}; std::atomic 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 fd9605fc31..f6f1400b2b 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 #include #include -#include +#include +#include #include -#include +#include #include -#include #include #include @@ -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 PostForm(const std::string&, + const std::unordered_map&, + const std::unordered_map&, + const ErrorHandler&, AuthSession&) override { + ++requests; + return response(); } + std::function()> response = [] { + return IOError("refresh failed"); + }; + std::atomic requests{0}; +}; - scheduler.Shutdown(); -} - -// Verifies that cancelling a task prevents it from executing -TEST(TokenRefreshSchedulerTest, CancelPreventsExecution) { - TokenRefreshScheduler scheduler; - std::atomic 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 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> RefreshingSession( + const std::shared_ptr& 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, 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(); + client->response = [&]() -> Result { + 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 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(); + client->response = [&]() -> Result { + 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 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(); + client->response = [&]() -> Result { + 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(); - 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 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(); + 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 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(); + std::promise entered; + std::promise release; + auto entered_future = entered.get_future(); + auto release_future = release.get_future().share(); + client->response = [&]() -> Result { + 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 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(); + 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 a3a801b9b2..236c2d0e88 100644 --- a/src/iceberg/test/executor.h +++ b/src/iceberg/test/executor.h @@ -20,15 +20,40 @@ #pragma once #include +#include +#include #include #include #include +#include + #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 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 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 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 86865a6638..75bea71407 100644 --- a/src/iceberg/test/rest_catalog_unit_test.cc +++ b/src/iceberg/test/rest_catalog_unit_test.cc @@ -29,6 +29,7 @@ #include #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, PostForm, + (const std::string& path, + (const std::unordered_map&)form_data, + (const std::unordered_map&)headers, + const ErrorHandler& error_handler, auth::AuthSession& session), + (override)); + MOCK_METHOD(Result, Delete, (const std::string& path, (const std::unordered_map&)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> { + config, [&](const RestCatalogProperties&) -> Result> { 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 { + 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 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 29bc3ab042..47d06f34cc 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& 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 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 {}; TEST_F(RestMetricsReporterTest, SuppressesHttpErrors) { - ManualExecutor executor; + test::MockExecutor executor; + std::optional 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 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 0000000000..64d9f7e8a6 --- /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 +#include +#include +#include +#include +#include + +#include + +#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 entered; + std::promise release; + auto entered_done = entered.get_future(); + auto release_done = release.get_future().share(); + std::promise earlier; + std::promise later; + auto earlier_done = earlier.get_future(); + auto later_done = later.get_future(); + std::atomic 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 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 dispatched; + auto done = dispatched.get_future(); + std::atomic notifications{0}; + ScheduledExecutor first(first_executor); + ScheduledExecutor second(second_executor); + auto capture = std::make_shared(42); + std::weak_ptr 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(7); + std::weak_ptr 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 scheduler; + std::promise rejected; + }; + auto context = std::make_shared(); + context->scheduler = std::make_unique(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 release; + std::promise 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 625c833f12..4b81ca44d0 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 8d2ed2a183..afd3500418 100644 --- a/src/iceberg/util/retry_util.h +++ b/src/iceberg/util/retry_util.h @@ -94,8 +94,8 @@ template 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; template using OnlyRetryOn = RetryPolicy; +/// An empty stop list retries all errors. template using StopRetryOn = RetryPolicy; diff --git a/src/iceberg/util/scheduled_executor.cc b/src/iceberg/util/scheduled_executor.cc new file mode 100644 index 0000000000..0c143920d1 --- /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 +#include +#include +#include +#include +#include +#include + +#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 owner, + std::chrono::steady_clock::time_point deadline, ExecutorTask task, + std::optional 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& 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 owner; + ExecutorTask task; + std::optional on_submission; + }; + + using Queue = std::multimap; + + static bool SameOwner(const std::weak_ptr& a, + const std::weak_ptr& 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 { + 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 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 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(executor)) {} + +ScheduledExecutor::~ScheduledExecutor() { Shutdown(); } + +void ScheduledExecutor::Schedule(std::chrono::milliseconds delay, ExecutorTask task, + std::optional 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 0000000000..15521ef3fe --- /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 +#include +#include + +#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; + +/// 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 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 state_; +}; + +} // namespace iceberg::internal