Fixes race condition in MojoMessagePump

I didn't realize ScheduleWork() can be invoked on any thread.

BUG=359973
TEST=none
R=darin@chromium.org

Review URL: https://codereview.chromium.org/239683002

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@264125 0039d316-1c4b-4281-b951-d872f2087c98

2 files changed, 66 insertions, 43 deletions

diff --git a/mojo/common/message_pump_mojo.cc b/mojo/common/message_pump_mojo.cc
index 2afa72f..5bc3462b 100644
--- a/mojo/common/message_pump_mojo.cc
+++ b/mojo/common/message_pump_mojo.cc
@@ -68,69 +68,84 @@ void MessagePumpMojo::RemoveHandler(const Handle& handle) {
 }
 
 void MessagePumpMojo::Run(Delegate* delegate) {
-  RunState* old_state = run_state_;
   RunState run_state;
   // TODO: better deal with error handling.
   CHECK(run_state.read_handle.is_valid());
   CHECK(run_state.write_handle.is_valid());
-  run_state_ = &run_state;
+  RunState* old_state = NULL;
+  {
+    base::AutoLock auto_lock(run_state_lock_);
+    old_state = run_state_;
+    run_state_ = &run_state;
+  }
+  DoRunLoop(&run_state, delegate);
+  {
+    base::AutoLock auto_lock(run_state_lock_);
+    run_state_ = old_state;
+  }
+}
+
+void MessagePumpMojo::Quit() {
+  base::AutoLock auto_lock(run_state_lock_);
+  if (run_state_)
+    run_state_->should_quit = true;
+}
+
+void MessagePumpMojo::ScheduleWork() {
+  base::AutoLock auto_lock(run_state_lock_);
+  if (run_state_)
+    SignalControlPipe(*run_state_);
+}
+
+void MessagePumpMojo::ScheduleDelayedWork(
+    const base::TimeTicks& delayed_work_time) {
+  base::AutoLock auto_lock(run_state_lock_);
+  if (!run_state_)
+    return;
+  run_state_->delayed_work_time = delayed_work_time;
+  SignalControlPipe(*run_state_);
+}
+
+void MessagePumpMojo::DoRunLoop(RunState* run_state, Delegate* delegate) {
   bool more_work_is_plausible = true;
   for (;;) {
     const bool block = !more_work_is_plausible;
-    DoInternalWork(block);
+    DoInternalWork(*run_state, block);
 
     // There isn't a good way to know if there are more handles ready, we assume
     // not.
     more_work_is_plausible = false;
 
-    if (run_state.should_quit)
+    if (run_state->should_quit)
       break;
 
     more_work_is_plausible |= delegate->DoWork();
-    if (run_state.should_quit)
+    if (run_state->should_quit)
       break;
 
     more_work_is_plausible |= delegate->DoDelayedWork(
-        &run_state.delayed_work_time);
-    if (run_state.should_quit)
+        &run_state->delayed_work_time);
+    if (run_state->should_quit)
       break;
 
     if (more_work_is_plausible)
       continue;
 
     more_work_is_plausible = delegate->DoIdleWork();
-    if (run_state.should_quit)
+    if (run_state->should_quit)
       break;
   }
-  run_state_ = old_state;
-}
-
-void MessagePumpMojo::Quit() {
-  if (run_state_)
-    run_state_->should_quit = true;
 }
 
