feat(bazel-diff): improve performance (#120)

Malinskiy · tinder-maxwellelliott · web-flow · commit bc35aee65c44 · 2022-04-27T23:21:12.000-07:00
* Use ByteBufferPool for effective pooling of memory for file hashing
* Use parallelStream() to speed up IO operations related to file hashing
* Use Maps.difference from guava for getImpactedTargets

Co-authored-by: Maxwell Elliott &lt;56700854+tinder-maxwellelliott@users.noreply.github.com&gt;
diff --git a/artifacts.bzl b/artifacts.bzl
@@ -7,5 +7,6 @@ BAZEL_DIFF_MAVEN_ARTIFACTS = [
     maven.artifact("org.mockito", "mockito-core", "3.5.15", testonly = True),
     "info.picocli:picocli:jar:4.3.2",
     "com.google.code.gson:gson:jar:2.8.6",
-    "com.google.guava:guava:29.0-jre"
+    "com.google.guava:guava:29.0-jre",
+    "org.apache.commons:commons-pool2:2.11.1",
 ]
diff --git a/src/main/java/com/bazel_diff/BUILD b/src/main/java/com/bazel_diff/BUILD
@@ -27,7 +27,8 @@ java_library(
         "@com_google_protobuf//:protobuf_java",
         "@bazel_diff_maven//:info_picocli_picocli",
         "@bazel_diff_maven//:com_google_code_gson_gson",
-        "@bazel_diff_maven//:com_google_guava_guava"
+        "@bazel_diff_maven//:com_google_guava_guava",
+        "@bazel_diff_maven//:org_apache_commons_commons_pool2"
     ],
     javacopts = select({
         ":enable_debug": ["-ADEBUG=true"],
diff --git a/src/main/java/com/bazel_diff/BazelClient.java b/src/main/java/com/bazel_diff/BazelClient.java
@@ -1,26 +1,27 @@
 package com.bazel_diff;
 
+import com.google.common.hash.Hasher;
+import com.google.common.hash.Hashing;
 import com.google.devtools.build.lib.query2.proto.proto2api.Build;
 
 import java.io.*;
 import java.nio.charset.StandardCharsets;
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
-import java.security.MessageDigest;
-import java.security.NoSuchAlgorithmException;
 import java.time.Duration;
 import java.time.Instant;
 import java.util.ArrayList;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 import java.util.Arrays;
 
 interface BazelClient {
     List<BazelTarget> queryAllTargets() throws IOException;
-    Map<String, BazelSourceFileTarget> queryAllSourcefileTargets() throws IOException, NoSuchAlgorithmException;
+
+    Map<String, BazelSourceFileTarget> queryAllSourcefileTargets() throws Exception;
 }
 
 class BazelClientImpl implements BazelClient {
@@ -43,8 +44,8 @@ class BazelClientImpl implements BazelClient {
     ) {
         this.workingDirectory = workingDirectory.normalize();
         this.bazelPath = bazelPath;
-        this.startupOptions = startupOptions != null ? Arrays.asList(startupOptions.split(" ")): new ArrayList<String>();
-        this.commandOptions = commandOptions != null ? Arrays.asList(commandOptions.split(" ")): new ArrayList<String>();
+        this.startupOptions = startupOptions != null ? Arrays.asList(startupOptions.split(" ")) : new ArrayList<String>();
+        this.commandOptions = commandOptions != null ? Arrays.asList(commandOptions.split(" ")) : new ArrayList<String>();
         this.verbose = verbose;
         this.keepGoing = keepGoing;
         this.debug = debug;
@@ -59,11 +60,11 @@ public List<BazelTarget> queryAllTargets() throws IOException {
             long querySeconds = Duration.between(queryStartTime, queryEndTime).getSeconds();
             System.out.printf("BazelDiff: All targets queried in %d seconds%n", querySeconds);
         }
-        return targets.stream().map( target -> new BazelTargetImpl(target)).collect(Collectors.toList());
+        return targets.stream().map(target -> new BazelTargetImpl(target)).collect(Collectors.toList());
     }
 
     @Override
-    public Map<String, BazelSourceFileTarget> queryAllSourcefileTargets() throws IOException, NoSuchAlgorithmException {
+    public Map<String, BazelSourceFileTarget> queryAllSourcefileTargets() throws Exception {
         Instant queryStartTime = Instant.now();
         List<Build.Target> targets = performBazelQuery("kind('source file', //...:all-targets)");
         Instant queryEndTime = Instant.now();
@@ -78,26 +79,59 @@ public Map<String, BazelSourceFileTarget> queryAllSourcefileTargets() throws IOE
         return sourceFileTargets;
     }
 
-    private Map<String, BazelSourceFileTarget> processBazelSourcefileTargets(List<Build.Target> targets, Boolean readSourcefileTargets) throws IOException, NoSuchAlgorithmException {
-        Map<String, BazelSourceFileTarget> sourceTargets = new HashMap<>();
-        for (Build.Target target : targets) {
-            Build.SourceFile sourceFile = target.getSourceFile();
-            if (sourceFile != null) {
-                MessageDigest digest = MessageDigest.getInstance("SHA-256");
-                digest.update(sourceFile.getNameBytes().toByteArray());
-                for (String subinclude : sourceFile.getSubincludeList()) {
-                    digest.update(subinclude.getBytes());
-                }
-                BazelSourceFileTargetImpl sourceFileTarget = new BazelSourceFileTargetImpl(
-                        sourceFile.getName(),
-                        digest.digest().clone(),
-                        readSourcefileTargets ? workingDirectory : null,
-                        verbose
-                );
-                sourceTargets.put(sourceFileTarget.getName(), sourceFileTarget);
-            }
+    private Map<String, BazelSourceFileTarget> processBazelSourcefileTargets(List<Build.Target> targets, Boolean readSourcefileTargets) throws Exception {
+        AtomicReference<Exception> exception = new AtomicReference(null);
+        Map<String, BazelSourceFileTarget> result = targets.parallelStream().map((target -> {
+                    Build.SourceFile sourceFile = target.getSourceFile();
+                    if (sourceFile != null) {
+                        Hasher hasher = Hashing.sha256().newHasher();
+                        hasher.putBytes(sourceFile.getNameBytes().toByteArray());
+                        for (String subinclude : sourceFile.getSubincludeList()) {
+                            hasher.putBytes(subinclude.getBytes());
+                        }
+                        BazelSourceFileTargetImpl sourceFileTarget = null;
+                        try {
+                            sourceFileTarget = new BazelSourceFileTargetImpl(
+                                    sourceFile.getName(),
+                                    hasher.hash().asBytes().clone(),
+                                    readSourcefileTargets ? workingDirectory : null,
+                                    verbose
+                            );
+                        } catch (Exception e) {
+                            exception.set(e);
+                        }
+                        return new SourceTargetEntry(sourceFileTarget.getName(), sourceFileTarget);
+                    }
+                    return null;
+                }))
+                .filter(pair -> pair != null)
+                .collect(Collectors.toMap(SourceTargetEntry::getKey, SourceTargetEntry::getValue));
+
+        //Rethrowing nested parallel exception
+        Exception nestedException = exception.get();
+        if (nestedException != null) {
+            throw nestedException;
+        }
+
+        return result;
+    }
+
+    private static class SourceTargetEntry<K extends String, V extends BazelSourceFileTargetImpl> {
+        private K key;
+        private V value;
+
+        public SourceTargetEntry(K key, V value) {
+            this.key = key;
+            this.value = value;
+        }
+
+        public K getKey() {
+            return key;
+        }
+
+        public V getValue() {
+            return value;
         }
-        return sourceTargets;
     }
 
     private List<Build.Target> performBazelQuery(String query) throws IOException {
@@ -134,18 +168,20 @@ private List<Build.Target> performBazelQuery(String query) throws IOException {
         BufferedReader stdError = new BufferedReader(new InputStreamReader(process.getErrorStream()));
         Thread tStdError = new Thread(new Runnable() {
             String line = null;
+
             public void run() {
                 try {
                     while ((line = stdError.readLine()) != null) {
                         if (verbose) {
                             System.out.println(line);
                         }
 
-                        if(Thread.currentThread().isInterrupted()) {
+                        if (Thread.currentThread().isInterrupted()) {
                             return;
                         }
                     }
-                } catch(IOException e) {}
+                } catch (IOException e) {
+                }
             }
         });
         tStdError.start();
diff --git a/src/main/java/com/bazel_diff/BazelRule.java b/src/main/java/com/bazel_diff/BazelRule.java
@@ -1,13 +1,14 @@
 package com.bazel_diff;
 
-import java.security.MessageDigest;
-import java.security.NoSuchAlgorithmException;
 import java.util.List;
+
+import com.google.common.hash.Hasher;
+import com.google.common.hash.Hashing;
 import com.google.devtools.build.lib.query2.proto.proto2api.Build;
 import java.util.stream.Collectors;
 
 interface BazelRule {
-    byte[] getDigest() throws NoSuchAlgorithmException;
+    byte[] getDigest();
     List<String> getRuleInputList();
     String getName();
 }
@@ -20,15 +21,15 @@ class BazelRuleImpl implements BazelRule {
     }
 
     @Override
-    public byte[] getDigest() throws NoSuchAlgorithmException {
-        MessageDigest digest = MessageDigest.getInstance("SHA-256");
-        digest.update(rule.getRuleClassBytes().toByteArray());
-        digest.update(rule.getNameBytes().toByteArray());
-        digest.update(rule.getSkylarkEnvironmentHashCodeBytes().toByteArray());
+    public byte[] getDigest() {
+        Hasher hasher = Hashing.sha256().newHasher();
+        hasher.putBytes(rule.getRuleClassBytes().toByteArray());
+        hasher.putBytes(rule.getNameBytes().toByteArray());
+        hasher.putBytes(rule.getSkylarkEnvironmentHashCodeBytes().toByteArray());
         for (Build.Attribute attribute : rule.getAttributeList()) {
-            digest.update(attribute.toByteArray());
+            hasher.putBytes(attribute.toByteArray());
         }
-        return digest.digest();
+        return hasher.hash().asBytes();
     }
 
     @Override
diff --git a/src/main/java/com/bazel_diff/BazelSourceFileTarget.java b/src/main/java/com/bazel_diff/BazelSourceFileTarget.java
@@ -1,13 +1,12 @@
 package com.bazel_diff;
 
+import com.google.common.hash.Hasher;
+import com.google.common.hash.Hashing;
+
 import java.io.*;
 import java.nio.ByteBuffer;
-import java.nio.channels.FileChannel;
 import java.nio.file.Path;
 import java.nio.file.Paths;
-import java.io.IOException;
-import java.security.MessageDigest;
-import java.security.NoSuchAlgorithmException;
 
 interface BazelSourceFileTarget {
     String getName();
@@ -17,51 +16,40 @@ interface BazelSourceFileTarget {
 class BazelSourceFileTargetImpl implements BazelSourceFileTarget {
     private String name;
     private byte[] digest;
+    static private ByteBufferPool pool = new ByteBufferPool(1024, 10240); //10kb
 
-    private void digestLargeFile(MessageDigest finalDigest, FileChannel inChannel) throws IOException {
-        int bufferSize = 10240; // 10kb
-        ByteBuffer buffer = ByteBuffer.allocate(bufferSize);
-        while (inChannel.read(buffer) != -1) {
+    private void digest(Hasher finalDigest, InputStream stream) throws Exception {
+        ByteBuffer buffer = pool.borrow();
+        byte[] array = buffer.array(); //Available for non-direct buffers
+        Integer length = 0;
+        while (true) {
+            if (!((length = stream.read(array)) != -1)) break;
             buffer.flip();
-            finalDigest.update(buffer);
+            finalDigest.putBytes(array, 0, length);
             buffer.clear();
         }
+        pool.recycle(buffer);
     }
 
-    private void digestSmallFile(MessageDigest finalDigest, FileChannel inChannel) throws IOException {
-        long fileSize = inChannel.size();
-        ByteBuffer buffer = ByteBuffer.allocate((int) fileSize);
-        inChannel.read(buffer);
-        buffer.flip();
-        finalDigest.update(buffer);
-    }
-
-    BazelSourceFileTargetImpl(String name, byte[] digest, Path workingDirectory, Boolean verbose)
-        throws IOException, NoSuchAlgorithmException {
+    BazelSourceFileTargetImpl(String name, byte[] digest, Path workingDirectory, Boolean verbose) throws Exception {
         this.name = name;
-        MessageDigest finalDigest = MessageDigest.getInstance("SHA-256");
+        Hasher hasher = Hashing.sha256().newHasher();
         if (workingDirectory != null && name.startsWith("//")) {
             String filenameSubstring = name.substring(2);
             String filenamePath = filenameSubstring.replaceFirst(":", "/");
             Path absoluteFilePath = Paths.get(workingDirectory.toString(), filenamePath);
-            try (RandomAccessFile sourceFile = new RandomAccessFile(absoluteFilePath.toString(), "r")) {
-                FileChannel inChannel = sourceFile.getChannel();
-                if (inChannel.size() > 1048576) { // 1mb
-                    digestLargeFile(finalDigest, inChannel);
-                } else {
-                    digestSmallFile(finalDigest, inChannel);
-                }
-                sourceFile.close();
-                inChannel.close();
+
+            try (InputStream stream = new BufferedInputStream(new FileInputStream(absoluteFilePath.toString()))) {
+                digest(hasher, stream);
             } catch (FileNotFoundException e) {
                 if (verbose) {
                     System.out.printf("BazelDiff: [Warning] file %s not found%n", absoluteFilePath);
                 }
             }
         }
-        finalDigest.update(digest);
-        finalDigest.update(name.getBytes());
-        this.digest = finalDigest.digest();
+        hasher.putBytes(digest);
+        hasher.putBytes(name.getBytes());
+        this.digest = hasher.hash().asBytes();
     }
 
     @Override
diff --git a/src/main/java/com/bazel_diff/ByteBufferPool.java b/src/main/java/com/bazel_diff/ByteBufferPool.java
@@ -0,0 +1,67 @@
+package com.bazel_diff;
+
+import com.google.common.base.Supplier;
+import com.google.common.base.Suppliers;
+import org.apache.commons.pool2.PooledObject;
+import org.apache.commons.pool2.PooledObjectFactory;
+import org.apache.commons.pool2.impl.DefaultPooledObject;
+import org.apache.commons.pool2.impl.GenericObjectPool;
+import org.apache.commons.pool2.impl.GenericObjectPoolConfig;
+
+import java.nio.ByteBuffer;
+
+public class ByteBufferPool {
+    private Integer poolSize;
+    private Integer bufferSize;
+    private Supplier<GenericObjectPool<ByteBuffer>> delegate = Suppliers.memoize(() -> {
+        GenericObjectPoolConfig<ByteBuffer> config = new GenericObjectPoolConfig<>();
+        config.setMaxTotal(poolSize);
+        return new GenericObjectPool(new ByteBufferObjectFactory(bufferSize), config);
+    });
+
+    public ByteBufferPool(Integer poolSize, Integer bufferSize) {
+        this.poolSize = poolSize;
+        this.bufferSize = bufferSize;
+    }
+
+    public ByteBuffer borrow() throws Exception {
+        return delegate.get().borrowObject();
+    }
+
+    public void recycle(ByteBuffer buffer) {
+        delegate.get().returnObject(buffer);
+    }
+}
+
+class ByteBufferObjectFactory implements PooledObjectFactory<ByteBuffer> {
+    private Integer bufferSize;
+
+    public ByteBufferObjectFactory(Integer bufferSize) {
+        this.bufferSize = bufferSize;
+    }
+
+    @Override
+    public void activateObject(PooledObject<ByteBuffer> p) throws Exception {
+        p.getObject().clear();
+    }
+
+    @Override
+    public void destroyObject(PooledObject<ByteBuffer> p) throws Exception {
+        p.getObject().clear();
+    }
+
+    @Override
+    public PooledObject<ByteBuffer> makeObject() throws Exception {
+        return new DefaultPooledObject(ByteBuffer.allocate(bufferSize));
+    }
+
+    @Override
+    public void passivateObject(PooledObject<ByteBuffer> p) throws Exception {
+        p.getObject().clear();
+    }
+
+    @Override
+    public boolean validateObject(PooledObject<ByteBuffer> p) {
+        return p.getObject().capacity() == bufferSize && !p.getObject().isDirect();
+    }
+}
diff --git a/src/main/java/com/bazel_diff/TargetHashingClient.java b/src/main/java/com/bazel_diff/TargetHashingClient.java
diff --git a/src/main/java/com/bazel_diff/main.java b/src/main/java/com/bazel_diff/main.java
diff --git a/test/java/com/bazel_diff/TargetHashingClientImplTests.java b/test/java/com/bazel_diff/TargetHashingClientImplTests.java

Original file line number	Diff line number	Diff line change
`@@ -7,5 +7,6 @@ BAZEL_DIFF_MAVEN_ARTIFACTS = [`
`7`	`7`	`maven.artifact("org.mockito", "mockito-core", "3.5.15", testonly = True),`
`8`	`8`	`"info.picocli:picocli:jar:4.3.2",`
`9`	`9`	`"com.google.code.gson:gson:jar:2.8.6",`
`10`		`- "com.google.guava:guava:29.0-jre"`
	`10`	`+ "com.google.guava:guava:29.0-jre",`
	`11`	`+ "org.apache.commons:commons-pool2:2.11.1",`
`11`	`12`	`]`