fix(auth): harden scheduled refresh lifecycle

Author: wgtmac

src/iceberg/catalog/rest/auth/auth_session_internal.h

@@ -224,49 +224,87 @@ 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 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] {
+          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);
             }
           }),
           ::iceberg::internal::ScheduledTaskRejection(
-              [weak_self = std::move(weak_for_rejection)](Status) {
-                if (auto self = weak_self.lock()) self->ScheduleRefresh();
+              [weak_self = std::move(weak_for_rejection), generation](Status) {
+                if (auto self = weak_self.lock()) {
+                  self->ScheduledTaskRejected(generation, 0);
+                }
               }));
-      StoreScheduledTask(std::move(retry_task));
+      StoreScheduledTask(std::move(retry_task), generation);
     }
   }
 
-  void ScheduleRefresh() {
+  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 = CalculateRefreshDelay();
+    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 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();
+        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)](Status) {
-              if (auto self = weak_self.lock()) self->ScheduleRefresh();
+            [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));
+    StoreScheduledTask(std::move(task), generation);
   }
 
-  void StoreScheduledTask(::iceberg::internal::ScheduledTask task) {
+  uint64_t BeginScheduledTask() {
     std::lock_guard lock(scheduled_mutex_);
-    if (closed_.load()) {
+    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);
+  }
+
+  void StoreScheduledTask(::iceberg::internal::ScheduledTask task, uint64_t generation) {
+    std::lock_guard lock(scheduled_mutex_);
+    if (closed_.load() || generation != schedule_generation_) {
       task.Cancel();
       return;
     }
-    scheduled_task_.Cancel();
     scheduled_task_ = std::move(task);
   }
 
@@ -294,6 +332,7 @@ class OAuth2Session final : public AuthSession,
   std::shared_ptr<HttpClient> client_;
   ::iceberg::internal::ScheduledTask scheduled_task_;
   std::mutex scheduled_mutex_;
+  uint64_t schedule_generation_ = 0;
   std::atomic<bool> closed_{false};
   std::mutex refresh_mutex_;
   std::condition_variable refresh_cv_;

src/iceberg/util/scheduled_executor.cc

@@ -17,27 +17,29 @@
  * 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;
-  }
+  if (auto state = state_.lock()) state->Cancel(id_);
+  state_.reset();
+  id_ = 0;
 }
 
-ScheduledExecutor::ScheduledExecutor(Executor& executor)
-    : executor_(executor), timer_thread_([this] { Run(); }) {}
+ScheduledExecutorState::ScheduledExecutorState(Executor& executor)
+    : executor_(executor) {}
 
-ScheduledExecutor::~ScheduledExecutor() { Shutdown(); }
+ScheduledExecutorState::~ScheduledExecutorState() { Shutdown(); }
 
-ScheduledTask ScheduledExecutor::Schedule(
+void ScheduledExecutorState::Start() {
+  auto self = shared_from_this();
+  timer_thread_ = std::thread([self = std::move(self)] { self->Run(); });
+}
+
+ScheduledTask ScheduledExecutorState::Schedule(
     std::chrono::steady_clock::time_point deadline, ExecutorTask task,
     std::optional<ScheduledTaskRejection> on_rejected) {
   std::lock_guard lock(mutex_);
@@ -49,17 +51,10 @@ ScheduledTask ScheduledExecutor::Schedule(
                    .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));
+  return ScheduledTask(weak_from_this(), id);
 }
 
