wgtmac commented on code in PR #1033:
URL: https://github.com/apache/iceberg-cpp/pull/1033#discussion_r4305367336
##########
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:
Fixed. Scheduling now uses a generation token: the callback marks its
generation started before running, and a late StoreScheduledTask rejects and
cancels stale handles. This prevents the original registration from overwriting
and cancelling a successor scheduled by the callback.
##########
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:
Fixed. The timer loop now runs in a shared self-owned state. If teardown
happens on the timer thread, shutdown detaches that thread while its
self-reference keeps the state alive until Run exits, avoiding self-join and
use-after-free.
##########
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:
Fixed. Executor rejection retries now use a positive exponential dispatch
backoff (200 ms up to 10 seconds) with a bounded number of attempts, instead of
recomputing an expired token's zero delay indefinitely.
--
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]
