Add Support to Get a COM STA Task Runner from the Task Scheduler API

base/test/scoped_task_scheduler.cc

Index: base/test/scoped_task_scheduler.cc
diff --git a/base/test/scoped_task_scheduler.cc b/base/test/scoped_task_scheduler.cc
index 61ed2c7730ffd0f30c79ff744ceebe122290db5d..001a4a3acbed4d783e31621ebb859ba56a6a64ef 100644
--- a/base/test/scoped_task_scheduler.cc
+++ b/base/test/scoped_task_scheduler.cc
@@ -23,6 +23,12 @@
 #include "base/task_scheduler/task_scheduler.h"
 #include "base/task_scheduler/task_tracker.h"
 #include "base/threading/thread_task_runner_handle.h"
+#include "base/time/time.h"
+#include "build/build_config.h"
+
+#if defined(OS_WIN)
+#include "base/win/scoped_com_initializer.h"
+#endif  // defined(OS_WIN)
 
 namespace base {
 namespace test {
@@ -50,6 +56,10 @@ class TestTaskScheduler : public TaskScheduler {
       const TaskTraits& traits) override;
   scoped_refptr<SingleThreadTaskRunner> CreateSingleThreadTaskRunnerWithTraits(
       const TaskTraits& traits) override;
+#if defined(OS_WIN)
+  scoped_refptr<SingleThreadTaskRunner> CreateCOMSTATaskRunnerWithTraits(
+      const TaskTraits& traits) override;
+#endif  // defined(OS_WIN)
   std::vector<const HistogramBase*> GetHistograms() const override;
   int GetMaxConcurrentTasksWithTraitsDeprecated(
       const TaskTraits& traits) const override;
@@ -80,6 +90,27 @@ class TestTaskScheduler : public TaskScheduler {
     return message_loop_->task_runner();
   }
 
+#if defined(OS_WIN)
+  void EnsureCOMSTA() {
+    if (!com_init_requested) {
+      com_init_requested = true;
+      DCHECK(!scoped_com_initializer_);
+      PostTask(MakeUnique<internal::Task>(
+                   FROM_HERE,
+                   Bind(
+                       [](std::unique_ptr<win::ScopedCOMInitializer>*
+                              scoped_com_initializer) {
+                         DCHECK(!*scoped_com_initializer);
+                         *scoped_com_initializer =
+                             MakeUnique<win::ScopedCOMInitializer>();
+                       },
+                       &scoped_com_initializer_),
+                   TaskTraits(), TimeDelta()),
+               SequenceToken());
+    }
+  }
+#endif  // defined(OS_WIN)
+
   // |message_loop_owned_| will be non-null if this TestTaskScheduler owns the
   // MessageLoop (wasn't provided an external one at construction).
   // |message_loop_| will always be set and is used by this TestTaskScheduler to
@@ -96,6 +127,14 @@ class TestTaskScheduler : public TaskScheduler {
   // TaskRunner set by external code.
   scoped_refptr<SingleThreadTaskRunner> saved_task_runner_;
 
+#if defined(OS_WIN)
+  // Maintains the lifetime of the COM Single-Threaded Apartment. Allocation and
+  // deallocation should be done in the |message_loop_| via PostTask.
+  std::unique_ptr<win::ScopedCOMInitializer> scoped_com_initializer_;
+
+  bool com_init_requested = false;
+#endif  // defined(OS_WIN)
+
   // Handles shutdown behaviors and sets up the environment to run a task.
   internal::TaskTracker task_tracker_;
 
@@ -138,6 +177,18 @@ TestTaskScheduler::~TestTaskScheduler() {
   // Shutdown if it hasn't already been done explicitly.
   if (!task_tracker_.HasShutdownStarted())
     Shutdown();
+
+#if defined(OS_WIN)
+  MessageLoopTaskRunner()->PostTask(
+      FROM_HERE,
+      Bind([](std::unique_ptr<win::ScopedCOMInitializer>*
+                  scoped_com_initializer) { scoped_com_initializer->reset(); },
+           &scoped_com_initializer_));
+
+  RunLoop().RunUntilIdle();

[fdoray, 2017/03/28 14:14:46]

RunLoop().RunUntilIdle() runs tasks on the current thread. Therefore, if the
MessageLoop runs on another thread, it is not guaranteed that the
scoped_com_initializer->reset() task has run when this returns (DCHECK fails and
use-after-free). Anyways, ScopedTaskScheduler is not intended to be used with a
MessageLoop from another thread (use ScopedAsyncTaskScheduler for async
execution).

[robliao, 2017/03/28 17:15:01]

Yeah, Shutdown() seemed pretty strange to me when I saw
RunLoop().RunUntilIdle(). We can simplify this path. Done.

+
+  DCHECK(!scoped_com_initializer_);
+#endif  // defined(OS_WIN)
 }
 
 void TestTaskScheduler::PostDelayedTaskWithTraits(
@@ -168,6 +219,15 @@ TestTaskScheduler::CreateSingleThreadTaskRunnerWithTraits(
       this, ExecutionMode::SINGLE_THREADED, traits));
 }
 
+#if defined(OS_WIN)
+scoped_refptr<SingleThreadTaskRunner>
+TestTaskScheduler::CreateCOMSTATaskRunnerWithTraits(const TaskTraits& traits) {
+  EnsureCOMSTA();
+  return make_scoped_refptr(new TestTaskSchedulerTaskRunner(
+      this, ExecutionMode::SINGLE_THREADED, traits));
+}
+#endif  // defined(OS_WIN)
+
 std::vector<const HistogramBase*> TestTaskScheduler::GetHistograms() const {
   NOTREACHED();
   return std::vector<const HistogramBase*>();