-void ScheduledExecutor::Cancel(uint64_t id) {
+void ScheduledExecutorState::Cancel(uint64_t id) {
   if (id == 0) return;
   {
     std::lock_guard lock(mutex_);
@@ -68,7 +63,7 @@ void ScheduledExecutor::Cancel(uint64_t id) {
   cv_.notify_one();
 }
 
-void ScheduledExecutor::Shutdown() {
+void ScheduledExecutorState::Shutdown() {
   {
     std::lock_guard lock(mutex_);
     if (shutdown_) return;
@@ -78,10 +73,15 @@ void ScheduledExecutor::Shutdown() {
     scheduled_.clear();
   }
   cv_.notify_one();
-  if (timer_thread_.joinable()) timer_thread_.join();
+  if (!timer_thread_.joinable()) return;
+  if (timer_thread_.get_id() == std::this_thread::get_id()) {
+    timer_thread_.detach();
+  } else {
+    timer_thread_.join();
+  }
 }
 
-void ScheduledExecutor::Run() {
+void ScheduledExecutorState::Run() {
   while (true) {
     std::optional<Task> task;
     {
@@ -108,11 +108,36 @@ void ScheduledExecutor::Run() {
       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));
     }
   }
 }
 
+ScheduledExecutor::ScheduledExecutor(Executor& executor)
+    : state_(std::make_shared<ScheduledExecutorState>(executor)) {
+  state_->Start();
+}
+
+ScheduledExecutor::~ScheduledExecutor() { Shutdown(); }
+
+ScheduledTask ScheduledExecutor::Schedule(
+    std::chrono::steady_clock::time_point deadline, ExecutorTask task,
+    std::optional<ScheduledTaskRejection> on_rejected) {
+  return state_->Schedule(deadline, std::move(task), std::move(on_rejected));
+}
+
+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::Shutdown() {
+  if (state_) state_->Shutdown();
+}
+
 }  // namespace iceberg::internal

src/iceberg/util/scheduled_executor_internal.h

@@ -22,6 +22,7 @@
 #include <chrono>
 #include <condition_variable>
 #include <cstdint>
+#include <memory>
 #include <mutex>
 #include <optional>
 #include <queue>
@@ -35,45 +36,42 @@
 
 namespace iceberg::internal {
 
-class ScheduledExecutor;
+class ScheduledExecutorState;
 
 using ScheduledTaskRejection = FnOnce<void(Status)>;
 
 class ICEBERG_EXPORT ScheduledTask {
  public:
   ScheduledTask() = default;
-  ScheduledTask(ScheduledExecutor* executor, uint64_t id)
-      : executor_(executor), id_(id) {}
+  ScheduledTask(std::weak_ptr<ScheduledExecutorState> state, uint64_t id)
+      : state_(std::move(state)), id_(id) {}
 
   void Cancel();
-  bool valid() const { return executor_ != nullptr && id_ != 0; }
+  bool valid() const { return id_ != 0 && !state_.expired(); }
 
  private:
-  ScheduledExecutor* executor_ = nullptr;
+  std::weak_ptr<ScheduledExecutorState> state_;
   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 ICEBERG_EXPORT ScheduledExecutor {
+class ICEBERG_EXPORT ScheduledExecutorState
+    : public std::enable_shared_from_this<ScheduledExecutorState> {
  public:
-  explicit ScheduledExecutor(Executor& executor);
-  ~ScheduledExecutor();
+  explicit ScheduledExecutorState(Executor& executor);
+  ~ScheduledExecutorState();
 
-  ScheduledExecutor(const ScheduledExecutor&) = delete;
-  ScheduledExecutor& operator=(const ScheduledExecutor&) = delete;
+  ScheduledExecutorState(const ScheduledExecutorState&) = delete;
+  ScheduledExecutorState& operator=(const ScheduledExecutorState&) = delete;
 
-  ScheduledTask Schedule(
-      std::chrono::steady_clock::time_point deadline, ExecutorTask task,
-      std::optional<ScheduledTaskRejection> on_rejected = std::nullopt);
-  ScheduledTask Schedule(
-      std::chrono::milliseconds delay, ExecutorTask task,
-      std::optional<ScheduledTaskRejection> on_rejected = std::nullopt);
+  ScheduledTask Schedule(std::chrono::steady_clock::time_point deadline,
+                         ExecutorTask task,
+                         std::optional<ScheduledTaskRejection> on_rejected);
   void Cancel(uint64_t id);
   void Shutdown();
 
  private:
+  friend class ScheduledExecutor;
+
   struct Task {
     uint64_t id;
     std::chrono::steady_clock::time_point deadline;
@@ -87,6 +85,7 @@ class ICEBERG_EXPORT ScheduledExecutor {
     }
   };
 
+  void Start();
   void Run();
 
   Executor& executor_;
@@ -100,4 +99,27 @@ class ICEBERG_EXPORT ScheduledExecutor {
   std::thread timer_thread_;
 };
 
+/// 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 ICEBERG_EXPORT 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,
+      std::optional<ScheduledTaskRejection> on_rejected = std::nullopt);
+  ScheduledTask Schedule(
+      std::chrono::milliseconds delay, ExecutorTask task,
+      std::optional<ScheduledTaskRejection> on_rejected = std::nullopt);
+  void Shutdown();
+
+ private:
+  std::shared_ptr<ScheduledExecutorState> state_;
+};
+
 }  // namespace iceberg::internal