libbrillo: Add ComputeDirectoryDiskUsage() for calculating disk usage.

BUG=chromium:1080554
TEST=FEATURES=test emerge-$BOARD libbrillo

Change-Id: I443165ebad8ebdfd06c346b127541799acfc0a7c
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform2/+/2195876
Tested-by: John L Chen <zuan@chromium.org>
Commit-Queue: John L Chen <zuan@chromium.org>
Reviewed-by: Gwendal Grignou <gwendal@chromium.org>
Reviewed-by: Eric Caruso <ejcaruso@chromium.org>
Reviewed-by: François Degros <fdegros@chromium.org>
Cr-Mirrored-From: https://chromium.googlesource.com/chromiumos/platform2
Cr-Mirrored-Commit: 226ab943e7e01181feef79c9caeb6b1f56db6878

3 files changed, 219 insertions, 0 deletions

diff --git a/brillo/file_utils.cc b/brillo/file_utils.cc
index 0f43402..921fb57 100644
--- a/brillo/file_utils.cc
+++ b/brillo/file_utils.cc
@@ -11,6 +11,7 @@
 #include <utility>
 #include <vector>
 
+#include <base/files/file_enumerator.h>
 #include <base/files/file_path.h>
 #include <base/files/file_util.h>
 #include <base/logging.h>
@@ -506,4 +507,19 @@ bool WriteToFileAtomic(const base::FilePath& path,
   return true;
 }
 
+int64_t ComputeDirectoryDiskUsage(const base::FilePath& root_path) {
+  int64_t running_blocks = 0;
+  base::FileEnumerator file_iter(root_path, true,
+                                 base::FileEnumerator::FILES |
+                                     base::FileEnumerator::DIRECTORIES |
+                                     base::FileEnumerator::SHOW_SYM_LINKS);
+  while (!file_iter.Next().empty()) {
+    // st_blocks in struct stat is the number of S_BLKSIZE (512) bytes sized
+    // blocks occupied by this file.
+    running_blocks += file_iter.GetInfo().stat().st_blocks;
+  }
+  // Each block is S_BLKSIZE (512) bytes so *S_BLKSIZE.
+  return running_blocks * S_BLKSIZE;
+}
+
 }  // namespace brillo
diff --git a/brillo/file_utils.h b/brillo/file_utils.h
index 3862a43..f328165 100644
--- a/brillo/file_utils.h
+++ b/brillo/file_utils.h
@@ -144,6 +144,48 @@ BRILLO_EXPORT bool WriteBlobToFileAtomic(const base::FilePath& path,
                            blob.size(), mode);
 }
 
+// ComputeDirectoryDiskUsage() is similar to base::ComputeDirectorySize() in
+// libbase, but it returns the actual disk usage instead of the apparent size.
+// In another word, ComputeDirectoryDiskUsage() behaves like "du -s
+// --apparent-size", and ComputeDirectorySize() behaves like "du -s". The
+// primary difference is that sparse file and files on filesystem with
+// transparent compression will report smaller file size than
+// ComputeDirectorySize(). Returns the total used bytes.
+// The following behaviours of this function is guaranteed and is verified by
+// unit tests:
+// - This function recursively processes directory down the tree, so disk space
+// used by files in all the subdirectories are counted.
+// - Symbolic links will not be followed (the size of link itself is counted,
+// the target is not)
+// - Hidden files are counted as well.
+// The following behaviours are not guaranteed, and it is recommended to avoid
+// them in the field. Their current behaviour is provided for reference only:
+// - This function doesn't care about filesystem boundaries, so it'll cross
+// filesystem boundary to count file size if there's one in the specified
+// directory.
+// - Hard links will be treated like normal files, so they could be
+// over-reported.
+// - Directories that the current user doesn't have permission to list/stat will
+// be ignored, and an error will be logged but the returned result could be
+// under-reported without error in the returned value.
+// - Deduplication (should the filesystem support it) is ignored, and the result
+// could be over-reported.
+// - Doesn't check if |root_path| exists, a non-existent directory will results
+// in 0 bytes without any error.
+// - There are no limit on the depth of file system tree, the program will crash
+// if it run out of memory to hold the entire depth of file system tree.
+// - If the directory is modified during this function call, there's no
+// guarantee on if the function will count the updated or original file system
+// state. The function could choose to count the updated state for one file and
+// original state for another file.
+// - Non-POSIX system is not supported.
+// - Disk space used by directory (and its subdirectories) itself is counted.
+//
+// Parameters
+//   root_path - The directory to compute the size for
+BRILLO_EXPORT int64_t
+ComputeDirectoryDiskUsage(const base::FilePath& root_path);
+
 }  // namespace brillo
 
 #endif  // LIBBRILLO_BRILLO_FILE_UTILS_H_
