[Chromoting] Introducing refcount-based life time management of the message loops in the service (d…

remoting/host/win/host_service.cc

 // Copyright (c) 2012 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.
 //
 // This file implements the Windows service controlling Me2Me host processes
 // running within user sessions.
 
 #include "remoting/host/win/host_service.h"
 
 #include <windows.h>
 #include <shellapi.h>
 #include <wtsapi32.h>
 
 #include "base/at_exit.h"
 #include "base/base_paths.h"
 #include "base/base_switches.h"
 #include "base/bind.h"
 #include "base/command_line.h"
 #include "base/file_path.h"
 #include "base/logging.h"
 #include "base/message_loop.h"
 #include "base/single_thread_task_runner.h"
 #include "base/stringprintf.h"
 #include "base/threading/thread.h"
 #include "base/utf_string_conversions.h"
 #include "base/win/wrapped_window_proc.h"
+#include "remoting/base/auto_message_loop.h"
 #include "remoting/base/breakpad.h"
 #include "remoting/base/scoped_sc_handle_win.h"
 #include "remoting/base/stoppable.h"
 #include "remoting/host/branding.h"
 
 #if defined(REMOTING_MULTI_PROCESS)
 #include "remoting/host/daemon_process.h"
 #endif  // defined(REMOTING_MULTI_PROCESS)
 
 #include "remoting/host/usage_stats_consent.h"
@@ skipping 48 matching lines @@
 // Exit codes:
 const int kSuccessExitCode = 0;
 const int kUsageExitCode = 1;
 const int kErrorExitCode = 2;
 
 void usage(const FilePath& program_name) {
   LOG(INFO) << StringPrintf(kUsageMessage,
                             UTF16ToWide(program_name.value()).c_str());
 }
 
+void QuitMessageLoop(MessageLoop* message_loop) {
+  message_loop->PostTask(FROM_HERE, MessageLoop::QuitClosure());
+}
+
 }  // namespace
 
 namespace remoting {
 
 HostService::HostService() :
   console_session_id_(kInvalidSessionId),
   run_routine_(&HostService::RunAsService),
   service_status_handle_(0),
   stopped_event_(true, false) {
 }
 
 HostService::~HostService() {
 }
 
 void HostService::AddWtsConsoleObserver(WtsConsoleObserver* observer) {
   DCHECK(main_task_runner_->BelongsToCurrentThread());
 
   console_observers_.AddObserver(observer);
 }
 
 void HostService::RemoveWtsConsoleObserver(WtsConsoleObserver* observer) {
   DCHECK(main_task_runner_->BelongsToCurrentThread());
 
   console_observers_.RemoveObserver(observer);
 }
 
 void HostService::OnChildStopped() {
   child_.reset(NULL);
-  main_task_runner_->PostTask(FROM_HERE, MessageLoop::QuitClosure());
+  main_task_runner_ = NULL;
 }
 
 void HostService::OnSessionChange() {
   // WTSGetActiveConsoleSessionId is a very cheap API. It basically reads
   // a single value from shared memory. Therefore it is better to check if
   // the console session is still the same every time a session change
   // notification event is posted. This also takes care of coalescing multiple
   // events into one since we look at the latest state.
   uint32 console_session_id = WTSGetActiveConsoleSessionId();
   if (console_session_id_ != console_session_id) {
@@ skipping 55 matching lines @@
     run_routine_ = &HostService::RunInConsole;
   }
 
   return true;
 }
 
 int HostService::Run() {
   return (this->*run_routine_)();
 }
 
-void HostService::RunMessageLoop(MessageLoop* message_loop) {
-  // Launch the I/O thread.
-  base::Thread io_thread(kIoThreadName);
-  base::Thread::Options io_thread_options(MessageLoop::TYPE_IO, 0);
-  if (!io_thread.StartWithOptions(io_thread_options)) {
-    LOG(ERROR) << "Failed to start the I/O thread";
-    stopped_event_.Signal();
-    return;
-  }
+void HostService::CreateLauncher(
+    scoped_refptr<base::SingleThreadTaskRunner> io_task_runner) {
 
 #if defined(REMOTING_MULTI_PROCESS)
 
   child_ = DaemonProcess::Create(
       main_task_runner_,
-      io_thread.message_loop_proxy(),
+      io_task_runner,
       base::Bind(&HostService::OnChildStopped,
                  base::Unretained(this))).PassAs<Stoppable>();
 
 #else  // !defined(REMOTING_MULTI_PROCESS)
 
   // Create the session process launcher.
   child_.reset(new WtsSessionProcessLauncher(
       base::Bind(&HostService::OnChildStopped, base::Unretained(this)),
       this,
       main_task_runner_,
-      io_thread.message_loop_proxy()));
+      io_task_runner));
 
 #endif  // !defined(REMOTING_MULTI_PROCESS)
+}
+
+void HostService::RunMessageLoop(MessageLoop* message_loop) {
+  // Launch the I/O thread.
+  base::Thread io_thread(kIoThreadName);
+  base::Thread::Options io_thread_options(MessageLoop::TYPE_IO, 0);
+  if (!io_thread.StartWithOptions(io_thread_options)) {
+    LOG(FATAL) << "Failed to start the I/O thread";
+    return;
+  }
+
+  CreateLauncher(new AutoMessageLoop(main_task_runner_,
+                                     io_thread.message_loop_proxy(),
+                                     base::Closure()));
 
   // Run the service.
   message_loop->Run();
-
-  // Release the control handler.
-  stopped_event_.Signal();
 }
 
 int HostService::Elevate() {
   // Get the name of the binary to launch.
   FilePath binary =
       CommandLine::ForCurrentProcess()->GetSwitchValuePath(kElevateSwitchName);
 
   // Create the child process command line by copying known switches from our
   // command line.
   CommandLine command_line(CommandLine::NO_PROGRAM);
@@ skipping 29 matching lines @@
         << "Failed to connect to the service control manager";
     return kErrorExitCode;
   }
 
   return kSuccessExitCode;
 }
 
 int HostService::RunInConsole() {
   MessageLoop message_loop(MessageLoop::TYPE_UI);
 
-  // Allow other threads to post to our message loop.
-  main_task_runner_ = message_loop.message_loop_proxy();
+  // Keep a reference to the main message loop while it is used. Once the last
+  // reference is dropped, QuitClosure() will be posted to the loop.
+  main_task_runner_ =
+      new AutoMessageLoop(NULL, message_loop.message_loop_proxy(),
+                          base::Bind(&QuitMessageLoop,
+                                     base::Unretained(&message_loop)));
 
   int result = kErrorExitCode;
 
   // Subscribe to Ctrl-C and other console events.
   if (!SetConsoleCtrlHandler(&HostService::ConsoleControlHandler, TRUE)) {
     LOG_GETLASTERROR(ERROR)
         << "Failed to set console control handler";
     return result;
   }
 
@@ skipping 26 matching lines @@
   // session.
   main_task_runner_->PostTask(FROM_HERE, base::Bind(
       &HostService::OnSessionChange, base::Unretained(this)));
 
   // Subscribe to session change notifications.
   if (WTSRegisterSessionNotification(window,
                                      NOTIFY_FOR_ALL_SESSIONS) != FALSE) {
     // Run the service.
     RunMessageLoop(&message_loop);
 
+    // Release the control handler.
+    stopped_event_.Signal();
+
     WTSUnRegisterSessionNotification(window);
     result = kSuccessExitCode;
   }
 
 cleanup:
   if (window != NULL) {
     DestroyWindow(window);
   }
 
   if (atom != 0) {
@@ skipping 28 matching lines @@
       self->main_task_runner_->PostTask(FROM_HERE, base::Bind(
           &HostService::OnSessionChange, base::Unretained(self)));
       return NO_ERROR;
 
     default:
       return ERROR_CALL_NOT_IMPLEMENTED;
   }
 }
 
 VOID WINAPI HostService::ServiceMain(DWORD argc, WCHAR* argv[]) {
-  MessageLoop message_loop;
+  MessageLoop message_loop(MessageLoop::TYPE_DEFAULT);
 
-  // Allow other threads to post to our message loop.
+  // Keep a reference to the main message loop while it is used. Once the last
+  // reference is dropped QuitClosure() will be posted to the loop.
   HostService* self = HostService::GetInstance();
-  self->main_task_runner_ = message_loop.message_loop_proxy();
+  self->main_task_runner_ =
+      new AutoMessageLoop(NULL, message_loop.message_loop_proxy(),
+                          base::Bind(&QuitMessageLoop,
+                                     base::Unretained(&message_loop)));
 
   // Register the service control handler.
   self->service_status_handle_ =
       RegisterServiceCtrlHandlerExW(kWindowsServiceName,
                                     &HostService::ServiceControlHandler,
                                     self);
   if (self->service_status_handle_ == 0) {
     LOG_GETLASTERROR(ERROR)
         << "Failed to register the service control handler";
     return;
@@ skipping 16 matching lines @@
   }
 
   // Post a dummy session change notification to peek up the current console
   // session.
   self->main_task_runner_->PostTask(FROM_HERE, base::Bind(
       &HostService::OnSessionChange, base::Unretained(self)));
 
   // Run the service.
   self->RunMessageLoop(&message_loop);
 
+  // Release the control handler.
+  self->stopped_event_.Signal();
+
   // Tell SCM that the service is stopped.
   service_status.dwCurrentState = SERVICE_STOPPED;
   service_status.dwControlsAccepted = 0;
 
   if (!SetServiceStatus(self->service_status_handle_, &service_status)) {
     LOG_GETLASTERROR(ERROR)
         << "Failed to report service status to the service control manager";
     return;
   }
 }
@@ skipping 51 matching lines @@
   }
 
   remoting::HostService* service = remoting::HostService::GetInstance();
   if (!service->InitWithCommandLine(command_line)) {
     usage(command_line->GetProgram());
     return kUsageExitCode;
   }
 
   return service->Run();
 }