FOR REFERENCE ONLY Task Scheduler COM Task Runner

base/task_scheduler/scheduler_single_thread_task_runner_manager.cc

Index: base/task_scheduler/scheduler_single_thread_task_runner_manager.cc
diff --git a/base/task_scheduler/scheduler_single_thread_task_runner_manager.cc b/base/task_scheduler/scheduler_single_thread_task_runner_manager.cc
index 6a2676fb49b9e20bb8f06ea8b52a4839a76a8922..379b03112b1fc6a2a229fdf3ad9a21b4d620effd 100644
--- a/base/task_scheduler/scheduler_single_thread_task_runner_manager.cc
+++ b/base/task_scheduler/scheduler_single_thread_task_runner_manager.cc
@@ -22,6 +22,13 @@
 #include "base/task_scheduler/task_traits.h"
 #include "base/threading/platform_thread.h"
 #include "base/time/time.h"
+#include "build/build_config.h"
+
+#if defined(OS_WIN)
+#include <windows.h>
+
+#include "base/win/scoped_com_initializer.h"
+#endif
 
 namespace base {
 namespace internal {
@@ -129,6 +136,65 @@ class SchedulerWorkerDelegate : public SchedulerWorker::Delegate {
   DISALLOW_COPY_AND_ASSIGN(SchedulerWorkerDelegate);
 };
 
+#if defined(OS_WIN)
+
+void PumpOneMessage() {
+  MSG msg;
+  if (PeekMessage(&msg, nullptr, 0, 0, PM_REMOVE) != FALSE) {
+    TranslateMessage(&msg);
+    DispatchMessage(&msg);
+  }
+}
+
+class SchedulerWorkerCOMDelegate : public SchedulerWorkerDelegate {
+ public:
+  SchedulerWorkerCOMDelegate(const std::string& thread_name)
+      : SchedulerWorkerDelegate(thread_name) {}
+
+  // SchedulerWorker::Delegate:
+  void OnMainEntry(SchedulerWorker* worker) override {
+    SchedulerWorkerDelegate::OnMainEntry(worker);
+
+    scoped_com_initializer_ = MakeUnique<win::ScopedCOMInitializer>();
+  }
+
+  scoped_refptr<Sequence> GetWork(SchedulerWorker* worker) override {
+    pump_one_message_ = !pump_one_message_;
+    if (pump_one_message_) {
+      message_pump_sequence_->PushTask(MakeUnique<Task>(
+          FROM_HERE, base::Bind(&PumpOneMessage), TaskTraits().MayBlock(),
+          TimeDelta(), Task::TaskType::INTERNAL));
+      return message_pump_sequence_;
+    }
+    return SchedulerWorkerDelegate::GetWork(worker);
+  }
+
+  void OnMainExit() override { scoped_com_initializer_.reset(); }
+
+  void WaitForWork(WaitableEvent* wake_up_event) override {
+    DCHECK(wake_up_event);
+    const TimeDelta sleep_time = GetSleepTimeout();
+    const DWORD milliseconds_wait =
+        sleep_time.is_max() ? INFINITE : sleep_time.InMilliseconds();
+    HANDLE wake_up_event_handle = wake_up_event->handle();
+    DWORD result = MsgWaitForMultipleObjectsEx(
+        1, &wake_up_event_handle, milliseconds_wait, QS_ALLEVENTS, 0);
+    if (result == WAIT_OBJECT_0) {
+      // Reset the event since we woke up due to it.
+      wake_up_event->Reset();
+    }
+  }
+
+ private:
+  bool pump_one_message_ = true;
+  scoped_refptr<Sequence> message_pump_sequence_ = new Sequence;
+  std::unique_ptr<win::ScopedCOMInitializer> scoped_com_initializer_;
+
+  DISALLOW_COPY_AND_ASSIGN(SchedulerWorkerCOMDelegate);
+};
+
+#endif  // defined(OS_WIN)
+
 }  // namespace
 
 class SchedulerSingleThreadTaskRunnerManager::SchedulerSingleThreadTaskRunner
@@ -241,7 +307,19 @@ SchedulerSingleThreadTaskRunnerManager::CreateSingleThreadTaskRunnerWithTraits(
   DCHECK_LT(index, worker_pool_params_vector_.size());
   return new SchedulerSingleThreadTaskRunner(
       this, traits,
-      CreateAndRegisterSchedulerWorker(worker_pool_params_vector_[index]));
+      CreateAndRegisterSchedulerWorker(worker_pool_params_vector_[index],
+                                       DelegateType::REGULAR));
+}
+
+scoped_refptr<SingleThreadTaskRunner>
+SchedulerSingleThreadTaskRunnerManager::CreateCOMSTATaskRunnerWithTraits(
+    const TaskTraits& traits) {
+  size_t index = worker_pool_index_for_traits_callback_.Run(traits);
+  DCHECK_LT(index, worker_pool_params_vector_.size());
+  return new SchedulerSingleThreadTaskRunner(
+      this, traits,
+      CreateAndRegisterSchedulerWorker(worker_pool_params_vector_[index],
+                                       DelegateType::COM_STA));
 }
 
 void SchedulerSingleThreadTaskRunnerManager::JoinForTesting() {
@@ -264,11 +342,19 @@ void SchedulerSingleThreadTaskRunnerManager::JoinForTesting() {
 
 SchedulerWorker*
 SchedulerSingleThreadTaskRunnerManager::CreateAndRegisterSchedulerWorker(
-    const SchedulerWorkerPoolParams& params) {
+    const SchedulerWorkerPoolParams& params,
+    DelegateType delegate_type) {
   AutoSchedulerLock auto_lock(workers_lock_);
   int id = next_worker_id_++;
-  auto delegate = MakeUnique<SchedulerWorkerDelegate>(base::StringPrintf(
-      "TaskSchedulerSingleThreadWorker%d%s", id, params.name().c_str()));
+
+  auto delegate =
+      delegate_type == DelegateType::COM_STA
+          ? MakeUnique<SchedulerWorkerCOMDelegate>(
+                base::StringPrintf("TaskSchedulerSingleThreadWorker%d%sCOMSTA",
+                                   id, params.name().c_str()))
+          : MakeUnique<SchedulerWorkerDelegate>(
+                base::StringPrintf("TaskSchedulerSingleThreadWorker%d%s", id,
+                                   params.name().c_str()));
   workers_.emplace_back(SchedulerWorker::Create(
       params.priority_hint(), std::move(delegate), task_tracker_,
       SchedulerWorker::InitialState::DETACHED));