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 (!scoped_com_initializer_)
+      scoped_com_initializer_ = MakeUnique<win::ScopedCOMInitializer>();
+  }
+#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_;
+#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 52 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