Make WorkerMultiplexerManager not key hash maps with integers, but with the actual WorkerKey object.
This allows disambiguation of objects that have the same hashcode. Since the hash code is expensive to
calculate, but the object is final, we cache the hash code in WorkerKey.
RELNOTES: none
PiperOrigin-RevId: 325268372
diff --git a/src/main/java/com/google/devtools/build/lib/worker/WorkerFactory.java b/src/main/java/com/google/devtools/build/lib/worker/WorkerFactory.java
index ff2e760..b59f899 100644
--- a/src/main/java/com/google/devtools/build/lib/worker/WorkerFactory.java
+++ b/src/main/java/com/google/devtools/build/lib/worker/WorkerFactory.java
@@ -69,11 +69,7 @@
} else if (key.getProxied()) {
worker =
new WorkerProxy(
- key,
- workerId,
- key.getExecRoot(),
- logFile,
- WorkerMultiplexerManager.getInstance(key.hashCode()));
+ key, workerId, key.getExecRoot(), logFile, WorkerMultiplexerManager.getInstance(key));
} else {
worker = new Worker(key, workerId, key.getExecRoot(), logFile);
}
diff --git a/src/main/java/com/google/devtools/build/lib/worker/WorkerKey.java b/src/main/java/com/google/devtools/build/lib/worker/WorkerKey.java
index b2ab4c9..3480596 100644
--- a/src/main/java/com/google/devtools/build/lib/worker/WorkerKey.java
+++ b/src/main/java/com/google/devtools/build/lib/worker/WorkerKey.java
@@ -26,6 +26,9 @@
/**
* Data container that uniquely identifies a kind of worker process and is used as the key for the
* {@link WorkerPool}.
+ *
+ * <p>We expect a small number of WorkerKeys per mnemonic. Unbounded creation of WorkerKeys will
+ * break various things as well as render the workers less useful.
*/
final class WorkerKey {
private final ImmutableList<String> args;
@@ -43,6 +46,11 @@
private final boolean mustBeSandboxed;
/** A WorkerProxy will be instantiated if true, instantiate a regular Worker if false. */
private final boolean proxied;
+ /**
+ * 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;
private final WorkerProtocolFormat protocolFormat;
@@ -74,6 +82,8 @@
this.proxied = proxied;
/** The format of the worker protocol sent to and read from the worker. */
this.protocolFormat = protocolFormat;
+
+ hash = calculateHashCode();
}
/** Getter function for variable args. */
@@ -154,8 +164,13 @@
}
+ /** Since all fields involved in the {@code hashCode} are final, we cache the result. */
@Override
public int hashCode() {
+ return hash;
+ }
+
+ private int calculateHashCode() {
int result = args.hashCode();
result = 31 * result + env.hashCode();
result = 31 * result + execRoot.hashCode();
diff --git a/src/main/java/com/google/devtools/build/lib/worker/WorkerMultiplexerManager.java b/src/main/java/com/google/devtools/build/lib/worker/WorkerMultiplexerManager.java
index 5387cd4..fefea03 100644
--- a/src/main/java/com/google/devtools/build/lib/worker/WorkerMultiplexerManager.java
+++ b/src/main/java/com/google/devtools/build/lib/worker/WorkerMultiplexerManager.java
@@ -14,6 +14,7 @@
package com.google.devtools.build.lib.worker;
+import com.google.common.annotations.VisibleForTesting;
import com.google.devtools.build.lib.actions.UserExecException;
import com.google.devtools.build.lib.server.FailureDetails;
import com.google.devtools.build.lib.server.FailureDetails.FailureDetail;
@@ -32,31 +33,23 @@
* A map from the hash of {@code WorkerKey} objects to the corresponding information about the
* multiplexer instance.
*/
- private static Map<Integer, InstanceInfo> multiplexerInstance;
+ private static final Map<WorkerKey, InstanceInfo> multiplexerInstance = new HashMap<>();
/** A semaphore to protect {@code multiplexerInstance} and {@code multiplexerRefCount} objects. */
- private static Semaphore semMultiplexer;
-
- static {
- multiplexerInstance = new HashMap<>();
- semMultiplexer = new Semaphore(1);
- }
+ private static final Semaphore semMultiplexer = new Semaphore(1);
private WorkerMultiplexerManager() {}
/**
* Returns a {@code WorkerMultiplexer} instance to {@code WorkerProxy}. {@code WorkerProxy}
- * objects with the same workerHash talk to the same {@code WorkerMultiplexer}. Also, record how
- * many {@code WorkerProxy} objects are talking to this {@code WorkerMultiplexer}.
+ * objects with the same {@code WorkerKey} talk to the same {@code WorkerMultiplexer}. Also,
+ * record how many {@code WorkerProxy} objects are talking to this {@code WorkerMultiplexer}.
*/
- public static WorkerMultiplexer getInstance(Integer workerHash) throws InterruptedException {
+ public static WorkerMultiplexer getInstance(WorkerKey key) throws InterruptedException {
semMultiplexer.acquire();
- if (!multiplexerInstance.containsKey(workerHash)) {
- multiplexerInstance.put(workerHash, new InstanceInfo());
- }
- multiplexerInstance.get(workerHash).increaseRefCount();
- WorkerMultiplexer workerMultiplexer =
- multiplexerInstance.get(workerHash).getWorkerMultiplexer();
+ multiplexerInstance.putIfAbsent(key, new InstanceInfo());
+ multiplexerInstance.get(key).increaseRefCount();
+ WorkerMultiplexer workerMultiplexer = multiplexerInstance.get(key).getWorkerMultiplexer();
semMultiplexer.release();
return workerMultiplexer;
}
@@ -65,15 +58,14 @@
* Removes the {@code WorkerMultiplexer} instance and reference count since it is no longer in
* use.
*/
- public static void removeInstance(Integer workerHash)
- throws InterruptedException, UserExecException {
+ public static void removeInstance(WorkerKey key) throws InterruptedException, UserExecException {
semMultiplexer.acquire();
try {
- multiplexerInstance.get(workerHash).decreaseRefCount();
- if (multiplexerInstance.get(workerHash).getRefCount() == 0) {
- multiplexerInstance.get(workerHash).getWorkerMultiplexer().interrupt();
- multiplexerInstance.get(workerHash).getWorkerMultiplexer().destroyMultiplexer();
- multiplexerInstance.remove(workerHash);
+ multiplexerInstance.get(key).decreaseRefCount();
+ if (multiplexerInstance.get(key).getRefCount() == 0) {
+ multiplexerInstance.get(key).getWorkerMultiplexer().interrupt();
+ multiplexerInstance.get(key).getWorkerMultiplexer().destroyMultiplexer();
+ multiplexerInstance.remove(key);
}
} catch (Exception e) {
String message = "NullPointerException while accessing non-existent multiplexer instance.";
@@ -83,25 +75,28 @@
}
}
- public static WorkerMultiplexer getMultiplexer(Integer workerHash) throws UserExecException {
+ @VisibleForTesting
+ static WorkerMultiplexer getMultiplexer(WorkerKey key) throws UserExecException {
try {
- return multiplexerInstance.get(workerHash).getWorkerMultiplexer();
+ return multiplexerInstance.get(key).getWorkerMultiplexer();
} catch (NullPointerException e) {
String message = "NullPointerException while accessing non-existent multiplexer instance.";
throw createUserExecException(e, message, Code.MULTIPLEXER_DOES_NOT_EXIST);
}
}
- public static Integer getRefCount(Integer workerHash) throws UserExecException {
+ @VisibleForTesting
+ static Integer getRefCount(WorkerKey key) throws UserExecException {
try {
- return multiplexerInstance.get(workerHash).getRefCount();
+ return multiplexerInstance.get(key).getRefCount();
} catch (NullPointerException e) {
String message = "NullPointerException while accessing non-existent multiplexer instance.";
throw createUserExecException(e, message, Code.MULTIPLEXER_DOES_NOT_EXIST);
}
}
- public static Integer getInstanceCount() {
+ @VisibleForTesting
+ static Integer getInstanceCount() {
return multiplexerInstance.keySet().size();
}
diff --git a/src/main/java/com/google/devtools/build/lib/worker/WorkerProxy.java b/src/main/java/com/google/devtools/build/lib/worker/WorkerProxy.java
index 03a353e..63de9f4 100644
--- a/src/main/java/com/google/devtools/build/lib/worker/WorkerProxy.java
+++ b/src/main/java/com/google/devtools/build/lib/worker/WorkerProxy.java
@@ -66,7 +66,7 @@
@Override
synchronized void destroy() throws IOException {
try {
- WorkerMultiplexerManager.removeInstance(workerKey.hashCode());
+ WorkerMultiplexerManager.removeInstance(workerKey);
} catch (InterruptedException e) {
logger.atWarning().withCause(e).log(
"InterruptedException was caught while destroying multiplexer. "
diff --git a/src/test/java/com/google/devtools/build/lib/BUILD b/src/test/java/com/google/devtools/build/lib/BUILD
index e07aaa8..ace04f8 100644
--- a/src/test/java/com/google/devtools/build/lib/BUILD
+++ b/src/test/java/com/google/devtools/build/lib/BUILD
@@ -538,6 +538,7 @@
":AllTests",
"//src/main/java/com/google/devtools/build/lib/actions",
"//src/main/java/com/google/devtools/build/lib/actions:execution_requirements",
+ "//src/main/java/com/google/devtools/build/lib/clock",
"//src/main/java/com/google/devtools/build/lib/sandbox",
"//src/main/java/com/google/devtools/build/lib/shell",
"//src/main/java/com/google/devtools/build/lib/util",
diff --git a/src/test/java/com/google/devtools/build/lib/worker/WorkerFactoryTest.java b/src/test/java/com/google/devtools/build/lib/worker/WorkerFactoryTest.java
index 6ca7731..5724afa 100644
--- a/src/test/java/com/google/devtools/build/lib/worker/WorkerFactoryTest.java
+++ b/src/test/java/com/google/devtools/build/lib/worker/WorkerFactoryTest.java
@@ -104,7 +104,7 @@
Worker proxiedWorker = workerFactory.create(proxiedWorkerKey);
// If proxied = true, WorkerProxy is created along with a WorkerMultiplexer.
// Destroy WorkerMultiplexer to avoid unexpected behavior in WorkerMultiplexerManagerTest.
- WorkerMultiplexerManager.removeInstance(proxiedWorkerKey.hashCode());
+ WorkerMultiplexerManager.removeInstance(proxiedWorkerKey);
assertThat(proxiedWorker.getClass()).isEqualTo(WorkerProxy.class);
}
}
diff --git a/src/test/java/com/google/devtools/build/lib/worker/WorkerMultiplexerManagerTest.java b/src/test/java/com/google/devtools/build/lib/worker/WorkerMultiplexerManagerTest.java
index 4957dad..03b3bfd 100644
--- a/src/test/java/com/google/devtools/build/lib/worker/WorkerMultiplexerManagerTest.java
+++ b/src/test/java/com/google/devtools/build/lib/worker/WorkerMultiplexerManagerTest.java
@@ -17,7 +17,16 @@
import static com.google.common.truth.Truth.assertThat;
import static org.junit.Assert.assertThrows;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSortedMap;
+import com.google.common.hash.HashCode;
+import com.google.devtools.build.lib.actions.ExecutionRequirements.WorkerProtocolFormat;
import com.google.devtools.build.lib.actions.UserExecException;
+import com.google.devtools.build.lib.clock.BlazeClock;
+import com.google.devtools.build.lib.vfs.FileSystem;
+import com.google.devtools.build.lib.vfs.inmemoryfs.InMemoryFileSystem;
+import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
@@ -26,50 +35,77 @@
@RunWith(JUnit4.class)
public class WorkerMultiplexerManagerTest {
+ private FileSystem fileSystem;
+
+ @Before
+ public void setUp() {
+ fileSystem = new InMemoryFileSystem(BlazeClock.instance());
+ }
+
@Test
public void instanceCreationRemovalTest() throws Exception {
// Create a WorkerProxy hash and request for a WorkerMultiplexer.
- Integer worker1Hash = "worker1".hashCode();
- WorkerMultiplexer wm1 = WorkerMultiplexerManager.getInstance(worker1Hash);
+ WorkerKey workerKey1 =
+ new WorkerKey(
+ ImmutableList.of(),
+ ImmutableMap.of(),
+ fileSystem.getPath("/execRoot"),
+ "mnemonic1",
+ HashCode.fromInt(1),
+ ImmutableSortedMap.of(),
+ false,
+ false,
+ WorkerProtocolFormat.PROTO);
+ WorkerMultiplexer wm1 = WorkerMultiplexerManager.getInstance(workerKey1);
- assertThat(WorkerMultiplexerManager.getMultiplexer(worker1Hash)).isEqualTo(wm1);
- assertThat(WorkerMultiplexerManager.getRefCount(worker1Hash)).isEqualTo(1);
+ assertThat(WorkerMultiplexerManager.getMultiplexer(workerKey1)).isEqualTo(wm1);
+ assertThat(WorkerMultiplexerManager.getRefCount(workerKey1)).isEqualTo(1);
assertThat(WorkerMultiplexerManager.getInstanceCount()).isEqualTo(1);
// Create another WorkerProxy hash and request for a WorkerMultiplexer.
- Integer worker2Hash = "worker2".hashCode();
- WorkerMultiplexer wm2 = WorkerMultiplexerManager.getInstance(worker2Hash);
+ WorkerKey workerKey2 =
+ new WorkerKey(
+ ImmutableList.of(),
+ ImmutableMap.of(),
+ fileSystem.getPath("/execRoot"),
+ "mnemonic2",
+ HashCode.fromInt(1),
+ ImmutableSortedMap.of(),
+ false,
+ false,
+ WorkerProtocolFormat.PROTO);
+ WorkerMultiplexer wm2 = WorkerMultiplexerManager.getInstance(workerKey2);
- assertThat(WorkerMultiplexerManager.getMultiplexer(worker2Hash)).isEqualTo(wm2);
- assertThat(WorkerMultiplexerManager.getRefCount(worker2Hash)).isEqualTo(1);
+ assertThat(WorkerMultiplexerManager.getMultiplexer(workerKey2)).isEqualTo(wm2);
+ assertThat(WorkerMultiplexerManager.getRefCount(workerKey2)).isEqualTo(1);
assertThat(WorkerMultiplexerManager.getInstanceCount()).isEqualTo(2);
// Use the same WorkerProxy hash, it shouldn't instantiate a new WorkerMultiplexer.
- WorkerMultiplexer wm2Annex = WorkerMultiplexerManager.getInstance(worker2Hash);
+ WorkerMultiplexer wm2Annex = WorkerMultiplexerManager.getInstance(workerKey2);
assertThat(wm2).isEqualTo(wm2Annex);
- assertThat(WorkerMultiplexerManager.getRefCount(worker2Hash)).isEqualTo(2);
+ assertThat(WorkerMultiplexerManager.getRefCount(workerKey2)).isEqualTo(2);
assertThat(WorkerMultiplexerManager.getInstanceCount()).isEqualTo(2);
// Remove an instance. If reference count is larger than 0, instance shouldn't be destroyed.
- WorkerMultiplexerManager.removeInstance(worker2Hash);
+ WorkerMultiplexerManager.removeInstance(workerKey2);
- assertThat(WorkerMultiplexerManager.getRefCount(worker2Hash)).isEqualTo(1);
+ assertThat(WorkerMultiplexerManager.getRefCount(workerKey2)).isEqualTo(1);
assertThat(WorkerMultiplexerManager.getInstanceCount()).isEqualTo(2);
// Remove an instance. Reference count is down to 0, instance should be destroyed.
- WorkerMultiplexerManager.removeInstance(worker2Hash);
+ WorkerMultiplexerManager.removeInstance(workerKey2);
assertThrows(
- UserExecException.class, () -> WorkerMultiplexerManager.getMultiplexer(worker2Hash));
+ UserExecException.class, () -> WorkerMultiplexerManager.getMultiplexer(workerKey2));
assertThat(WorkerMultiplexerManager.getInstanceCount()).isEqualTo(1);
// WorkerProxy hash not found.
assertThrows(
- UserExecException.class, () -> WorkerMultiplexerManager.removeInstance(worker2Hash));
+ UserExecException.class, () -> WorkerMultiplexerManager.removeInstance(workerKey2));
// Remove all the instances.
- WorkerMultiplexerManager.removeInstance(worker1Hash);
+ WorkerMultiplexerManager.removeInstance(workerKey1);
assertThat(WorkerMultiplexerManager.getInstanceCount()).isEqualTo(0);
}
