Refactor combined cache. (bazelbuild#16110)

Instead of guessing when to use remote/disk part, combined cache now uses read/write cache policy provided by the context. The call sites can change the policy based on the requirements. Fixes bazelbuild#15934. But unfortunately, I am not able to add an integration test for it because our own remote worker doesn't support the asset API. Closes bazelbuild#16039. PiperOrigin-RevId: 465577383 Change-Id: I99effab1cdcba0890671ea64c4660ea31b059ce7
databricks · Aug 16, 2022 · 83041b1 · 83041b1
1 parent e745468
commit 83041b1
Show file tree

Hide file tree

Showing 9 changed files with 232 additions and 256 deletions.
diff --git a/src/main/java/com/google/devtools/build/lib/remote/ByteStreamBuildEventArtifactUploader.java b/src/main/java/com/google/devtools/build/lib/remote/ByteStreamBuildEventArtifactUploader.java
@@ -30,7 +30,7 @@
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.events.ExtendedEventHandler;
 import com.google.devtools.build.lib.remote.common.RemoteActionExecutionContext;
-import com.google.devtools.build.lib.remote.common.RemoteActionExecutionContext.Step;
+import com.google.devtools.build.lib.remote.common.RemoteActionExecutionContext.CachePolicy;
 import com.google.devtools.build.lib.remote.util.DigestUtil;
 import com.google.devtools.build.lib.remote.util.TracingMetadataUtils;
 import com.google.devtools.build.lib.vfs.Path;
@@ -274,8 +274,9 @@ private Single<PathConverter> upload(Set<Path> files) {
 
     RequestMetadata metadata =
         TracingMetadataUtils.buildMetadata(buildRequestId, commandId, "bes-upload", null);
-    RemoteActionExecutionContext context = RemoteActionExecutionContext.create(metadata);
-    context.setStep(Step.UPLOAD_BES_FILES);
+    RemoteActionExecutionContext context =
+        RemoteActionExecutionContext.create(metadata)
+            .withWriteCachePolicy(CachePolicy.REMOTE_CACHE_ONLY);
 
     return Single.using(
         remoteCache::retain,

diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteCacheClientFactory.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteCacheClientFactory.java
@@ -44,12 +44,11 @@ public static RemoteCacheClient createDiskAndRemoteClient(
       PathFragment diskCachePath,
       boolean remoteVerifyDownloads,
       DigestUtil digestUtil,
-      RemoteCacheClient remoteCacheClient,
-      RemoteOptions options)
+      RemoteCacheClient remoteCacheClient)
       throws IOException {
     DiskCacheClient diskCacheClient =
         createDiskCache(workingDirectory, diskCachePath, remoteVerifyDownloads, digestUtil);
-    return new DiskAndRemoteCacheClient(diskCacheClient, remoteCacheClient, options);
+    return new DiskAndRemoteCacheClient(diskCacheClient, remoteCacheClient);
   }
 
   public static RemoteCacheClient create(
@@ -147,12 +146,7 @@ private static RemoteCacheClient createDiskAndHttpCache(
 
     RemoteCacheClient httpCache = createHttp(options, cred, digestUtil);
     return createDiskAndRemoteClient(
-        workingDirectory,
-        diskCachePath,
-        options.remoteVerifyDownloads,
-        digestUtil,
-        httpCache,
-        options);
+        workingDirectory, diskCachePath, options.remoteVerifyDownloads, digestUtil, httpCache);
   }
 
   public static boolean isDiskCache(RemoteOptions options) {

diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteExecutionService.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteExecutionService.java
diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteModule.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteModule.java
@@ -567,8 +567,7 @@ public void beforeCommand(CommandEnvironment env) throws AbruptExitException {
                   remoteOptions.diskCache,
                   remoteOptions.remoteVerifyDownloads,
                   digestUtil,
-                  cacheClient,
-                  remoteOptions);
+                  cacheClient);
         } catch (IOException e) {
           handleInitFailure(env, e, Code.CACHE_INIT_FAILURE);
           return;
@@ -626,8 +625,7 @@ public void beforeCommand(CommandEnvironment env) throws AbruptExitException {
                   remoteOptions.diskCache,
                   remoteOptions.remoteVerifyDownloads,
                   digestUtil,
-                  cacheClient,
-                  remoteOptions);
+                  cacheClient);
         } catch (IOException e) {
           handleInitFailure(env, e, Code.CACHE_INIT_FAILURE);
           return;

diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteSpawnCache.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteSpawnCache.java
@@ -82,8 +82,10 @@ RemoteExecutionService getRemoteExecutionService() {
   @Override
   public CacheHandle lookup(Spawn spawn, SpawnExecutionContext context)
       throws InterruptedException, IOException, ExecException, ForbiddenActionInputException {
-    boolean shouldAcceptCachedResult = remoteExecutionService.shouldAcceptCachedResult(spawn);
-    boolean shouldUploadLocalResults = remoteExecutionService.shouldUploadLocalResults(spawn);
+    boolean shouldAcceptCachedResult =
+        remoteExecutionService.getReadCachePolicy(spawn).allowAnyCache();
+    boolean shouldUploadLocalResults =
+        remoteExecutionService.getWriteCachePolicy(spawn).allowAnyCache();
     if (!shouldAcceptCachedResult && !shouldUploadLocalResults) {
       return SpawnCache.NO_RESULT_NO_STORE;
     }

diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteSpawnRunner.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteSpawnRunner.java
@@ -183,10 +183,11 @@ public SpawnResult exec(Spawn spawn, SpawnExecutionContext context)
         "Spawn can't be executed remotely. This is a bug.");
 
     Stopwatch totalTime = Stopwatch.createStarted();
-    boolean uploadLocalResults = remoteExecutionService.shouldUploadLocalResults(spawn);
-    boolean acceptCachedResult = remoteExecutionService.shouldAcceptCachedResult(spawn);
+    boolean acceptCachedResult = remoteExecutionService.getReadCachePolicy(spawn).allowAnyCache();
+    boolean uploadLocalResults = remoteExecutionService.getWriteCachePolicy(spawn).allowAnyCache();
 
     RemoteAction action = remoteExecutionService.buildRemoteAction(spawn, context);
+
     SpawnMetrics.Builder spawnMetrics =
         SpawnMetrics.Builder.forRemoteExec()
             .setInputBytes(action.getInputBytes())

diff --git a/src/main/java/com/google/devtools/build/lib/remote/common/RemoteActionExecutionContext.java b/src/main/java/com/google/devtools/build/lib/remote/common/RemoteActionExecutionContext.java
@@ -13,48 +13,89 @@
 // limitations under the License.
 package com.google.devtools.build.lib.remote.common;
 
-import build.bazel.remote.execution.v2.ExecuteResponse;
 import build.bazel.remote.execution.v2.RequestMetadata;
 import com.google.devtools.build.lib.actions.ActionExecutionMetadata;
 import com.google.devtools.build.lib.actions.Spawn;
 import javax.annotation.Nullable;
 
 /** A context that provide remote execution related information for executing an action remotely. */
 public class RemoteActionExecutionContext {
-  /** The current step of the context. */
-  public enum Step {
-    INIT,
-    CHECK_ACTION_CACHE,
-    UPLOAD_INPUTS,
-    EXECUTE_REMOTELY,
-    UPLOAD_OUTPUTS,
-    DOWNLOAD_OUTPUTS,
-    UPLOAD_BES_FILES,
+  /** Determines whether to read/write remote cache, disk cache or both. */
+  public enum CachePolicy {
+    NO_CACHE,
+    REMOTE_CACHE_ONLY,
+    DISK_CACHE_ONLY,
+    ANY_CACHE;
+
+    public static CachePolicy create(boolean allowRemoteCache, boolean allowDiskCache) {
+      if (allowRemoteCache && allowDiskCache) {
+        return ANY_CACHE;
+      } else if (allowRemoteCache) {
+        return REMOTE_CACHE_ONLY;
+      } else if (allowDiskCache) {
+        return DISK_CACHE_ONLY;
+      } else {
+        return NO_CACHE;
+      }
+    }
+
+    public boolean allowAnyCache() {
+      return this != NO_CACHE;
+    }
+
+    public boolean allowRemoteCache() {
+      return this == REMOTE_CACHE_ONLY || this == ANY_CACHE;
+    }
+
+    public boolean allowDiskCache() {
+      return this == DISK_CACHE_ONLY || this == ANY_CACHE;
+    }
+
+    public CachePolicy addRemoteCache() {
+      if (this == DISK_CACHE_ONLY || this == ANY_CACHE) {
+        return ANY_CACHE;
+      }
+
+      return REMOTE_CACHE_ONLY;
+    }
   }
 
   @Nullable private final Spawn spawn;
   private final RequestMetadata requestMetadata;
   private final NetworkTime networkTime;
-
-  @Nullable private ExecuteResponse executeResponse;
-  private Step step;
+  private final CachePolicy writeCachePolicy;
+  private final CachePolicy readCachePolicy;
 
   private RemoteActionExecutionContext(
       @Nullable Spawn spawn, RequestMetadata requestMetadata, NetworkTime networkTime) {
     this.spawn = spawn;
     this.requestMetadata = requestMetadata;
     this.networkTime = networkTime;
-    this.step = Step.INIT;
+    this.writeCachePolicy = CachePolicy.ANY_CACHE;
+    this.readCachePolicy = CachePolicy.ANY_CACHE;
   }
 
-  /** Returns current {@link Step} of the context. */
-  public Step getStep() {
-    return step;
+  private RemoteActionExecutionContext(
+      @Nullable Spawn spawn,
+      RequestMetadata requestMetadata,
+      NetworkTime networkTime,
+      CachePolicy writeCachePolicy,
+      CachePolicy readCachePolicy) {
+    this.spawn = spawn;
+    this.requestMetadata = requestMetadata;
+    this.networkTime = networkTime;
+    this.writeCachePolicy = writeCachePolicy;
+    this.readCachePolicy = readCachePolicy;
+  }
+
+  public RemoteActionExecutionContext withWriteCachePolicy(CachePolicy writeCachePolicy) {
+    return new RemoteActionExecutionContext(
+        spawn, requestMetadata, networkTime, writeCachePolicy, readCachePolicy);
   }
 
-  /** Sets current {@link Step} of the context. */
-  public void setStep(Step step) {
-    this.step = step;
+  public RemoteActionExecutionContext withReadCachePolicy(CachePolicy readCachePolicy) {
+    return new RemoteActionExecutionContext(
+        spawn, requestMetadata, networkTime, writeCachePolicy, readCachePolicy);
   }
 
   /** Returns the {@link Spawn} of the action being executed or {@code null}. */
@@ -86,13 +127,12 @@ public ActionExecutionMetadata getSpawnOwner() {
     return spawn.getResourceOwner();
   }
 
-  public void setExecuteResponse(@Nullable ExecuteResponse executeResponse) {
-    this.executeResponse = executeResponse;
+  public CachePolicy getWriteCachePolicy() {
+    return writeCachePolicy;
   }
 
-  @Nullable
-  public ExecuteResponse getExecuteResponse() {
-    return executeResponse;
+  public CachePolicy getReadCachePolicy() {
+    return readCachePolicy;
   }
 
   /** Creates a {@link RemoteActionExecutionContext} with given {@link RequestMetadata}. */
@@ -108,4 +148,13 @@ public static RemoteActionExecutionContext create(
       @Nullable Spawn spawn, RequestMetadata metadata) {
     return new RemoteActionExecutionContext(spawn, metadata, new NetworkTime());
   }
+
+  public static RemoteActionExecutionContext create(
+      @Nullable Spawn spawn,
+      RequestMetadata requestMetadata,
+      CachePolicy writeCachePolicy,
+      CachePolicy readCachePolicy) {
+    return new RemoteActionExecutionContext(
+        spawn, requestMetadata, new NetworkTime(), writeCachePolicy, readCachePolicy);
+  }
 }