[7.6.0] Fix race condition with multiplex sandboxed workers (#25491)

Prior to this change, multiplex sandboxed workers shared a working
directory per mnemonic. This caused a race condition when a new
multiplex sandboxed worker with the same mnemonic was launched because
when launching it cleaned the working directory. That could cause
problems for any actions executing in that directory.

This change makes it so each multiplex sandbox worker process has a
unique working directory. It does so while ensuring each
SandboxedWorkerProxy and the associated sandbox are still associated
with the correct multiplexer process and working directory.

Resolves #22589

I couldn't figure out if the Bazel repo has an autoformatter somewhere,
so I did my best to manually format the code with a style that follows
the existing code.

Closes #25400.

PiperOrigin-RevId: 732256101
Change-Id: If8deea240fda77780feaeac352cf099fb9bfcee3 (cherry picked from
commit d54fc62)

Fixes #25460

Co-authored-by: jjudd <[email protected]>

Author: fmeum

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

@@ -96,7 +96,8 @@ public Worker create(WorkerKey key) throws IOException {
     if (key.isSandboxed()) {
       if (key.isMultiplex()) {
         WorkerMultiplexer workerMultiplexer = WorkerMultiplexerManager.getInstance(key, logFile);
-        Path workDir = getSandboxedWorkerPath(key);
+        int multiplexerId = workerMultiplexer.getMultiplexerId();
+        Path workDir = getMultiplexSandboxedWorkerPath(key, multiplexerId);
         worker = new SandboxedWorkerProxy(key, workerId, logFile, workerMultiplexer, workDir);
       } else {
         Path workDir = getSandboxedWorkerPath(key, workerId);
@@ -136,10 +137,11 @@ Path getSandboxedWorkerPath(WorkerKey key, int workerId) {
         .getRelative(workspaceName);
   }
 
-  Path getSandboxedWorkerPath(WorkerKey key) {
+  Path getMultiplexSandboxedWorkerPath(WorkerKey key, int multiplexerId) {
     String workspaceName = key.getExecRoot().getBaseName();
     return workerBaseDir
-        .getRelative(key.getMnemonic() + "-" + key.getWorkerTypeName() + "-workdir")
+        .getRelative(
+            key.getMnemonic() + "-" + key.getWorkerTypeName() + "-" + multiplexerId + "-workdir")
         .getRelative(workspaceName);
   }
 

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

@@ -54,34 +54,47 @@
  * WorkerMultiplexer} wakes up the relevant {@code WorkerProxy} to retrieve the response.
  */
 public class WorkerMultiplexer {
+  /**
+   * An ID for this multiplexer that can be used by sandboxed multiplex workers to generate their
+   * workdir. The workdir needs to be the same for all {@code SandboxedWorkerProxy} instances
+   * associated with a {@code WorkerMultiplexer}, but needs to be unique across multiplexers for the
+   * same mnemonic. This is analogous to the @{code workerId} created in {@code WorkerFactory}.
+   */
+  private final int multiplexerId;
+
   /**
    * A queue of {@link WorkRequest} instances that need to be sent to the worker. {@link
    * WorkerProxy} instances add to this queue, while the requestSender subthread remove requests and
    * send them to the worker. This prevents dynamic execution interrupts from corrupting the {@code
    * stdin} of the worker process.
    */
   @VisibleForTesting final BlockingQueue<WorkRequest> pendingRequests = new LinkedBlockingQueue<>();
+
   /**
    * A map of {@code WorkResponse}s received from the worker process. They are stored in this map
    * keyed by the request id until the corresponding {@code WorkerProxy} picks them up.
    */
   private final ConcurrentMap<Integer, WorkResponse> workerProcessResponse =
       new ConcurrentHashMap<>();
+
   /**
    * A map of semaphores corresponding to {@code WorkRequest}s. After sending the {@code
    * WorkRequest}, {@code WorkerProxy} will wait on a semaphore to be released. {@code
    * WorkerMultiplexer} is responsible for releasing the corresponding semaphore in order to signal
    * {@code WorkerProxy} that the {@code WorkerResponse} has been received.
    */
   private final ConcurrentMap<Integer, Semaphore> responseChecker = new ConcurrentHashMap<>();
+
   /**
    * The worker process that this WorkerMultiplexer should be talking to. This should only be set
    * once, when creating a new process. If the process dies or its stdio streams get corrupted, the
    * {@code WorkerMultiplexer} gets discarded as well and a new one gets created as needed.
    */
   @VisibleForTesting Subprocess process;
+
   /** The implementation of the worker protocol (JSON or Proto). */
   private WorkerProtocolImpl workerProtocol;
+
   /** InputStream from the worker process. */
   @LazyInit private RecordingInputStream recordingStream;
   /** True if this multiplexer was explicitly destroyed. */
@@ -117,9 +130,10 @@ public class WorkerMultiplexer {
    */
   private Thread shutdownHook;
 
-  WorkerMultiplexer(Path logFile, WorkerKey workerKey) {
+  WorkerMultiplexer(Path logFile, WorkerKey workerKey, int multiplexerId) {
     this.logFile = logFile;
     this.workerKey = workerKey;
+    this.multiplexerId = multiplexerId;
   }
 
   /** Sets or clears the reporter for outputting verbose info. */
@@ -157,7 +171,7 @@ public synchronized void createSandboxedProcess(
           SandboxOutputs.getEmptyInstance());
       SandboxHelpers.cleanExisting(
           workDir.getParentDirectory(), inputFiles, inputsToCreate, dirsToCreate, workDir);
-      SandboxHelpers.createDirectories(dirsToCreate, workDir, /* strict=*/ false);
+      SandboxHelpers.createDirectories(dirsToCreate, workDir, /* strict= */ false);
       WorkerExecRoot.createInputs(inputsToCreate, inputFiles.limitedCopy(workerFiles), workDir);
       createProcess(workDir);
     }
@@ -455,4 +469,8 @@ public long getProcessId() {
     }
     return process.getProcessId();
   }
+
+  public int getMultiplexerId() {
+    return this.multiplexerId;
+  }
 }

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

@@ -24,6 +24,7 @@
 import com.google.devtools.build.lib.vfs.Path;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
 import javax.annotation.Nullable;
 
 /**
@@ -40,6 +41,15 @@ public class WorkerMultiplexerManager {
 
   private WorkerMultiplexerManager() {}
 
+  /**
+   * A counter used to provide unique IDs across sandboxed multiplexer instances. It is used in
+   * determining the workdir for the multiplexer process. This is analogous to the {@code
+   * pidCounter} in {@code WorkerFactory}. It is ok to use an {@code AtomicInteger} here for the
+   * same reasons as it is there: the counter is only incremented when spawning a new multiplexer,
+   * so even in the worst case of workers quitting after each action it shouldn't overflow.
+   */
+  private static final AtomicInteger multiplexerIdCounter = new AtomicInteger(1);
+
   /**
    * Returns a {@code WorkerMultiplexer} instance to {@code WorkerProxy}. {@code WorkerProxy}
    * objects with the same {@code WorkerKey} talk to the same {@code WorkerMultiplexer}. Also,
@@ -48,7 +58,10 @@ private WorkerMultiplexerManager() {}
   public static synchronized WorkerMultiplexer getInstance(WorkerKey key, Path logFile) {
     InstanceInfo instanceInfo =
         multiplexerInstance.computeIfAbsent(
-            key, k -> new InstanceInfo(new WorkerMultiplexer(logFile, k)));
+            key,
+            k ->
+                new InstanceInfo(
+                    new WorkerMultiplexer(logFile, k, multiplexerIdCounter.getAndIncrement())));
     instanceInfo.increaseRefCount();
     return instanceInfo.getWorkerMultiplexer();
   }

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

@@ -21,7 +21,9 @@
 
 import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Lists;
 import com.google.devtools.build.lib.actions.Spawn;
+import com.google.devtools.build.lib.actions.UserExecException;
 import com.google.devtools.build.lib.vfs.DigestHashFunction;
 import com.google.devtools.build.lib.vfs.FileSystem;
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
@@ -31,6 +33,7 @@
 import com.google.devtools.build.lib.worker.WorkerProtocol.WorkRequest;
 import java.io.IOException;
 import java.io.PipedInputStream;
+import java.util.ArrayList;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -58,8 +61,12 @@ public void setUp() throws IOException {
 
   @Test
   public void prepareExecution_createsFilesInSandbox() throws IOException, InterruptedException {
-    SandboxedWorkerProxy proxy = createSandboxedWorkerProxy();
-    Path workDir = workerBaseDir.getChild("Mnem-multiplex-worker-workdir").getChild("execroot");
+    SandboxedWorkerProxy proxy = createSandboxedWorkerProxies("Mnem", 1).get(0);
+    int multiplexerId = proxy.workerMultiplexer.getMultiplexerId();
+    Path workDir =
+        workerBaseDir
+            .getChild("Mnem-multiplex-worker-" + multiplexerId + "-workdir")
+            .getChild("execroot");
     Path sandboxDir =
         workDir
             .getChild("__sandbox")
@@ -94,7 +101,11 @@ public void prepareExecution_createsFilesInSandbox() throws IOException, Interru
   @Test
   public void putRequest_setsSandboxDir() throws IOException, InterruptedException {
     SandboxedWorkerProxy worker = createFakedSandboxedWorkerProxy();
-    Path workDir = workerBaseDir.getChild("Mnem-multiplex-worker-workdir").getChild("execroot");
+    int multiplexerId = worker.workerMultiplexer.getMultiplexerId();
+    Path workDir =
+        workerBaseDir
+            .getChild("Mnem-multiplex-worker-" + multiplexerId + "-workdir")
+            .getChild("execroot");
     SandboxHelper sandboxHelper =
         new SandboxHelper(globalExecRoot, workDir)
             .addAndCreateInputFile("anInputFile", "anInputFile", "Just stuff")
@@ -115,7 +126,11 @@ public void putRequest_setsSandboxDir() throws IOException, InterruptedException
   @Test
   public void finishExecution_copiesOutputs() throws IOException, InterruptedException {
     SandboxedWorkerProxy worker = createFakedSandboxedWorkerProxy();
-    Path workDir = workerBaseDir.getChild("Mnem-multiplex-worker-workdir").getChild("execroot");
+    int multiplexerId = worker.workerMultiplexer.getMultiplexerId();
+    Path workDir =
+        workerBaseDir
+            .getChild("Mnem-multiplex-worker-" + multiplexerId + "-workdir")
+            .getChild("execroot");
     SandboxHelper sandboxHelper =
         new SandboxHelper(globalExecRoot, workDir)
             .addAndCreateInputFile("anInputFile", "anInputFile", "Just stuff")
@@ -150,8 +165,37 @@ public void finishExecution_copiesOutputs() throws IOException, InterruptedExcep
         .isEqualTo("some output");
   }
 
-  private SandboxedWorkerProxy createSandboxedWorkerProxy() throws IOException {
-    ImmutableMap.Builder<String, String> req = TestUtils.execRequirementsBuilder("Mnem");
+  @Test
+  public void differentProxiesSameMultiplexerHaveSameWorkDir()
+      throws IOException, InterruptedException {
+    ArrayList<SandboxedWorkerProxy> proxies = createSandboxedWorkerProxies("Mnem", 2);
+    SandboxedWorkerProxy proxyOne = proxies.get(0);
+    SandboxedWorkerProxy proxyTwo = proxies.get(1);
+
+    int multiplexerIdProxyOne = proxyOne.workerMultiplexer.getMultiplexerId();
+    Path expectedWorkDirProxyOne =
+        workerBaseDir
+            .getChild("Mnem-multiplex-worker-" + multiplexerIdProxyOne + "-workdir")
+            .getChild("execroot");
+
+    int multiplexerIdProxyTwo = proxyTwo.workerMultiplexer.getMultiplexerId();
+    Path expectedWorkDirProxyTwo =
+        workerBaseDir
+            .getChild("Mnem-multiplex-worker-" + multiplexerIdProxyTwo + "-workdir")
+            .getChild("execroot");
+
+    assertThat(proxyOne.workDir).isEqualTo(proxyTwo.workDir);
+    assertThat(proxyOne.workDir).isEqualTo(expectedWorkDirProxyOne);
+    assertThat(proxyTwo.workDir).isEqualTo(expectedWorkDirProxyTwo);
+  }
+
+  @Test
+  public void differentProxiesDifferentMultiplexerSameMnemHaveDifferentWorkDirs()
+      throws IOException, InterruptedException, UserExecException {
+    String sharedMnemonic = "Mnem";
+
+    // Create a proxy on the first multiplexer
+    ImmutableMap.Builder<String, String> req = TestUtils.execRequirementsBuilder(sharedMnemonic);
     req.put(SUPPORTS_MULTIPLEX_SANDBOXING, "1");
     Spawn spawn = TestUtils.createSpawn(req.buildOrThrow());
 
@@ -163,7 +207,51 @@ private SandboxedWorkerProxy createSandboxedWorkerProxy() throws IOException {
         TestUtils.createWorkerKeyFromOptions(
             PROTO, globalOutputBase, options, true, spawn, "worker.sh");
     WorkerFactory factory = new WorkerFactory(workerBaseDir);
-    return (SandboxedWorkerProxy) factory.create(key);
+
+    SandboxedWorkerProxy proxyOneMultiplexerOne = (SandboxedWorkerProxy) factory.create(key);
+    int multiplexerIdOne = proxyOneMultiplexerOne.workerMultiplexer.getMultiplexerId();
+    Path expectedWorkDirOne =
+        workerBaseDir
+            .getChild(sharedMnemonic + "-multiplex-worker-" + multiplexerIdOne + "-workdir")
+            .getChild("execroot");
+
+    // Shut down the first multiplexer, so we get a different multiplexer for the next proxy
+    WorkerMultiplexerManager.removeInstance(key);
+
+    // Create a proxy on the second multiplexer
+    SandboxedWorkerProxy proxyOneMultiplexerTwo = (SandboxedWorkerProxy) factory.create(key);
+    int multiplexerIdTwo = proxyOneMultiplexerTwo.workerMultiplexer.getMultiplexerId();
+    Path expectedWorkDirTwo =
+        workerBaseDir
+            .getChild(sharedMnemonic + "-multiplex-worker-" + multiplexerIdTwo + "-workdir")
+            .getChild("execroot");
+
+    assertThat(proxyOneMultiplexerOne.workDir).isNotEqualTo(proxyOneMultiplexerTwo.workDir);
+    assertThat(proxyOneMultiplexerOne.workDir).isEqualTo(expectedWorkDirOne);
+    assertThat(proxyOneMultiplexerTwo.workDir).isEqualTo(expectedWorkDirTwo);
+  }
+
+  private ArrayList<SandboxedWorkerProxy> createSandboxedWorkerProxies(
+      String mnemonic, int numProxiesToCreate) throws IOException {
+    ImmutableMap.Builder<String, String> req = TestUtils.execRequirementsBuilder(mnemonic);
+    req.put(SUPPORTS_MULTIPLEX_SANDBOXING, "1");
+    Spawn spawn = TestUtils.createSpawn(req.buildOrThrow());
+
+    WorkerOptions options = new WorkerOptions();
+    options.workerMultiplex = true;
+    options.multiplexSandboxing = true;
+
+    WorkerKey key =
+        TestUtils.createWorkerKeyFromOptions(
+            PROTO, globalOutputBase, options, true, spawn, "worker.sh");
+    WorkerFactory factory = new WorkerFactory(workerBaseDir);
+
+    assertThat(numProxiesToCreate).isGreaterThan(0);
+    ArrayList<SandboxedWorkerProxy> proxies = Lists.newArrayListWithCapacity(numProxiesToCreate);
+    for (int i = 0; i < numProxiesToCreate; i++) {
+      proxies.add((SandboxedWorkerProxy) factory.create(key));
+    }
+    return proxies;
   }
 
   private SandboxedWorkerProxy createFakedSandboxedWorkerProxy() throws IOException {
@@ -180,7 +268,7 @@ private SandboxedWorkerProxy createFakedSandboxedWorkerProxy() throws IOExceptio
             PROTO, globalOutputBase, options, true, spawn, "worker.sh");
     WorkerMultiplexerManager.injectForTesting(
         key,
-        new WorkerMultiplexer(globalExecRoot.getChild("testWorker.log"), key) {
+        new WorkerMultiplexer(globalExecRoot.getChild("testWorker.log"), key, 0) {
           @Override
           public synchronized void createProcess(Path workDir) throws IOException {
             PipedInputStream serverInputStream = new PipedInputStream();