Use a window to catch WM_ENDSESSION.

chrome/chrome_watcher/chrome_watcher_main.cc

Index: chrome/chrome_watcher/chrome_watcher_main.cc
diff --git a/chrome/chrome_watcher/chrome_watcher_main.cc b/chrome/chrome_watcher/chrome_watcher_main.cc
index 8dab91a01e71e7c0b5450ae1b5c4a0eb5ce7d77d..f167daadd31ea41b9ae358729688ef04ef049cc9 100644
--- a/chrome/chrome_watcher/chrome_watcher_main.cc
+++ b/chrome/chrome_watcher/chrome_watcher_main.cc
@@ -5,11 +5,22 @@
 #include <windows.h>
 
 #include "base/at_exit.h"
+#include "base/bind.h"
+#include "base/bind_helpers.h"
 #include "base/command_line.h"
 #include "base/logging_win.h"
+#include "base/macros.h"
+#include "base/memory/ref_counted.h"
+#include "base/message_loop/message_loop.h"
+#include "base/message_loop/message_loop_proxy.h"
 #include "base/process/process.h"
+#include "base/run_loop.h"
+#include "base/sequenced_task_runner.h"
 #include "base/synchronization/waitable_event.h"
 #include "base/template_util.h"
+#include "base/threading/thread.h"
+#include "base/time/time.h"
+#include "components/browser_watcher/endsession_watcher_window_win.h"
 #include "components/browser_watcher/exit_code_watcher_win.h"
 #include "components/browser_watcher/exit_funnel_win.h"
 #include "components/browser_watcher/watcher_main_api_win.h"
@@ -22,32 +33,128 @@ const GUID kChromeWatcherTraceProviderName = {
     0x7fe69228, 0x633e, 0x4f06,
         { 0x80, 0xc1, 0x52, 0x7f, 0xea, 0x23, 0xe3, 0xa7 } };
 
