Add the COM STA Task Runner

base/task_scheduler/scheduler_single_thread_task_runner_manager.cc

 // Copyright 2017 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/task_scheduler/scheduler_single_thread_task_runner_manager.h"
 
 #include <algorithm>
 #include <memory>
 #include <string>
 
 #include "base/bind.h"
 #include "base/callback.h"
 #include "base/memory/ptr_util.h"
 #include "base/single_thread_task_runner.h"
 #include "base/strings/stringprintf.h"
 #include "base/synchronization/atomic_flag.h"
 #include "base/task_scheduler/delayed_task_manager.h"
 #include "base/task_scheduler/scheduler_worker.h"
 #include "base/task_scheduler/sequence.h"
 #include "base/task_scheduler/task.h"
 #include "base/task_scheduler/task_tracker.h"
 #include "base/task_scheduler/task_traits.h"
 #include "base/threading/platform_thread.h"
 #include "base/time/time.h"
+#include "build/build_config.h"

[fdoray, 2017/03/21 15:23:12]

Not needed if you include it in the .h.

[robliao, 2017/03/21 20:00:19]

Done.

+
+#if defined(OS_WIN)
+#include <windows.h>
+
+#include "base/win/scoped_com_initializer.h"
+#endif

[gab, 2017/03/21 21:09:34]

#endif  // defined(OS_WIN)

[robliao, 2017/03/21 22:25:32]

Ah, I thought we didn't want these for certain short scoped things.
Added it back in.

[gab, 2017/03/22 16:16:04]

Yeah, it's kind of an hand-wavy rule, for tight scopes yes, though I'm stricter
on #endif comments than closing namespace comments because of lack of
indentation making it quickly hard to read with preprocessor macros.

