Make ValidatesRunner faster: Improvements to file staging

runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/PackageUtil.java

@@ -425,9 +425,7 @@ public static PackageAttributes forFileToStage(
       switch (dest) {
         case "dataflow-worker.jar":
         case "windmill_main":
-          target =
-              Environments.createStagingFileName(
-                  file, Files.asByteSource(file).hash(Hashing.sha256()));
+          target = Environments.createStagingFileName(file, Environments.getFileHash(file));
           LOG.info("Staging custom {} as {}", dest, target);
           break;
         default:

sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/Environments.java

@@ -22,12 +22,18 @@
 import java.io.File;
 import java.io.FileOutputStream;
 import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
 import java.util.UUID;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutionException;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor;
@@ -55,10 +61,13 @@
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.Cache;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheBuilder;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.HashCode;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hasher;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hashing;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.Files;
 import org.checkerframework.checker.nullness.qual.Nullable;
@@ -98,6 +107,11 @@ public class Environments {
           .put(ENVIRONMENT_PROCESS, ImmutableSet.of(processCommandOption, processVariablesOption))
           .build();
 
+  private static final Cache<FileHashCacheKey, HashCode> FILE_HASH_CACHE =
+      CacheBuilder.newBuilder().maximumSize(10_000).build();
+  private static final ConcurrentHashMap<HashCode, File> DIRECTORY_ZIP_CACHE =
+      new ConcurrentHashMap<>();
+
   public enum JavaVersion {
     java11("java11", "11", 11),
     java17("java17", "17", 17),
@@ -434,7 +448,7 @@ public static List<ArtifactInformation> getArtifacts(List<String> stagingFiles)
         File zippedFile;
         try {
           zippedFile = zipDirectory(file);
-          hashCode = Files.asByteSource(zippedFile).hash(Hashing.sha256());
+          hashCode = getFileHash(zippedFile);
         } catch (IOException e) {
           throw new RuntimeException(e);
         }
@@ -448,7 +462,7 @@ public static List<ArtifactInformation> getArtifacts(List<String> stagingFiles)
 
       } else {
         try {
-          hashCode = Files.asByteSource(file).hash(Hashing.sha256());
+          hashCode = getFileHash(file);
         } catch (IOException e) {
           throw new RuntimeException(e);
         }
@@ -538,6 +552,53 @@ public static String createStagingFileName(File path, HashCode hash) {
     return String.format("%s-%s%s", fileName, encodedHash, suffix);
   }
 
+  /**
+   * Returns the SHA-256 {@link HashCode} for {@code file}, caching the result in memory keyed by
+   * the file's absolute path, length, and last-modified timestamp.
+   */
+  public static HashCode getFileHash(File file) throws IOException {
+    FileHashCacheKey key = new FileHashCacheKey(file);
+    try {
+      return FILE_HASH_CACHE.get(key, () -> Files.asByteSource(file).hash(Hashing.sha256()));
+    } catch (ExecutionException e) {
+      if (e.getCause() instanceof IOException) {
+        throw (IOException) e.getCause();
+      }
+      throw new RuntimeException(e.getCause());
+    }
+  }
+
+  private static final class FileHashCacheKey {
+    private final String absolutePath;
+    private final long length;
+    private final long lastModified;
+
+    FileHashCacheKey(File file) {
+      this.absolutePath = file.getAbsolutePath();
+      this.length = file.length();
+      this.lastModified = file.lastModified();
+    }
+
+    @Override
+    public boolean equals(@Nullable Object o) {
+      if (this == o) {
+        return true;
+      }
+      if (!(o instanceof FileHashCacheKey)) {
+        return false;
+      }
+      FileHashCacheKey that = (FileHashCacheKey) o;
+      return length == that.length
+          && lastModified == that.lastModified
+          && Objects.equals(absolutePath, that.absolutePath);
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hash(absolutePath, length, lastModified);
+    }
+  }
+
   public static String getExternalServiceAddress(PortablePipelineOptions options) {
     String environmentConfig = options.getDefaultEnvironmentConfig();
     String environmentOption =
@@ -549,11 +610,52 @@ public static String getExternalServiceAddress(PortablePipelineOptions options)
   }
 
   private static File zipDirectory(File directory) throws IOException {
-    File zipFile = File.createTempFile(directory.getName(), ".zip");
-    try (FileOutputStream fos = new FileOutputStream(zipFile)) {
-      ZipFiles.zipDirectory(directory, fos);
+    HashCode metadataHash = computeDirectoryMetadataHash(directory);
+    try {
+      return DIRECTORY_ZIP_CACHE.compute(
+          metadataHash,
+          (k, cachedZip) -> {
+            if (cachedZip != null && cachedZip.exists()) {
+              return cachedZip;
+            }
+            try {
+              File zipFile = File.createTempFile(directory.getName(), ".zip");
+              zipFile.deleteOnExit();
+              try (FileOutputStream fos = new FileOutputStream(zipFile)) {
+                ZipFiles.zipDirectory(directory, fos);
+              }
+              return zipFile;
+            } catch (IOException e) {
+              throw new UncheckedIOException(e);
+            }
+          });
+    } catch (UncheckedIOException e) {
+      throw e.getCause();
+    }
+  }
+
+  private static HashCode computeDirectoryMetadataHash(File directory) {
+    Hasher hasher = Hashing.sha256().newHasher();
+    hasher.putString(directory.getAbsolutePath(), StandardCharsets.UTF_8);
+    hashDirectoryMetadataRecursive(directory, "", hasher);
+    return hasher.hash();
+  }
+
+  private static void hashDirectoryMetadataRecursive(File inputFile, String prefix, Hasher hasher) {
+    String entryName = prefix + inputFile.getName();
+    hasher.putString(entryName, StandardCharsets.UTF_8);
+    if (inputFile.isDirectory()) {
+      File[] children = inputFile.listFiles();
+      if (children != null) {
+        Arrays.sort(children);
+        for (File child : children) {
+          hashDirectoryMetadataRecursive(child, entryName + "/", hasher);
+        }
+      }
+    } else {
+      hasher.putLong(inputFile.length());
+      hasher.putLong(inputFile.lastModified());
     }
-    return zipFile;
   }
 
   private static class ProcessPayloadReferenceJSON {

sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/resources/ClasspathScanningResourcesDetector.java

@@ -19,8 +19,13 @@
 
 import io.github.classgraph.ClassGraph;
 import java.io.File;
+import java.lang.ref.WeakReference;
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
+import java.util.Objects;
 import java.util.stream.Collectors;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 /**
  * Attempts to detect all the resources to be staged using classgraph library.
@@ -30,6 +35,27 @@
  */
 public class ClasspathScanningResourcesDetector implements PipelineResourcesDetector {
 
+  private static final class CachedClasspath {
+    private final WeakReference<ClassLoader> classLoader;
+    private final @Nullable String javaClassPath;
+    private final List<String> files;
+
+    CachedClasspath(ClassLoader classLoader, @Nullable String javaClassPath, List<String> files) {
+      this.classLoader = new WeakReference<>(classLoader);
+      this.javaClassPath = javaClassPath;
+      this.files = Collections.unmodifiableList(new ArrayList<>(files));
+    }
+
+    boolean matches(@Nullable ClassLoader loader, @Nullable String currentJavaClassPath) {
+      return loader != null
+          && classLoader.get() == loader
+          && Objects.equals(javaClassPath, currentJavaClassPath);
+    }
+  }
+
+  private static final Object LOCK = new Object();
+  private static volatile @Nullable CachedClasspath cachedClasspath;
+
   private transient ClassGraph classGraph;
 
   public ClasspathScanningResourcesDetector(ClassGraph classGraph) {
@@ -43,14 +69,34 @@ public ClasspathScanningResourcesDetector(ClassGraph classGraph) {
    * @return A list of absolute paths to the resources the class loader uses.
    */
   @Override
-  public List<String> detect(ClassLoader classLoader) {
-    List<File> classpathContents =
-        classGraph
-            .disableNestedJarScanning()
-            .addClassLoader(classLoader)
-            .scan(1)
-            .getClasspathFiles();
-
-    return classpathContents.stream().map(File::getAbsolutePath).collect(Collectors.toList());
+  public List<String> detect(@Nullable ClassLoader classLoader) {
+    String currentJavaClassPath = System.getProperty("java.class.path");
+    CachedClasspath snapshot = cachedClasspath;
+    if (snapshot != null && snapshot.matches(classLoader, currentJavaClassPath)) {
+      return new ArrayList<>(snapshot.files);
+    }
+
+    synchronized (LOCK) {
+      currentJavaClassPath = System.getProperty("java.class.path");
+      snapshot = cachedClasspath;
+      if (snapshot != null && snapshot.matches(classLoader, currentJavaClassPath)) {
+        return new ArrayList<>(snapshot.files);
+      }
+
+      List<File> classpathContents;
+      if (classLoader != null) {
+        classpathContents =
+            classGraph.disableNestedJarScanning().addClassLoader(classLoader).getClasspathFiles();
+      } else {
+        classpathContents = classGraph.disableNestedJarScanning().getClasspathFiles();
+      }
+
+      List<String> result =
+          classpathContents.stream().map(File::getAbsolutePath).collect(Collectors.toList());
+      if (classLoader != null) {
+        cachedClasspath = new CachedClasspath(classLoader, currentJavaClassPath, result);
+      }
+      return result;
+    }
   }
 }

sdks/java/core/src/test/java/org/apache/beam/sdk/util/construction/EnvironmentsTest.java

@@ -24,12 +24,15 @@
 import static org.hamcrest.Matchers.hasItem;
 import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.not;
 import static org.hamcrest.Matchers.notNullValue;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertSame;
 
 import java.io.File;
 import java.io.IOException;
 import java.io.Serializable;
+import java.nio.charset.StandardCharsets;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Optional;
@@ -58,6 +61,8 @@
 import org.apache.beam.sdk.values.TupleTagList;
 import org.apache.beam.sdk.values.WindowingStrategy;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.HashCode;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.Files;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
@@ -399,4 +404,52 @@ public void testResolveAnyOfEnvironment() {
             env, BeamUrns.getUrn(StandardEnvironments.Environments.EXTERNAL)),
         notNullValue());
   }
+
+  @Test
+  public void testGetFileHashCachingAndInvalidation() throws Exception {
+    File file = File.createTempFile("hash-cache-test-", ".txt");
+    file.deleteOnExit();
+    Files.asCharSink(file, StandardCharsets.UTF_8).write("initial-content");
+
+    HashCode hash1 = Environments.getFileHash(file);
+    HashCode hash2 = Environments.getFileHash(file);
+    assertSame(hash1, hash2);
+
+    // Modifying the file content and length invalidates the cache entry.
+    Files.asCharSink(file, StandardCharsets.UTF_8).write("updated-longer-content");
+    HashCode hash3 = Environments.getFileHash(file);
+    assertThat(hash3, not(equalTo(hash1)));
+  }
+
+  @Test
+  public void testGetArtifactsDirectoryZipCachingAndInvalidation() throws Exception {
+    File tempDir = Files.createTempDir();
+    tempDir.deleteOnExit();
+    File child = new File(tempDir, "entry.txt");
+    child.deleteOnExit();
+    Files.asCharSink(child, StandardCharsets.UTF_8).write("v1");
+
+    List<ArtifactInformation> firstArtifacts =
+        Environments.getArtifacts(ImmutableList.of(tempDir.getAbsolutePath()));
+    List<ArtifactInformation> secondArtifacts =
+        Environments.getArtifacts(ImmutableList.of(tempDir.getAbsolutePath()));
+
+    RunnerApi.ArtifactFilePayload firstPayload =
+        RunnerApi.ArtifactFilePayload.parseFrom(firstArtifacts.get(0).getTypePayload());
+    RunnerApi.ArtifactFilePayload secondPayload =
+        RunnerApi.ArtifactFilePayload.parseFrom(secondArtifacts.get(0).getTypePayload());
+
+    // Unchanged directory reuses the same cached zip file and SHA-256 hash.
+    assertEquals(firstPayload.getPath(), secondPayload.getPath());
+    assertEquals(firstPayload.getSha256(), secondPayload.getSha256());
+
+    // Modifying a file inside the directory produces a new zip and hash.
+    Files.asCharSink(child, StandardCharsets.UTF_8).write("v2-modified");
+    List<ArtifactInformation> thirdArtifacts =
+        Environments.getArtifacts(ImmutableList.of(tempDir.getAbsolutePath()));
+    RunnerApi.ArtifactFilePayload thirdPayload =
+        RunnerApi.ArtifactFilePayload.parseFrom(thirdArtifacts.get(0).getTypePayload());
+    assertThat(thirdPayload.getPath(), not(equalTo(firstPayload.getPath())));
+    assertThat(thirdPayload.getSha256(), not(equalTo(firstPayload.getSha256())));
+  }
 }