-// The data shared from the main function to the console handler.
-struct HandlerData {
-  HandlerData() : handler_done(true /* manual_reset */,
-                               false /* initially_signaled */) {
-  }
+// Takes care of monitoring a browser. This class watches for a browser's exit
+// code, as well as listening for WM_ENDSESSION messages. Events are recorded
+// in an exit funnel, for reporting the next time Chrome runs.
+class BrowserMonitor {
+ public:
+  BrowserMonitor(base::RunLoop* run_loop, const base::char16* registry_path);
+  ~BrowserMonitor();
+
+  // Starts the monitor, returns true on success.
+  bool StartWatching(const base::char16* registry_path,
+                     base::Process process);
+
+ private:
+  // Called from EndSessionWatcherWindow on a WM_ENDSESSION message.
+  void OnEndSession(LPARAM lparam);
+
+  // Blocking function that runs on |background_thread_|.
+  void Watch();
+
+  // Posted to main thread from Watch when browser exits.
+  void BrowserExited();
+
+  // True if BrowserExited has run.
+  bool browser_exited_;
+
+  // The funnel used to record events for this browser.
+  browser_watcher::ExitFunnel exit_funnel_;
+
+  browser_watcher::ExitCodeWatcher exit_code_watcher_;
+  browser_watcher::EndSessionWatcherWindow end_session_watcher_window_;
 
-  base::WaitableEvent handler_done;
-  base::ProcessHandle process;
-  const base::char16* registry_path;
+  // The thread that runs Watch().
+  base::Thread background_thread_;
+
+  // The run loop for the main thread and its task runner.
+  base::RunLoop* run_loop_;
+  scoped_refptr<base::SequencedTaskRunner> main_thread_;
+
+  DISALLOW_COPY_AND_ASSIGN(BrowserMonitor);
 };
 
-HandlerData *g_handler_data = nullptr;
+BrowserMonitor::BrowserMonitor(base::RunLoop* run_loop,
+                               const base::char16* registry_path) :
+    browser_exited_(false),
+    exit_code_watcher_(registry_path),
+    end_session_watcher_window_(
+        base::Bind(&BrowserMonitor::OnEndSession, base::Unretained(this))),
+    background_thread_("BrowserWatcherThread"),
+    run_loop_(run_loop),
+    main_thread_(base::MessageLoopProxy::current()) {
+}
+
+BrowserMonitor::~BrowserMonitor() {
+}
+
+bool BrowserMonitor::StartWatching(const base::char16* registry_path,
+                                   base::Process process) {
+  if (!exit_code_watcher_.Initialize(process.Pass()))
+    return false;
+
+  if (!exit_funnel_.Init(registry_path,
+                        exit_code_watcher_.process().Handle())) {
+    return false;
+  }
 
-BOOL CALLBACK ConsoleCtrlHandler(DWORD ctl_type) {
-  if (g_handler_data && ctl_type == CTRL_LOGOFF_EVENT) {
-    // Record the watcher logoff event in the browser's exit funnel.
-    browser_watcher::ExitFunnel funnel;
-    funnel.Init(g_handler_data->registry_path, g_handler_data->process);
-    funnel.RecordEvent(L"WatcherLogoff");
+  if (!background_thread_.StartWithOptions(
+        base::Thread::Options(base::MessageLoop::TYPE_IO, 0))) {
+    return false;
+  }
 
-    // Release the main function.
-    g_handler_data->handler_done.Signal();
+  if (!background_thread_.task_runner()->PostTask(FROM_HERE,
+        base::Bind(&BrowserMonitor::Watch, base::Unretained(this)))) {
+    background_thread_.Stop();
+    return false;
   }
 
-  // Fall through to the next handler in turn.
-  return FALSE;
+  return true;
+}
+
+void BrowserMonitor::OnEndSession(LPARAM lparam) {
+  DCHECK_EQ(main_thread_, base::MessageLoopProxy::current());
+
+  exit_funnel_.RecordEvent(L"WatcherLogoff");
+
+  // Belt-and-suspenders; make sure our message loop exits ASAP.
+  if (browser_exited_)
+    run_loop_->Quit();
+}
+
+void BrowserMonitor::Watch() {
+  // This needs to run on an IO thread.
+  DCHECK_NE(main_thread_, base::MessageLoopProxy::current());
+
+  exit_code_watcher_.WaitForExit();
+  exit_funnel_.RecordEvent(L"BrowserExit");
+
+  main_thread_->PostTask(FROM_HERE,
+      base::Bind(&BrowserMonitor::BrowserExited, base::Unretained(this)));
+}
+
+void BrowserMonitor::BrowserExited() {
+  // This runs in the main thread.
+  DCHECK_EQ(main_thread_, base::MessageLoopProxy::current());
+
+  // Note that the browser has exited.
+  browser_exited_ = true;
+
+  // Our background thread has served it's purpose.
+  background_thread_.Stop();
+
+  const int exit_code = exit_code_watcher_.exit_code();
+  if (exit_code >= 0 && exit_code <= 28) {
+    // The browser exited with a well-known exit code, quit this process
+    // immediately.
+    run_loop_->Quit();
+  } else {
+    // The browser exited abnormally, wait around for a little bit to see
+    // whether this instance will get a logoff message.
+    main_thread_->PostDelayedTask(FROM_HERE,
+                                  run_loop_->QuitClosure(),
+                                  base::TimeDelta::FromSeconds(30));
+  }
 }
 
 }  // namespace
@@ -57,6 +164,7 @@ BOOL CALLBACK ConsoleCtrlHandler(DWORD ctl_type) {
 extern "C" int WatcherMain(const base::char16* registry_path,
                            HANDLE process_handle) {
   base::Process process(process_handle);
+
   // The exit manager is in charge of calling the dtors of singletons.
   base::AtExitManager exit_manager;
   // Initialize the commandline singleton from the environment.
@@ -66,52 +174,23 @@ extern "C" int WatcherMain(const base::char16* registry_path,
 
   // Arrange to be shut down as late as possible, as we want to outlive
   // chrome.exe in order to report its exit status.
-  // TODO(siggi): Does this (windowless) process need to register a console
-  //     handler too, in order to get notice of logoff events?
   ::SetProcessShutdownParameters(0x100, SHUTDOWN_NORETRY);
 
-  browser_watcher::ExitCodeWatcher exit_code_watcher(registry_path);
-  int ret = 1;
-  // Attempt to wait on our parent process, and record its exit status.
-  if (exit_code_watcher.Initialize(process.Duplicate())) {
-    // Set up a console control handler, and provide it the data it needs
-    // to record into the browser's exit funnel.
-    HandlerData data;
-    data.process = process.Handle();
-    data.registry_path = registry_path;
-    g_handler_data = &data;
-    ::SetConsoleCtrlHandler(&ConsoleCtrlHandler, TRUE /* Add */);
-
-    // Wait on the process.
-    exit_code_watcher.WaitForExit();
-
-    browser_watcher::ExitFunnel funnel;
-    funnel.Init(registry_path, process.Handle());
-    funnel.RecordEvent(L"BrowserExit");
-
-    // Chrome/content exit codes are currently in the range of 0-28.
-    // Anything outside this range is strange, so watch harder.
-    int exit_code = exit_code_watcher.exit_code();
-    if (exit_code < 0 || exit_code > 28) {
-      // Wait for a max of 30 seconds to see whether we get notified of logoff.
-      data.handler_done.TimedWait(base::TimeDelta::FromSeconds(30));
-    }
-
-    // Remove the console control handler.
-    // There is a potential race here, where the handler might be executing
-    // already as we fall through here. Hopefully SetConsoleCtrlHandler is
-    // synchronizing against handler execution, for there's no other way to
-    // close the race. Thankfully we'll just get snuffed out, as this is logoff.
-    ::SetConsoleCtrlHandler(&ConsoleCtrlHandler, FALSE /* Add */);
-    g_handler_data = nullptr;
-
-    ret = 0;
-  }
+  // Run a UI message loop on the main thread.
+  base::MessageLoop msg_loop(base::MessageLoop::TYPE_UI);
+  msg_loop.set_thread_name("WatcherMainThread");
+
+  base::RunLoop run_loop;
+  BrowserMonitor monitor(&run_loop, registry_path);
+  if (!monitor.StartWatching(registry_path, process.Pass()))
+    return 1;
+
+  run_loop.Run();
 
   // Wind logging down.
   logging::LogEventProvider::Uninitialize();
 
-  return ret;
+  return 0;
 }
 
 static_assert(base::is_same<decltype(&WatcherMain),