Revert 196343 "Made the ChromotingHost class not ref-counted."

r196343 broke ChromotingHost shutdown when connecting to an it2me host:

> [0425/154053:FATAL:observer_list.h(197)] Check failed: ObserverListBase<ObserverType>::size() == 0U (2 vs. 0)

> Made the ChromotingHost class not ref-counted.
> 
> ChromotingHost becomes a non thread-safe class should should live on the network thread. The CL removes ChromotingHost::Shutdown() allowing ChromotingHost to be destroyed synchronously closing all existing connections. It also updates HostStatusObserver::OnShutdown() handler to not destroy the host since it is the host's owner responsibility now.
> 
> Review URL: https://chromiumcodereview.appspot.com/13466014

TBR=alexeypa@chromium.org

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

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

5 files changed, 217 insertions, 115 deletions

diff --git a/remoting/host/chromoting_host.cc b/remoting/host/chromoting_host.cc
index f722d8f..9f7de561 100644
--- a/remoting/host/chromoting_host.cc
+++ b/remoting/host/chromoting_host.cc
@@ -76,14 +76,14 @@ ChromotingHost::ChromotingHost(
       network_task_runner_(network_task_runner),
       ui_task_runner_(ui_task_runner),
       signal_strategy_(signal_strategy),
-      started_(false),
+      state_(kInitial),
       protocol_config_(protocol::CandidateSessionConfig::CreateDefault()),
       login_backoff_(&kDefaultBackoffPolicy),
       authenticating_client_(false),
       reject_authenticating_client_(false),
       ALLOW_THIS_IN_INITIALIZER_LIST(weak_factory_(this)) {
-  DCHECK(network_task_runner_->BelongsToCurrentThread());
   DCHECK(signal_strategy);
+  DCHECK(network_task_runner_->BelongsToCurrentThread());
 
   if (!desktop_environment_factory_->SupportsAudioCapture()) {
     protocol::CandidateSessionConfig::DisableAudioChannel(
@@ -92,25 +92,19 @@ ChromotingHost::ChromotingHost(
 }
 
 ChromotingHost::~ChromotingHost() {
-  DCHECK(CalledOnValidThread());
-
-  // Disconnect all of the clients.
-  while (!clients_.empty()) {
-    clients_.front()->DisconnectSession();
-  }
-
-  // Notify observers.
-  if (started_) {
-    FOR_EACH_OBSERVER(HostStatusObserver, status_observers_, OnShutdown());
-  }
+  DCHECK(clients_.empty());
 }
 
 void ChromotingHost::Start(const std::string& xmpp_login) {
-  DCHECK(CalledOnValidThread());
-  DCHECK(!started_);
+  DCHECK(network_task_runner_->BelongsToCurrentThread());
 
   LOG(INFO) << "Starting host";
-  started_ = true;
+
+  // Make sure this object is not started.
+  if (state_ != kInitial)
+    return;
+  state_ = kStarted;
+
   FOR_EACH_OBSERVER(HostStatusObserver, status_observers_,
                     OnStart(xmpp_login));
 
@@ -118,13 +112,54 @@ void ChromotingHost::Start(const std::string& xmpp_login) {
   session_manager_->Init(signal_strategy_, this);
 }
 
+// This method is called when we need to destroy the host process.
+void ChromotingHost::Shutdown(const base::Closure& shutdown_task) {
+  if (!network_task_runner_->BelongsToCurrentThread()) {
+    network_task_runner_->PostTask(
+        FROM_HERE, base::Bind(&ChromotingHost::Shutdown, this, shutdown_task));
+    return;
+  }
+
+  switch (state_) {
+    case kInitial:
+    case kStopped:
+      // Nothing to do if we are not started.
+      state_ = kStopped;
+      if (!shutdown_task.is_null())
+        network_task_runner_->PostTask(FROM_HERE, shutdown_task);
+      break;
+
+    case kStopping:
+      // We are already stopping. Just save the task.
+      if (!shutdown_task.is_null())
+        shutdown_tasks_.push_back(shutdown_task);
+      break;
+
+    case kStarted:
+      if (!shutdown_task.is_null())
+        shutdown_tasks_.push_back(shutdown_task);
+      state_ = kStopping;
+
+      // Disconnect all of the clients.
+      while (!clients_.empty()) {
+        clients_.front()->DisconnectSession();
+      }
+
+      // Run the remaining shutdown tasks.
+      if (state_ == kStopping)
+        ShutdownFinish();
+
+      break;
+  }
+}
+
 void ChromotingHost::AddStatusObserver(HostStatusObserver* observer) {
-  DCHECK(CalledOnValidThread());
+  DCHECK(network_task_runner_->BelongsToCurrentThread());
   status_observers_.AddObserver(observer);
 }
 
 void ChromotingHost::RemoveStatusObserver(HostStatusObserver* observer) {
-  DCHECK(CalledOnValidThread());
+  DCHECK(network_task_runner_->BelongsToCurrentThread());
   status_observers_.RemoveObserver(observer);
 }
 
@@ -135,7 +170,7 @@ void ChromotingHost::RejectAuthenticatingClient() {
 
 void ChromotingHost::SetAuthenticatorFactory(
     scoped_ptr<protocol::AuthenticatorFactory> authenticator_factory) {
-  DCHECK(CalledOnValidThread());
+  DCHECK(network_task_runner_->BelongsToCurrentThread());
   session_manager_->set_authenticator_factory(authenticator_factory.Pass());
 }
 
@@ -147,7 +182,7 @@ void ChromotingHost::SetMaximumSessionDuration(
 ////////////////////////////////////////////////////////////////////////////
 // protocol::ClientSession::EventHandler implementation.
 void ChromotingHost::OnSessionAuthenticated(ClientSession* client) {
-  DCHECK(CalledOnValidThread());
+  DCHECK(network_task_runner_->BelongsToCurrentThread());
 
   login_backoff_.Reset();
 
@@ -180,7 +215,7 @@ void ChromotingHost::OnSessionAuthenticated(ClientSession* client) {
 }
 
 void ChromotingHost::OnSessionChannelsConnected(ClientSession* client) {
-  DCHECK(CalledOnValidThread());
+  DCHECK(network_task_runner_->BelongsToCurrentThread());
 
   // Notify observers.
   FOR_EACH_OBSERVER(HostStatusObserver, status_observers_,
@@ -188,7 +223,7 @@ void ChromotingHost::OnSessionChannelsConnected(ClientSession* client) {
 }
 
 void ChromotingHost::OnSessionAuthenticationFailed(ClientSession* client) {
-  DCHECK(CalledOnValidThread());
+  DCHECK(network_task_runner_->BelongsToCurrentThread());
 
   // Notify observers.
   FOR_EACH_OBSERVER(HostStatusObserver, status_observers_,
@@ -196,7 +231,7 @@ void ChromotingHost::OnSessionAuthenticationFailed(ClientSession* client) {
 }
 
 void ChromotingHost::OnSessionClosed(ClientSession* client) {
-  DCHECK(CalledOnValidThread());
+  DCHECK(network_task_runner_->BelongsToCurrentThread());
 
   ClientList::iterator it = std::find(clients_.begin(), clients_.end(), client);
   CHECK(it != clients_.end());
@@ -208,25 +243,28 @@ void ChromotingHost::OnSessionClosed(ClientSession* client) {
 
   clients_.erase(it);
   delete client;
+
+  if (state_ == kStopping && clients_.empty())
+    ShutdownFinish();
 }
 
 void ChromotingHost::OnSessionSequenceNumber(ClientSession* session,
                                              int64 sequence_number) {
-  DCHECK(CalledOnValidThread());
+  DCHECK(network_task_runner_->BelongsToCurrentThread());
 }
 
 void ChromotingHost::OnSessionRouteChange(
     ClientSession* session,
     const std::string& channel_name,
     const protocol::TransportRoute& route) {
-  DCHECK(CalledOnValidThread());
+  DCHECK(network_task_runner_->BelongsToCurrentThread());
   FOR_EACH_OBSERVER(HostStatusObserver, status_observers_,
                     OnClientRouteChange(session->client_jid(), channel_name,
                                         route));
 }
 
 void ChromotingHost::OnSessionManagerReady() {
-  DCHECK(CalledOnValidThread());
+  DCHECK(network_task_runner_->BelongsToCurrentThread());
   // Don't need to do anything here, just wait for incoming
   // connections.
 }
@@ -234,9 +272,9 @@ void ChromotingHost::OnSessionManagerReady() {
 void ChromotingHost::OnIncomingSession(
       protocol::Session* session,
       protocol::SessionManager::IncomingSessionResponse* response) {
-  DCHECK(CalledOnValidThread());
+  DCHECK(network_task_runner_->BelongsToCurrentThread());
 
-  if (!started_) {
+  if (state_ != kStarted) {
     *response = protocol::SessionManager::DECLINE;
     return;
   }
@@ -285,14 +323,18 @@ void ChromotingHost::OnIncomingSession(
 
 void ChromotingHost::set_protocol_config(
     scoped_ptr<protocol::CandidateSessionConfig> config) {
-  DCHECK(CalledOnValidThread());
+  DCHECK(network_task_runner_->BelongsToCurrentThread());
   DCHECK(config.get());
-  DCHECK(!started_);
+  DCHECK_EQ(state_, kInitial);
   protocol_config_ = config.Pass();
 }
 
 void ChromotingHost::DisconnectAllClients() {
-  DCHECK(CalledOnValidThread());
+  if (!network_task_runner_->BelongsToCurrentThread()) {
+    network_task_runner_->PostTask(
+        FROM_HERE, base::Bind(&ChromotingHost::DisconnectAllClients, this));
+    return;
+  }
 
   while (!clients_.empty()) {
     size_t size = clients_.size();
@@ -301,4 +343,35 @@ void ChromotingHost::DisconnectAllClients() {
   }
 }
 
+void ChromotingHost::ShutdownFinish() {
+  DCHECK(network_task_runner_->BelongsToCurrentThread());
+  DCHECK_EQ(state_, kStopping);
+
+  state_ = kStopped;
+
+  // Destroy session manager.
+  session_manager_.reset();
+
+  // Clear |desktop_environment_factory_| and |signal_strategy_| to
+  // ensure we don't try to touch them after running shutdown tasks
+  desktop_environment_factory_ = NULL;
+  signal_strategy_ = NULL;
+
+  // Keep reference to |this|, so that we don't get destroyed while
+  // sending notifications.
+  scoped_refptr<ChromotingHost> self(this);
+
+  // Notify observers.
+  FOR_EACH_OBSERVER(HostStatusObserver, status_observers_,
+                    OnShutdown());
+
+  for (std::vector<base::Closure>::iterator it = shutdown_tasks_.begin();
+       it != shutdown_tasks_.end(); ++it) {
+    it->Run();
+  }
+  shutdown_tasks_.clear();
+
+  weak_factory_.InvalidateWeakPtrs();
+}
+
 }  // namespace remoting
diff --git a/remoting/host/chromoting_host.h b/remoting/host/chromoting_host.h
index 35ef373..b1cd128 100644
--- a/remoting/host/chromoting_host.h
+++ b/remoting/host/chromoting_host.h
@@ -12,7 +12,6 @@
 #include "base/memory/ref_counted.h"
 #include "base/memory/weak_ptr.h"
 #include "base/observer_list.h"
-#include "base/threading/non_thread_safe.h"
 #include "base/threading/thread.h"
 #include "net/base/backoff_entry.h"
 #include "remoting/host/client_session.h"
@@ -60,13 +59,14 @@ class DesktopEnvironmentFactory;
 //    all pending tasks to complete. After all of that completed we
 //    return to the idle state. We then go to step (2) if there a new
 //    incoming connection.
-class ChromotingHost : public base::NonThreadSafe,
+class ChromotingHost : public base::RefCountedThreadSafe<ChromotingHost>,
                        public ClientSession::EventHandler,
                        public protocol::SessionManager::Listener,
                        public HostStatusMonitor {
  public:
-  // Both |signal_strategy| and |desktop_environment_factory| should outlive
-  // this object.
+  // The caller must ensure that |signal_strategy| and
+  // |desktop_environment_factory| remain valid at least until the
+  // |shutdown_task| supplied to Shutdown() has been notified.
   ChromotingHost(
       SignalStrategy* signal_strategy,
       DesktopEnvironmentFactory* desktop_environment_factory,
@@ -77,9 +77,8 @@ class ChromotingHost : public base::NonThreadSafe,
       scoped_refptr<base::SingleThreadTaskRunner> video_encode_task_runner,
       scoped_refptr<base::SingleThreadTaskRunner> network_task_runner,
       scoped_refptr<base::SingleThreadTaskRunner> ui_task_runner);
-  virtual ~ChromotingHost();
 
-  // Asynchronously starts the host.
+  // Asynchronously start the host process.
   //
   // After this is invoked, the host process will connect to the talk
   // network and start listening for incoming connections.
@@ -87,6 +86,10 @@ class ChromotingHost : public base::NonThreadSafe,
   // This method can only be called once during the lifetime of this object.
   void Start(const std::string& xmpp_login);
 
+  // Asynchronously shutdown the host process. |shutdown_task| is
+  // called after shutdown is completed.
+  void Shutdown(const base::Closure& shutdown_task);
+
   // HostStatusMonitor interface.
   virtual void AddStatusObserver(HostStatusObserver* observer) OVERRIDE;
   virtual void RemoveStatusObserver(HostStatusObserver* observer) OVERRIDE;
@@ -131,9 +134,10 @@ class ChromotingHost : public base::NonThreadSafe,
   // Sets desired configuration for the protocol. Must be called before Start().
   void set_protocol_config(scoped_ptr<protocol::CandidateSessionConfig> config);
 
-  // Immediately disconnects all active clients. Host-internal components may
-  // shutdown asynchronously, but the caller is guaranteed not to receive
-  // callbacks for disconnected clients after this call returns.
+  // Disconnects all active clients. Clients are disconnected
+  // asynchronously when this method is called on a thread other than
+  // the network thread. Potentically this may cause disconnection of
+  // clients that were not connected when this method is called.
   void DisconnectAllClients();
 
   base::WeakPtr<ChromotingHost> AsWeakPtr() {
@@ -141,10 +145,23 @@ class ChromotingHost : public base::NonThreadSafe,
   }
 
  private:
+  friend class base::RefCountedThreadSafe<ChromotingHost>;
   friend class ChromotingHostTest;
 
   typedef std::list<ClientSession*> ClientList;
 
+  enum State {
+    kInitial,
+    kStarted,
+    kStopping,
+    kStopped,
+  };
+
+  virtual ~ChromotingHost();
+
+  // Called from Shutdown() to finish shutdown.
+  void ShutdownFinish();
+
   // Unless specified otherwise all members of this class must be
   // used on the network thread only.
 
@@ -167,8 +184,8 @@ class ChromotingHost : public base::NonThreadSafe,
   // The connections to remote clients.
   ClientList clients_;
 
-  // True if the host has been started.
-  bool started_;
+  // Tracks the internal state of the host.
+  State state_;
 
   // Configuration of the protocol.
   scoped_ptr<protocol::CandidateSessionConfig> protocol_config_;
@@ -180,6 +197,10 @@ class ChromotingHost : public base::NonThreadSafe,
   bool authenticating_client_;
   bool reject_authenticating_client_;
 
+  // Stores list of tasks that should be executed when we finish
+  // shutdown. Used only while |state_| is set to kStopping.
+  std::vector<base::Closure> shutdown_tasks_;
+
   // The maximum duration of any session.
   base::TimeDelta max_session_duration_;
 
diff --git a/remoting/host/chromoting_host_unittest.cc b/remoting/host/chromoting_host_unittest.cc
index 809f6f0..3b20bd4 100644
--- a/remoting/host/chromoting_host_unittest.cc
+++ b/remoting/host/chromoting_host_unittest.cc
@@ -68,7 +68,7 @@ class ChromotingHostTest : public testing::Test {
   }
 
   virtual void SetUp() OVERRIDE {
-    task_runner_ = new AutoThreadTaskRunner(
+    ui_task_runner_ = new AutoThreadTaskRunner(
         message_loop_.message_loop_proxy(),
         base::Bind(&ChromotingHostTest::QuitMainMessageLoop,
                    base::Unretained(this)));
@@ -84,16 +84,16 @@ class ChromotingHostTest : public testing::Test {
 
     session_manager_ = new protocol::MockSessionManager();
 
-    host_.reset(new ChromotingHost(
+    host_ = new ChromotingHost(
         &signal_strategy_,
         desktop_environment_factory_.get(),
         scoped_ptr<protocol::SessionManager>(session_manager_),
-        task_runner_,   // Audio
-        task_runner_,   // Input
-        task_runner_,   // Video capture
-        task_runner_,   // Video encode
-        task_runner_,   // Network
-        task_runner_)); // UI
+        ui_task_runner_,  // Audio
+        ui_task_runner_,  // Input
+        ui_task_runner_,  // Video capture
+        ui_task_runner_,  // Video encode
+        ui_task_runner_,  // Network
+        ui_task_runner_); // UI
     host_->AddStatusObserver(&host_status_observer_);
 
     xmpp_login_ = "host@domain";
@@ -183,40 +183,41 @@ class ChromotingHostTest : public testing::Test {
     protocol::ConnectionToClient* connection_ptr = connection.get();
     scoped_ptr<ClientSession> client(new ClientSession(
         host_.get(),
-        task_runner_, // Audio
-        task_runner_, // Input
-        task_runner_, // Video capture
-        task_runner_, // Video encode
-        task_runner_, // Network
-        task_runner_, // UI
+        ui_task_runner_,  // Audio
+        ui_task_runner_,  // Input
+        ui_task_runner_,  // Video capture
+        ui_task_runner_,  // Video encode
+        ui_task_runner_,  // Network
+        ui_task_runner_,  // UI
         connection.Pass(),
         desktop_environment_factory_.get(),
         base::TimeDelta()));
 
-    connection_ptr->set_host_stub(client.get());
+    ClientSession* client_raw = client.get();
+    connection_ptr->set_host_stub(client_raw);
+
+    ui_task_runner_->PostTask(
+        FROM_HERE, base::Bind(&ChromotingHostTest::AddClientToHost,
+                              host_, base::Passed(&client)));
 
     if (authenticate) {
-      task_runner_->PostTask(
-          FROM_HERE,
-          base::Bind(&ClientSession::OnConnectionAuthenticated,
-                     base::Unretained(client.get()), connection_ptr));
+      ui_task_runner_->PostTask(
+          FROM_HERE, base::Bind(&ClientSession::OnConnectionAuthenticated,
+                                base::Unretained(client_raw), connection_ptr));
       if (!reject) {
-        task_runner_->PostTask(
+        ui_task_runner_->PostTask(
             FROM_HERE,
             base::Bind(&ClientSession::OnConnectionChannelsConnected,
-                       base::Unretained(client.get()), connection_ptr));
+                       base::Unretained(client_raw), connection_ptr));
       }
     } else {
-      task_runner_->PostTask(
+      ui_task_runner_->PostTask(
           FROM_HERE, base::Bind(&ClientSession::OnConnectionClosed,
-                                base::Unretained(client.get()), connection_ptr,
+                                base::Unretained(client_raw), connection_ptr,
                                 protocol::AUTHENTICATION_FAILED));
     }
 
-    get_client(connection_index) = client.get();
-
-    // |host| is responsible for deleting |client| from now on.
-    host_->clients_.push_back(client.release());
+    get_client(connection_index) = client_raw;
   }
 
   virtual void TearDown() OVERRIDE {
@@ -296,16 +297,23 @@ class ChromotingHostTest : public testing::Test {
     }
   }
 
+  static void AddClientToHost(scoped_refptr<ChromotingHost> host,
+                              scoped_ptr<ClientSession> client) {
+    // |host| is responsible for deleting |client| from now on.
+    host->clients_.push_back(client.release());
+  }
+
   void ShutdownHost() {
-    task_runner_->PostTask(
+    ui_task_runner_->PostTask(
         FROM_HERE,
-        base::Bind(&ChromotingHostTest::StopAndReleaseTaskRunner,
-                   base::Unretained(this)));
+        base::Bind(&ChromotingHost::Shutdown, host_,
+                   base::Bind(&ChromotingHostTest::ReleaseUiTaskRunner,
+                              base::Unretained(this))));
   }
 
-  void StopAndReleaseTaskRunner() {
-    host_.reset();
-    task_runner_ = NULL;
+  void ReleaseUiTaskRunner() {
+    ui_task_runner_ = NULL;
+    host_ = NULL;
     desktop_environment_factory_.reset();
   }
 
@@ -391,11 +399,11 @@ class ChromotingHostTest : public testing::Test {
 
  protected:
   MessageLoop message_loop_;
-  scoped_refptr<AutoThreadTaskRunner> task_runner_;
+  scoped_refptr<AutoThreadTaskRunner> ui_task_runner_;
   MockConnectionToClientEventHandler handler_;
   MockSignalStrategy signal_strategy_;
   scoped_ptr<MockDesktopEnvironmentFactory> desktop_environment_factory_;
-  scoped_ptr<ChromotingHost> host_;
+  scoped_refptr<ChromotingHost> host_;
   MockHostStatusObserver host_status_observer_;
   protocol::MockSessionManager* session_manager_;
   std::string xmpp_login_;
diff --git a/remoting/host/plugin/host_script_object.cc b/remoting/host/plugin/host_script_object.cc
index 7ed4bbb..b3c8612a 100644
--- a/remoting/host/plugin/host_script_object.cc
+++ b/remoting/host/plugin/host_script_object.cc
@@ -142,13 +142,8 @@ class HostNPScriptObject::It2MeImpl
                            const std::string& support_id,
                            const base::TimeDelta& lifetime);
 
-  // Shuts down |host_| on the network thread and posts ShutdownOnUiThread()
-  // to shut down UI thread resources.
-  void ShutdownOnNetworkThread();
-
-  // Shuts down |desktop_environment_factory_| and |policy_watcher_| on
-  // the UI thread.
-  void ShutdownOnUiThread();
+  // Called when ChromotingHost::Shutdown() has completed.
+  void OnShutdownFinished();
 
   // Called when initial policies are read, and when they change.
   void OnPolicyUpdate(scoped_ptr<base::DictionaryValue> policies);
@@ -173,7 +168,7 @@ class HostNPScriptObject::It2MeImpl
   scoped_ptr<DesktopEnvironmentFactory> desktop_environment_factory_;
   scoped_ptr<HostEventLogger> host_event_logger_;
 
-  scoped_ptr<ChromotingHost> host_;
+  scoped_refptr<ChromotingHost> host_;
   int failed_login_attempts_;
 
   scoped_ptr<policy_hack::PolicyWatcher> policy_watcher_;
@@ -259,13 +254,13 @@ void HostNPScriptObject::It2MeImpl::Disconnect() {
 
   switch (state_) {
     case kDisconnected:
-      ShutdownOnNetworkThread();
+      OnShutdownFinished();
       return;
 
     case kStarting:
       SetState(kDisconnecting);
       SetState(kDisconnected);
-      ShutdownOnNetworkThread();
+      OnShutdownFinished();
       return;
 
     case kDisconnecting:
@@ -275,16 +270,18 @@ void HostNPScriptObject::It2MeImpl::Disconnect() {
       SetState(kDisconnecting);
 
       if (!host_) {
-        SetState(kDisconnected);
-        ShutdownOnNetworkThread();
+        OnShutdownFinished();
         return;
       }
 
-      // Deleting the host destroys SignalStrategy synchronously, but
-      // SignalStrategy::Listener handlers are not allowed to destroy
-      // SignalStrategy, so post task to destroy the host later.
+      // ChromotingHost::Shutdown() may destroy SignalStrategy
+      // synchronously, but SignalStrategy::Listener handlers are not
+      // allowed to destroy SignalStrategy, so post task to call
+      // Shutdown() later.
       host_context_->network_task_runner()->PostTask(
-          FROM_HERE, base::Bind(&It2MeImpl::ShutdownOnNetworkThread, this));
+          FROM_HERE, base::Bind(
+              &ChromotingHost::Shutdown, host_,
+              base::Bind(&It2MeImpl::OnShutdownFinished, this)));
       return;
   }
 }
@@ -372,7 +369,7 @@ void HostNPScriptObject::It2MeImpl::FinishConnect(
   }
 
   // Create the host.
-  host_.reset(new ChromotingHost(
+  host_ = new ChromotingHost(
       signal_strategy_.get(),
       desktop_environment_factory_.get(),
       CreateHostSessionManager(network_settings,
@@ -382,7 +379,7 @@ void HostNPScriptObject::It2MeImpl::FinishConnect(
       host_context_->video_capture_task_runner(),
       host_context_->video_encode_task_runner(),
       host_context_->network_task_runner(),
-      host_context_->ui_task_runner()));
+      host_context_->ui_task_runner());
   host_->AddStatusObserver(this);
   log_to_server_.reset(
       new LogToServer(host_->AsWeakPtr(), ServerLogEntry::IT2ME,
@@ -407,21 +404,14 @@ void HostNPScriptObject::It2MeImpl::FinishConnect(
   return;
 }
 
-void HostNPScriptObject::It2MeImpl::ShutdownOnNetworkThread() {
-  DCHECK(host_context_->network_task_runner()->BelongsToCurrentThread());
-  DCHECK(state_ == kDisconnecting || state_ == kDisconnected);
-
-  if (state_ == kDisconnecting) {
-    host_.reset();
-    SetState(kDisconnected);
+void HostNPScriptObject::It2MeImpl::OnShutdownFinished() {
+  if (!host_context_->ui_task_runner()->BelongsToCurrentThread()) {
+    host_context_->ui_task_runner()->PostTask(
+        FROM_HERE, base::Bind(&It2MeImpl::OnShutdownFinished, this));
+    return;
   }
 
-  host_context_->ui_task_runner()->PostTask(
-      FROM_HERE, base::Bind(&It2MeImpl::ShutdownOnUiThread, this));
-}
-
-void HostNPScriptObject::It2MeImpl::ShutdownOnUiThread() {
-  DCHECK(host_context_->ui_task_runner()->BelongsToCurrentThread());
+  // Note that OnShutdownFinished() may be called more than once.
 
   // Destroy the DesktopEnvironmentFactory, to free thread references.
   desktop_environment_factory_.reset();
@@ -496,6 +486,11 @@ void HostNPScriptObject::It2MeImpl::OnShutdown() {
   signal_strategy_.reset();
   host_event_logger_.reset();
   host_->RemoveStatusObserver(this);
+  host_ = NULL;
+
+  if (state_ != kDisconnected) {
+    SetState(kDisconnected);
+  }
 }
 
 void HostNPScriptObject::It2MeImpl::OnPolicyUpdate(
diff --git a/remoting/host/remoting_me2me_host.cc b/remoting/host/remoting_me2me_host.cc
index 4e31ea8..217eb57 100644
--- a/remoting/host/remoting_me2me_host.cc
+++ b/remoting/host/remoting_me2me_host.cc
@@ -282,7 +282,7 @@ class HostProcess
   scoped_ptr<LogToServer> log_to_server_;
   scoped_ptr<HostEventLogger> host_event_logger_;
 
-  scoped_ptr<ChromotingHost> host_;
+  scoped_refptr<ChromotingHost> host_;
 
   // Used to keep this HostProcess alive until it is shutdown.
   scoped_refptr<HostProcess> self_;
@@ -950,7 +950,7 @@ void HostProcess::StartHost() {
     network_settings.max_port = NetworkSettings::kDefaultMaxPort;
   }
 
-  host_.reset(new ChromotingHost(
+  host_ = new ChromotingHost(
       signal_strategy_.get(),
       desktop_environment_factory_.get(),
       CreateHostSessionManager(network_settings,
@@ -960,7 +960,7 @@ void HostProcess::StartHost() {
       context_->video_capture_task_runner(),
       context_->video_encode_task_runner(),
       context_->network_task_runner(),
-      context_->ui_task_runner()));
+      context_->ui_task_runner());
 
   // TODO(simonmorris): Get the maximum session duration from a policy.
 #if defined(OS_LINUX)
@@ -1041,7 +1041,7 @@ void HostProcess::RestartHost() {
   DCHECK_EQ(state_, HOST_STARTED);
 
   state_ = HOST_STOPPING_TO_RESTART;
-  ShutdownOnNetworkThread();
+  host_->Shutdown(base::Bind(&HostProcess::ShutdownOnNetworkThread, this));
 }
 
 void HostProcess::ShutdownHost(int exit_code) {
@@ -1051,11 +1051,15 @@ void HostProcess::ShutdownHost(int exit_code) {
 
   switch (state_) {
     case HOST_INITIALIZING:
-    case HOST_STARTED:
       state_ = HOST_STOPPING;
       ShutdownOnNetworkThread();
       break;
 
+    case HOST_STARTED:
+      state_ = HOST_STOPPING;
+      host_->Shutdown(base::Bind(&HostProcess::ShutdownOnNetworkThread, this));
+      break;
+
     case HOST_STOPPING_TO_RESTART:
       state_ = HOST_STOPPING;
       break;
@@ -1070,7 +1074,7 @@ void HostProcess::ShutdownHost(int exit_code) {
 void HostProcess::ShutdownOnNetworkThread() {
   DCHECK(context_->network_task_runner()->BelongsToCurrentThread());
 
-  host_.reset();
+  host_ = NULL;
   curtaining_host_observer_.reset();
   host_event_logger_.reset();
   log_to_server_.reset();
@@ -1099,7 +1103,8 @@ void HostProcess::ShutdownOnNetworkThread() {
         FROM_HERE,
         base::Bind(&HostProcess::ShutdownOnUiThread, this));
   } else {
-    // This method is only called in STOPPING_TO_RESTART and STOPPING states.
+    // This method is used as a callback for ChromotingHost::Shutdown() which is
+    // called only in STOPPING_TO_RESTART and STOPPING states.
     NOTREACHED();
   }
 }