Copilot commented on code in PR #1033: URL: https://github.com/apache/iceberg-cpp/pull/1033#discussion_r4305286406
########## src/iceberg/util/scheduled_executor.cc: ########## @@ -0,0 +1,118 @@ +/* + * 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::optional<ScheduledTaskRejection> on_rejected) { + std::lock_guard lock(mutex_); + if (shutdown_) return {}; + auto id = next_id_++; + tasks_.push(Task{.id = id, + .deadline = deadline, + .task = std::move(task), + .on_rejected = std::move(on_rejected)}); + scheduled_.insert(id); + cv_.notify_one(); + return ScheduledTask(this, id); +} + +ScheduledTask ScheduledExecutor::Schedule( + std::chrono::milliseconds delay, ExecutorTask task, + std::optional<ScheduledTaskRejection> on_rejected) { + return Schedule(std::chrono::steady_clock::now() + delay, std::move(task), + std::move(on_rejected)); +} + +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; + } + auto status = executor_.Submit(std::move(task->task)); + if (!status.has_value() && task->on_rejected.has_value()) { + std::move (*task->on_rejected)(std::move(status)); Review Comment: The rejection callback runs on the timer thread and temporarily owns an OAuth2Session through weak_self.lock(). If the manager and external session references are released during that callback, releasing its local self can destroy the session and its scheduler on this thread. Shutdown() then joins its own timer thread, throwing and terminating the process. Keep the running state alive independently and arrange teardown/joining off the timer thread; simply skipping the join leaves Run() accessing destroyed state. ########## src/iceberg/catalog/rest/auth/auth_session_internal.h: ########## @@ -218,28 +225,49 @@ 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 weak_for_rejection = weak_self; + 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); + }), + ::iceberg::internal::ScheduledTaskRejection( + [weak_self = std::move(weak_for_rejection)](Status) { + if (auto self = weak_self.lock()) self->ScheduleRefresh(); + })); + StoreScheduledTask(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)] { + auto weak_for_rejection = weak_self; + auto 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); + }), + ::iceberg::internal::ScheduledTaskRejection( + [weak_self = std::move(weak_for_rejection)](Status) { + if (auto self = weak_self.lock()) self->ScheduleRefresh(); + })); + StoreScheduledTask(std::move(task)); + } + + void StoreScheduledTask(::iceberg::internal::ScheduledTask task) { + std::lock_guard lock(scheduled_mutex_); + if (closed_.load()) { + task.Cancel(); + return; + } + scheduled_task_.Cancel(); + scheduled_task_ = std::move(task); Review Comment: A due task can run or be rejected before its caller reaches StoreScheduledTask(). If that callback schedules and stores successor B first, the original caller then stores already-dispatched handle A and cancels B here, stopping refresh. Serialize enqueueing and handle installation together, or reject stale registrations, in both the normal-refresh and retry paths. ########## src/iceberg/catalog/rest/auth/auth_session_internal.h: ########## @@ -218,28 +225,49 @@ 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 weak_for_rejection = weak_self; + 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); + }), + ::iceberg::internal::ScheduledTaskRejection( + [weak_self = std::move(weak_for_rejection)](Status) { + if (auto self = weak_self.lock()) self->ScheduleRefresh(); + })); + StoreScheduledTask(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)] { + auto weak_for_rejection = weak_self; + auto 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); + }), + ::iceberg::internal::ScheduledTaskRejection( + [weak_self = std::move(weak_for_rejection)](Status) { + if (auto self = weak_self.lock()) self->ScheduleRefresh(); + })); Review Comment: Once the token expires, CalculateRefreshDelay() returns zero. If the executor keeps rejecting submissions, this callback repeatedly schedules immediately due tasks, creating a CPU-consuming loop on the timer thread. Use a positive, bounded dispatch-retry backoff in both rejection callbacks rather than the unchanged token's expiration time, and define when to stop retrying permanent rejection. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