Here the empty line made me lean towards preferring the closing comment (I never
oppose to closing comment for preprocessor macros FWIW because I find them so
hard to read in general).

 
 namespace base {
 namespace internal {
 
 namespace {
 
 // Allows for checking the PlatformThread::CurrentRef() against a set
 // PlatformThreadRef atomically without using locks.
 class AtomicThreadRefChecker {
  public:
@@ skipping 87 matching lines @@
   // Synchronizes access to |sequence_| and |has_work_|.
   SchedulerLock sequence_lock_;
   scoped_refptr<Sequence> sequence_ = new Sequence;
   bool has_work_ = false;
 
   AtomicThreadRefChecker thread_ref_checker_;
 
   DISALLOW_COPY_AND_ASSIGN(SchedulerWorkerDelegate);
 };
 
+#if defined(OS_WIN)
+
+class SchedulerWorkerCOMDelegate : public SchedulerWorkerDelegate {
+ public:
+  SchedulerWorkerCOMDelegate(const std::string& thread_name,
+                             TaskTracker* task_tracker)
+      : SchedulerWorkerDelegate(thread_name), task_tracker_(task_tracker) {}

[fdoray, 2017/03/21 15:23:12]

~SchedulerWorkerCOMDelegate() {
  DCHECK(!scoped_comp_initializer_);
}

[robliao, 2017/03/21 20:00:20]

Done.

+
+  // SchedulerWorker::Delegate:
+  void OnMainEntry(SchedulerWorker* worker) override {
+    SchedulerWorkerDelegate::OnMainEntry(worker);
+
+    scoped_com_initializer_ = MakeUnique<win::ScopedCOMInitializer>();
+  }
+
+  scoped_refptr<Sequence> GetWork(SchedulerWorker* worker) override {
+    // This scheme below allows us to cover the following scenarios:
+    // * Tasks only come from SchedulerWorkerDelegate::GetWork():

[gab, 2017/03/21 21:09:34]

s/only come/are only coming/ 

here and below? Feels this tense works better for me?

[robliao, 2017/03/21 22:25:32]

Rephrased
    // This scheme below allows us to cover the following scenarios:
    // * Only SchedulerWorkerDelegate::GetWork() has work:
    //   Always return the sequence from GetWork().
    // * Only the Windows Message Queue has work:
    //   Always return the sequence from GetWorkFromWindowsMessageQueue();
    // * Both SchedulerWorkerDelegate::GetWork() and the Windows Message Queue
    //   have work:
    //   Process sequences from each source round-robin style.

+    //   Only return the sequence from GetWork().
+    // * Tasks only come from the Windows Message Queue:
+    //   Only return the sequence from GetWorkFromWindowsMessageQueue();
+    // * Tasks come from both SchedulerWorkerDelegate::GetWork() and
+    //   the Windows Message Queue:
+    //   Process sequences from each source round-robin style.
+    scoped_refptr<Sequence> sequence;
+    if (get_work_first_) {
+      sequence = SchedulerWorkerDelegate::GetWork(worker);
+      if (sequence)
+        get_work_first_ = false;
+    }
+
+    if (!sequence) {
+      sequence = GetWorkFromWindowsMessageQueue();
+      if (sequence)
+        get_work_first_ = true;
+    }
+
+    if (!sequence && !get_work_first_) {
+      // This case is important if we checked the Windows Message Queue first
+      // and found there was no work. We don't want to return null immediately
+      // as that could cause the thread to go to sleep while work is waiting via
+      // SchedulerWorkerDelegate::GetWork(). As the same time, we don't want to

[gab, 2017/03/21 21:09:33]

s/As/At/

(but overall I don't think this last sentence is even required, it's pretty
obvious, just like we don't comment when we set |get_work_first_| above)

[robliao, 2017/03/21 22:25:32]

Removed.

+      // mark |get_work_first_| to continue to check the message queue first
+      // after this sequence is returned.
+      sequence = SchedulerWorkerDelegate::GetWork(worker);
+    }
+    return sequence;
+  }
+
+  void OnMainExit() override { scoped_com_initializer_.reset(); }
+
+  void WaitForWork(WaitableEvent* wake_up_event) override {
+    DCHECK(wake_up_event);
+    const TimeDelta sleep_time = GetSleepTimeout();
+    const DWORD milliseconds_wait =
+        sleep_time.is_max() ? INFINITE : sleep_time.InMilliseconds();
+    HANDLE wake_up_event_handle = wake_up_event->handle();
+    DWORD result = MsgWaitForMultipleObjectsEx(
+        1, &wake_up_event_handle, milliseconds_wait, QS_ALLINPUT, 0);
+    if (result == WAIT_OBJECT_0) {
+      // Reset the event since we woke up due to it.
+      wake_up_event->Reset();
+    }
+  }

[fdoray, 2017/03/21 15:23:12]

void ReEnqueueSequence(scoped_refptr<Sequence> sequence) override {
  DCHECK_NE(sequence, message_pump_sequence_);
  SchedulerWorkerDelegate::ReEnqueueSequence(std::move(sequence));
}

[robliao, 2017/03/21 20:00:19]

This check is already covered by SchedulerWorkerDelegate::ReEnqueueSequence
sequence check.

+
+ private:
+  scoped_refptr<Sequence> GetWorkFromWindowsMessageQueue() {
+    MSG msg;
+    if (PeekMessage(&msg, nullptr, 0, 0, PM_REMOVE) != FALSE) {
+      auto pump_message_task =
+          MakeUnique<Task>(FROM_HERE,
+                           base::Bind(
+                               [](MSG msg_in) {
+                                 MSG msg = msg_in;

[gab, 2017/03/21 21:09:34]

Why do you need this extra variable?

[robliao, 2017/03/21 22:25:32]

Because in the shuffle I lost a "const MSG& msg_in" :-).

In the past, I received feedback we didn't like non "const MSG&" even though we
were going mutate them immediately (although I can't seem to find the CR from
years ago). 

I'd actually prefer to do just MSG msg, and if you're fine with that, we'll go
with that.
The C++ style guide doesn't have a preference either way other than consistency.
https://google.github.io/styleguide/cppguide.html#Function_Parameter_Ordering

[gab, 2017/03/22 16:16:04]

sgtm

+                                 TranslateMessage(&msg);
+                                 DispatchMessage(&msg);
+                               },
+                               msg),

[gab, 2017/03/21 21:09:34]

std::move(msg) ? Not sure if MSG is moveable but if it is (or becomes so later)
that's a plus and it doesn't hurt otherwise?

Looks like it's a struct (from MSDN) so maybe not but I'll let you decide.

[robliao, 2017/03/21 22:25:32]

sgtm. No qualms with that. It is just a struct underneath the covers.

+                           TaskTraits().MayBlock(), TimeDelta());
+      if (task_tracker_->WillPostTask(pump_message_task.get())) {

[gab, 2017/03/21 21:09:33]

Otherwise do we have to tell Windows we're dropping this MSG?

[robliao, 2017/03/21 22:25:32]

Generally, the only thing you can do with a message is either process it or drop
it. What you do with it is up to the app.

We actually do drop messages in Chrome here:
https://cs.chromium.org/chromium/src/ui/base/ime/input_method_win.cc?rcl=e739...)

This is basically the same as shutting down with messages remaining in the
queue. Those are effectively dropped. The impact is that callers who used
SendMessage (like our unit test) will hang forever, but they should have been
using SendMessageTimeout anyways.

[gab, 2017/03/22 16:16:04]

Got it, and I guess SendMessage would return with an error or something if the
receiving process/thread dies (which is what would happen soon after TaskTracker
shuts down)?

[robliao, 2017/03/22 17:56:52]

Hrm... that's a good question. One of the things about SendMessage is you have
to make sure the target hwnd's thread is alive and pumping messages. I haven't
encountered a case where the thread dies during the call. I might have to try
this out sometime.

In either case, the caller needs to be using SendMessageTimeout.

+        message_pump_sequence_->PushTask(std::move(pump_message_task));

[fdoray, 2017/03/21 15:23:12]

bool was_empty = message_pump_sequence_->PushTask(...);
DCHECK(was_empty);

[robliao, 2017/03/21 20:00:19]

Done.

+        return message_pump_sequence_;
+      }
+    }
+    return nullptr;
+  }
+
+  bool get_work_first_ = true;
+  scoped_refptr<Sequence> message_pump_sequence_ = new Sequence;

[robliao, 2017/03/21 20:00:19]

Done.

+  TaskTracker* const task_tracker_;
+  std::unique_ptr<win::ScopedCOMInitializer> scoped_com_initializer_;
+
+  DISALLOW_COPY_AND_ASSIGN(SchedulerWorkerCOMDelegate);
+};
+
+#endif  // defined(OS_WIN)
+
 }  // namespace
 
 class SchedulerSingleThreadTaskRunnerManager::SchedulerSingleThreadTaskRunner
     : public SingleThreadTaskRunner {
  public:
   // Constructs a SchedulerSingleThreadTaskRunner that indirectly controls the
   // lifetime of a dedicated |worker| for |traits|.
   SchedulerSingleThreadTaskRunner(
       SchedulerSingleThreadTaskRunnerManager* const outer,
       const TaskTraits& traits,
@@ skipping 92 matching lines @@
 #endif
 }
 
 scoped_refptr<SingleThreadTaskRunner>
 SchedulerSingleThreadTaskRunnerManager::CreateSingleThreadTaskRunnerWithTraits(
     const TaskTraits& traits) {
   size_t index = worker_pool_index_for_traits_callback_.Run(traits);
   DCHECK_LT(index, worker_pool_params_vector_.size());
   return new SchedulerSingleThreadTaskRunner(
       this, traits,
-      CreateAndRegisterSchedulerWorker(worker_pool_params_vector_[index]));
+      CreateAndRegisterSchedulerWorker(worker_pool_params_vector_[index],
+                                       DelegateType::REGULAR));
 }
 
+#if defined(OS_WIN)
+scoped_refptr<SingleThreadTaskRunner>
+SchedulerSingleThreadTaskRunnerManager::CreateCOMSTATaskRunnerWithTraits(
+    const TaskTraits& traits) {
+  size_t index = worker_pool_index_for_traits_callback_.Run(traits);
+  DCHECK_LT(index, worker_pool_params_vector_.size());
+  return new SchedulerSingleThreadTaskRunner(
+      this, traits,
+      CreateAndRegisterSchedulerWorker(worker_pool_params_vector_[index],
+                                       DelegateType::COM_STA));
+}
+#endif  // defined(OS_WIN)
+
 void SchedulerSingleThreadTaskRunnerManager::JoinForTesting() {
   decltype(workers_) local_workers;
   {
     AutoSchedulerLock auto_lock(workers_lock_);
     local_workers = std::move(workers_);
   }
 
   for (const auto& worker : local_workers)
     worker->JoinForTesting();
 
   {
     AutoSchedulerLock auto_lock(workers_lock_);
     DCHECK(workers_.empty())
         << "New worker(s) unexpectedly registered during join.";
     workers_ = std::move(local_workers);
   }
 }
 
 SchedulerWorker*
 SchedulerSingleThreadTaskRunnerManager::CreateAndRegisterSchedulerWorker(
-    const SchedulerWorkerPoolParams& params) {
+    const SchedulerWorkerPoolParams& params,
+    DelegateType delegate_type) {
   AutoSchedulerLock auto_lock(workers_lock_);
   int id = next_worker_id_++;
-  auto delegate = MakeUnique<SchedulerWorkerDelegate>(base::StringPrintf(
-      "TaskSchedulerSingleThreadWorker%d%s", id, params.name().c_str()));
+
+  std::unique_ptr<SchedulerWorkerDelegate> delegate;

[fdoray, 2017/03/21 15:23:12]

Optional:
- Instantiate the right type of delegate from
CreateSingleThreadTaskRunnerWithTraits/CreateCOMSTATaskRunnerWithTraits and pass
it to CreateAndRegisterSchedulerWorker in lieu of the the enum.
- Use an AtomicSequenceNumber for the id (to avoid having to lock when
instantiating the delegate from
CreateSingleThreadTaskRunnerWithTraits/CreateCOMSTATaskRunnerWithTraits).
- Get rid of the enum.

I like to avoid enums when possible.

[robliao, 2017/03/21 20:00:19]

I guess this comes down to how much duplication we want.

We either duplicate the increment and emplace_back calls or we use the enum to
drive what we want to do.

This, in the future will likely need COM_MTA and/or WINRT. The only thing that
differs here is the delegate, so keeping the current structure seems to make
more sense.

[gab, 2017/03/21 21:09:34]

Or we make CreateAndRegisterSchedulerWorker<T> a templated method?

[robliao, 2017/03/21 22:25:32]

That is an interesting idea, but SchedulerWorkerCOMDelegate takes a
task_tracker_ and SchedulerWorkerDelegate doesn't.
As a result, we would need to specialize (separate functions), create a factory
scheme to package the task_tracker_, or base::Bind a function that creates it.

[robliao, 2017/03/22 08:29:42]

I think I just came up with a scheme to make this work. I'll add it in the
morning.

+  switch (delegate_type) {
+    case DelegateType::REGULAR:
+      delegate = MakeUnique<SchedulerWorkerDelegate>(base::StringPrintf(
+          "TaskSchedulerSingleThreadWorker%d%s", id, params.name().c_str()));
+      break;
+#if defined(OS_WIN)
+    case DelegateType::COM_STA:
+      delegate = MakeUnique<SchedulerWorkerCOMDelegate>(
+          base::StringPrintf("TaskSchedulerSingleThreadWorker%d%sCOMSTA", id,
+                             params.name().c_str()),
+          task_tracker_);
+      break;
+#endif
+  }
+
   workers_.emplace_back(SchedulerWorker::Create(
       params.priority_hint(), std::move(delegate), task_tracker_,
       SchedulerWorker::InitialState::DETACHED));
   return workers_.back().get();
 }
 
 void SchedulerSingleThreadTaskRunnerManager::UnregisterSchedulerWorker(
     SchedulerWorker* worker) {
   // Cleanup uses a SchedulerLock, so call Cleanup() after releasing
   // |workers_lock_|.
@@ skipping 17 matching lines @@
                      });
     DCHECK(worker_iter != workers_.end());
     worker_to_destroy = std::move(*worker_iter);
     workers_.erase(worker_iter);
   }
   worker_to_destroy->Cleanup();
 }
 
 }  // namespace internal
 }  // namespace base