This is an automated email from the ASF dual-hosted git repository. wgtmac pushed a commit to branch codex/oauth-refresh-executor in repository https://gitbox.apache.org/repos/asf/iceberg-cpp.git
commit 643b6ccc45f0cb48dc9202a8617a5d1a71bd173f Author: Gang Wu <[email protected]> AuthorDate: Sun Oct 18 00:54:36 2026 +0800 refactor(auth): use injected executor for token refresh --- src/iceberg/CMakeLists.txt | 1 + src/iceberg/catalog/rest/CMakeLists.txt | 1 - src/iceberg/catalog/rest/auth/auth_manager.cc | 34 +++--- src/iceberg/catalog/rest/auth/auth_manager.h | 4 + src/iceberg/catalog/rest/auth/auth_session.cc | 5 +- src/iceberg/catalog/rest/auth/auth_session.h | 4 +- .../catalog/rest/auth/auth_session_internal.h | 34 +++--- .../catalog/rest/auth/token_refresh_scheduler.cc | 125 --------------------- .../catalog/rest/auth/token_refresh_scheduler.h | 112 ------------------ src/iceberg/catalog/rest/rest_catalog.cc | 3 +- src/iceberg/catalog/rest/rest_catalog.h | 6 +- src/iceberg/cpp20_compatibility_internal.h | 1 - src/iceberg/test/auth_manager_test.cc | 61 +++++----- src/iceberg/util/scheduled_executor.cc | 109 ++++++++++++++++++ src/iceberg/util/scheduled_executor_internal.h | 95 ++++++++++++++++ 15 files changed, 294 insertions(+), 301 deletions(-) diff --git a/src/iceberg/CMakeLists.txt b/src/iceberg/CMakeLists.txt index c283b8af..a02349de 100644 --- a/src/iceberg/CMakeLists.txt +++ b/src/iceberg/CMakeLists.txt @@ -155,6 +155,7 @@ set(ICEBERG_SOURCES util/murmurhash3_internal.cc util/property_util.cc util/retry_util.cc + util/scheduled_executor.cc util/snapshot_util.cc util/string_util.cc util/struct_like_set.cc diff --git a/src/iceberg/catalog/rest/CMakeLists.txt b/src/iceberg/catalog/rest/CMakeLists.txt index 0736290b..a8146d46 100644 --- a/src/iceberg/catalog/rest/CMakeLists.txt +++ b/src/iceberg/catalog/rest/CMakeLists.txt @@ -24,7 +24,6 @@ set(ICEBERG_REST_SOURCES auth/auth_session.cc auth/oauth2_util.cc auth/sigv4_manager.cc - auth/token_refresh_scheduler.cc catalog_properties.cc endpoint.cc error_handlers.cc diff --git a/src/iceberg/catalog/rest/auth/auth_manager.cc b/src/iceberg/catalog/rest/auth/auth_manager.cc index 5facfa3e..a090c9fa 100644 --- a/src/iceberg/catalog/rest/auth/auth_manager.cc +++ b/src/iceberg/catalog/rest/auth/auth_manager.cc @@ -33,6 +33,7 @@ #include "iceberg/catalog/session_context.h" #include "iceberg/util/base64.h" #include "iceberg/util/macros.h" +#include "iceberg/util/scheduled_executor_internal.h" namespace iceberg::rest::auth { @@ -140,6 +141,15 @@ Result<std::unique_ptr<AuthManager>> MakeBasicAuthManager( /// \brief OAuth2 authentication manager. class OAuth2Manager : public AuthManager { public: + OAuth2Manager() = default; + + void SetAuthRefreshExecutor(Executor* executor) override { + scheduler_.reset(); + if (executor != nullptr) { + scheduler_ = std::make_unique<::iceberg::internal::ScheduledExecutor>(*executor); + } + } + Result<std::shared_ptr<AuthSession>> InitSession( std::shared_ptr<HttpClient> init_client, const std::unordered_map<std::string, std::string>& properties) override { @@ -181,7 +191,7 @@ class OAuth2Manager : public AuthManager { 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_); + config.optional_oauth_params(), refresh_client_, start_time_, scheduler_.get()); } // TODO(lishuxu): Honor token-refresh-enabled for catalog bearer tokens, matching @@ -192,10 +202,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 +213,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); @@ -259,10 +263,10 @@ class OAuth2Manager : public AuthManager { bool keep_refreshed) 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_, std::nullopt, scheduler_.get()); } Result<std::shared_ptr<AuthSession>> MaybeCreateChildSession( @@ -319,12 +323,14 @@ class OAuth2Manager : public AuthManager { std::optional<OAuthTokenResponse> auth_response_; std::optional<std::chrono::steady_clock::time_point> start_time_; std::shared_ptr<HttpClient> refresh_client_; + std::unique_ptr<::iceberg::internal::ScheduledExecutor> scheduler_; }; Result<std::unique_ptr<AuthManager>> MakeOAuth2Manager( [[maybe_unused]] std::string_view name, [[maybe_unused]] const std::unordered_map<std::string, std::string>& properties) { - return std::make_unique<OAuth2Manager>(); + std::unique_ptr<AuthManager> manager = std::make_unique<OAuth2Manager>(); + return manager; } } // namespace iceberg::rest::auth diff --git a/src/iceberg/catalog/rest/auth/auth_manager.h b/src/iceberg/catalog/rest/auth/auth_manager.h index 6ba25c28..b8be8a2b 100644 --- a/src/iceberg/catalog/rest/auth/auth_manager.h +++ b/src/iceberg/catalog/rest/auth/auth_manager.h @@ -38,6 +38,10 @@ class ICEBERG_REST_EXPORT AuthManager { public: virtual ~AuthManager() = default; + /// Configure the borrowed executor used for background authentication refresh. + /// The executor must outlive this manager and all sessions it creates. + virtual void SetAuthRefreshExecutor(Executor* /*executor*/) {} + /// \brief Create a short-lived session used to contact the configuration endpoint. /// /// This session is used only during catalog initialization to fetch server diff --git a/src/iceberg/catalog/rest/auth/auth_session.cc b/src/iceberg/catalog/rest/auth/auth_session.cc index 60aa0a3b..e20e73c0 100644 --- a/src/iceberg/catalog/rest/auth/auth_session.cc +++ b/src/iceberg/catalog/rest/auth/auth_session.cc @@ -30,7 +30,6 @@ #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" @@ -79,7 +78,8 @@ Result<std::shared_ptr<AuthSession>> internal::MakeOAuth2Session( const std::string& scope, bool keep_refreshed, const std::unordered_map<std::string, std::string>& optional_oauth_params, std::shared_ptr<HttpClient> client, - std::optional<std::chrono::steady_clock::time_point> token_request_started_at) { + std::optional<std::chrono::steady_clock::time_point> token_request_started_at, + ::iceberg::internal::ScheduledExecutor* scheduler) { internal::OAuth2Session::Config config{ .token_endpoint = token_endpoint, .client_id = client_id, @@ -87,6 +87,7 @@ Result<std::shared_ptr<AuthSession>> internal::MakeOAuth2Session( .scope = scope, .optional_oauth_params = optional_oauth_params, .keep_refreshed = keep_refreshed, + .scheduler = scheduler, }; ICEBERG_ASSIGN_OR_RAISE(auto session, internal::OAuth2Session::Make( initial_token, std::move(config), diff --git a/src/iceberg/catalog/rest/auth/auth_session.h b/src/iceberg/catalog/rest/auth/auth_session.h index bdc77ebc..88b60c35 100644 --- a/src/iceberg/catalog/rest/auth/auth_session.h +++ b/src/iceberg/catalog/rest/auth/auth_session.h @@ -100,7 +100,9 @@ class ICEBERG_REST_EXPORT AuthSession { /// \param client_id OAuth2 client ID for refresh requests. /// \param client_secret OAuth2 client secret for re-fetch if refresh fails. /// \param scope OAuth2 scope for refresh requests. - /// \param keep_refreshed Whether to schedule automatic token refresh. + /// \param keep_refreshed Whether to schedule automatic token refresh. Automatic refresh + /// is enabled only when the session is created by an auth manager with a refresh + /// executor. /// \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. diff --git a/src/iceberg/catalog/rest/auth/auth_session_internal.h b/src/iceberg/catalog/rest/auth/auth_session_internal.h index b8db3836..e8662f4f 100644 --- a/src/iceberg/catalog/rest/auth/auth_session_internal.h +++ b/src/iceberg/catalog/rest/auth/auth_session_internal.h @@ -34,9 +34,9 @@ #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/util/macros.h" +#include "iceberg/util/scheduled_executor_internal.h" namespace iceberg::rest::auth::internal { @@ -68,6 +68,7 @@ class OAuth2Session final : public AuthSession, std::string scope; std::unordered_map<std::string, std::string> optional_oauth_params; bool keep_refreshed; + ::iceberg::internal::ScheduledExecutor* scheduler; }; static Result<std::shared_ptr<OAuth2Session>> Make( @@ -118,10 +119,10 @@ class OAuth2Session final : public AuthSession, if (!closed_.compare_exchange_strong(expected, true)) { return {}; } - TokenRefreshScheduler::Instance().Cancel(scheduled_task_id_.exchange(0)); + scheduled_task_.Cancel(); std::unique_lock lock(refresh_mutex_); refresh_cv_.wait(lock, [this] { return active_refresh_count_ == 0; }); - TokenRefreshScheduler::Instance().Cancel(scheduled_task_id_.exchange(0)); + scheduled_task_.Cancel(); return {}; } @@ -182,7 +183,7 @@ class OAuth2Session final : public AuthSession, const OAuthTokenResponse& token_response, std::optional<std::chrono::steady_clock::time_point> token_request_started_at) { UpdateTokenState(token_response, token_request_started_at); - if (config_.keep_refreshed && + if (config_.keep_refreshed && config_.scheduler != nullptr && expires_at_ != std::chrono::steady_clock::time_point{}) { ScheduleRefresh(); } @@ -218,28 +219,28 @@ class OAuth2Session final : public AuthSession, std::min(std::chrono::duration_cast<std::chrono::milliseconds>(backoff * 2), kMaxBackoff); std::weak_ptr<OAuth2Session> weak_self = shared_from_this(); - auto retry_id = TokenRefreshScheduler::Instance().Schedule( - backoff, [weak_self = std::move(weak_self), next_attempt = attempt + 1, - next_backoff, refresh_started_at] { + auto retry_task = config_.scheduler->Schedule( + backoff, + ExecutorTask([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); + })); + scheduled_task_ = std::move(retry_task); } } void ScheduleRefresh() { - if (!config_.keep_refreshed || closed_.load()) return; + if (!config_.keep_refreshed || config_.scheduler == nullptr || closed_.load()) return; auto delay = CalculateRefreshDelay(); if (delay < std::chrono::milliseconds::zero()) return; std::weak_ptr<OAuth2Session> weak_self = shared_from_this(); - auto new_id = TokenRefreshScheduler::Instance().Schedule( - delay, [weak_self = std::move(weak_self)] { + scheduled_task_ = config_.scheduler->Schedule( + delay, ExecutorTask([weak_self = std::move(weak_self)] { if (auto self = weak_self.lock()) self->DoRefresh(); - }); - scheduled_task_id_.store(new_id); + })); } std::chrono::milliseconds CalculateRefreshDelay() const { @@ -264,7 +265,7 @@ class OAuth2Session final : public AuthSession, Config config_; AuthProperties refresh_properties_; std::shared_ptr<HttpClient> client_; - std::atomic<uint64_t> scheduled_task_id_{0}; + ::iceberg::internal::ScheduledTask scheduled_task_; std::atomic<bool> closed_{false}; std::mutex refresh_mutex_; std::condition_variable refresh_cv_; @@ -277,6 +278,7 @@ Result<std::shared_ptr<AuthSession>> MakeOAuth2Session( const std::string& scope, bool keep_refreshed, const std::unordered_map<std::string, std::string>& optional_oauth_params, std::shared_ptr<HttpClient> client, - std::optional<std::chrono::steady_clock::time_point> token_request_started_at); + std::optional<std::chrono::steady_clock::time_point> token_request_started_at, + ::iceberg::internal::ScheduledExecutor* scheduler = nullptr); } // namespace iceberg::rest::auth::internal diff --git a/src/iceberg/catalog/rest/auth/token_refresh_scheduler.cc b/src/iceberg/catalog/rest/auth/token_refresh_scheduler.cc deleted file mode 100644 index d66ed7d9..00000000 --- a/src/iceberg/catalog/rest/auth/token_refresh_scheduler.cc +++ /dev/null @@ -1,125 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ - -#include "iceberg/catalog/rest/auth/token_refresh_scheduler.h" - -#include <algorithm> -#include <ranges> - -namespace iceberg::rest::auth { - -TokenRefreshScheduler& TokenRefreshScheduler::Instance() { - // Intentionally leaked to avoid destruction-order races at process exit. - static auto* instance = new TokenRefreshScheduler(); - return *instance; -} - -TokenRefreshScheduler::TokenRefreshScheduler() : worker_([this] { Run(); }) {} - -TokenRefreshScheduler::~TokenRefreshScheduler() { Shutdown(); } - -uint64_t TokenRefreshScheduler::Schedule(std::chrono::milliseconds delay, - std::function<void()> callback) { - std::lock_guard lock(mutex_); - if (shutdown_) { - return 0; - } - uint64_t id = next_id_++; - tasks_.push_back(Task{.id = id, - .fire_at = std::chrono::steady_clock::now() + delay, - .callback = std::move(callback)}); - cv_.notify_one(); - return id; -} - -void TokenRefreshScheduler::Cancel(uint64_t handle) { - if (handle == 0) return; - std::lock_guard lock(mutex_); - std::erase_if(tasks_, [handle](const Task& t) { return t.id == handle; }); -} - -void TokenRefreshScheduler::Shutdown() { - { - std::lock_guard lock(mutex_); - if (shutdown_) return; - shutdown_ = true; - tasks_.clear(); - } - cv_.notify_one(); - if (worker_.joinable()) { - worker_.join(); - } -} - -void TokenRefreshScheduler::Run() { - while (true) { - std::function<void()> callback; - - { - std::unique_lock lock(mutex_); - - if (tasks_.empty() && !shutdown_) { - // Wait until a task is added or shutdown is requested - cv_.wait(lock, [this] { return !tasks_.empty() || shutdown_; }); - } - - if (shutdown_) break; - if (tasks_.empty()) continue; - - // Find the task with the earliest fire_at - auto earliest_it = std::ranges::min_element( - tasks_, [](const Task& a, const Task& b) { return a.fire_at < b.fire_at; }); - - auto fire_at = earliest_it->fire_at; - auto target_id = earliest_it->id; - - // Wait until fire_at or until woken (new task, cancel, or shutdown). - // Note: The predicate does O(n) scan on each spurious wakeup. This is - // acceptable for the expected task count (< 10). If task count grows - // significantly, consider replacing vector with a priority queue. - cv_.wait_until(lock, fire_at, [&] { - // Wake up if: shutdown, task list changed, or time is up - if (shutdown_) return true; - if (tasks_.empty()) return true; - // Check if the earliest task has changed (new task added or cancelled) - auto new_earliest = std::ranges::min_element( - tasks_, [](const Task& a, const Task& b) { return a.fire_at < b.fire_at; }); - return new_earliest->id != target_id; - }); - - if (shutdown_) break; - - // If we were woken because the earliest task changed, loop again - auto now = std::chrono::steady_clock::now(); - auto due_it = - std::ranges::find_if(tasks_, [now](const Task& t) { return t.fire_at <= now; }); - if (due_it == tasks_.end()) continue; - - callback = std::move(due_it->callback); - tasks_.erase(due_it); - } - - // Execute callback outside the lock - if (callback) { - callback(); - } - } -} - -} // namespace iceberg::rest::auth diff --git a/src/iceberg/catalog/rest/auth/token_refresh_scheduler.h b/src/iceberg/catalog/rest/auth/token_refresh_scheduler.h deleted file mode 100644 index 5ef20ed9..00000000 --- a/src/iceberg/catalog/rest/auth/token_refresh_scheduler.h +++ /dev/null @@ -1,112 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ - -#pragma once - -#include <chrono> -#include <condition_variable> -#include <cstdint> -#include <functional> -#include <mutex> -#include <thread> -#include <vector> - -#include "iceberg/catalog/rest/iceberg_rest_export.h" - -/// \file iceberg/catalog/rest/auth/token_refresh_scheduler.h -/// \brief Global scheduler for OAuth2 token refresh tasks. - -namespace iceberg::rest::auth { - -/// \brief A process-global scheduler for delayed token refresh tasks. -/// -/// Uses a single background thread that sleeps until the next task is due. -/// All OAuth2Session instances share this scheduler. Tasks are lightweight -/// (a single HTTP POST to refresh a token), so one thread is sufficient. -/// -/// Thread safety: All public methods are thread-safe. -/// -/// TODO(lishuxu): Migrate to the shared thread pool abstraction once available -/// (see https://github.com/apache/iceberg-cpp/pull/646#discussion_r3304315308). -class ICEBERG_REST_EXPORT TokenRefreshScheduler { - public: - /// \brief Get the global singleton instance. - /// - /// Lazily created on first access and intentionally leaked: the worker - /// thread is reclaimed by the OS at process exit. Tests needing - /// deterministic shutdown should use a local instance. - static TokenRefreshScheduler& Instance(); - - /// \brief Schedule a callback to run after a delay. - /// - /// \param delay Time to wait before executing the callback. - /// \param callback Function to execute when the delay expires. - /// \return A unique handle that can be used to cancel the task. - uint64_t Schedule(std::chrono::milliseconds delay, std::function<void()> callback); - - /// \brief Cancel a previously scheduled task. - /// - /// If the task has already fired or does not exist, this is a no-op. - /// - /// \param handle The handle returned by Schedule(). - void Cancel(uint64_t handle); - - /// \brief Shutdown the scheduler, cancelling all pending tasks. - /// - /// After shutdown, Schedule() calls are no-ops (return 0). - /// This is called automatically on destruction. - /// - /// WARNING: Do not call this on the global Instance() unless you intend to - /// permanently stop all token refresh for the entire process. This is mainly - /// useful for testing with locally-constructed scheduler instances. - void Shutdown(); - - ~TokenRefreshScheduler(); - - // Non-copyable, non-movable - TokenRefreshScheduler(const TokenRefreshScheduler&) = delete; - TokenRefreshScheduler& operator=(const TokenRefreshScheduler&) = delete; - TokenRefreshScheduler(TokenRefreshScheduler&&) = delete; - TokenRefreshScheduler& operator=(TokenRefreshScheduler&&) = delete; - - /// \brief Construct a scheduler (prefer Instance() for production use). - /// - /// This constructor is public to allow testing with isolated instances. - /// In production code, use Instance() to get the global singleton. - TokenRefreshScheduler(); - - private: - /// \brief Worker loop that processes tasks. - void Run(); - - struct Task { - uint64_t id; - std::chrono::steady_clock::time_point fire_at; - std::function<void()> callback; - }; - - std::mutex mutex_; - std::condition_variable cv_; - std::vector<Task> tasks_; - uint64_t next_id_ = 1; // 0 is reserved as "invalid handle" - bool shutdown_ = false; - std::thread worker_; -}; - -} // namespace iceberg::rest::auth diff --git a/src/iceberg/catalog/rest/rest_catalog.cc b/src/iceberg/catalog/rest/rest_catalog.cc index c4506731..67a87d15 100644 --- a/src/iceberg/catalog/rest/rest_catalog.cc +++ b/src/iceberg/catalog/rest/rest_catalog.cc @@ -433,7 +433,7 @@ RestCatalog::~RestCatalog() { Result<std::shared_ptr<RestCatalog>> RestCatalog::Make( const RestCatalogProperties& config, Executor* metrics_executor, - ClientFactory client_factory) { + ClientFactory client_factory, Executor* auth_executor) { if (!client_factory) { client_factory = [](const RestCatalogProperties& properties) -> Result<std::shared_ptr<HttpClient>> { @@ -445,6 +445,7 @@ Result<std::shared_ptr<RestCatalog>> RestCatalog::Make( std::string catalog_name = config.Get(RestCatalogProperties::kName); ICEBERG_ASSIGN_OR_RAISE(auto auth_manager, auth::AuthManagers::Load(catalog_name, config.configs())); + auth_manager->SetAuthRefreshExecutor(auth_executor); ICEBERG_ASSIGN_OR_RAISE( auto paths, ResourcePaths::Make(std::string(TrimTrailingSlash(uri)), diff --git a/src/iceberg/catalog/rest/rest_catalog.h b/src/iceberg/catalog/rest/rest_catalog.h index 1e713a50..e1b1a7df 100644 --- a/src/iceberg/catalog/rest/rest_catalog.h +++ b/src/iceberg/catalog/rest/rest_catalog.h @@ -63,9 +63,13 @@ class ICEBERG_REST_EXPORT RestCatalog final /// 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. + /// \param auth_executor Optional non-owning executor for OAuth2 token refresh. When + /// null, background token refresh is disabled. When non-null, it must outlive the + /// catalog and all authentication sessions created by it. static Result<std::shared_ptr<RestCatalog>> Make(const RestCatalogProperties& config, Executor* metrics_executor = nullptr, - ClientFactory client_factory = {}); + ClientFactory client_factory = {}, + Executor* auth_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 359e5ddf..64b3d932 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/auth_manager_test.cc b/src/iceberg/test/auth_manager_test.cc index fd9605fc..2c5796eb 100644 --- a/src/iceberg/test/auth_manager_test.cc +++ b/src/iceberg/test/auth_manager_test.cc @@ -38,7 +38,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,8 +45,10 @@ #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/util/base64.h" +#include "iceberg/util/scheduled_executor_internal.h" namespace iceberg::rest::auth { @@ -692,22 +693,23 @@ TEST_F(AuthManagerTest, ExpiresAtMillisInvalidTokensReturnNullopt) { } } -// ---- TokenRefreshScheduler tests ---- +// ---- ScheduledExecutor tests ---- // Verifies that a scheduled task fires after the specified delay -TEST(TokenRefreshSchedulerTest, ScheduleFiresAfterDelay) { - TokenRefreshScheduler scheduler; +TEST(ScheduledExecutorTest, ScheduleFiresAfterDelay) { + test::ThreadExecutor executor; + ::iceberg::internal::ScheduledExecutor scheduler(executor); std::mutex mutex; std::condition_variable cv; bool fired = false; - scheduler.Schedule(std::chrono::milliseconds(50), [&] { - { - std::lock_guard lock(mutex); - fired = true; - } - cv.notify_one(); - }); + scheduler.Schedule(std::chrono::milliseconds(50), ExecutorTask([&] { + { + std::lock_guard lock(mutex); + fired = true; + } + cv.notify_one(); + })); { std::unique_lock lock(mutex); @@ -718,15 +720,16 @@ TEST(TokenRefreshSchedulerTest, ScheduleFiresAfterDelay) { } // Verifies that cancelling a task prevents it from executing -TEST(TokenRefreshSchedulerTest, CancelPreventsExecution) { - TokenRefreshScheduler scheduler; +TEST(ScheduledExecutorTest, CancelPreventsExecution) { + test::ThreadExecutor executor; + ::iceberg::internal::ScheduledExecutor scheduler(executor); std::atomic<bool> fired{false}; - auto handle = - scheduler.Schedule(std::chrono::milliseconds(100), [&] { fired.store(true); }); + auto handle = scheduler.Schedule(std::chrono::milliseconds(100), + ExecutorTask([&] { fired.store(true); })); // Cancel before it fires - scheduler.Cancel(handle); + handle.Cancel(); // Wait past the scheduled time std::this_thread::sleep_for(std::chrono::milliseconds(200)); @@ -736,11 +739,13 @@ TEST(TokenRefreshSchedulerTest, CancelPreventsExecution) { } // Verifies that shutdown with pending tasks does not crash -TEST(TokenRefreshSchedulerTest, ShutdownWithPendingTasks) { - TokenRefreshScheduler scheduler; +TEST(ScheduledExecutorTest, ShutdownWithPendingTasks) { + test::ThreadExecutor executor; + ::iceberg::internal::ScheduledExecutor scheduler(executor); std::atomic<bool> fired{false}; - scheduler.Schedule(std::chrono::milliseconds(5000), [&] { fired.store(true); }); + scheduler.Schedule(std::chrono::milliseconds(5000), + ExecutorTask([&] { fired.store(true); })); // Shutdown immediately — should not crash and task should not fire scheduler.Shutdown(); @@ -750,20 +755,22 @@ TEST(TokenRefreshSchedulerTest, ShutdownWithPendingTasks) { } // Verifies that Schedule after shutdown returns invalid handle (0) -TEST(TokenRefreshSchedulerTest, ScheduleAfterShutdownIsNoop) { - TokenRefreshScheduler scheduler; +TEST(ScheduledExecutorTest, ScheduleAfterShutdownIsNoop) { + test::ThreadExecutor executor; + ::iceberg::internal::ScheduledExecutor scheduler(executor); scheduler.Shutdown(); - auto handle = scheduler.Schedule(std::chrono::milliseconds(10), [] {}); - EXPECT_EQ(0u, handle); + auto handle = scheduler.Schedule(std::chrono::milliseconds(10), ExecutorTask([] {})); + EXPECT_FALSE(handle.valid()); } // Verifies that cancelling an invalid handle does not crash -TEST(TokenRefreshSchedulerTest, CancelInvalidHandleIsNoop) { - TokenRefreshScheduler scheduler; +TEST(ScheduledExecutorTest, CancelInvalidHandleIsNoop) { + test::ThreadExecutor executor; + ::iceberg::internal::ScheduledExecutor scheduler(executor); // Should not crash - scheduler.Cancel(0); - scheduler.Cancel(999); + ::iceberg::internal::ScheduledTask handle; + handle.Cancel(); scheduler.Shutdown(); } diff --git a/src/iceberg/util/scheduled_executor.cc b/src/iceberg/util/scheduled_executor.cc new file mode 100644 index 00000000..faeb7e25 --- /dev/null +++ b/src/iceberg/util/scheduled_executor.cc @@ -0,0 +1,109 @@ +/* + * 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 <optional> +#include <utility> + +#include "iceberg/util/scheduled_executor_internal.h" + +namespace iceberg::internal { + +void ScheduledTask::Cancel() { + if (valid()) { + executor_->Cancel(id_); + executor_ = nullptr; + id_ = 0; + } +} + +ScheduledExecutor::ScheduledExecutor(Executor& executor) + : executor_(executor), timer_thread_([this] { Run(); }) {} + +ScheduledExecutor::~ScheduledExecutor() { Shutdown(); } + +ScheduledTask ScheduledExecutor::Schedule(std::chrono::steady_clock::time_point deadline, + ExecutorTask task) { + std::lock_guard lock(mutex_); + if (shutdown_) return {}; + auto id = next_id_++; + tasks_.push(Task{.id = id, .deadline = deadline, .task = std::move(task)}); + scheduled_.insert(id); + cv_.notify_one(); + return ScheduledTask(this, id); +} + +ScheduledTask ScheduledExecutor::Schedule(std::chrono::milliseconds delay, + ExecutorTask task) { + return Schedule(std::chrono::steady_clock::now() + delay, std::move(task)); +} + +void ScheduledExecutor::Cancel(uint64_t id) { + if (id == 0) return; + { + std::lock_guard lock(mutex_); + if (scheduled_.erase(id) != 0) cancelled_.insert(id); + } + cv_.notify_one(); +} + +void ScheduledExecutor::Shutdown() { + { + std::lock_guard lock(mutex_); + if (shutdown_) return; + shutdown_ = true; + while (!tasks_.empty()) tasks_.pop(); + cancelled_.clear(); + scheduled_.clear(); + } + cv_.notify_one(); + if (timer_thread_.joinable()) timer_thread_.join(); +} + +void ScheduledExecutor::Run() { + while (true) { + std::optional<Task> task; + { + std::unique_lock lock(mutex_); + cv_.wait(lock, [this] { return shutdown_ || !tasks_.empty(); }); + if (shutdown_) return; + + while (!tasks_.empty() && cancelled_.contains(tasks_.top().id)) { + cancelled_.erase(tasks_.top().id); + tasks_.pop(); + } + if (tasks_.empty()) continue; + + auto deadline = tasks_.top().deadline; + if (cv_.wait_until(lock, deadline, [this, deadline] { + return shutdown_ || tasks_.empty() || tasks_.top().deadline != deadline || + cancelled_.contains(tasks_.top().id); + })) { + continue; + } + + task.emplace(std::move(const_cast<Task&>(tasks_.top()))); + tasks_.pop(); + scheduled_.erase(task->id); + if (cancelled_.erase(task->id) != 0) continue; + } + std::ignore = executor_.Submit(std::move(task->task)); + } +} + +} // namespace iceberg::internal diff --git a/src/iceberg/util/scheduled_executor_internal.h b/src/iceberg/util/scheduled_executor_internal.h new file mode 100644 index 00000000..c23433ee --- /dev/null +++ b/src/iceberg/util/scheduled_executor_internal.h @@ -0,0 +1,95 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#pragma once + +#include <chrono> +#include <condition_variable> +#include <cstdint> +#include <mutex> +#include <queue> +#include <thread> +#include <unordered_set> +#include <vector> + +#include "iceberg/result.h" +#include "iceberg/util/executor.h" + +namespace iceberg::internal { + +class ScheduledExecutor; + +class ScheduledTask { + public: + ScheduledTask() = default; + ScheduledTask(ScheduledExecutor* executor, uint64_t id) + : executor_(executor), id_(id) {} + + void Cancel(); + bool valid() const { return executor_ != nullptr && id_ != 0; } + + private: + ScheduledExecutor* executor_ = nullptr; + uint64_t id_ = 0; +}; + +/// A delay queue that dispatches due tasks to an existing Executor. +/// +/// The executor is borrowed and must outlive this object and all scheduled tasks. +class ScheduledExecutor { + public: + explicit ScheduledExecutor(Executor& executor); + ~ScheduledExecutor(); + + ScheduledExecutor(const ScheduledExecutor&) = delete; + ScheduledExecutor& operator=(const ScheduledExecutor&) = delete; + + ScheduledTask Schedule(std::chrono::steady_clock::time_point deadline, + ExecutorTask task); + ScheduledTask Schedule(std::chrono::milliseconds delay, ExecutorTask task); + void Cancel(uint64_t id); + void Shutdown(); + + private: + struct Task { + uint64_t id; + std::chrono::steady_clock::time_point deadline; + ExecutorTask task; + }; + + struct Earlier { + bool operator()(const Task& lhs, const Task& rhs) const { + return lhs.deadline > rhs.deadline; + } + }; + + void Run(); + + Executor& executor_; + std::mutex mutex_; + std::condition_variable cv_; + std::priority_queue<Task, std::vector<Task>, Earlier> tasks_; + std::unordered_set<uint64_t> cancelled_; + std::unordered_set<uint64_t> scheduled_; + uint64_t next_id_ = 1; + bool shutdown_ = false; + std::thread timer_thread_; +}; + +} // namespace iceberg::internal
