(flink) 03/06: [FLINK-34134] Add methods to determine local directory sizes to FileUtils.

[sr...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

srichter pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git

commit 210ab633350c77b026f883f72d3884eda20a03ae
Author: Stefan Richter <sr...@confluent.io>
AuthorDate: Tue Jan 16 15:36:09 2024 +0100

    [FLINK-34134] Add methods to determine local directory sizes to FileUtils.
---
 .../main/java/org/apache/flink/util/FileUtils.java | 24 ++++++++++++++++++++++
 .../java/org/apache/flink/util/FileUtilsTest.java  | 13 ++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/flink-core/src/main/java/org/apache/flink/util/FileUtils.java b/flink-core/src/main/java/org/apache/flink/util/FileUtils.java
index e33e69e9c8d..a75df3d643b 100644
--- a/flink-core/src/main/java/org/apache/flink/util/FileUtils.java
+++ b/flink-core/src/main/java/org/apache/flink/util/FileUtils.java
@@ -580,6 +580,30 @@ public final class FileUtils {
         return filterFileVisitor.getFiles();
     }
 
+    /**
+     * Computes the sum of sizes of all files in the directory and it's subdirectories.
+     *
+     * @param path the root path from which to start the calculation.
+     * @param options visitation options for the directory traversal.
+     * @return sum of sizes of all files in the directory and it's subdirectories.
+     * @throws IOException if the size cannot be determined.
+     */
+    public static long getDirectoryFilesSize(java.nio.file.Path path, FileVisitOption... options)
+            throws IOException {
+
+        if (path == null) {
+            return 0L;
+        }
+
+        try (Stream<java.nio.file.Path> pathStream = Files.walk(path, options)) {
+            return pathStream
+                    .map(java.nio.file.Path::toFile)
+                    .filter(File::isFile)
+                    .mapToLong(File::length)
+                    .sum();
+        }
+    }
+
     /**
      * Absolutize the given path if it is relative.
      *
diff --git a/flink-core/src/test/java/org/apache/flink/util/FileUtilsTest.java b/flink-core/src/test/java/org/apache/flink/util/FileUtilsTest.java
index afd1b70c941..0a9953c060b 100644
--- a/flink-core/src/test/java/org/apache/flink/util/FileUtilsTest.java
+++ b/flink-core/src/test/java/org/apache/flink/util/FileUtilsTest.java
@@ -24,6 +24,7 @@ import org.apache.flink.core.fs.Path;
 import org.apache.flink.core.testutils.CheckedThread;
 import org.apache.flink.testutils.junit.utils.TempDirUtils;
 
+import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Tag;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
@@ -363,6 +364,18 @@ public class FileUtilsTest {
         assertThat(targetPath).isEqualTo(dirInLinked2);
     }
 
+    @Test
+    void testGetDirectorySize() throws Exception {
+        final File parent = TempDirUtils.newFolder(temporaryFolder);
+
+        // Empty directory should have size 0
+        Assertions.assertEquals(0, FileUtils.getDirectoryFilesSize(parent.toPath()));
+
+        // Expected size: (20*5^0 + 20*5^1 + 20*5^2 + 20*5^3) * 1 byte = 3120 bytes
+        generateRandomDirs(parent, 20, 5, 3);
+        Assertions.assertEquals(3120, FileUtils.getDirectoryFilesSize(parent.toPath()));
+    }
+
     // ------------------------------------------------------------------------
     //  Utilities
     // ------------------------------------------------------------------------