Enable connection pooling for the remote REST cache

Connection pooling is a useful optimization for REST caches that are
far away as it avoids constantly redoing the TCP handshake. It also
prevents large builds from exhausting the local interface's source
ports through tens of thousands of one-transaction connections.

The default connection pool size of 20 is fairly arbitrary. Users
probably want to set this to something close to the value of --jobs.

We introduce some generic infrastructure for closing remote cache
instances and use it to cleanly shutdown the connection pool between
builds.

Change-Id: I73adc29ecae15cc10a1217ffbaa483892bcd4f9a
PiperOrigin-RevId: 160264681

9 files changed, 120 insertions, 42 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/remote/GrpcActionCache.java b/src/main/java/com/google/devtools/build/lib/remote/GrpcActionCache.java
index ba8a9d4257..0f3f6045d2 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/GrpcActionCache.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/GrpcActionCache.java
@@ -81,6 +81,9 @@ public class GrpcActionCache implements RemoteActionCache {
     this.channel = channel;
   }
 
+  @Override
+  public void close() {}
+
   // All gRPC stubs are reused.
   private final Supplier<ContentAddressableStorageBlockingStub> casBlockingStub =
       Suppliers.memoize(
diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteActionCache.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteActionCache.java
index 44a30f5745..994c23e1e6 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/RemoteActionCache.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/RemoteActionCache.java
@@ -74,4 +74,7 @@ interface RemoteActionCache {
    */
   void upload(ActionKey actionKey, Path execRoot, Collection<Path> files, FileOutErr outErr)
       throws IOException, InterruptedException;
+
+  /** Release resources associated with the cache. The cache may not be used after calling this. */
+  void close();
 }
diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteActionContextProvider.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteActionContextProvider.java
index a3a9d700ea..29789cea13 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/RemoteActionContextProvider.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/RemoteActionContextProvider.java
@@ -34,6 +34,7 @@ import com.google.devtools.build.lib.standalone.StandaloneSpawnStrategy;
 final class RemoteActionContextProvider extends ActionContextProvider {
   private final CommandEnvironment env;
   private ActionInputPrefetcher actionInputPrefetcher;
+  private RemoteSpawnStrategy spawnStrategy;
 
   RemoteActionContextProvider(CommandEnvironment env) {
     this.env = env;
@@ -43,10 +44,6 @@ final class RemoteActionContextProvider extends ActionContextProvider {
   public void init(
       ActionInputFileCache actionInputFileCache, ActionInputPrefetcher actionInputPrefetcher) {
     this.actionInputPrefetcher = Preconditions.checkNotNull(actionInputPrefetcher);
-  }
-
-  @Override
-  public Iterable<? extends ActionContext> getActionContexts() {
     ExecutionOptions executionOptions = env.getOptions().getOptions(ExecutionOptions.class);
     LocalExecutionOptions localExecutionOptions =
         env.getOptions().getOptions(LocalExecutionOptions.class);
@@ -58,12 +55,25 @@ final class RemoteActionContextProvider extends ActionContextProvider {
             executionOptions.verboseFailures,
             env.getRuntime().getProductName(),
             ResourceManager.instance());
-    return ImmutableList.of(
+    spawnStrategy =
         new RemoteSpawnStrategy(
             env.getExecRoot(),
             env.getOptions().getOptions(RemoteOptions.class),
             env.getOptions().getOptions(AuthAndTLSOptions.class),
             executionOptions.verboseFailures,
-            fallbackStrategy));
+            fallbackStrategy);
+  }
+
+  @Override
+  public Iterable<? extends ActionContext> getActionContexts() {
+    return ImmutableList.of(Preconditions.checkNotNull(spawnStrategy));
+  }
+
+  @Override
+  public void executionPhaseEnding() {
+    if (spawnStrategy != null) {
+      spawnStrategy.close();
+      spawnStrategy = null;
+    }
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteOptions.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteOptions.java
index 35dd9f5841..ebbf4fd487 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/RemoteOptions.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/RemoteOptions.java
@@ -30,6 +30,14 @@ public final class RemoteOptions extends OptionsBase {
   public String remoteRestCache;
 
   @Option(
+    name = "remote_rest_cache_pool_size",
+    defaultValue = "20",
+    category = "remote",
+    help = "Size of the HTTP pool for making requests to the REST cache."
+  )
+  public int restCachePoolSize;
+
+  @Option(
     name = "hazelcast_node",
     defaultValue = "null",
     category = "remote",
diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteSpawnStrategy.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteSpawnStrategy.java
index 43bb3770ab..1a8ab90247 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/RemoteSpawnStrategy.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/RemoteSpawnStrategy.java
@@ -125,6 +125,13 @@ final class RemoteSpawnStrategy implements SpawnActionContext {
     }
   }
 
+  /** Release resources associated with this spawn strategy. */
+  public void close() {
+    if (remoteCache != null) {
+      remoteCache.close();
+    }
+  }
+
   private Action buildAction(
       Collection<? extends ActionInput> outputs,
       Digest command,
diff --git a/src/main/java/com/google/devtools/build/lib/remote/SimpleBlobStore.java b/src/main/java/com/google/devtools/build/lib/remote/SimpleBlobStore.java
index 6c46407763..66b80139dc 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/SimpleBlobStore.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/SimpleBlobStore.java
@@ -33,4 +33,7 @@ public interface SimpleBlobStore {
    * indexed by the same {@param key} will be overwritten.
    */
   void put(String key, byte[] value);
+
+  /** Close resources associated with the blob store. */
+  void close();
 }
diff --git a/src/main/java/com/google/devtools/build/lib/remote/SimpleBlobStoreActionCache.java b/src/main/java/com/google/devtools/build/lib/remote/SimpleBlobStoreActionCache.java
index 55e1860561..6fd37b0a61 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/SimpleBlobStoreActionCache.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/SimpleBlobStoreActionCache.java
@@ -272,4 +272,9 @@ public final class SimpleBlobStoreActionCache implements RemoteActionCache {
       throws InterruptedException {
     blobStore.put(actionKey.getDigest().getHash(), result.toByteArray());
   }
+
+  @Override
+  public void close() {
+    blobStore.close();
+  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/remote/SimpleBlobStoreFactory.java b/src/main/java/com/google/devtools/build/lib/remote/SimpleBlobStoreFactory.java
index 9f210cf940..9e906dee08 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/SimpleBlobStoreFactory.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/SimpleBlobStoreFactory.java
@@ -24,16 +24,16 @@ import com.hazelcast.core.HazelcastInstance;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.util.concurrent.ConcurrentMap;
-import org.apache.http.HttpEntity;
 import org.apache.http.HttpResponse;
 import org.apache.http.HttpStatus;
 import org.apache.http.client.HttpClient;
+import org.apache.http.client.ResponseHandler;
 import org.apache.http.client.methods.HttpGet;
 import org.apache.http.client.methods.HttpHead;
 import org.apache.http.client.methods.HttpPut;
 import org.apache.http.entity.ByteArrayEntity;
 import org.apache.http.impl.client.HttpClientBuilder;
-import org.apache.http.util.EntityUtils;
+import org.apache.http.impl.conn.PoolingHttpClientConnectionManager;
 
 /**
  * A factory class for providing a {@link SimpleBlobStore} to be used with {@link
@@ -67,6 +67,9 @@ public final class SimpleBlobStoreFactory {
     public void put(String key, byte[] value) {
       map.put(key, value);
     }
+
+    @Override
+    public void close() {}
   }
 
   /** Construct a {@link SimpleBlobStore} using Hazelcast's version of {@link ConcurrentMap} */
@@ -120,19 +123,38 @@ public final class SimpleBlobStoreFactory {
   private static class RestBlobStore implements SimpleBlobStore {
 
     private final String baseUrl;
+    private final PoolingHttpClientConnectionManager connMan;
+    private final HttpClientBuilder clientFactory;
 
-    RestBlobStore(String baseUrl) {
+    RestBlobStore(String baseUrl, int poolSize) {
       this.baseUrl = baseUrl;
+      connMan = new PoolingHttpClientConnectionManager();
+      connMan.setDefaultMaxPerRoute(poolSize);
+      connMan.setMaxTotal(poolSize);
+      clientFactory = HttpClientBuilder.create();
+      clientFactory.setConnectionManager(connMan);
+      clientFactory.setConnectionManagerShared(true);
+    }
+
+    @Override
+    public void close() {
+      connMan.close();
     }
 
     @Override
     public boolean containsKey(String key) {
       try {
-        HttpClient client = HttpClientBuilder.create().build();
+        HttpClient client = clientFactory.build();
         HttpHead head = new HttpHead(baseUrl + "/" + key);
-        HttpResponse response = client.execute(head);
-        int statusCode = response.getStatusLine().getStatusCode();
-        return HttpStatus.SC_OK == statusCode;
+        return client.execute(
+            head,
+            new ResponseHandler<Boolean>() {
+              @Override
+              public Boolean handleResponse(HttpResponse response) {
+                int statusCode = response.getStatusLine().getStatusCode();
+                return HttpStatus.SC_OK == statusCode;
+              }
+            });
       } catch (IOException e) {
         throw new RuntimeException(e);
       }
@@ -141,24 +163,27 @@ public final class SimpleBlobStoreFactory {
     @Override
     public byte[] get(String key) {
       try {
-        HttpClient client = HttpClientBuilder.create().build();
+        HttpClient client = clientFactory.build();
         HttpGet get = new HttpGet(baseUrl + "/" + key);
-        HttpResponse response = client.execute(get);
-        int statusCode = response.getStatusLine().getStatusCode();
-        if (HttpStatus.SC_NOT_FOUND == statusCode || HttpStatus.SC_NO_CONTENT == statusCode) {
-          return null;
-        }
-        if (HttpStatus.SC_OK != statusCode) {
-          throw new RuntimeException("GET failed with status code " + statusCode);
-        }
-        ByteArrayOutputStream buffer = new ByteArrayOutputStream();
-        HttpEntity entity = response.getEntity();
-        entity.writeTo(buffer);
-        buffer.flush();
-        EntityUtils.consume(entity);
-
-        return buffer.toByteArray();
-
+        return client.execute(
+            get,
+            new ResponseHandler<byte[]>() {
+              @Override
+              public byte[] handleResponse(HttpResponse response) throws IOException {
+                int statusCode = response.getStatusLine().getStatusCode();
+                if (HttpStatus.SC_NOT_FOUND == statusCode
+                    || HttpStatus.SC_NO_CONTENT == statusCode) {
+                  return null;
+                }
+                if (HttpStatus.SC_OK != statusCode) {
+                  throw new RuntimeException("GET failed with status code " + statusCode);
+                }
+                ByteArrayOutputStream buffer = new ByteArrayOutputStream();
+                response.getEntity().writeTo(buffer);
+                buffer.flush();
+                return buffer.toByteArray();
+              }
+            });
       } catch (IOException e) {
         throw new RuntimeException(e);
       }
@@ -167,20 +192,27 @@ public final class SimpleBlobStoreFactory {
     @Override
     public void put(String key, byte[] value) {
       try {
-        HttpClient client = HttpClientBuilder.create().build();
+        HttpClient client = clientFactory.build();
         HttpPut put = new HttpPut(baseUrl + "/" + key);
         put.setEntity(new ByteArrayEntity(value));
         put.setHeader("Content-Type", "application/octet-stream");
-        HttpResponse response = client.execute(put);
-        int statusCode = response.getStatusLine().getStatusCode();
-
-        // Accept more than SC_OK to be compatible with Nginx WebDav module.
-        if (HttpStatus.SC_OK != statusCode
-            && HttpStatus.SC_ACCEPTED != statusCode
-            && HttpStatus.SC_CREATED != statusCode
-            && HttpStatus.SC_NO_CONTENT != statusCode) {
-          throw new RuntimeException("PUT failed with status code " + statusCode);
-        }
+        client.execute(
+            put,
+            new ResponseHandler<Void>() {
+              @Override
+              public Void handleResponse(HttpResponse response) {
+                int statusCode = response.getStatusLine().getStatusCode();
+
+                // Accept more than SC_OK to be compatible with Nginx WebDav module.
+                if (HttpStatus.SC_OK != statusCode
+                    && HttpStatus.SC_ACCEPTED != statusCode
+                    && HttpStatus.SC_CREATED != statusCode
+                    && HttpStatus.SC_NO_CONTENT != statusCode) {
+                  throw new RuntimeException("PUT failed with status code " + statusCode);
+                }
+                return null;
+              }
+            });
       } catch (IOException e) {
         throw new RuntimeException(e);
       }
@@ -188,7 +220,7 @@ public final class SimpleBlobStoreFactory {
   }
 
   public static SimpleBlobStore createRest(RemoteOptions options) {
-    return new RestBlobStore(options.remoteRestCache);
+    return new RestBlobStore(options.remoteRestCache, options.restCachePoolSize);
   }
 
   public static SimpleBlobStore create(RemoteOptions options) {
diff --git a/src/test/shell/bazel/remote_execution_test.sh b/src/test/shell/bazel/remote_execution_test.sh
index 6dfd0ae529..b119f885f7 100755
--- a/src/test/shell/bazel/remote_execution_test.sh
+++ b/src/test/shell/bazel/remote_execution_test.sh
@@ -215,6 +215,13 @@ EOF
       || fail "Failed to build //a:test with remote gRPC cache service"
   diff bazel-bin/a/test ${TEST_TMPDIR}/test_expected \
       || fail "Remote cache generated different result"
+  # Check that persistent connections are closed after the build. Is there a good cross-platform way
+  # to check this?
+  if [[ "$PLATFORM" = "linux" ]]; then
+    if netstat -tn | grep -qE ":${hazelcast_port}\\s+ESTABLISHED$"; then
+      fail "connections to to cache not closed"
+    fi
+  fi
 }
 
 function test_py_test() {