Use WeakPtr for NotificationImageRetainer

This makes NotificationTemplateBuilder own a WeakPtr rather than
a raw pointer to NotificationImageRetainer, which helps with tracking
potential memory issues.

Bug: 884224
Change-Id: I229ffd187ccd9b424987ab1473b240f898c06023
Reviewed-on: https://chromium-review.googlesource.com/1227717Reviewed-by: Finnur Thorarinsson <finnur@chromium.org>
Commit-Queue: Xi Cheng <chengx@chromium.org>
Cr-Commit-Position: refs/heads/master@{#592759}

chrome/browser/notifications/notification_platform_bridge_win.cc

@@ -16,6 +16,7 @@
 #include "base/feature_list.h"
 #include "base/hash.h"
 #include "base/logging.h"
+#include "base/memory/weak_ptr.h"
 #include "base/metrics/histogram_functions.h"
 #include "base/sequenced_task_runner.h"
 #include "base/strings/string_number_conversions.h"
@@ -348,8 +349,8 @@ class NotificationPlatformBridgeWinImpl
                                    profile_id, incognito,
                                    notification->origin_url());
     std::unique_ptr<NotificationTemplateBuilder> notification_template =
-        NotificationTemplateBuilder::Build(image_retainer_.get(), launch_id,
-                                           profile_id, *notification);
+        NotificationTemplateBuilder::Build(
+            image_retainer_->AsWeakPtr(), launch_id, profile_id, *notification);
     mswr::ComPtr<winui::Notifications::IToastNotification> toast;
     hr = GetToastNotification(*notification, *notification_template, profile_id,
                               incognito, &toast);

chrome/browser/notifications/notification_platform_bridge_win_unittest.cc

@@ -14,6 +14,7 @@
 #include <wrl/implements.h>
 
 #include "base/hash.h"
+#include "base/memory/weak_ptr.h"
 #include "base/strings/string16.h"
 #include "base/strings/string_number_conversions.h"
 #include "base/strings/stringprintf.h"
@@ -66,10 +67,10 @@ class NotificationPlatformBridgeWinTest : public testing::Test {
         message_center::NotifierId(origin),
         message_center::RichNotificationData(), nullptr /* delegate */);
     notification->set_renotify(renotify);
-    MockNotificationImageRetainer image_retainer;
+    auto image_retainer = std::make_unique<MockNotificationImageRetainer>();
     std::unique_ptr<NotificationTemplateBuilder> builder =
-        NotificationTemplateBuilder::Build(&image_retainer, launch_id,
-                                           profile_id, *notification);
+        NotificationTemplateBuilder::Build(
+            image_retainer->AsWeakPtr(), launch_id, profile_id, *notification);
 
     mswr::ComPtr<winui::Notifications::IToastNotification> toast;
     HRESULT hr =

chrome/browser/notifications/win/notification_image_retainer.h

@@ -27,7 +27,8 @@ class Image;
 // can potentially be large, this presents a problem because Chrome might then
 // be leaving chunks of dead bits lying around on user’s computers during
 // unclean shutdowns.
