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