Incorporate the in-memory tracking setting into the WorkerKey.

This ensures that the worker is restarted if the option changes. We could alternatively update the setting on already running instances, but this way is more obviously correct, and we don't expect users to switch back and forth all the time.

PiperOrigin-RevId: 661226777
Change-Id: Ide49b8b7fcef95117258df66504a9b585758871b

src/main/java/com/google/devtools/build/lib/worker/SandboxedWorker.java

@@ -39,7 +39,6 @@
 import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.devtools.build.lib.vfs.Symlinks;
 import java.io.IOException;
-import java.util.List;
 import java.util.Map.Entry;
 import java.util.Set;
 import java.util.SortedMap;
@@ -109,6 +108,7 @@ public static WorkerSandboxOptions create(
       WorkerOptions workerOptions,
       @Nullable WorkerSandboxOptions hardenedSandboxOptions,
       TreeDeleter treeDeleter,
+      boolean useInMemoryTracking,
       @Nullable VirtualCgroupFactory cgroupFactory) {
     super(workerKey, workerId, workDir, logFile, workerOptions, cgroupFactory);
     Path tmpDirPath = SandboxHelpers.getTmpDirPath(workDir);
@@ -118,17 +118,11 @@ public static WorkerSandboxOptions create(
             hardenedSandboxOptions != null
                 ? ImmutableList.of(PathFragment.create(tmpDirPath.getPathString()))
                 : ImmutableList.of(),
-            shouldUseInMemoryTracking(workerOptions, workerKey));
+            useInMemoryTracking);
     this.hardenedSandboxOptions = hardenedSandboxOptions;
     this.treeDeleter = treeDeleter;
   }
 
