Reland "[Paint Preview] Create total size budget"

This is a reland of 09cf456a

Failures were due to flakiness in FileManagerTest.OldestFileForCleanup

Timing of last modified time seems to be racy leading to inconsistent
behavior. The fix for this is manually setting the modified time to
something longer.

Original change's description:
> [Paint Preview] Create total size budget
>
> This CL introduces a mechanism to determine the oldest files and
> enforce a total disk size budget for captures. By default this is set to
> 25 MB for the tab service. This will ensure we don't use excessive disk
> space.
>
> Bug: 1086963
> Change-Id: I71bd8efa5985214b322a4dc3048afbe4d0b717f2
> Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2218548
> Commit-Queue: Calder Kitagawa <ckitagawa@chromium.org>
> Reviewed-by: Mehran Mahmoudi <mahmoudi@chromium.org>
> Cr-Commit-Position: refs/heads/master@{#772809}

Bug: 1086963, 1087592
Change-Id: Id95442afcdeb55e60cafc4a8bf6cebef84b3fc8d
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2220483Reviewed-by: Mehran Mahmoudi <mahmoudi@chromium.org>
Commit-Queue: Calder Kitagawa <ckitagawa@chromium.org>
Cr-Commit-Position: refs/heads/master@{#773764}

chrome/browser/paint_preview/services/paint_preview_tab_service.cc

@@ -26,7 +26,8 @@ namespace paint_preview {
 
 namespace {
 
-constexpr size_t kMaxPerCaptureSizeBytes = 5 * 1000L * 1000L;  // 5 MB.
+constexpr size_t kMaxPerCaptureSizeBytes = 5 * 1000L * 1000L;    // 5 MB.
+constexpr size_t kMaximumTotalCaptureSize = 25 * 1000L * 1000L;  // 25 MB.
 
 #if defined(OS_ANDROID)
 void JavaBooleanCallbackAdapter(base::OnceCallback<void(bool)> callback,
@@ -247,6 +248,33 @@ void PaintPreviewTabService::OnFinished(int tab_id,
     captured_tab_ids_.insert(tab_id);
   std::move(callback).Run(success ? Status::kOk
                                   : Status::kProtoSerializationFailed);
+  auto file_manager = GetFileManager();
+  GetTaskRunner()->PostTaskAndReplyWithResult(
+      FROM_HERE,
+      base::BindOnce(&FileManager::GetOldestArtifactsForCleanup, file_manager,
+                     kMaximumTotalCaptureSize),
+      base::BindOnce(&PaintPreviewTabService::CleanupOldestFiles,
+                     weak_ptr_factory_.GetWeakPtr(), tab_id));
+}
+
+void PaintPreviewTabService::CleanupOldestFiles(
+    int tab_id,
+    const std::vector<DirectoryKey>& keys) {
+  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
+  std::vector<DirectoryKey> keys_to_delete;
+  keys_to_delete.reserve(keys.size());
+  for (const auto& key : keys) {
+    auto id = TabIdFromDirectoryKey(key);
+    if (id == tab_id)
+      continue;
+
+    captured_tab_ids_.erase(id);
+    keys_to_delete.push_back(key);
+  }
+
+  GetTaskRunner()->PostTask(FROM_HERE,
+                            base::BindOnce(&FileManager::DeleteArtifactSets,
+                                           GetFileManager(), keys_to_delete));
 }
 
 void PaintPreviewTabService::RunAudit(

chrome/browser/paint_preview/services/paint_preview_tab_service.h

@@ -115,6 +115,8 @@ class PaintPreviewTabService : public PaintPreviewBaseService {
 
   void OnFinished(int tab_id, FinishedCallback callback, bool success);
 
+  void CleanupOldestFiles(int tab_id, const std::vector<DirectoryKey>& keys);
+
   void RunAudit(const std::vector<int>& active_tab_ids,
                 const base::flat_set<DirectoryKey>& in_use_keys);
 

chrome/browser/paint_preview/services/paint_preview_tab_service_unittest.cc

@@ -422,4 +422,55 @@ TEST_F(PaintPreviewTabServiceTest, EarlyCapture) {
   task_environment()->RunUntilIdle();
 }
 
+TEST_F(PaintPreviewTabServiceTest, CaptureTabAndCleanup) {
+  content::NavigationSimulator::NavigateAndCommitFromBrowser(
+      web_contents(), GURL("http://www.example.com"));
+  const int kTabId = 1U;
+
+  MockPaintPreviewRecorder recorder;
+  recorder.SetResponse(mojom::PaintPreviewStatus::kOk);
+  OverrideInterface(&recorder);
+
+  auto service = BuildServiceWithCache({kTabId + 1});
+  task_environment()->RunUntilIdle();
+  EXPECT_TRUE(service->CacheInitialized());
+  base::FilePath old_path = GetPath()
+                                .AppendASCII("paint_preview")
+                                .AppendASCII(kFeatureName)
+                                .AppendASCII(base::NumberToString(kTabId + 1));
+  // The threshold for cleanup is 25 MB.
+  std::string data(25 * 1000 * 1000, 'x');
+  EXPECT_TRUE(base::WriteFile(old_path.AppendASCII("foo.txt"), data.data(),
+                              data.size()));
+  EXPECT_TRUE(base::PathExists(old_path));
+  EXPECT_TRUE(service->HasCaptureForTab(kTabId + 1));
+
+  service->CaptureTab(kTabId, web_contents(),
+                      base::BindOnce([](PaintPreviewTabService::Status status) {
+                        EXPECT_EQ(status, PaintPreviewTabService::Status::kOk);
+                      }));
+  task_environment()->RunUntilIdle();
+  EXPECT_TRUE(service->HasCaptureForTab(kTabId));
+
+  auto file_manager = service->GetFileManager();
+  auto key = file_manager->CreateKey(kTabId);
+  service->GetTaskRunner()->PostTaskAndReplyWithResult(
+      FROM_HERE,
+      base::BindOnce(&FileManager::DirectoryExists, file_manager, key),
+      base::BindOnce([](bool exists) { EXPECT_TRUE(exists); }));
+  task_environment()->RunUntilIdle();
+  EXPECT_FALSE(base::PathExists(old_path));
+  EXPECT_FALSE(service->HasCaptureForTab(kTabId + 1));
+  EXPECT_TRUE(service->HasCaptureForTab(kTabId));
+
+  service->TabClosed(kTabId);
+  EXPECT_FALSE(service->HasCaptureForTab(kTabId));
+  task_environment()->RunUntilIdle();
+  service->GetTaskRunner()->PostTaskAndReplyWithResult(
+      FROM_HERE,
+      base::BindOnce(&FileManager::DirectoryExists, file_manager, key),
+      base::BindOnce([](bool exists) { EXPECT_FALSE(exists); }));
+  task_environment()->RunUntilIdle();
+}
+
 }  // namespace paint_preview

components/paint_preview/browser/directory_key.cc

@@ -10,4 +10,8 @@ bool operator<(const DirectoryKey& a, const DirectoryKey& b) {
   return a.AsciiDirname() < b.AsciiDirname();
 }
 
+bool operator==(const DirectoryKey& a, const DirectoryKey& b) {
+  return a.AsciiDirname() == b.AsciiDirname();
+}
+
 }  // namespace paint_preview

components/paint_preview/browser/directory_key.h

@@ -27,6 +27,8 @@ class DirectoryKey {
 
 bool operator<(const DirectoryKey& a, const DirectoryKey& b);
 
+bool operator==(const DirectoryKey& a, const DirectoryKey& b);
+
 }  // namespace paint_preview
 
 #endif  // COMPONENTS_PAINT_PREVIEW_BROWSER_DIRECTORY_KEY_H_

components/paint_preview/browser/file_manager.cc

@@ -4,6 +4,9 @@
 
 #include "components/paint_preview/browser/file_manager.h"
 
+#include <algorithm>
+#include <vector>
+
 #include "base/files/file_enumerator.h"
 #include "base/files/file_util.h"
 #include "base/hash/hash.h"
@@ -18,6 +21,11 @@ namespace {
 constexpr char kProtoName[] = "proto.pb";
 constexpr char kZipExt[] = ".zip";
 
+bool CompareByLastModified(const base::FileEnumerator::FileInfo& a,
+                           const base::FileEnumerator::FileInfo& b) {
+  return a.GetLastModifiedTime() < b.GetLastModifiedTime();
+}
+
 }  // namespace
 
 FileManager::FileManager(
@@ -212,6 +220,45 @@ base::flat_set<DirectoryKey> FileManager::ListUsedKeys() const {
   return base::flat_set<DirectoryKey>(std::move(keys));
 }
 
+std::vector<DirectoryKey> FileManager::GetOldestArtifactsForCleanup(
+    size_t max_size) {
+  // The rest of this function is expensive so cleanup should exit early if not
+  // required.
+  size_t size = base::ComputeDirectorySize(root_directory_);
+  if (size <= max_size)
+    return std::vector<DirectoryKey>();
+
+  base::FileEnumerator file_enum(
+      root_directory_, false,
+      base::FileEnumerator::FILES | base::FileEnumerator::DIRECTORIES);
+  std::vector<base::FileEnumerator::FileInfo> file_infos;
+  for (base::FilePath name = file_enum.Next(); !name.empty();
+       name = file_enum.Next()) {
+    file_infos.push_back(file_enum.GetInfo());
+  }
+
+  std::sort(file_infos.begin(), file_infos.end(), CompareByLastModified);
+
+  std::vector<DirectoryKey> keys_to_remove;
+  for (const auto& file_info : file_infos) {
+    base::FilePath full_path = root_directory_.Append(file_info.GetName());
+
+    size_t size_delta = file_info.GetSize();
+    // Computing a directory size is expensive. Most files should hopefully be
+    // compressed already.
+    if (file_info.IsDirectory())
+      size_delta = base::ComputeDirectorySize(full_path);
+
+    // Directory names should always be ASCII.
+    keys_to_remove.emplace_back(
+        file_info.GetName().RemoveExtension().MaybeAsASCII());
+    size -= size_delta;
+    if (size <= max_size)
+      break;
+  }
+  return keys_to_remove;
+}
+
 FileManager::StorageType FileManager::GetPathForKey(
     const DirectoryKey& key,
     base::FilePath* path) const {

components/paint_preview/browser/file_manager.h

@@ -87,6 +87,11 @@ class FileManager : public base::RefCountedThreadSafe<FileManager> {
   // Lists the current set of in-use DirectoryKeys.
   base::flat_set<DirectoryKey> ListUsedKeys() const;
 
+  // Returns a list of the least recently modified artifact sets until which
+  // when deleted would result in a total capture size on disk that is less than
+  // |max_size|.
+  std::vector<DirectoryKey> GetOldestArtifactsForCleanup(size_t max_size);
+
  private:
   friend class base::RefCountedThreadSafe<FileManager>;
   ~FileManager();

components/paint_preview/browser/file_manager_unittest.cc

@@ -138,6 +138,7 @@ TEST_F(FileManagerTest, TestCompression) {
   // Compress.
   base::FilePath zip_path = directory.AddExtensionASCII(".zip");
   EXPECT_TRUE(manager->CompressDirectory(key));
+  EXPECT_TRUE(manager->CompressDirectory(key));  // no-op
   EXPECT_GT(manager->GetSizeOfArtifacts(key), 0U);
   EXPECT_FALSE(base::PathExists(directory));
   EXPECT_FALSE(base::PathExists(test_file));
@@ -275,10 +276,66 @@ TEST_F(FileManagerTest, HandleProtoCompressed) {
   metadata->set_url(GURL("www.chromium.org").spec());
 
   EXPECT_TRUE(manager->SerializePaintPreviewProto(key, original_proto, true));
+  EXPECT_TRUE(manager->CaptureExists(key));
+
   EXPECT_TRUE(base::PathExists(path.AddExtensionASCII(".zip")));
   auto out_proto = manager->DeserializePaintPreviewProto(key);
   EXPECT_NE(out_proto, nullptr);
   EXPECT_THAT(*out_proto, EqualsProto(original_proto));
+
+  EXPECT_TRUE(manager->CaptureExists(key));
+}
+
+TEST_F(FileManagerTest, OldestFilesForCleanup) {
+  base::ScopedTempDir temp_dir;
+  ASSERT_TRUE(temp_dir.CreateUniqueTempDir());
+  auto manager =
+      base::MakeRefCounted<FileManager>(temp_dir.GetPath(), MainTaskRunner());
+
+  const std::string data = "Foobar";
+
+  auto key_0 = manager->CreateKey(0U);
+  base::FilePath path_0 =
+      manager->CreateOrGetDirectory(key_0, true).value_or(base::FilePath());
+  base::WriteFile(path_0.AppendASCII("0.txt"), data.data(), data.size());
+  {
+    auto to_delete = manager->GetOldestArtifactsForCleanup(0U);
+    ASSERT_EQ(to_delete.size(), 1U);
+    EXPECT_EQ(to_delete[0], key_0);
+  }
+  {
+    auto to_delete = manager->GetOldestArtifactsForCleanup(50U);
+    EXPECT_EQ(to_delete.size(), 0U);
+  }
+
+  auto key_1 = manager->CreateKey(1U);
+  base::FilePath path_1 =
+      manager->CreateOrGetDirectory(key_1, true).value_or(base::FilePath());
+  base::WriteFile(path_1.AppendASCII("1.txt"), data.data(), data.size());
+  manager->CompressDirectory(key_1);
+  base::Time modified_time = base::Time::Now();
+  auto path_1_zip = path_1.AddExtensionASCII(".zip");
+  base::TouchFile(path_0, modified_time - base::TimeDelta::FromSeconds(10),
+                  modified_time - base::TimeDelta::FromSeconds(10));
+  base::TouchFile(path_1_zip, modified_time, modified_time);
+
+  {
+    auto to_delete = manager->GetOldestArtifactsForCleanup(0U);
+    ASSERT_EQ(to_delete.size(), 2U);
+    // Elements should be ordered in oldest to newest.
+    EXPECT_EQ(to_delete[0], key_0);
+    EXPECT_EQ(to_delete[1], key_1);
+  }
+  {
+    // Zip is ~116 bytes.
+    auto to_delete = manager->GetOldestArtifactsForCleanup(120U);
+    ASSERT_EQ(to_delete.size(), 1U);
+    EXPECT_EQ(to_delete[0], key_0);
+  }
+  {
+    auto to_delete = manager->GetOldestArtifactsForCleanup(150U);
+    EXPECT_EQ(to_delete.size(), 0U);
+  }
 }
 
 }  // namespace paint_preview