Replace Direct WTS Registration Calls with the Windows Session Change Observer

Each WTSRegisterSessionNotification call has the potential to jank the UI
thread. This change reduces the jank to once per UI process by registering a
singleton observer bound to the lifetime of the process.

Any new UI thread component that needs to know about Windows Session Change
notifications registers against this component instead of calling
WTSRegisterSessionNotification to avoid janking the thread.

BUG=477462

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

Cr-Commit-Position: refs/heads/master@{#328406}

6 files changed, 170 insertions, 32 deletions

diff --git a/ui/gfx/win/msg_util.h b/ui/gfx/win/msg_util.h
index b973941..0eb5123 100644
--- a/ui/gfx/win/msg_util.h
+++ b/ui/gfx/win/msg_util.h
@@ -1750,17 +1750,6 @@
     }                                                      \
   }
 
-// void OnWTSSessionChange(WPARAM nStatusCode, PWTSSESSION_NOTIFICATION
-// nSessionID)
-#define CR_MSG_WM_WTSSESSION_CHANGE(func)           \
-  if (uMsg == WM_WTSSESSION_CHANGE) {               \
-    SetMsgHandled(TRUE);                            \
-    func(wParam, (PWTSSESSION_NOTIFICATION)lParam); \
-    lResult = 0;                                    \
-    if (IsMsgHandled())                             \
-      return TRUE;                                  \
-  }
-
 // OnThemeChanged()
 #define CR_MSG_WM_THEMECHANGED(func) \
   if (uMsg == WM_THEMECHANGED) {     \
diff --git a/ui/views/views.gyp b/ui/views/views.gyp
index 89afb9e..272c7a7 100644
--- a/ui/views/views.gyp
+++ b/ui/views/views.gyp
@@ -347,6 +347,8 @@
       'win/hwnd_util_aurawin.cc',
       'win/scoped_fullscreen_visibility.cc',
       'win/scoped_fullscreen_visibility.h',
+      'win/windows_session_change_observer.cc',
+      'win/windows_session_change_observer.h',
     ],
     'views_aura_sources': [
       'accessibility/ax_aura_obj_cache.cc',
diff --git a/ui/views/win/hwnd_message_handler.cc b/ui/views/win/hwnd_message_handler.cc
index 8499cbf..653fe3c 100644
--- a/ui/views/win/hwnd_message_handler.cc
+++ b/ui/views/win/hwnd_message_handler.cc
@@ -7,10 +7,9 @@
 #include <dwmapi.h>
 #include <oleacc.h>
 #include <shellapi.h>
-#include <wtsapi32.h>
-#pragma comment(lib, "wtsapi32.lib")
 
 #include "base/bind.h"
+#include "base/bind_helpers.h"
 #include "base/profiler/scoped_tracker.h"
 #include "base/trace_event/trace_event.h"
 #include "base/tracked_objects.h"
@@ -43,6 +42,7 @@
 #include "ui/views/win/fullscreen_handler.h"
 #include "ui/views/win/hwnd_message_handler_delegate.h"
 #include "ui/views/win/scoped_fullscreen_visibility.h"
+#include "ui/views/win/windows_session_change_observer.h"
 
 namespace views {
 namespace {
@@ -1464,7 +1464,9 @@ LRESULT HWNDMessageHandler::OnCreate(CREATESTRUCT* create_struct) {
   tracked_objects::ScopedTracker tracking_profile8(
       FROM_HERE_WITH_EXPLICIT_FUNCTION("440919 HWNDMessageHandler::OnCreate8"));
 
-  WTSRegisterSessionNotification(hwnd(), NOTIFY_FOR_THIS_SESSION);
+  windows_session_change_observer_.reset(new WindowsSessionChangeObserver(
+      base::Bind(&HWNDMessageHandler::OnSessionChange,
+                 base::Unretained(this))));
 
   // TODO(beng): move more of NWW::OnCreate here.
   return 0;
@@ -1475,7 +1477,7 @@ void HWNDMessageHandler::OnDestroy() {
   tracked_objects::ScopedTracker tracking_profile(
       FROM_HERE_WITH_EXPLICIT_FUNCTION("440919 HWNDMessageHandler::OnDestroy"));
 
-  WTSUnRegisterSessionNotification(hwnd());
+  windows_session_change_observer_.reset(nullptr);
   delegate_->HandleDestroying();
 }
 
@@ -2207,21 +2209,6 @@ LRESULT HWNDMessageHandler::OnScrollMessage(UINT message,
   return 0;
 }
 
-void HWNDMessageHandler::OnSessionChange(WPARAM status_code,
-                                         PWTSSESSION_NOTIFICATION session_id) {
-  // TODO(vadimt): Remove ScopedTracker below once crbug.com/440919 is fixed.
-  tracked_objects::ScopedTracker tracking_profile(
-      FROM_HERE_WITH_EXPLICIT_FUNCTION(
-          "440919 HWNDMessageHandler::OnSessionChange"));
-
-  // Direct3D presents are ignored while the screen is locked, so force the
-  // window to be redrawn on unlock.
-  if (status_code == WTS_SESSION_UNLOCK)
-    ForceRedrawWindow(10);
-
-  SetMsgHandled(FALSE);
-}
-
 LRESULT HWNDMessageHandler::OnSetCursor(UINT message,
                                         WPARAM w_param,
                                         LPARAM l_param) {
@@ -2634,6 +2621,18 @@ void HWNDMessageHandler::OnWindowPosChanged(WINDOWPOS* window_pos) {
   SetMsgHandled(FALSE);
 }
 
+void HWNDMessageHandler::OnSessionChange(WPARAM status_code) {
+  // TODO(vadimt): Remove ScopedTracker below once crbug.com/440919 is fixed.
+  tracked_objects::ScopedTracker tracking_profile(
+      FROM_HERE_WITH_EXPLICIT_FUNCTION(
+          "440919 HWNDMessageHandler::OnSessionChange"));
+
+  // Direct3D presents are ignored while the screen is locked, so force the
+  // window to be redrawn on unlock.
+  if (status_code == WTS_SESSION_UNLOCK)
+    ForceRedrawWindow(10);
+}
+
 void HWNDMessageHandler::HandleTouchEvents(const TouchEvents& touch_events) {
   base::WeakPtr<HWNDMessageHandler> ref(weak_factory_.GetWeakPtr());
   for (size_t i = 0; i < touch_events.size() && ref; ++i)
diff --git a/ui/views/win/hwnd_message_handler.h b/ui/views/win/hwnd_message_handler.h
index d2dc968..a011f76 100644
--- a/ui/views/win/hwnd_message_handler.h
+++ b/ui/views/win/hwnd_message_handler.h
@@ -42,6 +42,7 @@ namespace views {
 class FullscreenHandler;
 class HWNDMessageHandlerDelegate;
 class InputMethod;
+class WindowsSessionChangeObserver;
 
 // These two messages aren't defined in winuser.h, but they are sent to windows
 // with captions. They appear to paint the window caption and frame.
@@ -400,7 +401,6 @@ class VIEWS_EXPORT HWNDMessageHandler :
     CR_MSG_WM_THEMECHANGED(OnThemeChanged)
     CR_MSG_WM_WINDOWPOSCHANGED(OnWindowPosChanged)
     CR_MSG_WM_WINDOWPOSCHANGING(OnWindowPosChanging)
-    CR_MSG_WM_WTSSESSION_CHANGE(OnSessionChange)
   CR_END_MSG_MAP()
 
   // Message Handlers.
@@ -445,7 +445,6 @@ class VIEWS_EXPORT HWNDMessageHandler :
   void OnPaint(HDC dc);
   LRESULT OnReflectedMessage(UINT message, WPARAM w_param, LPARAM l_param);
   LRESULT OnScrollMessage(UINT message, WPARAM w_param, LPARAM l_param);
-  void OnSessionChange(WPARAM status_code, PWTSSESSION_NOTIFICATION session_id);
   LRESULT OnSetCursor(UINT message, WPARAM w_param, LPARAM l_param);
   void OnSetFocus(HWND last_focused_window);
   LRESULT OnSetIcon(UINT size_type, HICON new_icon);
@@ -458,6 +457,9 @@ class VIEWS_EXPORT HWNDMessageHandler :
   void OnWindowPosChanging(WINDOWPOS* window_pos);
   void OnWindowPosChanged(WINDOWPOS* window_pos);
 
+  // Receives Windows Session Change notifications.
+  void OnSessionChange(WPARAM status_code);
+
   typedef std::vector<ui::TouchEvent> TouchEvents;
   // Helper to handle the list of touch events passed in. We need this because
   // touch events on windows don't fire if we enter a modal loop in the context
@@ -629,6 +631,9 @@ class VIEWS_EXPORT HWNDMessageHandler :
   // The factory used with BEGIN_SAFE_MSG_MAP_EX.
   base::WeakPtrFactory<HWNDMessageHandler> weak_factory_;
 
+  // Manages observation of Windows Session Change messages.
+  scoped_ptr<WindowsSessionChangeObserver> windows_session_change_observer_;
+
   DISALLOW_COPY_AND_ASSIGN(HWNDMessageHandler);
 };
 
diff --git a/ui/views/win/windows_session_change_observer.cc b/ui/views/win/windows_session_change_observer.cc
new file mode 100644
index 0000000..2b831f2
--- /dev/null
+++ b/ui/views/win/windows_session_change_observer.cc
@@ -0,0 +1,107 @@
+// Copyright 2015 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 "ui/views/win/windows_session_change_observer.h"
+
+#include <wtsapi32.h>
+#pragma comment(lib, "wtsapi32.lib")
+
+#include "base/bind.h"
+#include "base/bind_helpers.h"
+#include "base/memory/scoped_ptr.h"
+#include "base/memory/singleton.h"
+#include "base/observer_list.h"
+#include "ui/gfx/win/singleton_hwnd.h"
+#include "ui/gfx/win/singleton_hwnd_observer.h"
+
+namespace views {
+
+class WindowsSessionChangeObserver::WtsRegistrationNotificationManager {
+ public:
+  static WtsRegistrationNotificationManager* GetInstance() {
+    return Singleton<WtsRegistrationNotificationManager>::get();
+  }
+
+  WtsRegistrationNotificationManager() {
+    AddSingletonHwndObserver();
+  }
+
+  ~WtsRegistrationNotificationManager() {
+    RemoveSingletonHwndObserver();
+  }
+
+  void AddObserver(WindowsSessionChangeObserver* observer) {
+    DCHECK(singleton_hwnd_observer_);
+    observer_list_.AddObserver(observer);
+  }
+
+  void RemoveObserver(WindowsSessionChangeObserver* observer) {
+    observer_list_.RemoveObserver(observer);
+  }
+
+ private:
+  friend struct DefaultSingletonTraits<WtsRegistrationNotificationManager>;
+
+  void OnWndProc(HWND hwnd, UINT message, WPARAM wparam, LPARAM lparam) {
+    switch (message) {
+      case WM_WTSSESSION_CHANGE:
+        FOR_EACH_OBSERVER(WindowsSessionChangeObserver,
+                          observer_list_,
+                          OnSessionChange(wparam));
+        break;
+      case WM_DESTROY:
+        RemoveSingletonHwndObserver();
+        break;
+    }
+  }
+
+  void AddSingletonHwndObserver() {
+    DCHECK(!singleton_hwnd_observer_);
+    singleton_hwnd_observer_.reset(new gfx::SingletonHwndObserver(
+        base::Bind(&WtsRegistrationNotificationManager::OnWndProc,
+                   base::Unretained(this))));
+    WTSRegisterSessionNotification(gfx::SingletonHwnd::GetInstance()->hwnd(),
+                                   NOTIFY_FOR_THIS_SESSION);
+  }
+
+  void RemoveSingletonHwndObserver() {
+    if (!singleton_hwnd_observer_)
+      return;
+
+    singleton_hwnd_observer_.reset(nullptr);
+    WTSUnRegisterSessionNotification(gfx::SingletonHwnd::GetInstance()->hwnd());
+    FOR_EACH_OBSERVER(WindowsSessionChangeObserver,
+                      observer_list_,
+                      ClearCallback());
+  }
+
+  ObserverList<WindowsSessionChangeObserver, true> observer_list_;
+  scoped_ptr<gfx::SingletonHwndObserver> singleton_hwnd_observer_;
+
+  DISALLOW_COPY_AND_ASSIGN(WtsRegistrationNotificationManager);
+};
+
+WindowsSessionChangeObserver::WindowsSessionChangeObserver(
+    const WtsCallback& callback)
+    : callback_(callback) {
+  DCHECK(!callback_.is_null());
+  WtsRegistrationNotificationManager::GetInstance()->AddObserver(this);
+}
+
+WindowsSessionChangeObserver::~WindowsSessionChangeObserver() {
+  ClearCallback();
+}
+
+void WindowsSessionChangeObserver::OnSessionChange(WPARAM wparam) {
+  callback_.Run(wparam);
+}
+
+void WindowsSessionChangeObserver::ClearCallback() {
+  if (!callback_.is_null()) {
+    callback_.Reset();
+    WtsRegistrationNotificationManager::GetInstance()->RemoveObserver(this);
+  }
+}
+
+}  // namespace views
diff --git a/ui/views/win/windows_session_change_observer.h b/ui/views/win/windows_session_change_observer.h
new file mode 100644
index 0000000..38a3e08
--- /dev/null
+++ b/ui/views/win/windows_session_change_observer.h
@@ -0,0 +1,36 @@
+// Copyright 2015 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 UI_VIEWS_WIN_WINDOWS_SESSION_CHANGE_OBSERVER_H_
+#define UI_VIEWS_WIN_WINDOWS_SESSION_CHANGE_OBSERVER_H_
+
+#include <windows.h>
+
+#include "base/callback.h"
+#include "base/macros.h"
+
+namespace views {
+
+// Calls the provided callback on WM_WTSSESSION_CHANGE messages along with
+// managing the tricky business of observing a singleton object.
+class WindowsSessionChangeObserver {
+ public:
+  typedef base::Callback<void(WPARAM)> WtsCallback;
+  explicit WindowsSessionChangeObserver(const WtsCallback& callback);
+  ~WindowsSessionChangeObserver();
+
+ private:
+  class WtsRegistrationNotificationManager;
+
+  void OnSessionChange(WPARAM wparam);
+  void ClearCallback();
+
+  WtsCallback callback_;
+
+  DISALLOW_COPY_AND_ASSIGN(WindowsSessionChangeObserver);
+};
+
+}  // namespace views
+
+#endif  // UI_VIEWS_WIN_WINDOWS_SESSION_CHANGE_OBSERVER_H_