-class NotificationImageRetainer {
+class NotificationImageRetainer
+    : public base::SupportsWeakPtr<NotificationImageRetainer> {
  public:
   explicit NotificationImageRetainer(
       scoped_refptr<base::SequencedTaskRunner> task_runner);

chrome/browser/notifications/win/notification_template_builder.cc

@@ -82,7 +82,7 @@ const char* NotificationTemplateBuilder::context_menu_label_override_ = nullptr;
 
 // static
 std::unique_ptr<NotificationTemplateBuilder> NotificationTemplateBuilder::Build(
-    NotificationImageRetainer* notification_image_retainer,
+    base::WeakPtr<NotificationImageRetainer> notification_image_retainer,
     const NotificationLaunchId& launch_id,
     const std::string& profile_id,
     const message_center::Notification& notification) {
@@ -144,7 +144,7 @@ std::unique_ptr<NotificationTemplateBuilder> NotificationTemplateBuilder::Build(
 }
 
 NotificationTemplateBuilder::NotificationTemplateBuilder(
-    NotificationImageRetainer* notification_image_retainer,
+    base::WeakPtr<NotificationImageRetainer> notification_image_retainer,
     const std::string& profile_id)
     : xml_writer_(std::make_unique<XmlWriter>()),
       image_retainer_(notification_image_retainer),
@@ -266,6 +266,9 @@ void NotificationTemplateBuilder::WriteImageElement(
     const GURL& origin,
     const std::string& placement,
     const std::string& hint_crop) {
+  // Although image_retainer_ is a WeakPtr, it should never be nullptr here.
+  DCHECK(image_retainer_);
+
   base::FilePath path =
       image_retainer_->RegisterTemporaryImage(image, profile_id_, origin);
   if (!path.empty()) {
@@ -357,6 +360,9 @@ void NotificationTemplateBuilder::WriteActionElement(
   xml_writer_->AddAttribute(kArguments, copied_launch_id.Serialize());
 
   if (!button.icon.IsEmpty()) {
+    // Although image_retainer_ is a WeakPtr, it should never be nullptr here.
+    DCHECK(image_retainer_);
+
     base::FilePath path = image_retainer_->RegisterTemporaryImage(
         button.icon, profile_id_, origin);
     if (!path.empty())

chrome/browser/notifications/win/notification_template_builder.h

@@ -10,6 +10,7 @@
 #include <vector>
 
 #include "base/macros.h"
+#include "base/memory/weak_ptr.h"
 #include "base/strings/string16.h"
 #include "chrome/browser/notifications/notification_common.h"
 
@@ -46,7 +47,7 @@ class NotificationTemplateBuilder {
  public:
   // Builds the notification template for the given |notification|.
   static std::unique_ptr<NotificationTemplateBuilder> Build(
-      NotificationImageRetainer* notification_image_retainer,
+      base::WeakPtr<NotificationImageRetainer> notification_image_retainer,
       const NotificationLaunchId& launch_id,
       const std::string& profile_id,
       const message_center::Notification& notification);
@@ -65,7 +66,7 @@ class NotificationTemplateBuilder {
   enum class TextType { NORMAL, ATTRIBUTION };
 
   NotificationTemplateBuilder(
-      NotificationImageRetainer* notification_image_retainer,
+      base::WeakPtr<NotificationImageRetainer> notification_image_retainer,
       const std::string& profile_id);
 
   // Formats the |origin| for display in the notification template.
@@ -143,8 +144,8 @@ class NotificationTemplateBuilder {
   // The XML writer to which the template will be written.
   std::unique_ptr<XmlWriter> xml_writer_;
 
-  // The image retainer. Weak, not owned by us.
-  NotificationImageRetainer* image_retainer_;
+  // The image retainer.
+  base::WeakPtr<NotificationImageRetainer> image_retainer_;
 
   // The id of the profile the notification is intended for.
   std::string profile_id_;

chrome/browser/notifications/win/notification_template_builder_unittest.cc

@@ -8,6 +8,7 @@
 #include <string>
 
 #include "base/macros.h"
+#include "base/memory/weak_ptr.h"
 #include "base/strings/string16.h"
 #include "base/strings/string_number_conversions.h"
 #include "base/strings/stringprintf.h"
@@ -82,10 +83,10 @@ class NotificationTemplateBuilderTest : public ::testing::Test {
   // must be wrapped in ASSERT_NO_FATAL_FAILURE().
   void VerifyXml(const message_center::Notification& notification,
                  const base::string16& xml_template) {
-    MockNotificationImageRetainer image_retainer;
+    auto image_retainer = std::make_unique<MockNotificationImageRetainer>();
     NotificationLaunchId launch_id(kEncodedId);
-    template_ = NotificationTemplateBuilder::Build(&image_retainer, launch_id,
-                                                   kProfileId, notification);
+    template_ = NotificationTemplateBuilder::Build(
+        image_retainer->AsWeakPtr(), launch_id, kProfileId, notification);
 
     ASSERT_TRUE(template_);