shangxinli commented on code in PR #1033:
URL: https://github.com/apache/iceberg-cpp/pull/1033#discussion_r4316572932
##########
src/iceberg/catalog/rest/auth/auth_session_internal.h:
##########
@@ -217,29 +224,88 @@ class OAuth2Session final : public AuthSession,
auto next_backoff =
std::min(std::chrono::duration_cast<std::chrono::milliseconds>(backoff * 2),
kMaxBackoff);
+ auto generation = BeginScheduledTask();
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), generation,
+ next_attempt = attempt + 1, next_backoff,
refresh_started_at] {
if (auto self = weak_self.lock()) {
+ self->ScheduledTaskStarted(generation);
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), generation](Status) {
+ if (auto self = weak_self.lock()) {
+ self->ScheduledTaskRejected(generation, 0);
+ }
+ }));
+ StoreScheduledTask(std::move(retry_task), generation);
}
}
- void ScheduleRefresh() {
- if (!config_.keep_refreshed || closed_.load()) return;
- auto delay = CalculateRefreshDelay();
+ void ScheduleRefresh() { ScheduleRefreshAfter(std::nullopt, 0); }
+
+ void ScheduleRefreshAfter(std::optional<std::chrono::milliseconds>
override_delay,
+ int dispatch_attempt) {
+ if (!config_.keep_refreshed || config_.scheduler == nullptr ||
closed_.load()) return;
+ auto delay = override_delay.value_or(CalculateRefreshDelay());
if (delay < std::chrono::milliseconds::zero()) return;
+ auto generation = BeginScheduledTask();
std::weak_ptr<OAuth2Session> weak_self = shared_from_this();
- auto new_id = TokenRefreshScheduler::Instance().Schedule(
- delay, [weak_self = std::move(weak_self)] {
- if (auto self = weak_self.lock()) self->DoRefresh();
- });
- scheduled_task_id_.store(new_id);
+ auto weak_for_rejection = weak_self;
+ auto task = config_.scheduler->Schedule(
+ delay, ExecutorTask([weak_self = std::move(weak_self), generation] {
+ if (auto self = weak_self.lock()) {
+ self->ScheduledTaskStarted(generation);
+ self->DoRefresh();
+ }
+ }),
+ ::iceberg::internal::ScheduledTaskRejection(
+ [weak_self = std::move(weak_for_rejection), generation,
+ dispatch_attempt](Status) {
+ if (auto self = weak_self.lock()) {
+ self->ScheduledTaskRejected(generation, dispatch_attempt);
+ }
+ }));
+ StoreScheduledTask(std::move(task), generation);
+ }
+
+ uint64_t BeginScheduledTask() {
+ std::lock_guard lock(scheduled_mutex_);
+ scheduled_task_.Cancel();
+ return ++schedule_generation_;
+ }
+
+ void ScheduledTaskStarted(uint64_t generation) {
+ std::lock_guard lock(scheduled_mutex_);
+ if (generation == schedule_generation_) {
+ ++schedule_generation_;
+ scheduled_task_.Cancel();
+ }
+ }
+
+ void ScheduledTaskRejected(uint64_t generation, int dispatch_attempt) {
+ ScheduledTaskStarted(generation);
+ static constexpr int kMaxDispatchRetries = 5;
+ static constexpr auto kInitialDispatchBackoff =
std::chrono::milliseconds(200);
+ if (dispatch_attempt >= kMaxDispatchRetries || closed_.load()) return;
+ auto multiplier = 1 << dispatch_attempt;
+ auto delay =
+ std::min(kInitialDispatchBackoff * multiplier,
std::chrono::milliseconds(10'000));
+ ScheduleRefreshAfter(delay, dispatch_attempt + 1);
Review Comment:
This reschedule is not gated on the generation still being current, so a
stale rejection cancels the live refresh and replaces it with a 200ms dispatch
retry, discarding whatever backoff it was carrying.
--
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]
