Fix memory leak issue around partial screenshot view on tests.

BUG=166995
TEST=run heapcheck and make sure the leak doesn't happen anymore


Review URL: https://chromiumcodereview.appspot.com/11645038

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

2 files changed, 94 insertions, 72 deletions

diff --git a/ash/wm/partial_screenshot_view.cc b/ash/wm/partial_screenshot_view.cc
index 2510d06..8680509 100644
--- a/ash/wm/partial_screenshot_view.cc
+++ b/ash/wm/partial_screenshot_view.cc
@@ -4,6 +4,8 @@
 
 #include "ash/wm/partial_screenshot_view.h"
 
+#include <algorithm>
+
 #include "ash/display/mouse_cursor_event_filter.h"
 #include "ash/screenshot_delegate.h"
 #include "ash/shell.h"
@@ -17,53 +19,27 @@
 #include "ui/gfx/rect.h"
 #include "ui/views/view.h"
 #include "ui/views/widget/widget.h"
+#include "ui/views/widget/widget_observer.h"
 
 namespace ash {
-namespace {
-class PartialScreenshotOverlayDelegate
-    : public internal::OverlayEventFilter::Delegate {
+
+// A self-owned object to handle the cancel and the finish of current partial
+// screenshot session.
+class PartialScreenshotView::OverlayDelegate
+    : public internal::OverlayEventFilter::Delegate,
+      public views::WidgetObserver {
  public:
-  PartialScreenshotOverlayDelegate() {
+  OverlayDelegate() {
     Shell::GetInstance()->overlay_filter()->Activate(this);
   }
 
-  virtual ~PartialScreenshotOverlayDelegate() {
-    Shell::GetInstance()->overlay_filter()->Deactivate();
-  }
-
-  void CreatePartialScreenshotUI(ScreenshotDelegate* screenshot_delegate,
-                                 aura::RootWindow* root_window) {
-    views::Widget* widget = new views::Widget;
-    PartialScreenshotView* view = new PartialScreenshotView(
-        this, screenshot_delegate);
-    views::Widget::InitParams params(
-        views::Widget::InitParams::TYPE_WINDOW_FRAMELESS);
-    params.transparent = true;
-    params.delegate = view;
-    // The partial screenshot rectangle has to be at the real top of
-    // the screen.
-    params.parent = Shell::GetContainer(
-        root_window,
-        internal::kShellWindowId_OverlayContainer);
-
-    widget->Init(params);
-    widget->SetContentsView(view);
-    widget->SetBounds(root_window->GetBoundsInScreen());
-    widget->GetNativeView()->SetName("PartialScreenshotView");
-    widget->StackAtTop();
-    widget->Show();
-    // Releases the mouse capture to let mouse events come to the view. This
-    // will close the context menu.
-    aura::client::CaptureClient* capture_client =
-        aura::client::GetCaptureClient(root_window);
-    if (capture_client->GetCaptureWindow())
-      capture_client->ReleaseCapture(capture_client->GetCaptureWindow());
-
+  void RegisterWidget(views::Widget* widget) {
     widgets_.push_back(widget);
+    widget->AddObserver(this);
   }
 
-  // OverlayEventFilter::Delegate overrides:
-  void Cancel() OVERRIDE {
+  // Overridden from OverlayEventFilter::Delegate:
+  virtual void Cancel() OVERRIDE {
     // Make sure the mouse_warp_mode allows warping. It can be stopped by a
     // partial screenshot view.
     internal::MouseCursorEventFilter* mouse_cursor_filter =
@@ -72,28 +48,51 @@ class PartialScreenshotOverlayDelegate
         internal::MouseCursorEventFilter::WARP_ALWAYS);
     for (size_t i = 0; i < widgets_.size(); ++i)
       widgets_[i]->Close();
-    delete this;
   }
 
-  bool IsCancelingKeyEvent(ui::KeyEvent* event) OVERRIDE {
+  virtual bool IsCancelingKeyEvent(ui::KeyEvent* event) OVERRIDE {
     return event->key_code() == ui::VKEY_ESCAPE;
   }
 
-  aura::Window* GetWindow() OVERRIDE {
+  virtual aura::Window* GetWindow() OVERRIDE {
     // Just returns NULL because this class does not handle key events in
     // OverlayEventFilter, except for cancel keys.
     return NULL;
   }
 
+  // Overridden from views::WidgetObserver:
+  virtual void OnWidgetClosing(views::Widget* widget) {
+    widget->RemoveObserver(this);
+    widgets_.erase(std::remove(widgets_.begin(), widgets_.end(), widget));
+    if (widgets_.empty())
+      delete this;
+  }
+
  private:
+  virtual ~OverlayDelegate() {
+    Shell::GetInstance()->overlay_filter()->Deactivate();
+  }
+
   std::vector<views::Widget*> widgets_;
 
-  DISALLOW_COPY_AND_ASSIGN(PartialScreenshotOverlayDelegate);
+  DISALLOW_COPY_AND_ASSIGN(OverlayDelegate);
 };
+
+// static
+void PartialScreenshotView::StartPartialScreenshot(
+    ScreenshotDelegate* screenshot_delegate) {
+  OverlayDelegate* overlay_delegate = new OverlayDelegate();
+  Shell::RootWindowList root_windows = Shell::GetAllRootWindows();
+  for (Shell::RootWindowList::iterator it = root_windows.begin();
+       it != root_windows.end(); ++it) {
+    PartialScreenshotView* new_view = new PartialScreenshotView(
+        overlay_delegate, screenshot_delegate);
+    new_view->Init(*it);
+  }
 }
 
 PartialScreenshotView::PartialScreenshotView(
-    internal::OverlayEventFilter::Delegate* overlay_delegate,
+    PartialScreenshotView::OverlayDelegate* overlay_delegate,
     ScreenshotDelegate* screenshot_delegate)
     : is_dragging_(false),
       overlay_delegate_(overlay_delegate),
@@ -101,19 +100,44 @@ PartialScreenshotView::PartialScreenshotView(
 }
 
 PartialScreenshotView::~PartialScreenshotView() {
+  overlay_delegate_ = NULL;
   screenshot_delegate_ = NULL;
 }
 
-// static
-void PartialScreenshotView::StartPartialScreenshot(
-    ScreenshotDelegate* screenshot_delegate) {
-  PartialScreenshotOverlayDelegate* overlay_delegate =
-      new PartialScreenshotOverlayDelegate();
-  Shell::RootWindowList root_windows = Shell::GetAllRootWindows();
-  for (Shell::RootWindowList::iterator it = root_windows.begin();
-       it != root_windows.end(); ++it) {
-    overlay_delegate->CreatePartialScreenshotUI(screenshot_delegate, *it);
-  }
+void PartialScreenshotView::Init(aura::RootWindow* root_window) {
+  views::Widget* widget = new views::Widget;
+  views::Widget::InitParams params(
+      views::Widget::InitParams::TYPE_WINDOW_FRAMELESS);
+  params.transparent = true;
+  params.delegate = this;
+  // The partial screenshot rectangle has to be at the real top of
+  // the screen.
+  params.parent = Shell::GetContainer(
+      root_window,
+      internal::kShellWindowId_OverlayContainer);
+
+  widget->Init(params);
+  widget->SetContentsView(this);
+  widget->SetBounds(root_window->GetBoundsInScreen());
+  widget->GetNativeView()->SetName("PartialScreenshotView");
+  widget->StackAtTop();
+  widget->Show();
+  // Releases the mouse capture to let mouse events come to the view. This
+  // will close the context menu.
+  aura::client::CaptureClient* capture_client =
+      aura::client::GetCaptureClient(root_window);
+  if (capture_client->GetCaptureWindow())
+    capture_client->ReleaseCapture(capture_client->GetCaptureWindow());
+
+  overlay_delegate_->RegisterWidget(widget);
+}
+
+gfx::Rect PartialScreenshotView::GetScreenshotRect() const {
+  int left = std::min(start_position_.x(), current_position_.x());
+  int top = std::min(start_position_.y(), current_position_.y());
+  int width = ::abs(start_position_.x() - current_position_.x());
+  int height = ::abs(start_position_.y() - current_position_.y());
+  return gfx::Rect(left, top, width, height);
 }
 
 gfx::NativeCursor PartialScreenshotView::GetCursor(
@@ -174,12 +198,4 @@ void PartialScreenshotView::OnMouseReleased(const ui::MouseEvent& event) {
   }
 }
 
-gfx::Rect PartialScreenshotView::GetScreenshotRect() const {
-  int left = std::min(start_position_.x(), current_position_.x());
-  int top = std::min(start_position_.y(), current_position_.y());
-  int width = ::abs(start_position_.x() - current_position_.x());
-  int height = ::abs(start_position_.y() - current_position_.y());
-  return gfx::Rect(left, top, width, height);
-}
-
 }  // namespace ash
diff --git a/ash/wm/partial_screenshot_view.h b/ash/wm/partial_screenshot_view.h
index 9ae3430..a2ca607 100644
--- a/ash/wm/partial_screenshot_view.h
+++ b/ash/wm/partial_screenshot_view.h
@@ -6,11 +6,14 @@
 #define ASH_WM_PARTIAL_SCREENSHOT_VIEW_H_
 
 #include "ash/ash_export.h"
-#include "ash/wm/overlay_event_filter.h"
 #include "base/compiler_specific.h"
 #include "ui/gfx/point.h"
 #include "ui/views/widget/widget_delegate.h"
 
+namespace aura {
+class RootWindow;
+}
+
 namespace ash {
 class ScreenshotDelegate;
 
@@ -19,21 +22,24 @@ class ScreenshotDelegate;
 // the current mode.
 class ASH_EXPORT PartialScreenshotView : public views::WidgetDelegateView {
  public:
-  PartialScreenshotView(
-      internal::OverlayEventFilter::Delegate* overlay_delegate,
-      ScreenshotDelegate* screenshot_delegate);
-  virtual ~PartialScreenshotView();
-
   // Starts the UI for taking partial screenshot; dragging to select a region.
   static void StartPartialScreenshot(ScreenshotDelegate* screenshot_delegate);
 
-  // Overriddden from View:
-  virtual gfx::NativeCursor GetCursor(const ui::MouseEvent& event) OVERRIDE;
-
  private:
+  class OverlayDelegate;
+
+  PartialScreenshotView(OverlayDelegate* overlay_delegate,
+                        ScreenshotDelegate* screenshot_delegate);
+  virtual ~PartialScreenshotView();
+
+  // Initializes partial screenshot UI widget for |root_window|.
+  void Init(aura::RootWindow* root_window);
+
+  // Returns the currently selected region.
   gfx::Rect GetScreenshotRect() const;
 
-  // Overridden from View:
+  // Overridden from views::View:
+  virtual gfx::NativeCursor GetCursor(const ui::MouseEvent& event) OVERRIDE;
   virtual void OnPaint(gfx::Canvas* canvas) OVERRIDE;
   virtual bool OnMousePressed(const ui::MouseEvent& event) OVERRIDE;
   virtual bool OnMouseDragged(const ui::MouseEvent& event) OVERRIDE;
@@ -45,7 +51,7 @@ class ASH_EXPORT PartialScreenshotView : public views::WidgetDelegateView {
   gfx::Point current_position_;
 
   // The delegate to receive Cancel. No ownership.
-  internal::OverlayEventFilter::Delegate* overlay_delegate_;
+  OverlayDelegate* overlay_delegate_;
 
   // ScreenshotDelegate to take the actual screenshot. No ownership.
   ScreenshotDelegate* screenshot_delegate_;