Fix thread usage in chromoting host

There are several things done in this patch:
1. Isloate thread start and stop to ChromotingHostContext
2. SessionManager now doesn't own capturer and encoder, ownership moved to ChromotingHost
3. Fix up the sequence of actions when ChromotingHost shuts down

TEST=remoting_unittests
BUG=none

Review URL: http://codereview.chromium.org/2829018

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

remoting/host/chromoting_host.cc

@@ -5,117 +5,63 @@
 #include "remoting/host/chromoting_host.h"
 
 #include "base/stl_util-inl.h"
-#include "base/waitable_event.h"
+#include "base/task.h"
 #include "build/build_config.h"
 #include "remoting/base/constants.h"
 #include "remoting/base/protocol_decoder.h"
+#include "remoting/host/chromoting_host_context.h"
 #include "remoting/host/host_config.h"
 #include "remoting/host/session_manager.h"
 #include "remoting/jingle_glue/jingle_channel.h"
 
 namespace remoting {
 
-ChromotingHost::ChromotingHost(MutableHostConfig* config,
+ChromotingHost::ChromotingHost(ChromotingHostContext* context,
+                               MutableHostConfig* config,
                                Capturer* capturer,
                                Encoder* encoder,
-                               EventExecutor* executor,
-                               base::WaitableEvent* host_done)
-      : main_thread_("MainThread"),
-        capture_thread_("CaptureThread"),
-        encode_thread_("EncodeThread"),
-        config_(config),
-        capturer_(capturer),
-        encoder_(encoder),
-        executor_(executor),
-        host_done_(host_done) {
-  // TODO(ajwong): The thread injection and object ownership is odd here.
-  // Fix so we do not start this thread in the constructor, so we only
-  // take in a session manager, don't let session manager own the
-  // capturer/encoder, and then associate the capturer and encoder threads with
-  // the capturer and encoder objects directly.  This will require a
-  // non-refcounted NewRunnableMethod.
-  main_thread_.StartWithOptions(
-      base::Thread::Options(MessageLoop::TYPE_UI, 0));
-  network_thread_.Start();
+                               EventExecutor* executor)
+    : context_(context),
+      config_(config),
+      capturer_(capturer),
+      encoder_(encoder),
+      executor_(executor),
+      state_(kInitial) {
 }
 
 ChromotingHost::~ChromotingHost() {
-  // TODO(ajwong): We really need to inject these threads and get rid of these
-  // start/stops.
-  main_thread_.Stop();
-  network_thread_.Stop();
-  DCHECK(!encode_thread_.IsRunning());
-  DCHECK(!capture_thread_.IsRunning());
 }
 
-void ChromotingHost::Run() {
+void ChromotingHost::Start(Task* shutdown_task) {
   // Submit a task to perform host registration. We'll also start
   // listening to connection if registration is done.
-  message_loop()->PostTask(
+  context_->main_message_loop()->PostTask(
       FROM_HERE,
-      NewRunnableMethod(this, &ChromotingHost::RegisterHost));
+      NewRunnableMethod(this, &ChromotingHost::DoStart, shutdown_task));
 }
 
 // This method is called when we need to destroy the host process.
-void ChromotingHost::DestroySession() {
-  DCHECK_EQ(message_loop(), MessageLoop::current());
-
-  // First we tell the session to pause and then we wait until all
-  // the tasks are done.
-  if (session_.get()) {
-    session_->Pause();
-
-    // TODO(hclam): Revise the order.
-    DCHECK(encode_thread_.IsRunning());
-    encode_thread_.Stop();
-
-    DCHECK(capture_thread_.IsRunning());
-    capture_thread_.Stop();
-  }
-}
-
-// This method talks to the cloud to register the host process. If
-// successful we will start listening to network requests.
-void ChromotingHost::RegisterHost() {
-  DCHECK_EQ(message_loop(), MessageLoop::current());
-  DCHECK(!jingle_client_);
-
-  std::string xmpp_login;
-  std::string xmpp_auth_token;
-  if (!config_->GetString(kXmppLoginConfigPath, &xmpp_login) ||
-      !config_->GetString(kXmppAuthTokenConfigPath, &xmpp_auth_token)) {
-    LOG(ERROR) << "XMMP credentials are not defined in config.";
-    return;
-  }
-
-  // Connect to the talk network with a JingleClient.
-  jingle_client_ = new JingleClient(&network_thread_);
-  jingle_client_->Init(xmpp_login, xmpp_auth_token,
-                       kChromotingTokenServiceName, this);
+void ChromotingHost::Shutdown() {
+  context_->main_message_loop()->PostTask(
+      FROM_HERE,
+      NewRunnableMethod(this, &ChromotingHost::DoShutdown));
 }
 
 // This method is called if a client is connected to this object.
 void ChromotingHost::OnClientConnected(ClientConnection* client) {
-  DCHECK_EQ(message_loop(), MessageLoop::current());
+  DCHECK_EQ(context_->main_message_loop(), MessageLoop::current());
 
   // Create a new RecordSession if there was none.
   if (!session_.get()) {
-    // The first we need to make sure capture and encode thread are
-    // running.
-    capture_thread_.Start();
-    encode_thread_.Start();
-
     // Then we create a SessionManager passing the message loops that
     // it should run on.
-    // Note that we pass the ownership of the capturer and encoder to
-    // the session manager.
     DCHECK(capturer_.get());
     DCHECK(encoder_.get());
-    session_ = new SessionManager(capture_thread_.message_loop(),
-                                  encode_thread_.message_loop(),
-                                  message_loop(),
-                                  capturer_.release(),
-                                  encoder_.release());
+    session_ = new SessionManager(context_->capture_message_loop(),
+                                  context_->encode_message_loop(),
+                                  context_->main_message_loop(),
+                                  capturer_.get(),
+                                  encoder_.get());
 
     // Immediately add the client and start the session.
     session_->AddClient(client);
@@ -128,26 +74,24 @@ void ChromotingHost::OnClientConnected(ClientConnection* client) {
 }
 
 void ChromotingHost::OnClientDisconnected(ClientConnection* client) {
-  DCHECK_EQ(message_loop(), MessageLoop::current());
+  DCHECK_EQ(context_->main_message_loop(), MessageLoop::current());
 
-  // Remove the client from the session manager.
-  if (session_.get())
+  // Remove the client from the session manager and pause the session.
+  // TODO(hclam): Pause only if the last client disconnected.
+  if (session_.get()) {
     session_->RemoveClient(client);
+    session_->Pause();
+  }
 
   // Also remove reference to ClientConnection from this object.
   client_ = NULL;
-
-  // TODO(hclam): If the last client has disconnected we need to destroy
-  // the session manager and shutdown the capture and encode threads.
-  // Right now we assume that there's only one client.
-  DestroySession();
 }
 
 ////////////////////////////////////////////////////////////////////////////
 // ClientConnection::EventHandler implementations
 void ChromotingHost::HandleMessages(ClientConnection* client,
                                 ClientMessageList* messages) {
-  DCHECK_EQ(message_loop(), MessageLoop::current());
+  DCHECK_EQ(context_->main_message_loop(), MessageLoop::current());
 
   // Delegate the messages to EventExecutor and delete the unhandled
   // messages.
@@ -157,7 +101,7 @@ void ChromotingHost::HandleMessages(ClientConnection* client,
 }
 
 void ChromotingHost::OnConnectionOpened(ClientConnection* client) {
-  DCHECK_EQ(message_loop(), MessageLoop::current());
+  DCHECK_EQ(context_->main_message_loop(), MessageLoop::current());
 
   // Completes the client connection.
   LOG(INFO) << "Connection to client established.";
@@ -165,7 +109,7 @@ void ChromotingHost::OnConnectionOpened(ClientConnection* client) {
 }
 
 void ChromotingHost::OnConnectionClosed(ClientConnection* client) {
-  DCHECK_EQ(message_loop(), MessageLoop::current());
+  DCHECK_EQ(context_->main_message_loop(), MessageLoop::current());
 
   // Completes the client connection.
   LOG(INFO) << "Connection to client closed.";
@@ -173,7 +117,7 @@ void ChromotingHost::OnConnectionClosed(ClientConnection* client) {
 }
 
 void ChromotingHost::OnConnectionFailed(ClientConnection* client) {
-  DCHECK_EQ(message_loop(), MessageLoop::current());
+  DCHECK_EQ(context_->main_message_loop(), MessageLoop::current());
 
   // The client has disconnected.
   LOG(ERROR) << "Connection failed unexpectedly.";
@@ -183,10 +127,9 @@ void ChromotingHost::OnConnectionFailed(ClientConnection* client) {
 ////////////////////////////////////////////////////////////////////////////
 // JingleClient::Callback implementations
 void ChromotingHost::OnStateChange(JingleClient* jingle_client,
-                               JingleClient::State state) {
-  DCHECK_EQ(jingle_client_.get(), jingle_client);
-
+                                   JingleClient::State state) {
   if (state == JingleClient::CONNECTED) {
+    DCHECK_EQ(jingle_client_.get(), jingle_client);
     LOG(INFO) << "Host connected as "
               << jingle_client->GetFullJid() << "." << std::endl;
 
@@ -197,17 +140,21 @@ void ChromotingHost::OnStateChange(JingleClient* jingle_client,
     LOG(INFO) << "Host disconnected from talk network." << std::endl;
     heartbeat_sender_ = NULL;
 
-    // Quit the message loop if disconected.
     // TODO(sergeyu): We should try reconnecting here instead of terminating
     // the host.
-    message_loop()->PostTask(FROM_HERE, new MessageLoop::QuitTask());
-    host_done_->Signal();
+    // Post a shutdown task to properly shutdown the chromoting host.
+    context_->main_message_loop()->PostTask(
+        FROM_HERE, NewRunnableMethod(this, &ChromotingHost::DoShutdown));
   }
 }
 
 bool ChromotingHost::OnAcceptConnection(
     JingleClient* jingle_client, const std::string& jid,
     JingleChannel::Callback** channel_callback) {
+  AutoLock auto_lock(lock_);
+  if (state_ != kStarted)
+    return false;
+
   DCHECK_EQ(jingle_client_.get(), jingle_client);
 
   // TODO(hclam): Allow multiple clients to connect to the host.
@@ -218,13 +165,18 @@ bool ChromotingHost::OnAcceptConnection(
 
   // If we accept the connected then create a client object and set the
   // callback.
-  client_ = new ClientConnection(message_loop(), new ProtocolDecoder(), this);
+  client_ = new ClientConnection(context_->main_message_loop(),
+                                 new ProtocolDecoder(), this);
   *channel_callback = client_.get();
   return true;
 }
 
 void ChromotingHost::OnNewConnection(JingleClient* jingle_client,
-                                 scoped_refptr<JingleChannel> channel) {
+                                     scoped_refptr<JingleChannel> channel) {
+  AutoLock auto_lock(lock_);
+  if (state_ != kStarted)
+    return;
+
   DCHECK_EQ(jingle_client_.get(), jingle_client);
 
   // Since the session manager has not started, it is still safe to access
@@ -233,8 +185,67 @@ void ChromotingHost::OnNewConnection(JingleClient* jingle_client,
   client_->set_jingle_channel(channel);
 }
 
-MessageLoop* ChromotingHost::message_loop() {
-  return main_thread_.message_loop();
+void ChromotingHost::DoStart(Task* shutdown_task) {
+  DCHECK_EQ(context_->main_message_loop(), MessageLoop::current());
+  DCHECK(!jingle_client_);
+  DCHECK(shutdown_task);
+
+  // Make sure this object is not started.
+  {
+    AutoLock auto_lock(lock_);
+    if (state_ != kInitial)
+      return;
+    state_ = kStarted;
+  }
+
+  // Save the shutdown task.
+  shutdown_task_.reset(shutdown_task);
+
+  std::string xmpp_login;
+  std::string xmpp_auth_token;
+  if (!config_->GetString(kXmppLoginConfigPath, &xmpp_login) ||
+      !config_->GetString(kXmppAuthTokenConfigPath, &xmpp_auth_token)) {
+    LOG(ERROR) << "XMMP credentials are not defined in config.";
+    return;
+  }
+
+  // Connect to the talk network with a JingleClient.
+  jingle_client_ = new JingleClient(context_->jingle_thread());
+  jingle_client_->Init(xmpp_login, xmpp_auth_token,
+                       kChromotingTokenServiceName, this);
+}
+
+void ChromotingHost::DoShutdown() {
+  DCHECK_EQ(context_->main_message_loop(), MessageLoop::current());
+
+  // No-op if this object is not started yet.
+  {
+    AutoLock auto_lock(lock_);
+    if (state_ != kStarted)
+      return;
+    state_ = kStopped;
+  }
+
+  // Tell the session to pause and then disconnect all clients.
+  if (session_.get()) {
+    session_->Pause();
+    session_->RemoveAllClients();
+  }
+
+  // Disconnect all clients.
+  if (client_) {
+    client_->Disconnect();
+  }
+
+  // Disconnect from the talk network.
+  if (jingle_client_) {
+    jingle_client_->Close();
+  }
+
+  // Lastly call the shutdown task.
+  if (shutdown_task_.get()) {
+    shutdown_task_->Run();
+  }
 }
 
 }  // namespace remoting

remoting/host/chromoting_host.h

@@ -17,12 +17,11 @@
 #include "remoting/jingle_glue/jingle_client.h"
 #include "remoting/jingle_glue/jingle_thread.h"
 
-namespace base {
-class WaitableEvent;
-}  // namespace base
+class Task;
 
 namespace remoting {
 
+class ChromotingHostContext;
 class MutableHostConfig;
 
 // A class to implement the functionality of a host process.
@@ -51,24 +50,24 @@ class MutableHostConfig;
 //    return to the idle state. We then go to step (2) if there a new
 //    incoming connection.
 class ChromotingHost : public base::RefCountedThreadSafe<ChromotingHost>,
-                   public ClientConnection::EventHandler,
-                   public JingleClient::Callback {
+                       public ClientConnection::EventHandler,
+                       public JingleClient::Callback {
  public:
-  ChromotingHost(MutableHostConfig* config, Capturer* capturer,
-                 Encoder* encoder, EventExecutor* executor,
-                 base::WaitableEvent* host_done);
+  ChromotingHost(ChromotingHostContext* context, MutableHostConfig* config,
+                 Capturer* capturer, Encoder* encoder, EventExecutor* executor);
   virtual ~ChromotingHost();
 
-  // Run the host porcess. This method returns only after the message loop
-  // of the host process exits.
-  void Run();
+  // Start the host porcess. This methods starts the chromoting host
+  // asynchronously.
+  //
+  // |shutdown_task| is called if Start() has failed ot Shutdown() is called
+  // and all related operations are completed.
+  //
+  // This method can only be called once during the lifetime of this object.
+  void Start(Task* shutdown_task);
 
   // This method is called when we need to the host process.
-  void DestroySession();
-
-  // This method talks to the cloud to register the host process. If
-  // successful we will start listening to network requests.
-  void RegisterHost();
+  void Shutdown();
 
   // This method is called if a client is connected to this object.
   void OnClientConnected(ClientConnection* client);
@@ -95,20 +94,21 @@ class ChromotingHost : public base::RefCountedThreadSafe<ChromotingHost>,
       scoped_refptr<JingleChannel> channel);
 
  private:
-  // The message loop that this class runs on.
-  MessageLoop* message_loop();
-
-  // The main thread that this object runs on.
-  base::Thread main_thread_;
+  enum State {
+    kInitial,
+    kStarted,
+    kStopped,
+  };
 
-  // Used to handle the Jingle connection.
-  JingleThread network_thread_;
+  // This method connects to the talk network and start listening for incoming
+  // connections.
+  void DoStart(Task* shutdown_task);
 
-  // A thread that hosts capture operations.
-  base::Thread capture_thread_;
+  // This method shuts down the host process.
+  void DoShutdown();
 
-  // A thread that hosts encode operations.
-  base::Thread encode_thread_;
+  // The context that the chromoting host runs on.
+  ChromotingHostContext* context_;
 
   scoped_refptr<MutableHostConfig> config_;
 
@@ -137,8 +137,17 @@ class ChromotingHost : public base::RefCountedThreadSafe<ChromotingHost>,
   // Session manager for the host process.
   scoped_refptr<SessionManager> session_;
 
-  // Signals the host is ready to be destroyed.
-  base::WaitableEvent* host_done_;
+  // This task gets executed when this object fails to connect to the
+  // talk network or Shutdown() is called.
+  scoped_ptr<Task> shutdown_task_;
+
+  // Tracks the internal state of the host.
+  // This variable is written on the main thread of ChromotingHostContext
+  // and read by jingle thread.
+  State state_;
+
+  // Lock is to lock the access to |state_|.
+  Lock lock_;
 
   DISALLOW_COPY_AND_ASSIGN(ChromotingHost);
 };

remoting/host/chromoting_host_context.cc

+// Copyright (c) 2010 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 "remoting/host/chromoting_host_context.h"
+
+#include <string>
+
+#include "base/thread.h"
+#include "remoting/jingle_glue/jingle_thread.h"
+
+namespace remoting {
+
+ChromotingHostContext::ChromotingHostContext()
+    : main_thread_("ChromotingMainThread"),
+      capture_thread_("ChromotingCaptureThread"),
+      encode_thread_("ChromotingEncodeThread") {
+}
+
+ChromotingHostContext::~ChromotingHostContext() {
+}
+
+void ChromotingHostContext::Start() {
+  // Start all the threads.
+  main_thread_.StartWithOptions(
+      base::Thread::Options(MessageLoop::TYPE_UI, 0));
+  capture_thread_.Start();
+  encode_thread_.Start();
+  jingle_thread_.Start();
+}
+
+void ChromotingHostContext::Stop() {
+  // Stop all the threads.
+  jingle_thread_.Stop();
+  encode_thread_.Stop();
+  capture_thread_.Stop();
+  main_thread_.Stop();
+}
+
+JingleThread* ChromotingHostContext::jingle_thread() {
+  return &jingle_thread_;
+}
+
+MessageLoop* ChromotingHostContext::main_message_loop() {
+  return main_thread_.message_loop();
+}
+
+MessageLoop* ChromotingHostContext::capture_message_loop() {
+  return capture_thread_.message_loop();
+}
+
+MessageLoop* ChromotingHostContext::encode_message_loop() {
+  return encode_thread_.message_loop();
+}
+
+}  // namespace remoting

remoting/host/chromoting_host_context.h

+// Copyright (c) 2010 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.
+
+#ifndef REMOTING_CHROMOTING_HOST_CONTEXT_H_
+#define REMOTING_CHROMOTING_HOST_CONTEXT_H_
+
+#include <string>
+
+#include "base/thread.h"
+#include "remoting/jingle_glue/jingle_thread.h"
+#include "testing/gtest/include/gtest/gtest_prod.h"
+
+namespace remoting {
+
+// A class that manages threads and running context for the chromoting host
+// process.
+class ChromotingHostContext {
+ public:
+  ChromotingHostContext();
+  virtual ~ChromotingHostContext();
+
+  virtual void Start();
+  virtual void Stop();
+
+  virtual JingleThread* jingle_thread();
+  virtual MessageLoop* main_message_loop();
+  virtual MessageLoop* capture_message_loop();
+  virtual MessageLoop* encode_message_loop();
+
+ private:
+  FRIEND_TEST(ChromotingHostContextTest, StartAndStop);
+
+  // A thread that host network operations.
+  JingleThread jingle_thread_;
+
+  // A thread that host ChromotingHost.
+  base::Thread main_thread_;
+
+  // A thread that host all capture operations.
+  base::Thread capture_thread_;
+
+  // A thread that host all encode operations.
+  base::Thread encode_thread_;
+
+  DISALLOW_COPY_AND_ASSIGN(ChromotingHostContext);
+};
+
+}  // namespace remoting
+
+#endif  // REMOTING_HOST_CHROMOTING_HOST_CONTEXT_H_

remoting/host/chromoting_host_context_unittest.cc

+// Copyright (c) 2010 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/message_loop.h"
+#include "remoting/host/chromoting_host_context.h"
+#include "testing/gtest/include/gtest/gtest.h"
+
+namespace remoting {
+
+// A simple test that starts and stop the context. This tests the context
+// operates properly and all threads and message loops are valid.
+TEST(ChromotingHostContextTest, StartAndStop) {
+  ChromotingHostContext context;
+  context.Start();
+  EXPECT_TRUE(context.jingle_thread());
+  EXPECT_TRUE(context.main_message_loop());
+  EXPECT_TRUE(context.capture_message_loop());
+  EXPECT_TRUE(context.encode_message_loop());
+  context.Stop();
+
+  // Expect all the threads are stopped.
+  EXPECT_FALSE(context.main_thread_.IsRunning());
+  EXPECT_FALSE(context.capture_thread_.IsRunning());
+  EXPECT_FALSE(context.encode_thread_.IsRunning());
+}
+
+}  // namespace remoting

remoting/host/client_connection.cc

@@ -38,7 +38,10 @@ ClientConnection::~ClientConnection() {
 void ClientConnection::SendInitClientMessage(int width, int height) {
   DCHECK_EQ(loop_, MessageLoop::current());
   DCHECK(!update_stream_size_);
-  DCHECK(channel_.get());
+
+  // If we are disconnected then return.
+  if (!channel_)
+    return;
 
   HostMessage msg;
   msg.mutable_init_client()->set_width(width);
@@ -49,7 +52,10 @@ void ClientConnection::SendInitClientMessage(int width, int height) {
 
 void ClientConnection::SendBeginUpdateStreamMessage() {
   DCHECK_EQ(loop_, MessageLoop::current());
-  DCHECK(channel_.get());
+
+  // If we are disconnected then return.
+  if (!channel_)
+    return;
 
   HostMessage msg;
   msg.mutable_begin_update_stream();
@@ -65,7 +71,10 @@ void ClientConnection::SendUpdateStreamPacketMessage(
     const UpdateStreamPacketHeader* header,
     scoped_refptr<DataBuffer> data) {
   DCHECK_EQ(loop_, MessageLoop::current());
-  DCHECK(channel_.get());
+
+  // If we are disconnected then return.
+  if (!channel_)
+    return;
 
   HostMessage msg;
   msg.mutable_update_stream_packet()->mutable_header()->CopyFrom(*header);
@@ -81,7 +90,10 @@ void ClientConnection::SendUpdateStreamPacketMessage(
 
 void ClientConnection::SendEndUpdateStreamMessage() {
   DCHECK_EQ(loop_, MessageLoop::current());
-  DCHECK(channel_.get());
+
+  // If we are disconnected then return.
+  if (!channel_)
+    return;
 
   HostMessage msg;
   msg.mutable_end_update_stream();
@@ -116,8 +128,11 @@ int ClientConnection::GetPendingUpdateStreamMessages() {
 void ClientConnection::Disconnect() {
   DCHECK_EQ(loop_, MessageLoop::current());
 
-  DCHECK(channel_.get());
-  channel_->Close();
+  // If there is a channel then close it and release the reference.
+  if (channel_) {
+    channel_->Close();
+    channel_ = NULL;
+  }
 }
 
 void ClientConnection::OnStateChange(JingleChannel* channel,

remoting/host/client_connection.h

@@ -90,7 +90,10 @@ class ClientConnection : public base::RefCountedThreadSafe<ClientConnection>,
   // TODO(hclam): Report this number accurately.
   virtual int GetPendingUpdateStreamMessages();
 
-  // Disconnect the remote viewer.
+  // Disconnect the client connection. This method is allowed to be called
+  // more than once and calls after the first one will be ignored.
+  //
+  // After this method is called all the send method calls will be ignored.
   virtual void Disconnect();
 
   /////////////////////////////////////////////////////////////////////////////

remoting/host/client_connection_unittest.cc

@@ -12,6 +12,7 @@
 
 using ::testing::_;
 using ::testing::NotNull;
+using ::testing::StrictMock;
 
 namespace remoting {
 
@@ -23,13 +24,13 @@ class ClientConnectionTest : public testing::Test {
  protected:
   virtual void SetUp() {
     decoder_ = new MockProtocolDecoder();
-    channel_ = new MockJingleChannel();
+    channel_ = new StrictMock<MockJingleChannel>();
 
     // Allocate a ClientConnection object with the mock objects. we give the
     // ownership of decoder to the viewer.
     viewer_ = new ClientConnection(&message_loop_,
-                         decoder_,
-                         &handler_);
+                                   decoder_,
+                                   &handler_);
 
     viewer_->set_jingle_channel(channel_.get());
   }
@@ -37,9 +38,12 @@ class ClientConnectionTest : public testing::Test {
   MessageLoop message_loop_;
   MockProtocolDecoder* decoder_;
   MockClientConnectionEventHandler handler_;
-  scoped_refptr<MockJingleChannel> channel_;
   scoped_refptr<ClientConnection> viewer_;
 
+  // |channel_| is wrapped with StrictMock because we limit strictly the calls
+  // to it.
+  scoped_refptr<StrictMock<MockJingleChannel> > channel_;
+
  private:
   DISALLOW_COPY_AND_ASSIGN(ClientConnectionTest);
 };
@@ -95,4 +99,23 @@ TEST_F(ClientConnectionTest, ParseMessages) {
   message_loop_.RunAllPending();
 }
 
+// Test that we can close client connection more than once and operations
+// after the connection is closed has no effect.
+TEST_F(ClientConnectionTest, Close) {
+  EXPECT_CALL(*channel_, Close());
+  viewer_->Disconnect();
+
+  viewer_->SendBeginUpdateStreamMessage();
+  scoped_ptr<UpdateStreamPacketHeader> header(new UpdateStreamPacketHeader);
+  header->set_x(0);
+  header->set_y(0);
+  header->set_width(640);
+  header->set_height(480);
+
+  scoped_refptr<media::DataBuffer> data = new media::DataBuffer(10);
+  viewer_->SendUpdateStreamPacketMessage(header.get(), data);
+  viewer_->SendEndUpdateStreamMessage();
+  viewer_->Disconnect();
+}
+
 }  // namespace remoting

remoting/host/encoder_verbatim.cc

@@ -30,7 +30,8 @@ void EncoderVerbatim::Encode(const DirtyRects& dirty_rects,
       EncodingState state = EncodingInProgress;
       if (i == 0) {
         state |= EncodingStarting;
-      } else if (i == num_rects - 1) {
+      }
+      if (i == num_rects - 1) {
         state |= EncodingEnded;
       }
       data_available_callback->Run(header.release(),

remoting/host/session_manager.cc

@@ -139,6 +139,12 @@ void SessionManager::RemoveClient(scoped_refptr<ClientConnection> client) {
       NewRunnableMethod(this, &SessionManager::DoRemoveClient, client));
 }
 
+void SessionManager::RemoveAllClients() {
+  network_loop_->PostTask(
+      FROM_HERE,
+      NewRunnableMethod(this, &SessionManager::DoRemoveAllClients));
+}
+
 void SessionManager::DoCapture() {
   DCHECK_EQ(capture_loop_, MessageLoop::current());
 
@@ -166,8 +172,7 @@ void SessionManager::DoCapture() {
   ScheduleNextCapture();
 
   // And finally perform one capture.
-  DCHECK(capturer_.get());
-  capturer_->CaptureDirtyRects(
+  capturer()->CaptureDirtyRects(
       NewRunnableMethod(this, &SessionManager::CaptureDoneTask));
 }
 
@@ -190,13 +195,11 @@ void SessionManager::DoEncode(const CaptureData *capture_data) {
 
   DCHECK_EQ(encode_loop_, MessageLoop::current());
 
-  DCHECK(encoder_.get());
-
   // TODO(hclam): Enable |force_refresh| if a new client was
   // added.
-  encoder_->SetSize(capture_data->width_, capture_data->height_);
-  encoder_->SetPixelFormat(capture_data->pixel_format_);
-  encoder_->Encode(
+  encoder()->SetSize(capture_data->width_, capture_data->height_);
+  encoder()->SetPixelFormat(capture_data->pixel_format_);
+  encoder()->Encode(
       capture_data->dirty_rects_,
       capture_data->data_,
       capture_data->data_strides_,
@@ -239,7 +242,7 @@ void SessionManager::DoGetInitInfo(scoped_refptr<ClientConnection> client) {
   network_loop_->PostTask(
       FROM_HERE,
       NewRunnableMethod(this, &SessionManager::DoSendInit, client,
-                        capturer_->GetWidth(), capturer_->GetHeight()));
+                        capturer()->GetWidth(), capturer()->GetHeight()));
 
   // And then add the client to the list so it can receive update stream.
   // It is important we do so in such order or the client will receive
@@ -288,8 +291,16 @@ void SessionManager::DoRemoveClient(scoped_refptr<ClientConnection> client) {
   // TODO(hclam): Is it correct to do to a scoped_refptr?
   ClientConnectionList::iterator it
       = std::find(clients_.begin(), clients_.end(), client);
-  if (it != clients_.end())
+  if (it != clients_.end()) {
     clients_.erase(it);
+  }
+}
+
+void SessionManager::DoRemoveAllClients() {
+  DCHECK_EQ(network_loop_, MessageLoop::current());
+
+  // Clear the list of clients.
+  clients_.clear();
 }
 
 void SessionManager::DoRateControl() {
@@ -354,12 +365,12 @@ void SessionManager::CaptureDoneTask() {
   scoped_ptr<CaptureData> data(new CaptureData);
 
   // Save results of the capture.
-  capturer_->GetData(data->data_);
-  capturer_->GetDataStride(data->data_strides_);
-  capturer_->GetDirtyRects(&data->dirty_rects_);
-  data->pixel_format_ = capturer_->GetPixelFormat();
-  data->width_ = capturer_->GetWidth();
-  data->height_ = capturer_->GetHeight();
+  capturer()->GetData(data->data_);
+  capturer()->GetDataStride(data->data_strides_);
+  capturer()->GetDirtyRects(&data->dirty_rects_);
+  data->pixel_format_ = capturer()->GetPixelFormat();
+  data->width_ = capturer()->GetWidth();
+  data->height_ = capturer()->GetHeight();
 
   encode_loop_->PostTask(
       FROM_HERE,
@@ -384,10 +395,20 @@ void SessionManager::EncodeDataAvailableTask(
                         data,
                         state));
 
-  if (state == Encoder::EncodingEnded) {
+  if (state & Encoder::EncodingEnded) {
     capture_loop_->PostTask(
         FROM_HERE, NewRunnableMethod(this, &SessionManager::DoFinishEncode));
   }
 }
 
+Capturer* SessionManager::capturer() {
+  DCHECK_EQ(capture_loop_, MessageLoop::current());
+  return capturer_.get();
+}
+
+Encoder* SessionManager::encoder() {
+  DCHECK_EQ(encode_loop_, MessageLoop::current());
+  return encoder_.get();
+}
+
 }  // namespace remoting

remoting/host/session_manager.h

@@ -67,7 +67,7 @@ class SessionManager : public base::RefCountedThreadSafe<SessionManager> {
  public:
 
   // Construct a SessionManager. Message loops and threads are provided.
-  // Ownership of Capturer and Encoder are given to this object.
+  // This object does not own capturer and encoder.
   SessionManager(MessageLoop* capture_loop,
                  MessageLoop* encode_loop,
                  MessageLoop* network_loop,
@@ -95,6 +95,9 @@ class SessionManager : public base::RefCountedThreadSafe<SessionManager> {
   // Remove a client from receiving screen updates.
   void RemoveClient(scoped_refptr<ClientConnection> client);
 
+  // Remove all clients.
+  void RemoveAllClients();
+
  private:
 
   // Stores the data and information of a capture to pass off to the
@@ -130,6 +133,7 @@ class SessionManager : public base::RefCountedThreadSafe<SessionManager> {
   void DoSetMaxRate(double max_rate);
   void DoAddClient(scoped_refptr<ClientConnection> client);
   void DoRemoveClient(scoped_refptr<ClientConnection> client);
+  void DoRemoveAllClients();
   void DoRateControl();
 
   // Hepler method to schedule next capture using the current rate.
@@ -145,6 +149,10 @@ class SessionManager : public base::RefCountedThreadSafe<SessionManager> {
                                const scoped_refptr<media::DataBuffer>& data,
                                Encoder::EncodingState state);
 
+  // Getters for capturer and encoder.
+  Capturer* capturer();
+  Encoder* encoder();
+
   // Message loops used by this class.
   MessageLoop* capture_loop_;
   MessageLoop* encode_loop_;

remoting/host/simple_host_process.cc

@@ -27,6 +27,7 @@
 #include "base/waitable_event.h"
 #include "remoting/host/capturer_fake.h"
 #include "remoting/host/chromoting_host.h"
+#include "remoting/host/chromoting_host_context.h"
 #include "remoting/host/encoder_verbatim.h"
 #include "remoting/host/json_host_config.h"
 
@@ -52,6 +53,10 @@ const char kHomePath[] = "HOME";
 static char* GetEnvironmentVar(const char* x) { return getenv(x); }
 #endif
 
+void ShutdownTask(MessageLoop* message_loop) {
+  message_loop->PostTask(FROM_HERE, new MessageLoop::QuitTask());
+}
+
 const std::string kFakeSwitchName = "fake";
 const std::string kConfigSwitchName = "config";
 
@@ -96,6 +101,7 @@ int main(int argc, char** argv) {
 
   if (fake) {
     // Inject a fake capturer.
+    LOG(INFO) << "Usage a fake capturer.";
     capturer.reset(new remoting::CapturerFake());
   }
 
@@ -111,16 +117,25 @@ int main(int argc, char** argv) {
     return 1;
   }
 
-  base::WaitableEvent host_done(false, false);
+  // Allocate a chromoting context and starts it.
+  remoting::ChromotingHostContext context;
+  context.Start();
+
+  // Construct a chromoting host.
   scoped_refptr<remoting::ChromotingHost> host =
-      new remoting::ChromotingHost(config,
+      new remoting::ChromotingHost(&context,
+                                   config,
                                    capturer.release(),
                                    encoder.release(),
-                                   executor.release(),
-                                   &host_done);
-  host->Run();
-  host_done.Wait();
+                                   executor.release());
+
+  // Let the chromoting host runs until the shutdown task is executed.
+  MessageLoop message_loop;
+  host->Start(NewRunnableFunction(&ShutdownTask, &message_loop));
+  message_loop.Run();
 
+  // And then stop the chromoting context.
+  context.Stop();
   file_io_thread.Stop();
   return 0;
 }

remoting/remoting.gyp

@@ -145,6 +145,8 @@
         'host/capturer.h',
         'host/chromoting_host.cc',
         'host/chromoting_host.h',
+        'host/chromoting_host_context.cc',
+        'host/chromoting_hsot_context.h',
         'host/client_connection.cc',
         'host/client_connection.h',
         'host/differ.h',
@@ -323,6 +325,7 @@
         'base/protocol_decoder_unittest.cc',
         'client/mock_objects.h',
         'client/decoder_verbatim_unittest.cc',
+        'host/chromoting_host_context_unittest.cc',
         'host/differ_unittest.cc',
         'host/differ_block_unittest.cc',
         'host/json_host_config_unittest.cc',