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

base/test/scoped_task_scheduler.cc

 // Copyright 2016 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "base/test/scoped_task_scheduler.h"
 
 #include <memory>
 #include <utility>
 #include <vector>
 
 #include "base/bind.h"
 #include "base/logging.h"
 #include "base/memory/ptr_util.h"
 #include "base/memory/ref_counted.h"
 #include "base/message_loop/message_loop.h"
 #include "base/metrics/histogram_base.h"
 #include "base/run_loop.h"
 #include "base/sequence_token.h"
 #include "base/sequenced_task_runner.h"
 #include "base/single_thread_task_runner.h"
 #include "base/task_runner.h"
 #include "base/task_scheduler/task.h"
 #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 {
 
 namespace {
 
 enum class ExecutionMode { PARALLEL, SEQUENCED, SINGLE_THREADED };
 
 class TestTaskScheduler : public TaskScheduler {
  public:
   // |external_message_loop| is an externally provided MessageLoop on which to
   // run tasks. A MessageLoop will be created by TestTaskScheduler if
   // |external_message_loop| is nullptr.
   explicit TestTaskScheduler(MessageLoop* external_message_loop);
   ~TestTaskScheduler() override;
 
   // TaskScheduler:
   void PostDelayedTaskWithTraits(const tracked_objects::Location& from_here,
                                  const TaskTraits& traits,
                                  const Closure& task,
                                  TimeDelta delay) override;
   scoped_refptr<TaskRunner> CreateTaskRunnerWithTraits(
       const TaskTraits& traits) override;
   scoped_refptr<SequencedTaskRunner> CreateSequencedTaskRunnerWithTraits(
       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;
   void Shutdown() override;
   void FlushForTesting() override;
   void JoinForTesting() override;
 
   // Posts |task| to this TaskScheduler with |sequence_token|. Returns true on
   // success.
   bool PostTask(std::unique_ptr<internal::Task> task,
@@ skipping 10 matching lines @@
   // Returns the TaskRunner to which this TaskScheduler forwards tasks. It may
   // be |message_loop_->task_runner()| or a reference to it saved on entry to
   // RunTask().
   scoped_refptr<SingleThreadTaskRunner> MessageLoopTaskRunner() const {
     if (saved_task_runner_)
       return saved_task_runner_;
     DCHECK(message_loop_->task_runner());
     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
   // run tasks.
   std::unique_ptr<MessageLoop> message_loop_owned_;
   MessageLoop* message_loop_;
 
   // A reference to |message_loop_->task_runner()| saved on entry to RunTask().
   // This is required because RunTask() overrides
   // |message_loop_->task_runner()|.
   //
   // Note: |message_loop_->task_runner()| is accessed directly outside of
   // RunTask() to guarantee that ScopedTaskScheduler always uses the latest
   // 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_;
 
   DISALLOW_COPY_AND_ASSIGN(TestTaskScheduler);
 };
 
 class TestTaskSchedulerTaskRunner : public SingleThreadTaskRunner {
  public:
   TestTaskSchedulerTaskRunner(TestTaskScheduler* task_scheduler,
                               ExecutionMode execution_mode,
@@ skipping 19 matching lines @@
   DISALLOW_COPY_AND_ASSIGN(TestTaskSchedulerTaskRunner);
 };
 
 TestTaskScheduler::TestTaskScheduler(MessageLoop* external_message_loop)
     : message_loop_owned_(external_message_loop ? nullptr
                                                 : MakeUnique<MessageLoop>()),
       message_loop_(message_loop_owned_ ? message_loop_owned_.get()
                                         : 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(
     const tracked_objects::Location& from_here,
     const TaskTraits& traits,
     const Closure& task,
     TimeDelta delay) {
   CreateTaskRunnerWithTraits(traits)->PostDelayedTask(from_here, task, delay);
 }
 
@@ skipping 10 matching lines @@
       new TestTaskSchedulerTaskRunner(this, ExecutionMode::SEQUENCED, traits));
 }
 
 scoped_refptr<SingleThreadTaskRunner>
 TestTaskScheduler::CreateSingleThreadTaskRunnerWithTraits(
     const TaskTraits& traits) {
   return make_scoped_refptr(new TestTaskSchedulerTaskRunner(
       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*>();
 }
 
 int TestTaskScheduler::GetMaxConcurrentTasksWithTraitsDeprecated(
     const TaskTraits& traits) const {
   return 1;
 }
 
@@ skipping 106 matching lines @@
   DCHECK_EQ(task_scheduler_, TaskScheduler::GetInstance());
 
   // Per contract, call JoinForTesting() before deleting the TaskScheduler.
   TaskScheduler::GetInstance()->JoinForTesting();
 
   TaskScheduler::SetInstance(nullptr);
 }
 
 }  // namespace test
 }  // namespace base