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..cdd7fc040eabb10e3972bb883e33ec477d7e249f 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_;
 
@@ -135,9 +174,28 @@ TestTaskScheduler::TestTaskScheduler(MessageLoop* external_message_loop)
                                         : external_message_loop) {}
 
 TestTaskScheduler::~TestTaskScheduler() {
+#if defined(OS_WIN)
+  PostTask(MakeUnique<internal::Task>(

[fdoray, 2017/03/27 20:11:50]

If new tasks are posted during shutdown to a COM TaskRunner, they will run after
COM uninitialization. Do we need to explicitly reset |scoped_com_inializer_|? We
could let the member be destroyed automatically.

[robliao, 2017/03/27 20:28:05]

For correctness, this task does need to be posted to the message loop (someone
could have sent us a message loop running on a different thread, right?). Maybe
a message loop destruction observer is the way to go here?

[robliao, 2017/03/27 21:32:32]

Simpler: Run another task after shutdown. That's what I went with here.

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

ScopedTaskScheduler is not supposed to be used with a MessageLoop that belongs
to another thread
https://cs.chromium.org/chromium/src/base/test/scoped_task_scheduler.h?sq=pac...
The description could probably be clearer about that.

+               FROM_HERE,
+               Bind(
+                   [](std::unique_ptr<win::ScopedCOMInitializer>*
+                          scoped_com_initializer) {
+                     scoped_com_initializer->reset();
+                   },
+                   &scoped_com_initializer_),
+               TaskTraits().WithShutdownBehavior(
+                   TaskShutdownBehavior::BLOCK_SHUTDOWN),
+               TimeDelta()),
+           SequenceToken());
+#endif  // defined(OS_WIN)
+
   // Shutdown if it hasn't already been done explicitly.
   if (!task_tracker_.HasShutdownStarted())
     Shutdown();
+
+#if defined(OS_WIN)
+  DCHECK(!scoped_com_initializer_);
+#endif  // defined(OS_WIN)
 }
 
 void TestTaskScheduler::PostDelayedTaskWithTraits(
@@ -168,6 +226,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*>();