-void MessagePumpMojo::ScheduleWork() {
-  SignalControlPipe();
-}
-
-void MessagePumpMojo::ScheduleDelayedWork(
-    const base::TimeTicks& delayed_work_time) {
-  if (!run_state_)
-    return;
-  run_state_->delayed_work_time = delayed_work_time;
-  SignalControlPipe();
-}
-
-void MessagePumpMojo::DoInternalWork(bool block) {
-  const MojoDeadline deadline = block ? GetDeadlineForWait() : 0;
-  const WaitState wait_state = GetWaitState();
+void MessagePumpMojo::DoInternalWork(const RunState& run_state, bool block) {
+  const MojoDeadline deadline = block ? GetDeadlineForWait(run_state) : 0;
+  const WaitState wait_state = GetWaitState(run_state);
   const MojoResult result =
       WaitMany(wait_state.handles, wait_state.wait_flags, deadline);
   if (result == 0) {
     // Control pipe was written to.
     uint32_t num_bytes = 0;
-    ReadMessageRaw(run_state_->read_handle.get(), NULL, &num_bytes, NULL, NULL,
+    ReadMessageRaw(run_state.read_handle.get(), NULL, &num_bytes, NULL, NULL,
                    MOJO_READ_MESSAGE_FLAG_MAY_DISCARD);
   } else if (result > 0) {
     const size_t index = static_cast<size_t>(result);
@@ -187,18 +202,16 @@ void MessagePumpMojo::RemoveFirstInvalidHandle(const WaitState& wait_state) {
   }
 }
 
-void MessagePumpMojo::SignalControlPipe() {
-  if (!run_state_)
-    return;
-
+void MessagePumpMojo::SignalControlPipe(const RunState& run_state) {
   // TODO(sky): deal with error?
-  WriteMessageRaw(run_state_->write_handle.get(), NULL, 0, NULL, 0,
+  WriteMessageRaw(run_state.write_handle.get(), NULL, 0, NULL, 0,
                   MOJO_WRITE_MESSAGE_FLAG_NONE);
 }
 
-MessagePumpMojo::WaitState MessagePumpMojo::GetWaitState() const {
+MessagePumpMojo::WaitState MessagePumpMojo::GetWaitState(
+    const RunState& run_state) const {
   WaitState wait_state;
-  wait_state.handles.push_back(run_state_->read_handle.get());
+  wait_state.handles.push_back(run_state.read_handle.get());
   wait_state.wait_flags.push_back(MOJO_WAIT_FLAG_READABLE);
 
   for (HandleToHandler::const_iterator i = handlers_.begin();
@@ -209,8 +222,9 @@ MessagePumpMojo::WaitState MessagePumpMojo::GetWaitState() const {
   return wait_state;
 }
 
-MojoDeadline MessagePumpMojo::GetDeadlineForWait() const {
-  base::TimeTicks min_time = run_state_->delayed_work_time;
+MojoDeadline MessagePumpMojo::GetDeadlineForWait(
+    const RunState& run_state) const {
+  base::TimeTicks min_time = run_state.delayed_work_time;
   for (HandleToHandler::const_iterator i = handlers_.begin();
        i != handlers_.end(); ++i) {
     if (min_time.is_null() && i->second.deadline < min_time)
diff --git a/mojo/common/message_pump_mojo.h b/mojo/common/message_pump_mojo.h
index cc21565..c761106 100644
--- a/mojo/common/message_pump_mojo.h
+++ b/mojo/common/message_pump_mojo.h
@@ -9,6 +9,7 @@
 
 #include "base/memory/scoped_ptr.h"
 #include "base/message_loop/message_pump.h"
+#include "base/synchronization/lock.h"
 #include "base/time/time.h"
 #include "mojo/common/mojo_common_export.h"
 #include "mojo/public/cpp/system/core.h"
@@ -61,25 +62,33 @@ class MOJO_COMMON_EXPORT MessagePumpMojo : public base::MessagePump {
 
   typedef std::map<Handle, Handler> HandleToHandler;
 
+  // Implementation of Run().
+  void DoRunLoop(RunState* run_state, Delegate* delegate);
+
   // Services the set of handles ready. If |block| is true this waits for a
   // handle to become ready, otherwise this does not block.
-  void DoInternalWork(bool block);
+  void DoInternalWork(const RunState& run_state, bool block);
 
   // Removes the first invalid handle. This is called if MojoWaitMany finds an
   // invalid handle.
   void RemoveFirstInvalidHandle(const WaitState& wait_state);
 
-  void SignalControlPipe();
+  void SignalControlPipe(const RunState& run_state);
 
-  WaitState GetWaitState() const;
+  WaitState GetWaitState(const RunState& run_state) const;
 
   // Returns the deadline for the call to MojoWaitMany().
-  MojoDeadline GetDeadlineForWait() const;
+  MojoDeadline GetDeadlineForWait(const RunState& run_state) const;
 
   // If non-NULL we're running (inside Run()). Member is reference to value on
   // stack.
   RunState* run_state_;
 
+  // Lock for accessing |run_state_|. In general the only method that we have to
+  // worry about is ScheduleWork(). All other methods are invoked on the same
+  // thread.
+  base::Lock run_state_lock_;
+
   HandleToHandler handlers_;
 
   // An ever increasing value assigned to each Handler::id. Used to detect