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 22 matching lines @@
 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() {
   // 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(
     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