From 672cb4d6c9589295b15072c9804edfc80dd38dad Mon Sep 17 00:00:00 2001 From: "shuxu.li" Date: Mon, 5 Oct 2026 21:50:42 +0800 Subject: [PATCH] feat(auth): implement AuthSessionCache for OAuth2 child sessions --- src/iceberg/catalog/rest/CMakeLists.txt | 1 + src/iceberg/catalog/rest/auth/auth_manager.cc | 119 +++++- .../catalog/rest/auth/auth_properties.h | 4 + src/iceberg/catalog/rest/auth/auth_session.h | 2 + .../catalog/rest/auth/auth_session_cache.cc | 252 +++++++++++ .../rest/auth/auth_session_cache_internal.h | 129 ++++++ .../catalog/rest/auth/auth_session_internal.h | 59 ++- .../rest/auth/token_refresh_scheduler.cc | 4 +- .../rest/auth/token_refresh_scheduler.h | 3 + src/iceberg/catalog/session_context.h | 3 + src/iceberg/test/CMakeLists.txt | 1 + src/iceberg/test/auth_manager_test.cc | 63 +++ src/iceberg/test/auth_session_cache_test.cc | 403 ++++++++++++++++++ 13 files changed, 1008 insertions(+), 35 deletions(-) create mode 100644 src/iceberg/catalog/rest/auth/auth_session_cache.cc create mode 100644 src/iceberg/catalog/rest/auth/auth_session_cache_internal.h create mode 100644 src/iceberg/test/auth_session_cache_test.cc diff --git a/src/iceberg/catalog/rest/CMakeLists.txt b/src/iceberg/catalog/rest/CMakeLists.txt index f64860ff4..2ea291b72 100644 --- a/src/iceberg/catalog/rest/CMakeLists.txt +++ b/src/iceberg/catalog/rest/CMakeLists.txt @@ -22,6 +22,7 @@ set(ICEBERG_REST_SOURCES auth/auth_managers.cc auth/auth_properties.cc auth/auth_session.cc + auth/auth_session_cache.cc auth/oauth2_util.cc auth/sigv4_manager.cc auth/token_refresh_scheduler.cc diff --git a/src/iceberg/catalog/rest/auth/auth_manager.cc b/src/iceberg/catalog/rest/auth/auth_manager.cc index 5facfa3eb..d2305de7d 100644 --- a/src/iceberg/catalog/rest/auth/auth_manager.cc +++ b/src/iceberg/catalog/rest/auth/auth_manager.cc @@ -19,6 +19,7 @@ #include "iceberg/catalog/rest/auth/auth_manager.h" +#include #include #include #include @@ -28,8 +29,10 @@ #include "iceberg/catalog/rest/auth/auth_manager_internal.h" #include "iceberg/catalog/rest/auth/auth_properties.h" #include "iceberg/catalog/rest/auth/auth_session.h" +#include "iceberg/catalog/rest/auth/auth_session_cache_internal.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/session_context.h" #include "iceberg/util/base64.h" #include "iceberg/util/macros.h" @@ -156,13 +159,13 @@ class OAuth2Manager : public AuthManager { start_time_ = std::chrono::steady_clock::now(); ICEBERG_ASSIGN_OR_RAISE( auth_response_, OAuth2Util::FetchToken(*init_client, *init_session, config)); - // TODO(lishuxu): Match Java OAuth2Util.AuthSession.fromTokenResponse here. + // TODO(lishuxu): Build the initialization session from the full token response. return AuthSession::MakeDefault( OAuth2Util::AuthHeaders(auth_response_->access_token)); } if (!config.token().empty()) { - // TODO(lishuxu): Match Java OAuth2Util.AuthSession.fromAccessToken here. + // TODO(lishuxu): Preserve configured access-token expiration in the init session. return AuthSession::MakeDefault(OAuth2Util::AuthHeaders(config.token())); } @@ -176,6 +179,31 @@ class OAuth2Manager : public AuthManager { ICEBERG_PRECHECK(shared_client != nullptr, "OAuth2 catalog session HTTP client must not be null"); refresh_client_ = std::move(shared_client); + + // Initialize catalog-level configuration + keep_refreshed_ = config.keep_refreshed(); + exchange_enabled_ = config.exchange_enabled(); + session_timeout_ = + std::chrono::milliseconds(config.Get(AuthProperties::kSessionTimeoutMs)); + + // Create session cache + ICEBERG_ASSIGN_OR_RAISE( + session_cache_, + internal::AuthSessionCache::Make( + session_timeout_, [](std::shared_ptr session) { + if (auto oauth2 = + std::dynamic_pointer_cast(session)) { + oauth2->StopRefreshing(); + } + })); + + // Periodically reclaim idle sessions. Sweep at most once a minute and at least + // once a second, so a zero timeout still gets periodic cleanup. + auto sweep_interval = std::clamp(session_timeout_, std::chrono::milliseconds(1000), + std::chrono::milliseconds(60000)); + ICEBERG_RETURN_UNEXPECTED(session_cache_->StartPeriodicSweep( + TokenRefreshScheduler::Instance(), sweep_interval)); + // Reuse the token response and start time from the init phase. if (auth_response_.has_value()) { return internal::MakeOAuth2Session( @@ -184,8 +212,8 @@ class OAuth2Manager : public AuthManager { config.optional_oauth_params(), refresh_client_, start_time_); } - // TODO(lishuxu): Honor token-refresh-enabled for catalog bearer tokens, matching - // Java. If token is provided, use it directly. + // TODO(lishuxu): Honor token-refresh-enabled for catalog bearer tokens. + // If token is provided, use it directly. if (!config.token().empty()) { OAuthTokenResponse token_response{ .access_token = config.token(), @@ -217,21 +245,31 @@ class OAuth2Manager : public AuthManager { Result> ContextualSession( const SessionContext& context, std::shared_ptr parent) override { - // TODO(lishuxu): Add child-session caching and refresh, matching Java - // AuthSessionCache. + // Use session_id as cache key for contextual sessions + std::string cache_key = context.session_id.empty() ? "" : "ctx:" + context.session_id; return MaybeCreateChildSession(context.credentials, /*allow_credential=*/true, - std::move(parent)); + std::move(parent), cache_key); } Result> TableSession( [[maybe_unused]] const TableIdentifier& table, const std::unordered_map& properties, std::shared_ptr parent) override { + // Use token value as cache key for table sessions + auto token_it = properties.find(AuthProperties::kToken.key()); + std::string cache_key; + if (token_it != properties.end() && !token_it->second.empty()) { + cache_key = "tbl:" + token_it->second; + } return MaybeCreateChildSession(FilterTableSessionProperties(properties), - /*allow_credential=*/false, std::move(parent)); + /*allow_credential=*/false, std::move(parent), + cache_key); } Status Close() override { + if (session_cache_) { + session_cache_->Close(); + } refresh_client_.reset(); return {}; } @@ -267,7 +305,8 @@ class OAuth2Manager : public AuthManager { Result> MaybeCreateChildSession( const std::unordered_map& credentials, - bool allow_credential, std::shared_ptr parent) { + bool allow_credential, std::shared_ptr parent, + const std::string& cache_key) { auto token_it = credentials.find(AuthProperties::kToken.key()); auto credential_it = credentials.find(AuthProperties::kCredential.key()); auto typed_token = FindPreferredTypedToken(credentials); @@ -283,35 +322,65 @@ class OAuth2Manager : public AuthManager { ICEBERG_PRECHECK(parent_info.has_value(), "OAuth2 child session requires OAuth2 parent metadata"); + // Use cache if we have a valid cache_key + if (!cache_key.empty() && session_cache_) { + return session_cache_->Get(cache_key, + [&]() -> Result> { + return CreateChildSessionUncached( + credentials, allow_credential, *parent_info); + }); + } + + // No cache, create session directly + return CreateChildSessionUncached(credentials, allow_credential, *parent_info); + } + + Result> CreateChildSessionUncached( + const std::unordered_map& credentials, + bool allow_credential, const OAuth2SessionInfo& parent_info) { + auto token_it = credentials.find(AuthProperties::kToken.key()); + auto credential_it = credentials.find(AuthProperties::kCredential.key()); + auto typed_token = FindPreferredTypedToken(credentials); + if (token_it != credentials.end()) { - ICEBERG_ASSIGN_OR_RAISE(auto config, - ChildConfig(*parent_info, parent_info->credential)); + // Token child session: don't include parent credential in config. + // Token-based sessions are not refreshed; when token exchange is enabled, + // refresh will use the token itself rather than the parent credential. + auto properties = parent_info.optional_oauth_params; + properties[AuthProperties::kScope.key()] = parent_info.scope; + properties[AuthProperties::kOAuth2ServerUri.key()] = parent_info.oauth2_server_uri; + ICEBERG_ASSIGN_OR_RAISE(auto config, AuthProperties::FromProperties(properties)); return MakeSession(AccessTokenResponse(token_it->second), config, /*keep_refreshed=*/false); } if (allow_credential && credential_it != credentials.end()) { ICEBERG_ASSIGN_OR_RAISE(auto config, - ChildConfig(*parent_info, credential_it->second)); - ICEBERG_ASSIGN_OR_RAISE(auto response, - OAuth2Util::FetchToken(*refresh_client_, *parent, config)); - return MakeSession(response, config, /*keep_refreshed=*/false); + ChildConfig(parent_info, credential_it->second)); + // Use parent session headers for FetchToken authentication + auto temp_parent = AuthSession::MakeDefault(parent_info.headers); + ICEBERG_ASSIGN_OR_RAISE( + auto response, OAuth2Util::FetchToken(*refresh_client_, *temp_parent, config)); + // Credential child sessions use keep_refreshed from catalog level + return MakeSession(response, config, keep_refreshed_); } std::optional actor_token; std::optional actor_token_type; - if (!parent_info->token.empty()) { - actor_token = parent_info->token; - actor_token_type = parent_info->issued_token_type; + if (!parent_info.token.empty()) { + actor_token = parent_info.token; + actor_token_type = parent_info.issued_token_type; } + // Use parent session headers for ExchangeToken authentication + auto temp_parent = AuthSession::MakeDefault(parent_info.headers); ICEBERG_ASSIGN_OR_RAISE( auto response, - OAuth2Util::ExchangeToken(*refresh_client_, *parent, {}, typed_token->second, + OAuth2Util::ExchangeToken(*refresh_client_, *temp_parent, {}, typed_token->second, typed_token->first, actor_token, actor_token_type, - parent_info->scope, parent_info->oauth2_server_uri, - parent_info->optional_oauth_params)); + parent_info.scope, parent_info.oauth2_server_uri, + parent_info.optional_oauth_params)); ICEBERG_ASSIGN_OR_RAISE(auto config, - ChildConfig(*parent_info, parent_info->credential)); + ChildConfig(parent_info, parent_info.credential)); return MakeSession(response, config, /*keep_refreshed=*/false); } @@ -319,6 +388,12 @@ class OAuth2Manager : public AuthManager { std::optional auth_response_; std::optional start_time_; std::shared_ptr refresh_client_; + + // Catalog-level configuration initialized in CatalogSession + bool keep_refreshed_ = true; + bool exchange_enabled_ = true; + std::chrono::milliseconds session_timeout_{3'600'000}; + std::shared_ptr session_cache_; }; Result> MakeOAuth2Manager( diff --git a/src/iceberg/catalog/rest/auth/auth_properties.h b/src/iceberg/catalog/rest/auth/auth_properties.h index 8784194cc..cadb5913c 100644 --- a/src/iceberg/catalog/rest/auth/auth_properties.h +++ b/src/iceberg/catalog/rest/auth/auth_properties.h @@ -82,6 +82,10 @@ class ICEBERG_REST_EXPORT AuthProperties : public ConfigBase { inline static Entry kAudience{"audience", ""}; inline static Entry kResource{"resource", ""}; + /// Session cache timeout in milliseconds. Sessions will be eligible for eviction + /// after this duration of inactivity. Default is 1 hour (3,600,000 ms). + inline static Entry kSessionTimeoutMs{"auth.session-timeout-ms", 3'600'000}; + // ---- OAuth2 token type constants ---- inline static const std::string kAccessTokenType = diff --git a/src/iceberg/catalog/rest/auth/auth_session.h b/src/iceberg/catalog/rest/auth/auth_session.h index bdc77ebc4..35f0a88ea 100644 --- a/src/iceberg/catalog/rest/auth/auth_session.h +++ b/src/iceberg/catalog/rest/auth/auth_session.h @@ -37,11 +37,13 @@ namespace iceberg::rest::auth { /// \brief OAuth2 metadata used to derive child authentication sessions. struct ICEBERG_REST_EXPORT OAuth2SessionInfo { std::string token; + std::string token_type; std::string issued_token_type; std::string credential; std::string scope; std::string oauth2_server_uri; std::unordered_map optional_oauth_params; + std::unordered_map headers; }; /// \brief An authentication session that can authenticate outgoing HTTP requests. diff --git a/src/iceberg/catalog/rest/auth/auth_session_cache.cc b/src/iceberg/catalog/rest/auth/auth_session_cache.cc new file mode 100644 index 000000000..cb2d6d560 --- /dev/null +++ b/src/iceberg/catalog/rest/auth/auth_session_cache.cc @@ -0,0 +1,252 @@ +/* + * 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 "iceberg/catalog/rest/auth/auth_session_cache_internal.h" +#include "iceberg/catalog/rest/auth/auth_session_internal.h" +#include "iceberg/catalog/rest/auth/token_refresh_scheduler.h" +#include "iceberg/util/macros.h" + +namespace iceberg::rest::auth::internal { + +Result> AuthSessionCache::Make( + std::chrono::milliseconds session_timeout, RemovalListener on_removal, Clock clock) { + ICEBERG_PRECHECK(session_timeout >= std::chrono::milliseconds::zero(), + "Auth session timeout must not be negative: {} ms", + session_timeout.count()); + ICEBERG_PRECHECK(clock != nullptr, "Auth session cache clock must not be null"); + return std::shared_ptr( + new AuthSessionCache(session_timeout, std::move(on_removal), std::move(clock))); +} + +AuthSessionCache::AuthSessionCache(std::chrono::milliseconds session_timeout, + RemovalListener on_removal, Clock clock) + : session_timeout_(session_timeout), + on_removal_(std::move(on_removal)), + clock_(std::move(clock)) {} + +AuthSessionCache::~AuthSessionCache() { Close(); } + +Result> AuthSessionCache::Get(const std::string& key, + const Loader& loader) { + std::shared_ptr slot; + std::vector> removed; + { + std::unique_lock lock(mutex_); + while (true) { + if (closed_) { + return Invalid("Auth session cache is closed"); + } + auto it = slots_.find(key); + if (it == slots_.end()) { + // Cache miss: break out to create a new slot and call loader. + break; + } + auto existing = it->second; + if (existing->loading) { + // Another thread is loading this key. Wait for it to finish. + // A failed load removes the slot, so look up the key again after waking. + load_finished_.wait(lock, [&] { return closed_ || !existing->loading; }); + continue; + } + auto now = clock_(); + if (!IsExpired(*existing, now)) { + // Cache hit: update last access time and return the cached session. + existing->last_access = now; + return existing->session; + } + // Expired: remove from cache and let this thread reload it. + removed.push_back(std::move(existing->session)); + slots_.erase(it); + break; + } + // Create a new slot in loading state. Other threads calling Get(key) will + // see loading=true and wait on load_finished_. + slot = std::make_shared(); + slots_.emplace(key, slot); + } // Lock released here — loader runs without holding the cache lock. + + // Call removal listener outside the lock to avoid blocking other Get() calls. + NotifyRemoved(removed); + + // Call the loader outside the lock. This may take seconds (HTTP requests). + auto result = loader(); + if (result.has_value() && result.value() == nullptr) { + result = Invalid("Auth session loader returned a null session"); + } + + // Re-acquire the lock to update the slot state. + std::shared_ptr orphan; + { + std::lock_guard lock(mutex_); + slot->loading = false; + auto it = slots_.find(key); + bool current = it != slots_.end() && it->second == slot; + + if (!result.has_value()) { + // Load failed: remove the slot so the next Get() can retry. + // Failed loads are not cached. + if (current) { + slots_.erase(it); + } + } else if (closed_ || !current) { + // Cache was closed while loading, or another thread replaced this slot. + // Mark the session as orphan — it will be stopped outside the lock. + orphan = result.value(); + } else { + // Success: transition slot to Ready state and cache the session. + slot->session = result.value(); + slot->last_access = clock_(); + } + } + // Wake up all threads waiting for this key. They will re-check the map: + // - If load succeeded, they'll find loading=false and return the cached session. + // - If load failed, they'll find the slot removed and will try loading again. + load_finished_.notify_all(); + + if (orphan != nullptr) { + // Session loaded after cache closed or slot replaced: stop it immediately. + if (on_removal_) { + on_removal_(orphan); + } + return Invalid("Auth session cache is closed"); + } + return result; +} + +void AuthSessionCache::Sweep() { + std::vector> removed; + { + std::lock_guard lock(mutex_); + auto now = clock_(); + for (auto it = slots_.begin(); it != slots_.end();) { + if (!it->second->loading && IsExpired(*it->second, now)) { + removed.push_back(std::move(it->second->session)); + it = slots_.erase(it); + } else { + ++it; + } + } + } + NotifyRemoved(removed); +} + +Status AuthSessionCache::StartPeriodicSweep(TokenRefreshScheduler& scheduler, + std::chrono::milliseconds interval) { + ICEBERG_PRECHECK(interval > std::chrono::milliseconds::zero(), + "Auth session sweep interval must be positive: {} ms", + interval.count()); + std::lock_guard lock(mutex_); + if (closed_ || sweep_scheduler_ != nullptr) { + return {}; + } + sweep_scheduler_ = &scheduler; + sweep_interval_ = interval; + ScheduleSweepLocked(); + return {}; +} + +void AuthSessionCache::Close() { + std::vector> removed; + TokenRefreshScheduler* scheduler = nullptr; + uint64_t sweep_task_id = 0; + { + std::lock_guard lock(mutex_); + if (closed_) { + return; + } + closed_ = true; + for (auto& [key, slot] : slots_) { + // Sessions still loading are handed to the listener when their load ends. + if (!slot->loading) { + removed.push_back(std::move(slot->session)); + } + } + slots_.clear(); + scheduler = sweep_scheduler_; + sweep_task_id = std::exchange(sweep_task_id_, 0); + } + load_finished_.notify_all(); + if (scheduler != nullptr) { + scheduler->Cancel(sweep_task_id); + } + NotifyRemoved(removed); +} + +size_t AuthSessionCache::size() const { + std::lock_guard lock(mutex_); + return slots_.size(); +} + +bool AuthSessionCache::IsExpired(const Slot& slot, + std::chrono::steady_clock::time_point now) const { + // Check 1: Idle timeout (expireAfterAccess semantics) + if (now - slot.last_access >= session_timeout_) { + return true; + } + + // Check 2: Token expiration + // Even for refreshable sessions, check token expiration to handle cases where: + // - Refresh has permanently failed (after max retries) + // - Refresh is delayed or not scheduled + // - Token was obtained from external source (table token) + auto oauth2 = std::dynamic_pointer_cast(slot.session); + if (oauth2) { + auto expires_at = oauth2->ExpiresAt(); + if (expires_at.has_value() && now >= *expires_at) { + return true; + } + } + + return false; +} + +void AuthSessionCache::ScheduleSweepLocked() { + // Scheduling under the cache lock ensures Close() either sees this task's id + // or prevents it from being scheduled. + sweep_task_id_ = + sweep_scheduler_->Schedule(sweep_interval_, [weak_self = weak_from_this()] { + if (auto self = weak_self.lock()) { + self->RunScheduledSweep(); + } + }); +} + +void AuthSessionCache::RunScheduledSweep() { + Sweep(); + std::lock_guard lock(mutex_); + if (!closed_) { + ScheduleSweepLocked(); + } +} + +void AuthSessionCache::NotifyRemoved( + std::vector>& sessions) const { + if (!on_removal_) { + return; + } + for (auto& session : sessions) { + if (session != nullptr) { + on_removal_(session); + } + } +} + +} // namespace iceberg::rest::auth::internal diff --git a/src/iceberg/catalog/rest/auth/auth_session_cache_internal.h b/src/iceberg/catalog/rest/auth/auth_session_cache_internal.h new file mode 100644 index 000000000..3bdcffd79 --- /dev/null +++ b/src/iceberg/catalog/rest/auth/auth_session_cache_internal.h @@ -0,0 +1,129 @@ +/* + * 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 +#include +#include + +#include "iceberg/catalog/rest/iceberg_rest_export.h" +#include "iceberg/catalog/rest/type_fwd.h" +#include "iceberg/result.h" + +/// \file iceberg/catalog/rest/auth/auth_session_cache_internal.h +/// \brief Cache of child authentication sessions, keyed by caller identity. + +namespace iceberg::rest::auth { +class TokenRefreshScheduler; +} // namespace iceberg::rest::auth + +namespace iceberg::rest::auth::internal { + +/// \brief Caches child auth sessions and expires them after a period of inactivity. +/// +/// Entries expire after `session_timeout` without access, and each session that +/// leaves the cache is handed to the removal listener exactly once. Idle entries +/// are reclaimed when accessed or when Sweep() runs; StartPeriodicSweep() can +/// schedule Sweep() to run periodically. +/// +/// Thread safety: all public methods are thread-safe. Loaders and the removal +/// listener are always invoked without holding the cache lock. +class ICEBERG_REST_EXPORT AuthSessionCache + : public std::enable_shared_from_this { + public: + using Clock = std::function; + using Loader = std::function>()>; + /// Called once for every session that leaves the cache. Must not block. + using RemovalListener = std::function&)>; + + /// \brief Create a cache. + /// + /// \param session_timeout Idle time after which an entry expires; must not be + /// negative. Zero disables caching in effect. + /// \param on_removal Listener invoked for each session that leaves the cache. + /// \param clock Time source; tests can inject a fake clock. + static Result> Make( + std::chrono::milliseconds session_timeout, RemovalListener on_removal, + Clock clock = std::chrono::steady_clock::now); + + ~AuthSessionCache(); + + AuthSessionCache(const AuthSessionCache&) = delete; + AuthSessionCache& operator=(const AuthSessionCache&) = delete; + + /// \brief Return the cached session for `key`, loading it on a miss. + /// + /// Concurrent calls for the same key run `loader` at most once; the others wait + /// for its result. Failed loads are not cached. Returns an error once the cache + /// is closed; a session loaded after Close() is handed to the removal listener. + Result> Get(const std::string& key, const Loader& loader); + + /// \brief Remove all expired entries and notify the removal listener. + void Sweep(); + + /// \brief Run Sweep() every `interval` on `scheduler` until Close(). + /// + /// `scheduler` must outlive this cache or the call to Close(). `interval` must + /// be positive. Calling this more than once, or after Close(), has no effect. + Status StartPeriodicSweep(TokenRefreshScheduler& scheduler, + std::chrono::milliseconds interval); + + /// \brief Stop periodic sweeping and remove every entry. Idempotent. + void Close(); + + /// \brief Number of entries, including loads in progress. + size_t size() const; + + private: + struct Slot { + bool loading = true; + std::shared_ptr session; + std::chrono::steady_clock::time_point last_access; + }; + + AuthSessionCache(std::chrono::milliseconds session_timeout, RemovalListener on_removal, + Clock clock); + + bool IsExpired(const Slot& slot, std::chrono::steady_clock::time_point now) const; + void ScheduleSweepLocked(); + void RunScheduledSweep(); + void NotifyRemoved(std::vector>& sessions) const; + + const std::chrono::milliseconds session_timeout_; + const RemovalListener on_removal_; + const Clock clock_; + + mutable std::mutex mutex_; + std::condition_variable load_finished_; + std::unordered_map> slots_; + bool closed_ = false; + TokenRefreshScheduler* sweep_scheduler_ = nullptr; + std::chrono::milliseconds sweep_interval_{0}; + uint64_t sweep_task_id_ = 0; +}; + +} // namespace iceberg::rest::auth::internal diff --git a/src/iceberg/catalog/rest/auth/auth_session_internal.h b/src/iceberg/catalog/rest/auth/auth_session_internal.h index b8db3836e..de0c42fd2 100644 --- a/src/iceberg/catalog/rest/auth/auth_session_internal.h +++ b/src/iceberg/catalog/rest/auth/auth_session_internal.h @@ -94,16 +94,42 @@ class OAuth2Session final : public AuthSession, std::shared_lock lock(mutex_); return OAuth2SessionInfo{ .token = token_, + .token_type = "bearer", .issued_token_type = issued_token_type_, .credential = Credential(config_), .scope = config_.scope, .oauth2_server_uri = config_.token_endpoint, .optional_oauth_params = config_.optional_oauth_params, + .headers = headers_, }; } Status Close() override { return CloseImpl(); } + /// \brief Stop scheduling token refreshes without waiting. + /// + /// Unlike Close(), this does not wait for an in-flight refresh to finish, so it + /// is safe to call from any thread, including the refresh scheduler thread. The + /// session keeps its current token and remains usable; an in-flight refresh may + /// still update the token but will not schedule further refreshes. + void StopRefreshing() { + closed_.store(true); + TokenRefreshScheduler::Instance().Cancel( + scheduled_task_id_.exchange(kInvalidTaskHandle)); + } + + /// \brief Expiration time of the current token, or std::nullopt if unknown. + std::optional ExpiresAt() const { + std::shared_lock lock(mutex_); + if (expires_at_ == std::chrono::steady_clock::time_point{}) { + return std::nullopt; + } + return expires_at_; + } + + /// \brief Whether this session can refresh its token automatically. + bool IsRefreshable() const { return config_.keep_refreshed; } + ~OAuth2Session() override { std::ignore = CloseImpl(); } private: @@ -114,14 +140,16 @@ class OAuth2Session final : public AuthSession, client_(std::move(client)) {} Status CloseImpl() { - bool expected = false; - if (!closed_.compare_exchange_strong(expected, true)) { - return {}; - } - TokenRefreshScheduler::Instance().Cancel(scheduled_task_id_.exchange(0)); + closed_.store(true); + TokenRefreshScheduler::Instance().Cancel( + scheduled_task_id_.exchange(kInvalidTaskHandle)); std::unique_lock lock(refresh_mutex_); refresh_cv_.wait(lock, [this] { return active_refresh_count_ == 0; }); - TokenRefreshScheduler::Instance().Cancel(scheduled_task_id_.exchange(0)); + // Defensive: cancel again in case a refresh callback scheduled a retry between + // the first cancel and acquiring refresh_mutex_. With StoreScheduledTask()'s + // closed_ recheck, this should be a no-op, but it's cheap insurance. + TokenRefreshScheduler::Instance().Cancel( + scheduled_task_id_.exchange(kInvalidTaskHandle)); return {}; } @@ -200,8 +228,7 @@ class OAuth2Session final : public AuthSession, if (closed_.load()) return; auto empty_session = AuthSession::MakeDefault({}); - // TODO(lishuxu): Honor token-exchange-enabled and refresh via token exchange, - // matching Java. + // TODO(lishuxu): Use token exchange to refresh exchanged tokens when enabled. auto result = OAuth2Util::FetchToken(*client_, *empty_session, refresh_properties_); if (result.has_value()) { auto& response = result.value(); @@ -225,7 +252,7 @@ class OAuth2Session final : public AuthSession, self->DoRefreshAttempt(next_attempt, next_backoff, refresh_started_at); } }); - scheduled_task_id_.store(retry_id); + StoreScheduledTask(retry_id); } } @@ -239,7 +266,17 @@ class OAuth2Session final : public AuthSession, delay, [weak_self = std::move(weak_self)] { if (auto self = weak_self.lock()) self->DoRefresh(); }); - scheduled_task_id_.store(new_id); + StoreScheduledTask(new_id); + } + + void StoreScheduledTask(uint64_t task_id) { + scheduled_task_id_.store(task_id); + // StopRefreshing() may run between the caller's closed_ check and the store + // above; cancel the task it could not see. + if (closed_.load()) { + TokenRefreshScheduler::Instance().Cancel( + scheduled_task_id_.exchange(kInvalidTaskHandle)); + } } std::chrono::milliseconds CalculateRefreshDelay() const { @@ -264,7 +301,7 @@ class OAuth2Session final : public AuthSession, Config config_; AuthProperties refresh_properties_; std::shared_ptr client_; - std::atomic scheduled_task_id_{0}; + std::atomic scheduled_task_id_{kInvalidTaskHandle}; std::atomic closed_{false}; std::mutex refresh_mutex_; std::condition_variable refresh_cv_; diff --git a/src/iceberg/catalog/rest/auth/token_refresh_scheduler.cc b/src/iceberg/catalog/rest/auth/token_refresh_scheduler.cc index d66ed7d90..3c9da34a0 100644 --- a/src/iceberg/catalog/rest/auth/token_refresh_scheduler.cc +++ b/src/iceberg/catalog/rest/auth/token_refresh_scheduler.cc @@ -38,7 +38,7 @@ uint64_t TokenRefreshScheduler::Schedule(std::chrono::milliseconds delay, std::function callback) { std::lock_guard lock(mutex_); if (shutdown_) { - return 0; + return kInvalidTaskHandle; } uint64_t id = next_id_++; tasks_.push_back(Task{.id = id, @@ -49,7 +49,7 @@ uint64_t TokenRefreshScheduler::Schedule(std::chrono::milliseconds delay, } void TokenRefreshScheduler::Cancel(uint64_t handle) { - if (handle == 0) return; + if (handle == kInvalidTaskHandle) return; std::lock_guard lock(mutex_); std::erase_if(tasks_, [handle](const Task& t) { return t.id == handle; }); } diff --git a/src/iceberg/catalog/rest/auth/token_refresh_scheduler.h b/src/iceberg/catalog/rest/auth/token_refresh_scheduler.h index 5ef20ed94..3d3841c0a 100644 --- a/src/iceberg/catalog/rest/auth/token_refresh_scheduler.h +++ b/src/iceberg/catalog/rest/auth/token_refresh_scheduler.h @@ -44,6 +44,9 @@ namespace iceberg::rest::auth { /// /// TODO(lishuxu): Migrate to the shared thread pool abstraction once available /// (see https://github.com/apache/iceberg-cpp/pull/646#discussion_r3304315308). +/// Invalid task handle returned when scheduling fails or used to represent "no task". +constexpr uint64_t kInvalidTaskHandle = 0; + class ICEBERG_REST_EXPORT TokenRefreshScheduler { public: /// \brief Get the global singleton instance. diff --git a/src/iceberg/catalog/session_context.h b/src/iceberg/catalog/session_context.h index 19eb8fc5f..1a6bc819c 100644 --- a/src/iceberg/catalog/session_context.h +++ b/src/iceberg/catalog/session_context.h @@ -41,6 +41,9 @@ namespace iceberg { /// silently overriding credentials. struct ICEBERG_EXPORT SessionContext { /// Unique session identifier. Explicit contexts must provide a non-empty ID. + /// + /// A session ID identifies a stable caller context. Its identity and credentials + /// must not change while the ID is in use. std::string session_id; /// Caller identity. Empty means no identity was supplied. This value is diff --git a/src/iceberg/test/CMakeLists.txt b/src/iceberg/test/CMakeLists.txt index 6f6ff7603..74f1960f4 100644 --- a/src/iceberg/test/CMakeLists.txt +++ b/src/iceberg/test/CMakeLists.txt @@ -319,6 +319,7 @@ if(ICEBERG_BUILD_REST) add_rest_iceberg_test(rest_catalog_test SOURCES auth_manager_test.cc + auth_session_cache_test.cc error_handlers_test.cc endpoint_test.cc rest_file_io_test.cc diff --git a/src/iceberg/test/auth_manager_test.cc b/src/iceberg/test/auth_manager_test.cc index 61962d1c3..52e189c6a 100644 --- a/src/iceberg/test/auth_manager_test.cc +++ b/src/iceberg/test/auth_manager_test.cc @@ -432,6 +432,20 @@ TEST_F(AuthManagerTest, OAuth2StaticToken) { } // Verifies OAuth2 type is inferred from token property +TEST_F(AuthManagerTest, OAuth2ZeroSessionTimeoutCreatesCatalogSession) { + std::unordered_map properties = { + {AuthProperties::kAuthType, "oauth2"}, + {AuthProperties::kToken.key(), "my-static-token"}, + {AuthProperties::kSessionTimeoutMs.key(), "0"}, + }; + + ICEBERG_UNWRAP_OR_FAIL(auto manager, AuthManagers::Load("test-catalog", properties)); + ICEBERG_UNWRAP_OR_FAIL(auto session, manager->CatalogSession(client_, properties)); + + EXPECT_NE(session, nullptr); + EXPECT_THAT(manager->Close(), IsOk()); +} + TEST_F(AuthManagerTest, OAuth2InferredFromToken) { std::unordered_map properties = { {AuthProperties::kToken.key(), "inferred-token"}, @@ -791,6 +805,55 @@ TEST(OAuth2SessionTest, InitialTokenIsUsed) { session->Close(); } +TEST(OAuth2SessionTest, StopRefreshingKeepsSessionUsable) { + OAuthTokenResponse token_response{ + .access_token = "token", + .token_type = "bearer", + .expires_in_secs = 3600, + }; + ICEBERG_UNWRAP_OR_FAIL( + auto session, + internal::OAuth2Session::Make(token_response, + {.token_endpoint = "http://localhost/oauth/tokens", + .client_id = "client_id", + .client_secret = "client_secret", + .scope = "catalog", + .keep_refreshed = true}, + std::make_shared(), std::nullopt)); + + session->StopRefreshing(); + session->StopRefreshing(); + + ICEBERG_UNWRAP_OR_FAIL(auto request, session->Authenticate({})); + EXPECT_EQ(request.headers.at("Authorization"), "Bearer token"); + EXPECT_THAT(session->Close(), IsOk()); +} + +TEST(OAuth2SessionTest, ExpiresAtReflectsTokenLifetime) { + auto make_session = [](std::optional expires_in_secs) { + OAuthTokenResponse token_response{ + .access_token = "opaque-token", + .token_type = "bearer", + .expires_in_secs = expires_in_secs, + }; + return internal::OAuth2Session::Make( + token_response, + {.token_endpoint = "http://localhost/oauth/tokens", .keep_refreshed = false}, + std::make_shared(), std::nullopt); + }; + + auto before = std::chrono::steady_clock::now(); + ICEBERG_UNWRAP_OR_FAIL(auto expiring, make_session(60)); + auto after = std::chrono::steady_clock::now(); + auto expires_at = expiring->ExpiresAt(); + ASSERT_TRUE(expires_at.has_value()); + EXPECT_GE(*expires_at, before + std::chrono::seconds(60)); + EXPECT_LE(*expires_at, after + std::chrono::seconds(60)); + + ICEBERG_UNWRAP_OR_FAIL(auto unknown, make_session(std::nullopt)); + EXPECT_FALSE(unknown->ExpiresAt().has_value()); +} + TEST(OAuth2SessionTest, InitTokenExpirationUsesRequestStartTime) { OAuthTokenResponse token_response{ .access_token = "opaque-token", diff --git a/src/iceberg/test/auth_session_cache_test.cc b/src/iceberg/test/auth_session_cache_test.cc new file mode 100644 index 000000000..0ec2e55a5 --- /dev/null +++ b/src/iceberg/test/auth_session_cache_test.cc @@ -0,0 +1,403 @@ +/* + * 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 +#include +#include + +#include +#include + +#include "iceberg/catalog/rest/auth/auth_session.h" +#include "iceberg/catalog/rest/auth/auth_session_cache_internal.h" +#include "iceberg/catalog/rest/auth/token_refresh_scheduler.h" +#include "iceberg/test/matchers.h" + +namespace iceberg::rest::auth { + +namespace { + +using internal::AuthSessionCache; +using namespace std::chrono_literals; + +constexpr auto kTimeout = std::chrono::milliseconds(1000); + +class FakeClock { + public: + AuthSessionCache::Clock AsClock() { + return [this] { + return std::chrono::steady_clock::time_point(std::chrono::milliseconds(now_ms_)); + }; + } + + void Advance(std::chrono::milliseconds delta) { now_ms_ += delta.count(); } + + private: + std::atomic now_ms_{0}; +}; + +class RemovalRecorder { + public: + AuthSessionCache::RemovalListener Listener() { + return [this](const std::shared_ptr& session) { + std::lock_guard lock(mutex_); + removed_.push_back(session); + }; + } + + std::vector> removed() const { + std::lock_guard lock(mutex_); + return removed_; + } + + private: + mutable std::mutex mutex_; + std::vector> removed_; +}; + +/// One-shot latch used to hold a loader inside Get(). +class Gate { + public: + void Open() { + { + std::lock_guard lock(mutex_); + open_ = true; + } + cv_.notify_all(); + } + + void Wait() { + std::unique_lock lock(mutex_); + cv_.wait(lock, [this] { return open_; }); + } + + private: + std::mutex mutex_; + std::condition_variable cv_; + bool open_ = false; +}; + +AuthSessionCache::Loader CountingLoader(std::atomic& calls) { + return [&calls]() -> Result> { + ++calls; + return AuthSession::MakeDefault({}); + }; +} + +} // namespace + +class AuthSessionCacheTest : public ::testing::Test { + protected: + std::shared_ptr MakeCache( + std::chrono::milliseconds timeout = kTimeout) { + auto cache = AuthSessionCache::Make(timeout, recorder_.Listener(), clock_.AsClock()); + EXPECT_THAT(cache, IsOk()); + return cache.value(); + } + + // Declared before any cache so they outlive it. + FakeClock clock_; + RemovalRecorder recorder_; +}; + +TEST_F(AuthSessionCacheTest, RejectsNegativeTimeout) { + EXPECT_THAT(AuthSessionCache::Make(-1ms, recorder_.Listener()), + IsError(ErrorKind::kInvalidArgument)); +} + +TEST_F(AuthSessionCacheTest, ReturnsCachedSessionOnHit) { + auto cache = MakeCache(); + std::atomic calls = 0; + + ICEBERG_UNWRAP_OR_FAIL(auto first, cache->Get("tenant-a", CountingLoader(calls))); + ICEBERG_UNWRAP_OR_FAIL(auto second, cache->Get("tenant-a", CountingLoader(calls))); + + EXPECT_EQ(first, second); + EXPECT_EQ(calls, 1); + EXPECT_EQ(cache->size(), 1); +} + +TEST_F(AuthSessionCacheTest, KeysAreIsolated) { + auto cache = MakeCache(); + std::atomic calls = 0; + + ICEBERG_UNWRAP_OR_FAIL(auto a, cache->Get("tenant-a", CountingLoader(calls))); + ICEBERG_UNWRAP_OR_FAIL(auto b, cache->Get("tenant-b", CountingLoader(calls))); + + EXPECT_NE(a, b); + EXPECT_EQ(calls, 2); +} + +TEST_F(AuthSessionCacheTest, FailedLoadIsNotCached) { + auto cache = MakeCache(); + auto failing = []() -> Result> { + return AuthenticationFailed("token endpoint unavailable"); + }; + std::atomic calls = 0; + + EXPECT_THAT(cache->Get("tenant-a", failing), IsError(ErrorKind::kAuthenticationFailed)); + EXPECT_EQ(cache->size(), 0); + EXPECT_THAT(cache->Get("tenant-a", CountingLoader(calls)), IsOk()); + EXPECT_EQ(calls, 1); +} + +TEST_F(AuthSessionCacheTest, NullSessionIsRejected) { + auto cache = MakeCache(); + auto null_loader = []() -> Result> { return nullptr; }; + + EXPECT_THAT(cache->Get("tenant-a", null_loader), IsError(ErrorKind::kInvalid)); + EXPECT_EQ(cache->size(), 0); +} + +TEST_F(AuthSessionCacheTest, ExpiresAfterIdleTimeout) { + auto cache = MakeCache(); + std::atomic calls = 0; + ICEBERG_UNWRAP_OR_FAIL(auto first, cache->Get("tenant-a", CountingLoader(calls))); + + // Each access resets the idle timer, so the entry outlives the timeout in total. + clock_.Advance(kTimeout - 1ms); + ICEBERG_UNWRAP_OR_FAIL(auto hit, cache->Get("tenant-a", CountingLoader(calls))); + clock_.Advance(kTimeout - 1ms); + ICEBERG_UNWRAP_OR_FAIL(auto still_hit, cache->Get("tenant-a", CountingLoader(calls))); + EXPECT_EQ(hit, first); + EXPECT_EQ(still_hit, first); + EXPECT_EQ(calls, 1); + + clock_.Advance(kTimeout); + ICEBERG_UNWRAP_OR_FAIL(auto reloaded, cache->Get("tenant-a", CountingLoader(calls))); + EXPECT_NE(reloaded, first); + EXPECT_EQ(calls, 2); + EXPECT_THAT(recorder_.removed(), ::testing::ElementsAre(first)); +} + +TEST_F(AuthSessionCacheTest, SweepRemovesOnlyExpiredEntries) { + auto cache = MakeCache(); + std::atomic calls = 0; + ICEBERG_UNWRAP_OR_FAIL(auto idle, cache->Get("idle", CountingLoader(calls))); + ICEBERG_UNWRAP_OR_FAIL(auto active, cache->Get("active", CountingLoader(calls))); + + clock_.Advance(kTimeout - 1ms); + ASSERT_THAT(cache->Get("active", CountingLoader(calls)), IsOk()); + clock_.Advance(1ms); + cache->Sweep(); + + EXPECT_EQ(cache->size(), 1); + EXPECT_THAT(recorder_.removed(), ::testing::ElementsAre(idle)); + ICEBERG_UNWRAP_OR_FAIL(auto still_active, cache->Get("active", CountingLoader(calls))); + EXPECT_EQ(still_active, active); +} + +TEST_F(AuthSessionCacheTest, ConcurrentGetsForSameKeyLoadOnce) { + auto cache = MakeCache(); + std::atomic calls = 0; + Gate entered; + Gate release; + auto blocking_loader = [&]() -> Result> { + ++calls; + entered.Open(); + release.Wait(); + return AuthSession::MakeDefault({}); + }; + + constexpr int kThreads = 8; + std::vector> results(kThreads); + std::vector threads; + for (int i = 0; i < kThreads; ++i) { + threads.emplace_back([&, i] { + auto session = cache->Get("tenant-a", blocking_loader); + if (session.has_value()) { + results[i] = session.value(); + } + }); + } + entered.Wait(); + std::this_thread::sleep_for(20ms); // let the other callers queue up + release.Open(); + for (auto& thread : threads) { + thread.join(); + } + + EXPECT_EQ(calls, 1); + ASSERT_NE(results[0], nullptr); + for (const auto& session : results) { + EXPECT_EQ(session, results[0]); + } +} + +TEST_F(AuthSessionCacheTest, WaitersRetryAfterFailedLoad) { + auto cache = MakeCache(); + std::atomic calls = 0; + Gate entered; + Gate release; + auto loader = [&]() -> Result> { + if (++calls == 1) { + entered.Open(); + release.Wait(); + return AuthenticationFailed("first attempt fails"); + } + return AuthSession::MakeDefault({}); + }; + + Result> first_result = nullptr; + std::thread first([&] { first_result = cache->Get("tenant-a", loader); }); + entered.Wait(); + Result> waiter_result = nullptr; + std::thread waiter([&] { waiter_result = cache->Get("tenant-a", loader); }); + std::this_thread::sleep_for(20ms); + release.Open(); + first.join(); + waiter.join(); + + EXPECT_THAT(first_result, IsError(ErrorKind::kAuthenticationFailed)); + EXPECT_THAT(waiter_result, IsOk()); + EXPECT_EQ(calls, 2); + EXPECT_EQ(cache->size(), 1); +} + +TEST_F(AuthSessionCacheTest, CloseRemovesAllEntriesOnce) { + auto cache = MakeCache(); + std::atomic calls = 0; + ICEBERG_UNWRAP_OR_FAIL(auto a, cache->Get("tenant-a", CountingLoader(calls))); + ICEBERG_UNWRAP_OR_FAIL(auto b, cache->Get("tenant-b", CountingLoader(calls))); + + cache->Close(); + cache->Close(); + + EXPECT_EQ(cache->size(), 0); + EXPECT_THAT(recorder_.removed(), ::testing::UnorderedElementsAre(a, b)); + EXPECT_THAT(cache->Get("tenant-a", CountingLoader(calls)), + IsError(ErrorKind::kInvalid)); + EXPECT_EQ(calls, 2); +} + +TEST_F(AuthSessionCacheTest, LoadFinishingAfterCloseIsRemovedAndWaitersFail) { + auto cache = MakeCache(); + Gate entered; + Gate release; + std::shared_ptr loaded; + auto loader = [&]() -> Result> { + entered.Open(); + release.Wait(); + loaded = AuthSession::MakeDefault({}); + return loaded; + }; + + Result> loader_result = nullptr; + std::thread loading([&] { loader_result = cache->Get("tenant-a", loader); }); + entered.Wait(); + Result> waiter_result = nullptr; + std::thread waiter([&] { waiter_result = cache->Get("tenant-a", loader); }); + std::this_thread::sleep_for(20ms); + + cache->Close(); + waiter.join(); // Close() wakes waiters before the load finishes. + EXPECT_THAT(waiter_result, IsError(ErrorKind::kInvalid)); + + release.Open(); + loading.join(); + EXPECT_THAT(loader_result, IsError(ErrorKind::kInvalid)); + EXPECT_EQ(cache->size(), 0); + EXPECT_THAT(recorder_.removed(), ::testing::ElementsAre(loaded)); +} + +TEST_F(AuthSessionCacheTest, EverySessionIsRemovedExactlyOnceUnderContention) { + auto cache = MakeCache(); + std::mutex loaded_mutex; + std::set> loaded; + auto loader = [&]() -> Result> { + auto session = AuthSession::MakeDefault({}); + std::lock_guard lock(loaded_mutex); + loaded.insert(session); + return session; + }; + + std::atomic stop = false; + std::thread sweeper([&] { + while (!stop) { + cache->Sweep(); + std::this_thread::yield(); + } + }); + std::vector workers; + for (int t = 0; t < 4; ++t) { + workers.emplace_back([&, t] { + for (int i = 0; i < 500; ++i) { + EXPECT_THAT(cache->Get("key-" + std::to_string((i + t) % 5), loader), IsOk()); + if (i % 10 == 0) { + clock_.Advance(kTimeout / 3); + } + } + }); + } + for (auto& worker : workers) { + worker.join(); + } + stop = true; + sweeper.join(); + cache->Close(); + + auto removed = recorder_.removed(); + std::set> unique_removed(removed.begin(), removed.end()); + EXPECT_EQ(unique_removed.size(), removed.size()) << "a session was removed twice"; + EXPECT_EQ(unique_removed, loaded); +} + +TEST_F(AuthSessionCacheTest, PeriodicSweepRemovesExpiredEntries) { + TokenRefreshScheduler scheduler; // must outlive the cache + auto cache = MakeCache(); + std::atomic calls = 0; + ICEBERG_UNWRAP_OR_FAIL(auto idle, cache->Get("idle", CountingLoader(calls))); + clock_.Advance(kTimeout); + + ASSERT_THAT(cache->StartPeriodicSweep(scheduler, 10ms), IsOk()); + + auto deadline = std::chrono::steady_clock::now() + 5s; + while (cache->size() != 0 && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(5ms); + } + EXPECT_EQ(cache->size(), 0); + EXPECT_THAT(recorder_.removed(), ::testing::ElementsAre(idle)); +} + +TEST_F(AuthSessionCacheTest, PeriodicSweepArgumentsAndLifecycle) { + TokenRefreshScheduler scheduler; + auto cache = MakeCache(); + + EXPECT_THAT(cache->StartPeriodicSweep(scheduler, 0ms), + IsError(ErrorKind::kInvalidArgument)); + EXPECT_THAT(cache->StartPeriodicSweep(scheduler, 10ms), IsOk()); + EXPECT_THAT(cache->StartPeriodicSweep(scheduler, 10ms), IsOk()); // no-op + + // Destroying the cache cancels the sweep; the scheduler keeps running safely. + cache.reset(); + std::this_thread::sleep_for(30ms); + + auto closed = MakeCache(); + closed->Close(); + EXPECT_THAT(closed->StartPeriodicSweep(scheduler, 10ms), IsOk()); // no-op +} + +} // namespace iceberg::rest::auth