diff --git a/brillo/file_utils_test.cc b/brillo/file_utils_test.cc
index f4a5f55..3407cd1 100644
--- a/brillo/file_utils_test.cc
+++ b/brillo/file_utils_test.cc
@@ -344,4 +344,165 @@ TEST_F(FileUtilsTest,
   umask(old_mask);
 }
 
+TEST_F(FileUtilsTest, ComputeDirectoryDiskUsageNormalRandomFile) {
+  // 2MB test file.
+  constexpr size_t kFileSize = 2 * 1024 * 1024;
+
+  const base::FilePath dirname(GetTempName());
+  EXPECT_TRUE(base::CreateDirectory(dirname));
+  const base::FilePath filename = dirname.Append("test.temp");
+
+  std::string file_content = base::RandBytesAsString(kFileSize);
+  EXPECT_TRUE(WriteStringToFile(filename, file_content));
+
+  int64_t result_usage = ComputeDirectoryDiskUsage(dirname);
+  int64_t result_size = base::ComputeDirectorySize(dirname);
+
+  // result_usage (what we are testing here) should be within +/-10% of ground
+  // truth. The variation is to account for filesystem overhead variations.
+  EXPECT_GT(result_usage, kFileSize / 10 * 9);
+  EXPECT_LT(result_usage, kFileSize / 10 * 11);
+
+  // result_usage should be close to result_size, because the test file is
+  // random so it's disk usage should be similar to apparent size.
+  EXPECT_GT(result_usage, result_size / 10 * 9);
+  EXPECT_LT(result_usage, result_size / 10 * 11);
+}
+
+TEST_F(FileUtilsTest, ComputeDirectoryDiskUsageDeepRandomFile) {
+  // 2MB test file.
+  constexpr size_t kFileSize = 2 * 1024 * 1024;
+
+  const base::FilePath dirname(GetTempName());
+  EXPECT_TRUE(base::CreateDirectory(dirname));
+  base::FilePath currentlevel = dirname;
+  for (int i = 0; i < 10; i++) {
+    base::FilePath nextlevel = currentlevel.Append("test.dir");
+    EXPECT_TRUE(base::CreateDirectory(nextlevel));
+    currentlevel = nextlevel;
+  }
+  const base::FilePath filename = currentlevel.Append("test.temp");
+
+  std::string file_content = base::RandBytesAsString(kFileSize);
+  EXPECT_TRUE(WriteStringToFile(filename, file_content));
+
+  int64_t result_usage = ComputeDirectoryDiskUsage(dirname);
+  int64_t result_size = base::ComputeDirectorySize(dirname);
+
+  // result_usage (what we are testing here) should be within +/-10% of ground
+  // truth. The variation is to account for filesystem overhead variations.
+  EXPECT_GT(result_usage, kFileSize / 10 * 9);
+  EXPECT_LT(result_usage, kFileSize / 10 * 11);
+
+  // result_usage should be close to result_size, because the test file is
+  // random so it's disk usage should be similar to apparent size.
+  EXPECT_GT(result_usage, result_size / 10 * 9);
+  EXPECT_LT(result_usage, result_size / 10 * 11);
+}
+
+TEST_F(FileUtilsTest, ComputeDirectoryDiskUsageHiddenRandomFile) {
+  // 2MB test file.
+  constexpr size_t kFileSize = 2 * 1024 * 1024;
+
+  const base::FilePath dirname(GetTempName());
+  EXPECT_TRUE(base::CreateDirectory(dirname));
+  // File name starts with a dot, so it's a hidden file.
+  const base::FilePath filename = dirname.Append(".test.temp");
+
+  std::string file_content = base::RandBytesAsString(kFileSize);
+  EXPECT_TRUE(WriteStringToFile(filename, file_content));
+
+  int64_t result_usage = ComputeDirectoryDiskUsage(dirname);
+  int64_t result_size = base::ComputeDirectorySize(dirname);
+
+  // result_usage (what we are testing here) should be within +/-10% of ground
+  // truth. The variation is to account for filesystem overhead variations.
+  EXPECT_GT(result_usage, kFileSize / 10 * 9);
+  EXPECT_LT(result_usage, kFileSize / 10 * 11);
+
+  // result_usage should be close to result_size, because the test file is
+  // random so it's disk usage should be similar to apparent size.
+  EXPECT_GT(result_usage, result_size / 10 * 9);
+  EXPECT_LT(result_usage, result_size / 10 * 11);
+}
+
+TEST_F(FileUtilsTest, ComputeDirectoryDiskUsageSparseFile) {
+  // 128MB sparse test file.
+  constexpr size_t kFileSize = 128 * 1024 * 1024;
+  constexpr size_t kFileSizeThreshold = 64 * 1024;
+
+  const base::FilePath dirname(GetTempName());
+  EXPECT_TRUE(base::CreateDirectory(dirname));
+  const base::FilePath filename = dirname.Append("test.temp");
+
+  int fd =
+      open(filename.value().c_str(), O_CREAT | O_WRONLY, S_IRUSR | S_IWUSR);
+  EXPECT_NE(fd, -1);
+  // Calling ftruncate on an empty file will create a sparse file.
+  EXPECT_EQ(0, ftruncate(fd, kFileSize));
+
+  int64_t result_usage = ComputeDirectoryDiskUsage(dirname);
+  int64_t result_size = base::ComputeDirectorySize(dirname);
+
+  // result_usage (what we are testing here) should be less than
+  // kFileSizeThreshold, the threshold is to account for filesystem overhead
+  // variations.
+  EXPECT_LT(result_usage, kFileSizeThreshold);
+
+  // Since we are dealing with sparse files here, the apparent size should be
+  // much much larger than the actual disk usage.
+  EXPECT_LT(result_usage, result_size / 1000);
+}
+
+TEST_F(FileUtilsTest, ComputeDirectoryDiskUsageSymlinkFile) {
+  // 2MB test file.
+  constexpr size_t kFileSize = 2 * 1024 * 1024;
+
+  const base::FilePath dirname(GetTempName());
+  EXPECT_TRUE(base::CreateDirectory(dirname));
+  const base::FilePath filename = dirname.Append("test.temp");
+  const base::FilePath linkname = dirname.Append("test.link");
+
+  std::string file_content = base::RandBytesAsString(kFileSize);
+  EXPECT_TRUE(WriteStringToFile(filename, file_content));
+
+  // Create a symlink.
+  EXPECT_TRUE(base::CreateSymbolicLink(filename, linkname));
+
+  int64_t result_usage = ComputeDirectoryDiskUsage(dirname);
+
+  // result_usage (what we are testing here) should be within +/-10% of ground
+  // truth. The variation is to account for filesystem overhead variations.
+  // Note that it's not 2x kFileSize because symblink is not counted twice.
+  EXPECT_GT(result_usage, kFileSize / 10 * 9);
+  EXPECT_LT(result_usage, kFileSize / 10 * 11);
+}
+
+TEST_F(FileUtilsTest, ComputeDirectoryDiskUsageSymlinkDir) {
+  // 2MB test file.
+  constexpr size_t kFileSize = 2 * 1024 * 1024;
+
+  const base::FilePath parentname(GetTempName());
+  EXPECT_TRUE(base::CreateDirectory(parentname));
+  const base::FilePath dirname = parentname.Append("target.dir");
+  EXPECT_TRUE(base::CreateDirectory(dirname));
+  const base::FilePath linkname = parentname.Append("link.dir");
+
+  const base::FilePath filename = dirname.Append("test.temp");
+
+  std::string file_content = base::RandBytesAsString(kFileSize);
+  EXPECT_TRUE(WriteStringToFile(filename, file_content));
+
+  // Create a symlink.
+  EXPECT_TRUE(base::CreateSymbolicLink(dirname, linkname));
+
+  int64_t result_usage = ComputeDirectoryDiskUsage(dirname);
+
+  // result_usage (what we are testing here) should be within +/-10% of ground
+  // truth. The variation is to account for filesystem overhead variations.
+  // Note that it's not 2x kFileSize because symblink is not counted twice.
+  EXPECT_GT(result_usage, kFileSize / 10 * 9);
+  EXPECT_LT(result_usage, kFileSize / 10 * 11);
+}
+
 }  // namespace brillo