[ios] Restore snapshot for undo close all tabs

When closing tabs, an observer immediately deletes snapshots from the
snapshot cache. This prevents snapshots from reappearing in an undo
close all tabs operation.

This CL enables snapshot restoration on undo close all tabs by extending
the life of a snapshot when closed with a "close all" operation. Note
this is only extended for regular tabs, not incognito tabs.

The extended life is accomplished by marking tabs for deletion during
a close all operation, rather than immediately deleting the image
from the cache. Later, the marked images are either purged from the
cache, or the marked list is cleared (therefore leaving the images in
the cache).

Bug: 804567
Cq-Include-Trybots: master.tryserver.chromium.mac:ios-simulator-cronet;master.tryserver.chromium.mac:ios-simulator-full-configs
Change-Id: I4710bf46d3ebd4ad347563e5f6a6208a68a58209
Reviewed-on: https://chromium-review.googlesource.com/1041545
Commit-Queue: edchin <edchin@chromium.org>
Reviewed-by: Eugene But <eugenebut@chromium.org>
Reviewed-by: edchin <edchin@chromium.org>
Cr-Commit-Position: refs/heads/master@{#556162}

ios/chrome/browser/snapshots/snapshot_cache.h

@@ -38,7 +38,24 @@
                              callback:(void (^)(UIImage*))callback;
 
 - (void)setImage:(UIImage*)img withSessionID:(NSString*)sessionID;
+
+// Removes the image from both the LRU and disk cache, unless it is marked for
+// deferred deletion. Images marked for deferred deletion can only be removed by
+// calling |-removeMarkedImages|.
 - (void)removeImageWithSessionID:(NSString*)sessionID;
+
+// Marks an image for deferred deletion. The image will not be immediately
+// deleted when |-removeImageWithSessionID:| is called. Images marked for
+// deferred deletion can only be removed by calling |-removeMarkedImages|.
+- (void)markImageWithSessionID:(NSString*)sessionID;
+
+// Removes all marked images from both the LRU and disk cache.
+- (void)removeMarkedImages;
+
+// Unmarks all images, so they remain in the cache. They are no longer marked
+// for deferred deletion.
+- (void)unmarkAllImages;
+
 // Purge the cache of snapshots that are older than |date|. The snapshots for
 // the sessions given in |liveSessionIds| will be kept. This will be done
 // asynchronously on a background thread.

ios/chrome/browser/snapshots/snapshot_cache.mm

@@ -49,6 +49,8 @@
 @interface SnapshotCache ()
 // List of observers to be notified of changes to the snapshot cache.
 @property(nonatomic, strong) SnapshotCacheObservers* observers;
+// Marked set of identifiers for which images should not be immediately deleted.
+@property(nonatomic, strong) NSMutableSet* markedIDs;
 
 // Remove all UIImages from |lruCache_|.
 - (void)handleEnterBackground;
@@ -259,8 +261,9 @@ void ConvertAndSaveGreyImage(NSString* session_id,
   SEQUENCE_CHECKER(sequenceChecker_);
 }
 
-@synthesize pinnedIDs = pinnedIDs_;
-@synthesize observers = observers_;
+@synthesize pinnedIDs = _pinnedIDs;
+@synthesize observers = _observers;
+@synthesize markedIDs = _markedIDs;
 
 - (instancetype)init {
   base::FilePath cacheDirectory;
@@ -282,7 +285,8 @@ void ConvertAndSaveGreyImage(NSString* session_id,
     taskRunner_ = base::CreateSequencedTaskRunnerWithTraits(
         {base::MayBlock(), base::TaskPriority::USER_VISIBLE});
 
-    observers_ = [SnapshotCacheObservers observers];
+    _observers = [SnapshotCacheObservers observers];
+    _markedIDs = [[NSMutableSet alloc] init];
 
     [[NSNotificationCenter defaultCenter]
         addObserver:self
@@ -382,6 +386,10 @@ void ConvertAndSaveGreyImage(NSString* session_id,
 
 - (void)removeImageWithSessionID:(NSString*)sessionID {
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequenceChecker_);
+  // Do not immediately delete if the ID is marked.
+  if ([self.markedIDs containsObject:sessionID])
+    return;
+
   [lruCache_ removeObjectForKey:sessionID];
 
   if (!taskRunner_)
@@ -401,6 +409,22 @@ void ConvertAndSaveGreyImage(NSString* session_id,
       }));
 }
 
+- (void)markImageWithSessionID:(NSString*)sessionID {
+  [self.markedIDs addObject:sessionID];
+}
+
+- (void)removeMarkedImages {
+  while (self.markedIDs.count > 0) {
+    NSString* sessionID = [self.markedIDs anyObject];
+    [self.markedIDs removeObject:sessionID];
+    [self removeImageWithSessionID:sessionID];
+  }
+}
+
+- (void)unmarkAllImages {
+  [self.markedIDs removeAllObjects];
+}
+
 - (base::FilePath)imagePathForSessionID:(NSString*)sessionID {
   return ImagePath(sessionID, IMAGE_TYPE_COLOR, snapshotsScale_,
                    cacheDirectory_);
@@ -465,13 +489,13 @@ void ConvertAndSaveGreyImage(NSString* session_id,
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequenceChecker_);
   NSMutableDictionary<NSString*, UIImage*>* dictionary =
       [NSMutableDictionary dictionaryWithCapacity:2];
-  for (NSString* sessionID in pinnedIDs_) {
+  for (NSString* sessionID in self.pinnedIDs) {
     UIImage* image = [lruCache_ objectForKey:sessionID];
     if (image)
       [dictionary setObject:image forKey:sessionID];
   }
   [lruCache_ removeAllObjects];
-  for (NSString* sessionID in pinnedIDs_)
+  for (NSString* sessionID in self.pinnedIDs)
     [lruCache_ setObject:[dictionary objectForKey:sessionID] forKey:sessionID];
 }
 
@@ -482,7 +506,7 @@ void ConvertAndSaveGreyImage(NSString* session_id,
 
 - (void)handleBecomeActive {
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequenceChecker_);
-  for (NSString* sessionID in pinnedIDs_)
+  for (NSString* sessionID in self.pinnedIDs)
     [self retrieveImageForSessionID:sessionID
                            callback:^(UIImage*){
                            }];

ios/chrome/browser/snapshots/snapshot_cache_unittest.mm

@@ -80,6 +80,16 @@ class SnapshotCacheTest : public PlatformTest {
     return UIGraphicsGetImageFromCurrentImageContext();
   }
 
+  // Generates an image of |size|, filled with a random color.
+  UIImage* GenerateRandomImage(CGSize size) {
+    UIGraphicsBeginImageContextWithOptions(size, /*opaque=*/NO,
+                                           UIScreen.mainScreen.scale);
+    CGContextRef context = UIGraphicsGetCurrentContext();
+    UIImage* image = GenerateRandomImage(context);
+    UIGraphicsEndImageContext();
+    return image;
+  }
+
   // Flushes all the runloops internally used by the snapshot cache.
   void FlushRunLoops() {
     base::TaskScheduler::GetInstance()->FlushForTesting();
@@ -569,4 +579,39 @@ TEST_F(SnapshotCacheTest, DeleteRetinaImages) {
   EXPECT_FALSE(base::PathExists(retinaFile));
 }
 
+// Tests that a marked image does not immediately delete when calling
+// |-removeImageWithSessionID:|. Calling |-removeMarkedImages| immediately
+// deletes the marked image.
+TEST_F(SnapshotCacheTest, MarkedImageNotImmediatelyDeleted) {
+  SnapshotCache* cache = GetSnapshotCache();
+  UIImage* image =
+      GenerateRandomImage(CGSizeMake(kSnapshotPixelSize, kSnapshotPixelSize));
+  [cache setImage:image withSessionID:@"sessionID"];
+  base::FilePath image_path = [cache imagePathForSessionID:@"sessionID"];
+  [cache markImageWithSessionID:@"sessionID"];
+  [cache removeImageWithSessionID:@"sessionID"];
+  // Give enough time for deletion.
+  FlushRunLoops();
+  EXPECT_TRUE(base::PathExists(image_path));
+  [cache removeMarkedImages];
+  FlushRunLoops();
+  EXPECT_FALSE(base::PathExists(image_path));
+}
+
+// Tests that unmarked images are not deleted when calling
+// |-removeMarkedImages|.
+TEST_F(SnapshotCacheTest, UnmarkedImageNotDeleted) {
+  SnapshotCache* cache = GetSnapshotCache();
+  UIImage* image =
+      GenerateRandomImage(CGSizeMake(kSnapshotPixelSize, kSnapshotPixelSize));
+  [cache setImage:image withSessionID:@"sessionID"];
+  base::FilePath image_path = [cache imagePathForSessionID:@"sessionID"];
+  [cache markImageWithSessionID:@"sessionID"];
+  [cache unmarkAllImages];
+  [cache removeMarkedImages];
+  // Give enough time for deletion.
+  FlushRunLoops();
+  EXPECT_TRUE(base::PathExists(image_path));
+}
+
 }  // namespace

ios/chrome/browser/ui/tab_grid/tab_grid_mediator.mm

@@ -11,6 +11,8 @@
 #include "components/favicon/ios/web_favicon_driver.h"
 #include "ios/chrome/browser/browser_state/chrome_browser_state.h"
 #include "ios/chrome/browser/chrome_url_constants.h"
+#import "ios/chrome/browser/snapshots/snapshot_cache.h"
+#import "ios/chrome/browser/snapshots/snapshot_cache_factory.h"
 #import "ios/chrome/browser/snapshots/snapshot_tab_helper.h"
 #import "ios/chrome/browser/tabs/tab_model.h"
 #import "ios/chrome/browser/ui/tab_grid/grid/grid_consumer.h"
@@ -201,6 +203,7 @@ int GetIndexOfTabWithId(WebStateList* webStateList, NSString* identifier) {
 }
 
 - (void)insertNewItemAtIndex:(NSUInteger)index {
+  DCHECK(self.tabModel.browserState);
   web::WebState::CreateParams params(self.tabModel.browserState);
   std::unique_ptr<web::WebState> webState = web::WebState::Create(params);
   self.webStateList->InsertWebState(
@@ -239,6 +242,16 @@ int GetIndexOfTabWithId(WebStateList* webStateList, NSString* identifier) {
 - (void)saveAndCloseAllItems {
   if (self.webStateList->empty())
     return;
+  // Tell the cache to mark these images for deletion, rather than immediately
+  // deleting them.
+  DCHECK(self.tabModel.browserState);
+  SnapshotCache* cache =
+      SnapshotCacheFactory::GetForBrowserState(self.tabModel.browserState);
+  for (int i = 0; i < self.webStateList->count(); i++) {
+    web::WebState* webState = self.webStateList->GetWebStateAt(i);
+    TabIdTabHelper* tabHelper = TabIdTabHelper::FromWebState(webState);
+    [cache markImageWithSessionID:tabHelper->tab_id()];
+  }
   self.closedSessionWindow = SerializeWebStateList(self.webStateList);
   self.webStateList->CloseAllWebStates(WebStateList::CLOSE_USER_ACTION);
 }
@@ -246,16 +259,26 @@ int GetIndexOfTabWithId(WebStateList* webStateList, NSString* identifier) {
 - (void)undoCloseAllItems {
   if (!self.closedSessionWindow)
     return;
+  DCHECK(self.tabModel.browserState);
   web::WebState::CreateParams createParams(self.tabModel.browserState);
   DeserializeWebStateList(
       self.webStateList, self.closedSessionWindow,
       base::BindRepeating(&web::WebState::CreateWithStorageSession,
                           createParams));
   self.closedSessionWindow = nil;
+  // Unmark all images for deletion since they are now active tabs again.
+  ios::ChromeBrowserState* browserState = self.tabModel.browserState;
+  [SnapshotCacheFactory::GetForBrowserState(browserState) unmarkAllImages];
 }
 
 - (void)discardSavedClosedItems {
+  if (!self.closedSessionWindow)
+    return;
   self.closedSessionWindow = nil;
+  // Delete all marked images from the cache.
+  DCHECK(self.tabModel.browserState);
+  ios::ChromeBrowserState* browserState = self.tabModel.browserState;
+  [SnapshotCacheFactory::GetForBrowserState(browserState) removeMarkedImages];
 }
 
 #pragma mark - GridImageDataSource