Refactor thread pool

PiperOrigin-RevId: 742919232

internal/platform/implementation/windows/thread_pool.cc

@@ -16,33 +16,25 @@
 
 #include <windows.h>
 
+#include <cstdint>
 #include <memory>
+#include <optional>
 #include <queue>
 #include <utility>
 
 #include "absl/memory/memory.h"
 #include "absl/synchronization/mutex.h"
-#include "internal/platform/implementation/shared/count_down_latch.h"
+#include "absl/time/time.h"
 #include "internal/platform/logging.h"
 #include "internal/platform/runnable.h"
 
 namespace nearby {
 namespace windows {
 
-VOID CALLBACK WorkCallback(PTP_CALLBACK_INSTANCE instance, PVOID parameter,
-                           PTP_WORK work) {
-  // Instance is not used in thread pools.
-  UNREFERENCED_PARAMETER(instance);
-
-  ThreadPool* thread_pool = reinterpret_cast<ThreadPool*>(parameter);
-  thread_pool->RunNextTask();
-  CloseThreadpoolWork(work);
-}
-
-std::unique_ptr<ThreadPool> ThreadPool::Create(int max_pool_size) {
+std::unique_ptr<ThreadPool> ThreadPool::Create(uint32_t max_pool_size) {
   PTP_POOL thread_pool = nullptr;
   TP_CALLBACK_ENVIRON thread_pool_environ;
-  InitializeThreadpoolEnvironment(&thread_pool_environ);
+  PTP_CLEANUP_GROUP cleanup_group = nullptr;
 
   if (max_pool_size <= 0) {
     LOG(ERROR) << __func__
@@ -50,162 +42,332 @@ std::unique_ptr<ThreadPool> ThreadPool::Create(int max_pool_size) {
     return nullptr;
   }
 
-  thread_pool = CreateThreadpool(NULL);
-  if (thread_pool == nullptr) {
-    LOG(ERROR) << __func__ << ": failed to create thread pool. LastError: "
+  InitializeThreadpoolEnvironment(&thread_pool_environ);
+  cleanup_group = CreateThreadpoolCleanupGroup();
+  if (cleanup_group == nullptr) {
+    LOG(ERROR) << __func__
+               << ": Failed to create thread pool cleanup group. LastError: "
                << GetLastError();
     return nullptr;
   }
 
+  thread_pool = CreateThreadpool(nullptr);
+  if (thread_pool == nullptr) {
+    LOG(ERROR) << __func__ << ": Failed to create thread pool. LastError: "
+               << GetLastError();
+    CloseThreadpoolCleanupGroup(/*ptpcg=*/cleanup_group);
+    return nullptr;
+  }
+
   // Sets thread pool maximum value. In order to release all threads,
   // it will keep at least one thread.
-  SetThreadpoolThreadMaximum(thread_pool, max_pool_size);
-  if (!SetThreadpoolThreadMinimum(thread_pool, 1)) {
+  SetThreadpoolThreadMaximum(/*ptpp=*/thread_pool, /*cthrdMost=*/max_pool_size);
+  if (!SetThreadpoolThreadMinimum(/*ptpp=*/thread_pool, /*cthrdMic=*/1)) {
     LOG(ERROR) << __func__
                << ": failed to set minimum thread pool size. LastError: "
                << GetLastError();
-    CloseThreadpool(thread_pool);
+    CloseThreadpoolCleanupGroup(/*ptpcg=*/cleanup_group);
+    CloseThreadpool(/*ptpp=*/thread_pool);
     return nullptr;
   }
 
   //
   // Associate the callback environment with our thread pool.
   //
-  SetThreadpoolCallbackPool(&thread_pool_environ, thread_pool);
+  SetThreadpoolCallbackPool(/*pcbe=*/&thread_pool_environ,
+                            /*ptpp=*/thread_pool);
+  SetThreadpoolCallbackCleanupGroup(/*pcbe=*/&thread_pool_environ,
+                                    /*ptpcg=*/cleanup_group,
+                                    /*pfng=*/nullptr);
 
-  return absl::WrapUnique(
-      new ThreadPool(thread_pool, thread_pool_environ, max_pool_size));
-}
-
-ThreadPool::ThreadPool(PTP_POOL thread_pool,
-                       TP_CALLBACK_ENVIRON thread_pool_environ,
-                       int max_pool_size)
-    : thread_pool_(thread_pool),
-      thread_pool_environ_(thread_pool_environ),
-      max_pool_size_(max_pool_size) {
-  VLOG(1) << __func__ << ": Thread pool(" << this
-          << ") is created with size:" << max_pool_size_;
-}
-
-ThreadPool::~ThreadPool() {
-  VLOG(1) << __func__ << ": Thread pool(" << this << ") is released.";
-
-  if (thread_pool_ == nullptr) {
-    return;
-  }
-
-  ShutDown();
+  return absl::WrapUnique(new ThreadPool(thread_pool, thread_pool_environ,
+                                         cleanup_group, max_pool_size));
 }
 
 bool ThreadPool::Run(Runnable task) {
   absl::MutexLock lock(&mutex_);
 
-  if (thread_pool_ == nullptr) {
+  if (task == nullptr) {
+    LOG(WARNING) << __func__ << ": Invalid task.";
     return false;
   }
 
-  if (shutdown_latch_ != nullptr) {
-    LOG(WARNING) << __func__ << ": Thread pool is in shutting down.";
+  if (is_shut_down_) {
+    LOG(WARNING) << __func__ << ": Thread pool is shut down.";
     return false;
   }
 
   PTP_WORK work;
-  tasks_.push(std::move(task));
-  VLOG(1) << __func__ << ": Scheduled to run task(" << &tasks_.back() << ").";
 
-  work = CreateThreadpoolWork(WorkCallback, this, &thread_pool_environ_);
+  work = CreateThreadpoolWork(/*pfnwk=*/WorkCallback, /*pv=*/this,
+                              /*pcbe=*/&thread_pool_environ_);
   if (work == nullptr) {
     LOG(ERROR) << __func__ << ": failed to create thread pool work. LastError: "
                << GetLastError();
     return false;
   }
 
-  ++running_tasks_count_;
+  VLOG(1) << __func__ << ": Scheduled to run work(" << work << ").";
 
   //
   // Submit the work to the pool. Because this was a pre-allocated
   // work item (using CreateThreadpoolWork), it is guaranteed to execute.
   //
-  SubmitThreadpoolWork(work);
+  task_queue_.enqueue(std::move(task));
+  SubmitThreadpoolWork(/*pwk=*/work);
   return true;
 }
 
+std::optional<uint64_t> ThreadPool::Run(Runnable task, absl::Duration delay) {
+  return Run(std::move(task), delay, absl::ZeroDuration());
+}
+
+std::optional<uint64_t> ThreadPool::Run(Runnable task, absl::Duration delay,
+                                        absl::Duration period) {
+  absl::MutexLock lock(&mutex_);
+
+  if (task == nullptr) {
+    LOG(WARNING) << __func__ << ": Invalid task.";
+    return false;
+  }
+
+  if (is_shut_down_) {
+    LOG(WARNING) << __func__ << ": Thread pool is shut down.";
+    return false;
+  }
+
+  // Closing the timer within its callback is prohibited. The thread pool
+  // ensures the callback remains active until its completion. When a new
+  // delayed task arrives, all completed tasks are promptly cleaned. Developers
+  // should cancel timers as soon as they are no longer needed.
+  delayed_task_map_.clean_completed_tasks();
+
+  PTP_TIMER timer = CreateThreadpoolTimer(
+      /*pfnti=*/TimerCallback,
+      /*pv=*/this,
+      /*pcbe=*/&thread_pool_environ_);
+
+  if (timer == nullptr) {
+    LOG(WARNING) << __func__
+                 << ": failed to create thread pool timer. LastError: "
+                 << GetLastError();
+    return std::nullopt;
+  }
+
+  VLOG(1) << __func__ << ": Scheduled to run timer(" << timer << ").";
+
+  delayed_task_map_.put(
+      timer, std::make_unique<DelayedTaskInfo>(std::move(task), delay, period));
+
+  FILETIME file_due_time;
+  ULARGE_INTEGER due_time;
+  due_time.QuadPart = (ULONGLONG) - (absl::ToInt64Milliseconds(delay) * 10000);
+  file_due_time.dwLowDateTime = due_time.LowPart;
+  file_due_time.dwHighDateTime = due_time.HighPart;
+
+  SetThreadpoolTimerEx(
+      /*pti=*/timer, /*pftDueTime=*/&file_due_time,
+      /*msPeriod=*/static_cast<DWORD>(absl::ToInt64Milliseconds(period)),
+      /*msWindowLength=*/10);
+
+  return reinterpret_cast<uint64_t>(timer);
+}
+
 void ThreadPool::ShutDown() {
-  {
-    absl::MutexLock lock(&mutex_);
-
-    if (thread_pool_ == nullptr) {
-      LOG(WARNING) << __func__ << ": Shutdown on closed thread pool(" << this
-                   << ").";
-      return;
-    }
-
-    if (running_tasks_count_ == 0) {
-      CloseThreadpool(thread_pool_);
-      thread_pool_ = nullptr;
-      VLOG(1) << __func__ << ": Thread pool(" << this << ") is shut down.";
-      return;
-    }
-
-    if (shutdown_latch_ != nullptr) {
-      VLOG(1) << __func__ << ": Thread pool(" << this
-              << ") is already in shutting down.";
-      return;
-    }
-
-    VLOG(1) << __func__ << ": Thread pool(" << this << ") is shutting down.";
-
-    shutdown_latch_ = std::make_unique<shared::CountDownLatch>(1);
+  absl::MutexLock lock(&mutex_);
+  VLOG(1) << __func__ << ": Thread pool(" << this << ") is shutting down.";
+  if (is_shut_down_) {
+    LOG(INFO) << __func__ << ": Thread pool(" << this
+              << ") is already shut down.";
+    return;
   }
 
-  // Wait for all tasks to complete.
-  shutdown_latch_->Await();
+  is_shut_down_ = true;
 
-  {
-    absl::MutexLock lock(&mutex_);
-    CloseThreadpool(thread_pool_);
-    thread_pool_ = nullptr;
-    VLOG(1) << __func__ << ": Thread pool(" << this << ") is shut down.";
-  }
+  CloseThreadpoolCleanupGroupMembers(/*ptpcg=*/cleanup_group_,
+                                     /*fCancelPendingCallbacks=*/FALSE,
+                                     /*pvCleanupContext=*/nullptr);
+  CloseThreadpoolCleanupGroup(/*ptpcg=*/cleanup_group_);
+  CloseThreadpool(/*ptpp=*/thread_pool_);
+  task_queue_.clear();
+  delayed_task_map_.clear();
+  VLOG(1) << __func__ << ": Thread pool(" << this << ") is shut down.";
+}
+
+ThreadPool::ThreadPool(PTP_POOL thread_pool,
+                       TP_CALLBACK_ENVIRON thread_pool_environ,
+                       PTP_CLEANUP_GROUP cleanup_group, int max_pool_size)
+    : thread_pool_(thread_pool),
+      thread_pool_environ_(thread_pool_environ),
+      cleanup_group_(cleanup_group),
+      max_pool_size_(max_pool_size) {
+  VLOG(1) << __func__ << ": Thread pool(" << this
+          << ") is created with size:" << max_pool_size_;
+}
+
+ThreadPool::~ThreadPool() {
+  VLOG(1) << __func__ << ": Thread pool(" << this << ") is releasing.";
+  ShutDown();
+}
+
+VOID CALLBACK ThreadPool::WorkCallback(PTP_CALLBACK_INSTANCE instance,
+                                       PVOID parameter, PTP_WORK work) {
+  // Instance is not used in work callbacks.
+  UNREFERENCED_PARAMETER(instance);
+
+  VLOG(1) << __func__ << ": Start to run work(" << work << ").";
+  ThreadPool* thread_pool = static_cast<ThreadPool*>(parameter);
+  thread_pool->RunNextTask();
+  CloseThreadpoolWork(work);
+  VLOG(1) << __func__ << ": Completed to run work(" << work << ").";
+}
+
+VOID CALLBACK ThreadPool::TimerCallback(PTP_CALLBACK_INSTANCE instance,
+                                        PVOID context, PTP_TIMER timer) {
+  // Instance is not used in tiemr callbacks.
+  UNREFERENCED_PARAMETER(instance);
+
+  VLOG(1) << __func__ << ": Start to run timer(" << timer << ") callback.";
+  ThreadPool* thread_pool = static_cast<ThreadPool*>(context);
+  thread_pool->RunTimerCallback(timer);
+  VLOG(1) << __func__ << ": Completed to run timer(" << timer << ") callback.";
 }
 
 void ThreadPool::RunNextTask() {
-  Runnable task = nullptr;
+  std::optional<Runnable> task = task_queue_.dequeue();
+  if (!task.has_value()) {
+    return;
+  }
 
-  {
-    absl::MutexLock lock(&mutex_);
+  task.value()();
+}
 
-    if (thread_pool_ == nullptr) {
-      return;
-    }
-    if (!tasks_.empty()) {
-      VLOG(1) << __func__ << ": Run task(" << &tasks_.front() << ").";
+void ThreadPool::RunTimerCallback(PTP_TIMER timer) {
+  DelayedTaskInfo* task_info = delayed_task_map_.get(timer);
+  if (task_info == nullptr) {
+    return;
+  }
 
-      task = std::move(tasks_.front());
-      tasks_.pop();
+  DWORD thread_id = GetCurrentThreadId();
+  task_info->thread_id = thread_id;
 
-      if (task == nullptr) {
-        LOG(WARNING) << __func__
-                     << ": Tried to run task in an empty thread pool.";
-        --running_tasks_count_;
-        if (running_tasks_count_ == 0 && shutdown_latch_ != nullptr) {
-          shutdown_latch_->CountDown();
-        }
-        return;
-      }
+  task_info->task();
+
+  if (task_info->is_canceled) {
+    delayed_task_map_.erase(timer);
+    return;
+  }
+
+  if (task_info->period == absl::ZeroDuration()) {
+    task_info->is_done = true;
+  }
+}
+
+bool ThreadPool::CancelDelayedTask(uint64_t delayed_task_id) {
+  absl::MutexLock lock(&mutex_);
+
+  if (is_shut_down_) {
+    LOG(WARNING)
+        << __func__
+        << ": Ignore to cancel delayed task because thread pool is shut down.";
+    return false;
+  }
+
+  PTP_TIMER timer = reinterpret_cast<PTP_TIMER>(delayed_task_id);
+  VLOG(1) << __func__ << ": Start to cancel timer(" << timer << ").";
+  DelayedTaskInfo* task_info = delayed_task_map_.get(timer);
+
+  if (task_info == nullptr) {
+    return true;
+  }
+
+  BOOL is_canceled = SetThreadpoolTimerEx(/*pti=*/timer, /*pftDueTime=*/nullptr,
+                                          /*msPeriod=*/0, /*msWindowLength=*/0);
+  bool clean_timer = false;
+  if (!is_canceled) {
+    DWORD thread_id = GetCurrentThreadId();
+    if (thread_id == task_info->thread_id) {
+      // The timer is cancelled in the callback.
+      VLOG(1) << __func__ << ": The timer is cancelled in the same thread."
+              << task_info->thread_id;
+      // The task_info is deleted upon completion of the running timer callback.
+      // Developers must ensure proper resource management, considering access
+      // by the callback and other threads.
+      task_info->is_canceled = true;
+    } else {
+      WaitForThreadpoolTimerCallbacks(/*pti=*/timer,
+                                      /*fCancelPendingCallbacks=*/TRUE);
+      clean_timer = true;
     }
   }
 
-  task();
+  CloseThreadpoolTimer(/*pti=*/timer);
+  if (is_canceled || clean_timer) {
+    delayed_task_map_.erase(timer);
+  }
 
-  {
-    absl::MutexLock lock(&mutex_);
-    --running_tasks_count_;
-    if (running_tasks_count_ == 0 && shutdown_latch_ != nullptr) {
-      shutdown_latch_->CountDown();
+  VLOG(1) << __func__ << ": Completed to cancel timer(" << timer << ").";
+  return true;
+}
+
+void ThreadPool::TaskQueue::enqueue(Runnable task) {
+  absl::MutexLock lock(&mutex_);
+  queue_.push(std::move(task));
+}
+
+std::optional<Runnable> ThreadPool::TaskQueue::dequeue() {
+  absl::MutexLock lock(&mutex_);
+  if (queue_.empty()) {
+    return std::nullopt;
+  }
+
+  Runnable value = std::move(queue_.front());
+  queue_.pop();
+  return std::move(value);
+}
+
+void ThreadPool::TaskQueue::clear() {
+  absl::MutexLock lock(&mutex_);
+  queue_ = {};
+}
+
+void ThreadPool::DelayedTaskMap::put(
+    PTP_TIMER timer, std::unique_ptr<DelayedTaskInfo> task_info) {
+  absl::MutexLock lock(&mutex_);
+  task_map_[timer] = std::move(task_info);
+}
+
+ThreadPool::DelayedTaskInfo* ThreadPool::DelayedTaskMap::get(PTP_TIMER timer) {
+  absl::MutexLock lock(&mutex_);
+  auto it = task_map_.find(timer);
+  if (it == task_map_.end()) {
+    return nullptr;
+  }
+
+  return it->second.get();
+}
+
+void ThreadPool::DelayedTaskMap::erase(PTP_TIMER timer) {
+  absl::MutexLock lock(&mutex_);
+  task_map_.erase(timer);
+}
+
+void ThreadPool::DelayedTaskMap::clean_completed_tasks() {
+  absl::MutexLock lock(&mutex_);
+  for (auto it = task_map_.begin(); it != task_map_.end();) {
+    if (it->second->is_done) {
+      CloseThreadpoolTimer(it->first);
+      task_map_.erase(it++);
+    } else {
+      ++it;
     }
   }
 }
 
+void ThreadPool::DelayedTaskMap::clear() {
+  absl::MutexLock lock(&mutex_);
+  task_map_.clear();
+}
+
 }  // namespace windows
 }  // namespace nearby

internal/platform/implementation/windows/thread_pool.h

@@ -17,13 +17,17 @@
 
 #include <windows.h>
 
+#include <atomic>
+#include <cstdint>
 #include <memory>
+#include <optional>
 #include <queue>
 #include <utility>
 
 #include "absl/base/thread_annotations.h"
+#include "absl/container/flat_hash_map.h"
 #include "absl/synchronization/mutex.h"
-#include "internal/platform/implementation/shared/count_down_latch.h"
+#include "absl/time/time.h"
 #include "internal/platform/runnable.h"
 
 namespace nearby {
@@ -32,27 +36,98 @@ namespace windows {
 class ThreadPool {
  public:
   virtual ~ThreadPool();
-  static std::unique_ptr<ThreadPool> Create(int max_pool_size);
+  static std::unique_ptr<ThreadPool> Create(uint32_t max_pool_size);
 
   // Runs a task on thread pool. The result indicates whether the task is put
   // into the thread pool.
   bool Run(Runnable task) ABSL_LOCKS_EXCLUDED(mutex_);
 
+  // Runs a task on thread pool with a delay. A delayed task ID is returned if
+  // it is put into the thread pool. The ID can be used to cancel the task.
+  std::optional<uint64_t> Run(Runnable task, absl::Duration delay)
+      ABSL_LOCKS_EXCLUDED(mutex_);
+
+  // Runs a task on thread pool with a delay and period. A delayed task ID is
+  // returned if it is put into the thread pool. The ID can be used to cancel
+  // the task.
+  std::optional<uint64_t> Run(Runnable task, absl::Duration delay,
+                              absl::Duration period)
+      ABSL_LOCKS_EXCLUDED(mutex_);
+
+  // Cancels a delayed task. It will return after running callback is completed.
+  bool CancelDelayedTask(uint64_t delayed_task_id) ABSL_LOCKS_EXCLUDED(mutex_);
+
   // In Nearby platform, thread pool should make sure all queued tasks completed
   // in shut down.
-  void ShutDown();
+  void ShutDown() ABSL_LOCKS_EXCLUDED(mutex_);
 
  private:
+  struct DelayedTaskInfo {
+    Runnable task;
+    absl::Duration delay;
+    absl::Duration period;
+    // Thread id 0 is an invalid value in windows.
+    std::atomic_uint32_t thread_id = 0;
+    std::atomic_bool is_done = false;
+    std::atomic_bool is_canceled = false;
+
+    DelayedTaskInfo(Runnable task, absl::Duration delay,
+                    absl::Duration period) {
+      this->task = std::move(task);
+      this->delay = delay;
+      this->period = period;
+    }
+
+    DelayedTaskInfo(const DelayedTaskInfo&) = delete;
+    DelayedTaskInfo& operator=(const DelayedTaskInfo&) = delete;
+  };
+
+  class TaskQueue {
+   public:
+    TaskQueue() = default;
+    ~TaskQueue() = default;
+
+    void enqueue(Runnable task) ABSL_LOCKS_EXCLUDED(mutex_);
+    std::optional<Runnable> dequeue() ABSL_LOCKS_EXCLUDED(mutex_);
+    void clear() ABSL_LOCKS_EXCLUDED(mutex_);
+
+   private:
+    absl::Mutex mutex_;
+    std::queue<Runnable> queue_ ABSL_GUARDED_BY(mutex_);
+  };
+
+  class DelayedTaskMap {
+   public:
+    DelayedTaskMap() = default;
+    ~DelayedTaskMap() = default;
+
+    void put(PTP_TIMER timer, std::unique_ptr<DelayedTaskInfo> task_info)
+        ABSL_LOCKS_EXCLUDED(mutex_);
+    DelayedTaskInfo* get(PTP_TIMER timer) ABSL_LOCKS_EXCLUDED(mutex_);
+    void erase(PTP_TIMER timer) ABSL_LOCKS_EXCLUDED(mutex_);
+    void clean_completed_tasks() ABSL_LOCKS_EXCLUDED(mutex_);
+    void clear() ABSL_LOCKS_EXCLUDED(mutex_);
+
+   private:
+    absl::Mutex mutex_;
+    absl::flat_hash_map<PTP_TIMER, std::unique_ptr<DelayedTaskInfo>> task_map_
+        ABSL_GUARDED_BY(mutex_);
+  };
+
   ThreadPool(PTP_POOL thread_pool, TP_CALLBACK_ENVIRON thread_pool_environ,
-             int max_pool_size);
+             PTP_CLEANUP_GROUP cleanup_group, int max_pool_size);
+
+  static VOID CALLBACK WorkCallback(PTP_CALLBACK_INSTANCE instance,
+                                    PVOID parameter, PTP_WORK work);
+
+  static VOID CALLBACK TimerCallback(PTP_CALLBACK_INSTANCE instance,
+                                     PVOID context, PTP_TIMER Timer);
+
   void RunNextTask();
+  void RunTimerCallback(PTP_TIMER timer);
 
-  // Protects the access to tasks of the thread pool.
-  mutable absl::Mutex mutex_;
-
-  // The task queue of the thread pool. Thread pool will pick up task to run
-  // when it is idle.
-  std::queue<Runnable> tasks_ ABSL_GUARDED_BY(mutex_);
+  // Protects the access to methods of the thread pool.
+  absl::Mutex mutex_;
 
   // Keeps the pointer of the thread pool. It is created when constructing the
   // thread pool.
@@ -61,17 +136,21 @@ class ThreadPool {
   // Keeps the environment of the thread pool.
   TP_CALLBACK_ENVIRON thread_pool_environ_ ABSL_GUARDED_BY(mutex_);
 
+  // The cleanup group of the thread pool.
+  PTP_CLEANUP_GROUP cleanup_group_ ABSL_GUARDED_BY(mutex_) = nullptr;
+
   // The maximum thread count in the thread pool
   int max_pool_size_ ABSL_GUARDED_BY(mutex_) = 0;
 
-  // Current running task count
-  int running_tasks_count_ ABSL_GUARDED_BY(mutex_) = 0;
+  std::atomic_bool is_shut_down_ = false;
 
-  // The latch is used to wait for running tasks
-  std::unique_ptr<shared::CountDownLatch> shutdown_latch_ = nullptr;
+  // The task queue of the thread pool. Thread pool will pick up task to run
+  // when it is idle.
+  TaskQueue task_queue_;
 
-  friend VOID CALLBACK WorkCallback(PTP_CALLBACK_INSTANCE instance,
-                                    PVOID parameter, PTP_WORK work);
+  // The delayed task map of the thread pool. Thread pool will pick up delayed
+  // task to run when it is idle.
+  DelayedTaskMap delayed_task_map_;
 };
 
 }  // namespace windows

internal/platform/implementation/windows/thread_pool_test.cc

@@ -15,6 +15,8 @@
 #include "internal/platform/implementation/windows/thread_pool.h"
 
 #include <atomic>
+#include <cstdint>
+#include <optional>
 #include <vector>
 
 #include "gtest/gtest.h"
@@ -79,6 +81,114 @@ TEST(ThreadPool, ShutdownWaitsForRunningTasks) {
   EXPECT_EQ(value, 1);
 }
 
+TEST(ThreadPool, RunNoDelayedTask) {
+  auto pool = ThreadPool::Create(1);
+  std::atomic_int value = 0;
+  std::optional<uint64_t> delayed_task_id =
+      pool->Run([&]() { value += 1; }, absl::ZeroDuration());
+  EXPECT_TRUE(delayed_task_id.has_value());
+  absl::SleepFor(absl::Milliseconds(200));
+  EXPECT_EQ(value, 1);
+  pool->ShutDown();
+}
+
+TEST(ThreadPool, RunDelayedTask) {
+  auto pool = ThreadPool::Create(1);
+  std::atomic_int value = 0;
+  std::optional<uint64_t> delayed_task_id =
+      pool->Run([&]() { value += 1; }, absl::Seconds(1));
+  EXPECT_TRUE(delayed_task_id.has_value());
+  absl::SleepFor(absl::Milliseconds(200));
+  EXPECT_EQ(value, 0);
+  absl::SleepFor(absl::Milliseconds(1000));
+  EXPECT_EQ(value, 1);
+  pool->ShutDown();
+}
+
+TEST(ThreadPool, CancelDelayedTask) {
+  auto pool = ThreadPool::Create(1);
+  std::atomic_int value = 0;
+  std::optional<uint64_t> delayed_task_id =
+      pool->Run([&]() { value += 1; }, absl::Seconds(1));
+  EXPECT_TRUE(delayed_task_id.has_value());
+  absl::SleepFor(absl::Milliseconds(200));
+  EXPECT_TRUE(pool->CancelDelayedTask(delayed_task_id.value()));
+  absl::SleepFor(absl::Milliseconds(1000));
+  EXPECT_EQ(value, 0);
+  pool->ShutDown();
+}
+
+TEST(ThreadPool, CancelRunningDelayedTask) {
+  auto pool = ThreadPool::Create(1);
+  std::atomic_int value = 0;
+  std::optional<uint64_t> delayed_task_id = pool->Run(
+      [&]() {
+        value += 1;
+        absl::SleepFor(absl::Milliseconds(1000));
+      },
+      absl::Milliseconds(100));
+  EXPECT_TRUE(delayed_task_id.has_value());
+  absl::SleepFor(absl::Milliseconds(300));
+  EXPECT_TRUE(pool->CancelDelayedTask(delayed_task_id.value()));
+  EXPECT_EQ(value, 1);
+  pool->ShutDown();
+}
+
+TEST(ThreadPool, CancelDelayedTaskAfterShutDown) {
+  auto pool = ThreadPool::Create(1);
+  std::atomic_int value = 0;
+  std::optional<uint64_t> delayed_task_id =
+      pool->Run([&]() { value += 1; }, absl::Seconds(1));
+  absl::SleepFor(absl::Milliseconds(200));
+  pool->ShutDown();
+  EXPECT_FALSE(pool->CancelDelayedTask(delayed_task_id.value()));
+}
+
+TEST(ThreadPool, CancelDelayedTaskInCallback) {
+  auto pool = ThreadPool::Create(1);
+  std::atomic_int value = 0;
+  std::optional<uint64_t> delayed_task_id;
+  delayed_task_id = pool->Run(
+      [&]() {
+        value += 1;
+        pool->CancelDelayedTask(delayed_task_id.value());
+      },
+      absl::Milliseconds(100));
+  EXPECT_TRUE(delayed_task_id.has_value());
+  absl::SleepFor(absl::Milliseconds(500));
+  EXPECT_EQ(value, 1);
+  pool->ShutDown();
+}
+
+TEST(ThreadPool, RunRepeatedTask) {
+  auto pool = ThreadPool::Create(1);
+  std::atomic_int value = 0;
+  std::optional<uint64_t> delayed_task_id;
+  delayed_task_id = pool->Run([&]() { value += 1; }, absl::Milliseconds(500),
+                              absl::Milliseconds(500));
+  EXPECT_TRUE(delayed_task_id.has_value());
+  absl::SleepFor(absl::Milliseconds(1200));
+  pool->CancelDelayedTask(delayed_task_id.value());
+  EXPECT_EQ(value, 2);
+  pool->ShutDown();
+}
+
+TEST(ThreadPool, CancelRepeatedTaskInCallback) {
+  auto pool = ThreadPool::Create(1);
+  std::atomic_int value = 0;
+  std::optional<uint64_t> delayed_task_id;
+  delayed_task_id = pool->Run(
+      [&]() {
+        value += 1;
+        pool->CancelDelayedTask(delayed_task_id.value());
+      },
+      absl::Milliseconds(500), absl::Milliseconds(500));
+  EXPECT_TRUE(delayed_task_id.has_value());
+  absl::SleepFor(absl::Milliseconds(2000));
+  EXPECT_EQ(value, 1);
+  pool->ShutDown();
+}
+
 }  // namespace
 }  // namespace windows
 }  // namespace nearby