[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}

chrome/browser/paint_preview/services/paint_preview_tab_service.cc

@@ -27,6 +27,7 @@ namespace paint_preview {
 namespace {
 
 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,59 @@ 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);
+  {
+    auto to_delete = manager->GetOldestArtifactsForCleanup(0U);
+    ASSERT_EQ(to_delete.size(), 2U);
+    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