-  private static boolean shouldUseInMemoryTracking(
-      WorkerOptions workerOptions, WorkerKey workerKey) {
-    List<String> mnemonics = workerOptions.workerSandboxInMemoryTracking;
-    return mnemonics != null && mnemonics.contains(workerKey.getMnemonic());
-  }
-
   @Override
   public boolean isSandboxed() {
     return true;

src/main/java/com/google/devtools/build/lib/worker/WorkerFactory.java

@@ -112,6 +112,7 @@ public Worker create(WorkerKey key) throws IOException {
                 workerOptions,
                 hardenedSandboxOptions,
                 treeDeleter,
+                key.useInMemoryTracking(),
                 cgroupFactory);
       }
     } else if (key.isMultiplex()) {

src/main/java/com/google/devtools/build/lib/worker/WorkerKey.java

@@ -35,10 +35,13 @@
 public final class WorkerKey {
   /** Build options. */
   private final ImmutableList<String> args;
+
   /** Environment variables. */
   private final ImmutableMap<String, String> env;
+
   /** Execution root of Bazel process. */
   private final Path execRoot;
+
   /** Mnemonic of the worker. */
   private final String mnemonic;
 
@@ -48,19 +51,28 @@ public final class WorkerKey {
    * methods.
    */
   private final HashCode workerFilesCombinedHash;
+
   /** Worker files with the corresponding digest. */
   private final SortedMap<PathFragment, byte[]> workerFilesWithDigests;
+
   /** If true, the workers run inside a sandbox. */
   private final boolean sandboxed;
+
+  /** If true, the sandbox contents are tracked in memory to speed up cleanup. */
+  private final boolean useInMemoryTracking;
+
   /** A WorkerProxy will be instantiated if true, instantiate a regular Worker if false. */
   private final boolean multiplex;
+
   /** If true, the workers for this key are able to cancel work requests. */
   private final boolean cancellable;
+
   /**
    * Cached value for the hash of this key, because the value is expensive to calculate
    * (ImmutableMap and ImmutableList do not cache their hashcodes.
    */
   private final int hash;
+
   /** The format of the worker protocol sent to and read from the worker. */
   private final WorkerProtocolFormat protocolFormat;
 
@@ -72,6 +84,7 @@ public WorkerKey(
       HashCode workerFilesCombinedHash,
       SortedMap<PathFragment, byte[]> workerFilesWithDigests,
       boolean sandboxed,
+      boolean useInMemoryTracking,
       boolean multiplex,
       boolean cancellable,
       WorkerProtocolFormat protocolFormat) {
@@ -82,6 +95,7 @@ public WorkerKey(
     this.workerFilesCombinedHash = Preconditions.checkNotNull(workerFilesCombinedHash);
     this.workerFilesWithDigests = Preconditions.checkNotNull(workerFilesWithDigests);
     this.sandboxed = sandboxed;
+    this.useInMemoryTracking = useInMemoryTracking;
     this.multiplex = multiplex;
     this.cancellable = cancellable;
     this.protocolFormat = protocolFormat;
@@ -117,6 +131,10 @@ public boolean isSandboxed() {
     return sandboxed;
   }
 
+  public boolean useInMemoryTracking() {
+    return useInMemoryTracking;
+  }
+
   public boolean isMultiplex() {
     return multiplex;
   }
@@ -158,13 +176,16 @@ public boolean equals(Object o) {
     if (!args.equals(workerKey.args)) {
       return false;
     }
-    if (!multiplex == workerKey.multiplex) {
+    if (multiplex != workerKey.multiplex) {
+      return false;
+    }
+    if (cancellable != workerKey.cancellable) {
       return false;
     }
-    if (!cancellable == workerKey.cancellable) {
+    if (sandboxed != workerKey.sandboxed) {
       return false;
     }
-    if (!sandboxed == workerKey.sandboxed) {
+    if (useInMemoryTracking != workerKey.useInMemoryTracking) {
       return false;
     }
     if (!env.equals(workerKey.env)) {
@@ -197,6 +218,7 @@ private int calculateHashCode() {
         multiplex,
         cancellable,
         sandboxed,
+        useInMemoryTracking,
         protocolFormat.toString());
   }
 

src/main/java/com/google/devtools/build/lib/worker/WorkerParser.java

@@ -131,6 +131,7 @@ static WorkerKey createWorkerKey(
       WorkerOptions options,
       boolean dynamic,
       WorkerProtocolFormat protocolFormat) {
+    String workerKeyMnemonic = Spawns.getWorkerKeyMnemonic(spawn);
     boolean multiplex = options.workerMultiplex && Spawns.supportsMultiplexWorkers(spawn);
     if (dynamic && !(Spawns.supportsMultiplexSandboxing(spawn) && options.multiplexSandboxing)) {
       multiplex = false;
@@ -142,14 +143,20 @@ static WorkerKey createWorkerKey(
     } else {
       sandboxed = options.workerSandboxing || dynamic;
     }
+    boolean useInMemoryTracking = false;
+    if (sandboxed) {
+      List<String> mnemonics = options.workerSandboxInMemoryTracking;
+      useInMemoryTracking = mnemonics != null && mnemonics.contains(workerKeyMnemonic);
+    }
     return new WorkerKey(
         workerArgs,
         env,
         execRoot,
-        Spawns.getWorkerKeyMnemonic(spawn),
+        workerKeyMnemonic,
         workerFilesCombinedHash,
         workerFiles,
         sandboxed,
+        useInMemoryTracking,
         multiplex,
         Spawns.supportsWorkerCancellation(spawn),
         protocolFormat);

src/test/java/com/google/devtools/build/lib/actions/ResourceManagerTest.java

@@ -154,6 +154,7 @@ private WorkerKey createWorkerKey(String mnemonic) {
         /* workerFilesCombinedHash= */ HashCode.fromInt(0),
         /* workerFilesWithDigests= */ ImmutableSortedMap.of(),
         /* sandboxed= */ false,
+        /* useInMemoryTracking= */ false,
         /* multiplex= */ false,
         /* cancellable= */ false,
         WorkerProtocolFormat.PROTO);

src/test/java/com/google/devtools/build/lib/worker/WorkerFactoryTest.java

@@ -141,6 +141,7 @@ protected WorkerKey createWorkerKey(boolean mustBeSandboxed, boolean multiplex,
         /* workerFilesCombinedHash= */ HashCode.fromInt(0),
         /* workerFilesWithDigests= */ ImmutableSortedMap.of(),
         /* sandboxed= */ mustBeSandboxed,
+        /* useInMemoryTracking= */ false,
         /* multiplex= */ multiplex,
         /* cancellable= */ false,
         WorkerProtocolFormat.PROTO);

src/test/java/com/google/devtools/build/lib/worker/WorkerMultiplexerManagerTest.java

@@ -65,6 +65,7 @@ public void instanceCreationRemovalTest() throws Exception {
             ImmutableSortedMap.of(),
             false,
             false,
+            false,
             /* cancellable= */ false,
             WorkerProtocolFormat.PROTO);
     WorkerMultiplexer wm1 = WorkerMultiplexerManager.getInstance(workerKey1, logFile);
@@ -84,6 +85,7 @@ public void instanceCreationRemovalTest() throws Exception {
             ImmutableSortedMap.of(),
             false,
             false,
+            false,
             /* cancellable= */ false,
             WorkerProtocolFormat.PROTO);
     WorkerMultiplexer wm2 = WorkerMultiplexerManager.getInstance(workerKey2, logFile);