Impl lease extension

bazelbuild · Mar 31, 2023 · 66d476b · 66d476b
1 parent 6146e4a
commit 66d476b
Show file tree

Hide file tree

Showing 7 changed files with 325 additions and 15 deletions.
diff --git a/src/main/java/com/google/devtools/build/lib/actions/ActionCacheUtils.java b/src/main/java/com/google/devtools/build/lib/actions/ActionCacheUtils.java
@@ -14,24 +14,36 @@
 package com.google.devtools.build.lib.actions;
 
 import com.google.devtools.build.lib.actions.cache.ActionCache;
+import java.util.AbstractMap.SimpleEntry;
+import java.util.Map.Entry;
 import javax.annotation.Nullable;
 
 /** Utility functions for {@link ActionCache}. */
 public class ActionCacheUtils {
   private ActionCacheUtils() {}
 
-  /** Checks whether one of existing output paths is already used as a key. */
   @Nullable
-  public static ActionCache.Entry getCacheEntry(ActionCache actionCache, Action action) {
+  public static Entry<String, ActionCache.Entry> getCacheEntryWithKey(
+      ActionCache actionCache, Action action) {
     for (Artifact output : action.getOutputs()) {
       ActionCache.Entry entry = actionCache.get(output.getExecPathString());
       if (entry != null) {
-        return entry;
+        return new SimpleEntry<>(output.getExecPathString(), entry);
       }
     }
     return null;
   }
 
+  /** Checks whether one of existing output paths is already used as a key. */
+  @Nullable
+  public static ActionCache.Entry getCacheEntry(ActionCache actionCache, Action action) {
+    var entry = getCacheEntryWithKey(actionCache, action);
+    if (entry != null) {
+      return entry.getValue();
+    }
+    return null;
+  }
+
   public static void removeCacheEntry(ActionCache actionCache, Action action) {
     for (Artifact output : action.getOutputs()) {
       actionCache.remove(output.getExecPathString());

diff --git a/src/main/java/com/google/devtools/build/lib/actions/FileArtifactValue.java b/src/main/java/com/google/devtools/build/lib/actions/FileArtifactValue.java
@@ -641,6 +641,12 @@ public long getExpireAtEpochMilli() {
       return -1;
     }
 
+    /**
+     * Extends the expiration time for this metadata. If it was constructed without known expiration
+     * time (i.e. expireAtEpochMilli < 0), this extension does nothing.
+     */
+    public void extendExpireAtEpochMilli(long expireAtEpochMilli) {}
+
     public boolean isAlive(Instant now) {
       return true;
     }
@@ -669,7 +675,7 @@ public final String toString() {
 
   /** A remote artifact that expires at a particular time. */
   private static final class RemoteFileArtifactValueWithExpiration extends RemoteFileArtifactValue {
-    private final long expireAtEpochMilli;
+    private long expireAtEpochMilli;
 
     private RemoteFileArtifactValueWithExpiration(
         byte[] digest,
@@ -686,6 +692,11 @@ public long getExpireAtEpochMilli() {
       return expireAtEpochMilli;
     }
 
+    @Override
+    public void extendExpireAtEpochMilli(long expireAtEpochMilli) {
+      this.expireAtEpochMilli = expireAtEpochMilli;
+    }
+
     @Override
     public boolean isAlive(Instant now) {
       return now.toEpochMilli() < expireAtEpochMilli;

diff --git a/src/main/java/com/google/devtools/build/lib/remote/BUILD b/src/main/java/com/google/devtools/build/lib/remote/BUILD
@@ -55,6 +55,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib:runtime/command_line_path_factory",
         "//src/main/java/com/google/devtools/build/lib/actions",
         "//src/main/java/com/google/devtools/build/lib/actions:action_input_helper",
+        "//src/main/java/com/google/devtools/build/lib/actions:action_lookup_data",
         "//src/main/java/com/google/devtools/build/lib/actions:artifacts",
         "//src/main/java/com/google/devtools/build/lib/actions:execution_requirements",
         "//src/main/java/com/google/devtools/build/lib/actions:file_metadata",
@@ -97,7 +98,9 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/remote/options",
         "//src/main/java/com/google/devtools/build/lib/remote/util",
         "//src/main/java/com/google/devtools/build/lib/remote/zstd",
+        "//src/main/java/com/google/devtools/build/lib/skyframe:action_execution_value",
         "//src/main/java/com/google/devtools/build/lib/skyframe:mutable_supplier",
+        "//src/main/java/com/google/devtools/build/lib/skyframe:sky_functions",
         "//src/main/java/com/google/devtools/build/lib/skyframe:tree_artifact_value",
         "//src/main/java/com/google/devtools/build/lib/util",
         "//src/main/java/com/google/devtools/build/lib/util:abrupt_exit_exception",
@@ -114,6 +117,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/worker:worker_options",
         "//src/main/java/com/google/devtools/build/lib/worker:worker_spawn_runner",
         "//src/main/java/com/google/devtools/build/skyframe",
+        "//src/main/java/com/google/devtools/build/skyframe:skyframe-objects",
         "//src/main/java/com/google/devtools/common/options",
         "//src/main/protobuf:failure_details_java_proto",
         "//src/main/protobuf:spawn_java_proto",

diff --git a/src/main/java/com/google/devtools/build/lib/remote/LeaseService.java b/src/main/java/com/google/devtools/build/lib/remote/LeaseService.java
@@ -25,27 +25,38 @@
 
 /** A lease service that manages the lease of remote blobs. */
 public class LeaseService {
-  private final MemoizingEvaluator memoizingEvaluator;
-  @Nullable private final ActionCache actionCache;
+  protected final MemoizingEvaluator memoizingEvaluator;
+  @Nullable protected final ActionCache actionCache;
 
   public LeaseService(MemoizingEvaluator memoizingEvaluator, @Nullable ActionCache actionCache) {
     this.memoizingEvaluator = memoizingEvaluator;
     this.actionCache = actionCache;
   }
 
-  /** Clean up internal state when files are evicted from remote CAS. */
-  public void handleMissingInputs(Set<ActionInput> missingActionInputs) {
-    if (missingActionInputs.isEmpty()) {
-      return;
+  protected void startLeaseExtensionIfNot() {}
+
+  protected void stopLeaseExtension() {}
+
+  public void finalizeAction() {
+    startLeaseExtensionIfNot();
+  }
+
+  public void finalizeExecution(Set<ActionInput> missingActionInputs) {
+    stopLeaseExtension();
+
+    if (!missingActionInputs.isEmpty()) {
+      handleMissingInputs();
     }
+  }
 
+  /** Clean up internal state when files are evicted from remote CAS. */
+  private void handleMissingInputs() {
     // If any outputs are evicted, remove all remote metadata from skyframe and local action cache.
     //
     // With TTL based discarding and lease extension, remote cache eviction error won't happen if
     // remote cache can guarantee the TTL. However, if it happens, it usually means the remote cache
     // is under high load and it could possibly evict more blobs that Bazel wouldn't aware of.
     // Following builds could still fail for the same error (caused by different blobs).
-
     memoizingEvaluator.delete(
         key -> {
           if (key.functionName().equals(SkyFunctions.ACTION_EXECUTION)) {