Implement new WorkerPoolImpl guarded by --use_new_worker_pool without using the apache.commons.pool2 library. The implementation includes an additional mechanism to handle waiting borrowers when quota is unavailable (blocks) through the `waitingQueue` of `WorkerLatch(s)` and unblocks by signaling the latch in the event that a worker is returned or invalidated (w/o shrinking). While this mechanism is currently unused due to synchronization in the ResourceManager (as part of workers as a resource, causing it to only `#borrowObject` when resources are actually available), it is still included so for correctness sake, so that we don't need to assume that the API is called correctly. PiperOrigin-RevId: 613525949 Change-Id: I2bb6508bdb144e0489a6453ca9ccf78af5332b66
diff --git a/src/main/java/com/google/devtools/build/lib/worker/BUILD b/src/main/java/com/google/devtools/build/lib/worker/BUILD index 2c87983..8d64243 100644 --- a/src/main/java/com/google/devtools/build/lib/worker/BUILD +++ b/src/main/java/com/google/devtools/build/lib/worker/BUILD
@@ -95,6 +95,7 @@ ":singleplex_worker", ":worker_factory", ":worker_options", + ":worker_pool", ":worker_pool_config", ":worker_pool_impl", ":worker_process_metrics", @@ -223,6 +224,7 @@ ], deps = [ ":worker_factory", + "//third_party:guava", ], ) @@ -230,6 +232,7 @@ name = "worker_pool_impl", srcs = [ "SimpleWorkerPool.java", + "WorkerPoolImpl.java", "WorkerPoolImplLegacy.java", ], deps = [
diff --git a/src/main/java/com/google/devtools/build/lib/worker/WorkerModule.java b/src/main/java/com/google/devtools/build/lib/worker/WorkerModule.java index adf4e89..86b38bf 100644 --- a/src/main/java/com/google/devtools/build/lib/worker/WorkerModule.java +++ b/src/main/java/com/google/devtools/build/lib/worker/WorkerModule.java
@@ -50,9 +50,9 @@ private WorkerFactory workerFactory; private AsynchronousTreeDeleter treeDeleter; - @VisibleForTesting WorkerPoolImplLegacy workerPool; WorkerPoolConfig config; + @VisibleForTesting WorkerPool workerPool; @Nullable private WorkerLifecycleManager workerLifecycleManager; @Override @@ -160,7 +160,10 @@ WorkerPoolConfig newConfig = new WorkerPoolConfig( - workerFactory, options.workerMaxInstances, options.workerMaxMultiplexInstances); + workerFactory, + options.useNewWorkerPool, + options.workerMaxInstances, + options.workerMaxMultiplexInstances); // If the config changed compared to the last run, we have to create a new pool. if (workerPool == null || !newConfig.equals(config)) { @@ -171,7 +174,11 @@ } if (workerPool == null) { - workerPool = new WorkerPoolImplLegacy(newConfig); + if (options.useNewWorkerPool) { + workerPool = new WorkerPoolImpl(newConfig); + } else { + workerPool = new WorkerPoolImplLegacy(newConfig); + } config = newConfig; // If workerPool is restarted then we should recreate metrics. WorkerProcessMetricsCollector.instance().clear();
diff --git a/src/main/java/com/google/devtools/build/lib/worker/WorkerOptions.java b/src/main/java/com/google/devtools/build/lib/worker/WorkerOptions.java index 105fe49..af00a87 100644 --- a/src/main/java/com/google/devtools/build/lib/worker/WorkerOptions.java +++ b/src/main/java/com/google/devtools/build/lib/worker/WorkerOptions.java
@@ -75,6 +75,16 @@ } @Option( + name = "experimental_use_new_worker_pool", + defaultValue = "false", + documentationCategory = OptionDocumentationCategory.UNDOCUMENTED, + effectTags = {OptionEffectTag.EXECUTION}, + help = + "Uses a new worker pool implementation (no change in behavior, reimplementation of " + + " the worker pool in order to deprecate the use of a third party tool).") + public boolean useNewWorkerPool; + + @Option( name = "worker_max_instances", converter = MultiResourceConverter.class, defaultValue = "null",
diff --git a/src/main/java/com/google/devtools/build/lib/worker/WorkerPool.java b/src/main/java/com/google/devtools/build/lib/worker/WorkerPool.java index 97e23c1..ceaffb8 100644 --- a/src/main/java/com/google/devtools/build/lib/worker/WorkerPool.java +++ b/src/main/java/com/google/devtools/build/lib/worker/WorkerPool.java
@@ -28,19 +28,55 @@ */ public interface WorkerPool { + /** + * Returns the current maximum number of workers allowed in the pool. + * + * @param key the worker key. + */ int getMaxTotalPerKey(WorkerKey key); + /** + * Returns the number of active workers. + * + * @param key the worker key. + */ int getNumActive(WorkerKey key); + /** + * Evicts specified workers from the pool, destroying them. + * + * <p>It is possible that not all specified workers get evicted if they become active. + * + * @param workerIdsToEvict the worker ids to attempt to evict. + * @return a set of worker ids that were successfully evicted. + */ ImmutableSet<Integer> evictWorkers(ImmutableSet<Integer> workerIdsToEvict) throws InterruptedException; + /** Returns the idle workers in the pool (note that these workers can still become active). */ ImmutableSet<Integer> getIdleWorkers() throws InterruptedException; + /** + * Borrows a persistent worker from the pool, creating if necessary and blocking if unavailable. + * + * @param key the worker key. + */ Worker borrowObject(WorkerKey key) throws IOException, InterruptedException; + /** + * Returns an active worker back to the pool. + * + * @param key the worker key. + * @param obj the worker to be returned. + */ void returnObject(WorkerKey key, Worker obj); + /** + * Invalidates the worker, thus destroying it. + * + * @param key the worker key. + * @param obj the worker to be invalidated. + */ void invalidateObject(WorkerKey key, Worker obj) throws InterruptedException; void reset();
diff --git a/src/main/java/com/google/devtools/build/lib/worker/WorkerPoolConfig.java b/src/main/java/com/google/devtools/build/lib/worker/WorkerPoolConfig.java index 85b59e6..224e1b0 100644 --- a/src/main/java/com/google/devtools/build/lib/worker/WorkerPoolConfig.java +++ b/src/main/java/com/google/devtools/build/lib/worker/WorkerPoolConfig.java
@@ -14,6 +14,7 @@ package com.google.devtools.build.lib.worker; +import com.google.common.annotations.VisibleForTesting; import java.util.List; import java.util.Map.Entry; import java.util.Objects; @@ -24,18 +25,33 @@ */ public class WorkerPoolConfig { private final WorkerFactory workerFactory; + private final boolean useNewWorkerPool; private final List<Entry<String, Integer>> workerMaxInstances; private final List<Entry<String, Integer>> workerMaxMultiplexInstances; public WorkerPoolConfig( WorkerFactory workerFactory, + boolean useNewWorkerPool, List<Entry<String, Integer>> workerMaxInstances, List<Entry<String, Integer>> workerMaxMultiplexInstances) { this.workerFactory = workerFactory; + this.useNewWorkerPool = useNewWorkerPool; this.workerMaxInstances = workerMaxInstances; this.workerMaxMultiplexInstances = workerMaxMultiplexInstances; } + @VisibleForTesting + public WorkerPoolConfig( + WorkerFactory workerFactory, + List<Entry<String, Integer>> workerMaxInstances, + List<Entry<String, Integer>> workerMaxMultiplexInstances) { + this( + workerFactory, + /* useNewWorkerPool= */ false, + workerMaxInstances, + workerMaxMultiplexInstances); + } + public WorkerFactory getWorkerFactory() { return workerFactory; } @@ -58,12 +74,14 @@ } WorkerPoolConfig that = (WorkerPoolConfig) o; return workerFactory.equals(that.workerFactory) + && useNewWorkerPool == that.useNewWorkerPool && workerMaxInstances.equals(that.workerMaxInstances) && workerMaxMultiplexInstances.equals(that.workerMaxMultiplexInstances); } @Override public int hashCode() { - return Objects.hash(workerFactory, workerMaxInstances, workerMaxMultiplexInstances); + return Objects.hash( + workerFactory, useNewWorkerPool, workerMaxInstances, workerMaxMultiplexInstances); } }
diff --git a/src/main/java/com/google/devtools/build/lib/worker/WorkerPoolImpl.java b/src/main/java/com/google/devtools/build/lib/worker/WorkerPoolImpl.java new file mode 100644 index 0000000..6b90aa1 --- /dev/null +++ b/src/main/java/com/google/devtools/build/lib/worker/WorkerPoolImpl.java
@@ -0,0 +1,383 @@ +// Copyright 2024 The Bazel Authors. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.devtools.build.lib.worker; + +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.collect.ImmutableSet.toImmutableSet; + +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import com.google.common.flogger.GoogleLogger; +import com.google.devtools.build.lib.worker.WorkerProcessStatus.Status; +import java.io.IOException; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Map.Entry; +import java.util.Set; +import java.util.concurrent.BlockingDeque; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.LinkedBlockingDeque; +import java.util.concurrent.atomic.AtomicInteger; +import javax.annotation.Nullable; + +/** + * Implementation of the WorkerPool. + * + * <p>TODO(b/323880131): Remove documentation once we completely remove the legacy implementation. + * + * <p>Difference in internal implementation from {@code WorkerPoolLegacy}: + * + * <ul> + * <li>Legacy: WorkerPoolLegacy wraps multiple {@code SimpleWorkerPool} for each mnemonic. Each + * SimpleWorkerPool contains {@code Worker} instances capped per {@code WorkerKey}. + * <li>Current: This implementation flattens this to have a single {@code WorkerKeyPool} for each + * worker key (we don't need the indirection in referencing both mnemonic and worker key since + * the mnemonic is part of the key). + * <li>Legacy: SimpleWorkerPool extends {@code GenericKeyedObjectPool} that handles the logic to + * concurrent calls to borrow, return, invalidate and evict workers. + * <li>Current: WorkerKeyPool replaces this functionality directly, but can only handle pool logic + * for a single key (as compared to SimpleWorkerPool that handles multiple worker keys of the + * same mnemonic). Additionally, it bakes in pool shrinking logic so that we can handle + * concurrent calls. + * </ul> + */ +public class WorkerPoolImpl implements WorkerPool { + + private static final GoogleLogger logger = GoogleLogger.forEnclosingClass(); + + /** Unless otherwise specified, the max number of workers per WorkerKey. */ + private static final int DEFAULT_MAX_SINGLEPLEX_WORKERS = 4; + + /** Unless otherwise specified, the max number of multiplex workers per WorkerKey. */ + private static final int DEFAULT_MAX_MULTIPLEX_WORKERS = 8; + + private final WorkerFactory factory; + + private final ImmutableMap<String, Integer> singleplexMaxInstances; + private final ImmutableMap<String, Integer> multiplexMaxInstances; + private final ConcurrentHashMap<WorkerKey, WorkerKeyPool> pools = new ConcurrentHashMap<>(); + + public WorkerPoolImpl(WorkerPoolConfig config) { + this.factory = config.getWorkerFactory(); + this.singleplexMaxInstances = + getMaxInstances(config.getWorkerMaxInstances(), DEFAULT_MAX_SINGLEPLEX_WORKERS); + this.multiplexMaxInstances = + getMaxInstances(config.getWorkerMaxMultiplexInstances(), DEFAULT_MAX_MULTIPLEX_WORKERS); + } + + private static ImmutableMap<String, Integer> getMaxInstances( + List<Entry<String, Integer>> maxInstances, int defaultMaxWorkers) { + LinkedHashMap<String, Integer> newConfigBuilder = new LinkedHashMap<>(); + for (Map.Entry<String, Integer> entry : maxInstances) { + if (entry.getValue() != null) { + newConfigBuilder.put(entry.getKey(), entry.getValue()); + } else if (entry.getKey() != null) { + newConfigBuilder.put(entry.getKey(), defaultMaxWorkers); + } + } + return ImmutableMap.copyOf(newConfigBuilder); + } + + @Override + public int getMaxTotalPerKey(WorkerKey key) { + return getPool(key).getAvailableQuota(); + } + + @Override + public int getNumActive(WorkerKey key) { + return getPool(key).getNumActive(); + } + + @Override + public ImmutableSet<Integer> evictWorkers(ImmutableSet<Integer> workerIdsToEvict) + throws InterruptedException { + // TODO: Without having the Worker objects themselves, we can't directly pass the worker to the + // pool to be evicted. + ImmutableSet<Integer> evictedWorkerIds = + pools.values().stream() + .flatMap(p -> p.evictWorkers(workerIdsToEvict).stream()) + .collect(toImmutableSet()); + return workerIdsToEvict.stream().filter(evictedWorkerIds::contains).collect(toImmutableSet()); + } + + @Override + public ImmutableSet<Integer> getIdleWorkers() throws InterruptedException { + return pools.values().stream() + .flatMap(p -> p.getIdleWorkers().stream()) + .collect(toImmutableSet()); + } + + @Override + public Worker borrowObject(WorkerKey key) throws IOException, InterruptedException { + return getPool(key).borrowWorker(); + } + + @Override + public void returnObject(WorkerKey key, Worker obj) { + getPool(key).returnWorker(/* worker= */ obj); + } + + @Override + public void invalidateObject(WorkerKey key, Worker obj) throws InterruptedException { + invalidateWorker( + /* worker= */ obj, /* shouldShrinkPool= */ obj.getStatus().isPendingEviction()); + } + + /** + * TODO(b/323880131): This should be the main interface once the we remove the legacy worker pool + * implementation. + */ + private void invalidateWorker(Worker worker, boolean shouldShrinkPool) { + getPool(worker.getWorkerKey()).invalidateWorker(worker, shouldShrinkPool); + } + + @Override + public void reset() { + for (WorkerKeyPool pool : pools.values()) { + pool.reset(); + } + } + + @Override + public void close() { + for (WorkerKeyPool pool : pools.values()) { + pool.close(); + } + } + + private WorkerKeyPool getPool(WorkerKey key) { + return pools.computeIfAbsent(key, this::createPool); + } + + private WorkerKeyPool createPool(WorkerKey key) { + if (key.isMultiplex()) { + return new WorkerKeyPool( + key, + multiplexMaxInstances.getOrDefault(key.getMnemonic(), DEFAULT_MAX_MULTIPLEX_WORKERS)); + } + return new WorkerKeyPool( + key, + singleplexMaxInstances.getOrDefault(key.getMnemonic(), DEFAULT_MAX_SINGLEPLEX_WORKERS)); + } + + /** + * Actual pool implementation that handles the borrowing, returning and invalidation of workers of + * a single worker key. + */ + private class WorkerKeyPool { + + private final WorkerKey key; + private final int max; + // We maintain this as a separate counter from the activeSet so that we can create workers + // without locking the pool, while maintaining atomicity on + private final AtomicInteger acquired = new AtomicInteger(0); + private final AtomicInteger shrunk = new AtomicInteger(0); + + private final BlockingDeque<Worker> idleQueue = new LinkedBlockingDeque<>(); + + /** + * The waiting queue is meant to block borrowers when there are no workers available. With + * workers as a resource, workers are only borrowed when they get they are available, so this + * doesn't get used, i.e. there shouldn't be any borrowers waiting here, where the {@code + * ResourceManager} handles the proper synchronization to ensure that workers are borrowed + * together with its allocated resources. + * + * <p>Regardless, this implementation is still included to ensure correctness such that that + * multiple threads can still borrow concurrently, without needing to check how many workers are + * actually available (blocking if unavailable). + */ + private final BlockingDeque<WorkerLatch> waitingQueue = new LinkedBlockingDeque<>(); + + private final Set<Worker> activeSet = new HashSet<>(); + + public WorkerKeyPool(WorkerKey key, int max) { + this.key = key; + this.max = max; + } + + private synchronized Set<Integer> evictWorkers(Set<Integer> workerIdsToEvict) { + Set<Integer> evictedWorkerIds = new HashSet<>(); + for (Worker worker : idleQueue) { + if (workerIdsToEvict.contains(worker.getWorkerId())) { + evictedWorkerIds.add(worker.getWorkerId()); + worker.getStatus().maybeUpdateStatus(Status.PENDING_KILL_DUE_TO_MEMORY_PRESSURE); + // Currently when evicting idle workers, we do not shrink the pool. The pool is only + // shrunk when we have to postpone invalidation of the worker. + invalidateWorker(worker, /* shouldShrinkPool= */ false); + idleQueue.remove(worker); + logger.atInfo().log( + "Evicted %s worker (id %d, key hash %d).", + worker.getWorkerKey().getMnemonic(), + worker.getWorkerId(), + worker.getWorkerKey().hashCode()); + } + // TODO(b/323880131): Move postponing of invalidation from {@code WorkerLifecycleManager} + // here, since all we need to do is to update the statuses. We keep it like this for now + // to preserve the existing behavior. + } + return evictedWorkerIds; + } + + private synchronized int getNumActive() { + return acquired.get(); + } + + private synchronized int getAvailableQuota() { + return max - shrunk.get(); + } + + // Callers should atomically check to confirm that workers are available before calling this + // method or risk being blocked waiting for a worker to be available. + private Worker borrowWorker() throws IOException, InterruptedException { + Worker worker = null; + WorkerLatch latch = null; + // We don't want to hold the lock on the pool while creating or waiting for a worker or quota + // to be available. + synchronized (this) { + while (!idleQueue.isEmpty()) { + worker = idleQueue.takeLast(); + if (factory.validateWorker(worker.getWorkerKey(), worker)) { + acquired.incrementAndGet(); + break; + } + invalidateWorker(worker, /* shouldShrinkPool= */ false); + worker = null; + } + + if (worker == null) { + // If we were unable to get an idle worker, then either create or wait for one. + if (getAvailableQuota() - getNumActive() > 0) { + // No idle workers, but we have space to create another. + acquired.incrementAndGet(); + } else { + latch = new WorkerLatch(); + waitingQueue.add(latch); + } + } + } + + if (latch != null) { + // Wait until the resources are available. We cannot do this why synchronized because that + // would deadlock by blocking other threads from returning and thus freeing up quota for + // this to proceed. + worker = latch.await(); + } + + if (worker == null) { + worker = factory.create(key); + } + + activeSet.add(worker); + + checkArgument( + getAvailableQuota() - getNumActive() >= 0, + "Worker pool (mnemonic %s) does not have space to create another worker.", + key.getMnemonic()); + return worker; + } + + private synchronized void returnWorker(Worker worker) { + if (!factory.validateWorker(worker.getWorkerKey(), worker)) { + invalidateWorker(worker, true); + return; + } + + activeSet.remove(worker); + + WorkerLatch latch = waitingQueue.poll(); + if (latch != null) { + // Pass the worker directly to the waiting thread. + latch.countDown(worker); + } else { + acquired.decrementAndGet(); + idleQueue.addLast(worker); + } + } + + private synchronized void invalidateWorker(Worker worker, boolean shouldShrinkPool) { + factory.destroyWorker(worker.getWorkerKey(), worker); + + if (activeSet.remove(worker)) { + acquired.decrementAndGet(); + } else { + idleQueue.remove(worker); + return; + } + + // We don't want to shrink the pool to 0. + if (shouldShrinkPool && getAvailableQuota() > 1) { + shrunk.incrementAndGet(); + return; + } + + // If invalidating the worker has resulted in the freeing up of effective quota for waiting + // threads (also taking into account whether it was shrunk), then we signal for the next + // waiting thread proceed. + WorkerLatch latch = waitingQueue.poll(); + if (latch != null) { + // We signal to the waiting thread that it can proceed, but it has to create a worker + // for itself. + latch.countDown(null); + // We need to increment while synchronized here, so that we don't race with another + // thread that might borrow (thus taking up the quota) before the waiting thread actually + // manages to proceed (and increment on its own). + acquired.getAndIncrement(); + } + } + + /** + * It is not important that we synchronize here, the {@code WorkerLifecycleManager} takes the + * idle workers and figures out (non-atomically with respect to this instance) which workers to + * evict. So it is possible that an idle worker gets acquired before it decides to evict a + * previously idle worker. + */ + public Set<Integer> getIdleWorkers() { + return idleQueue.stream().map(Worker::getWorkerId).collect(toImmutableSet()); + } + + private void reset() { + shrunk.set(0); + logger.atInfo().log( + "clearing shrunk values for %s (key hash %d) worker pool", + key.getMnemonic(), key.hashCode()); + } + + private void close() {} + } + + /** + * Used to pass workers from threads that are returning the worker to the pool, bypassing the + * queue. + */ + private static class WorkerLatch { + final CountDownLatch latch = new CountDownLatch(1); + @Nullable volatile Worker worker = null; + + /** Returns a worker instance that has been freed, or null if the worker needs to be created. */ + @Nullable + public Worker await() throws InterruptedException { + latch.await(); + return worker; + } + + public void countDown(@Nullable Worker worker) { + this.worker = worker; + latch.countDown(); + } + } +}
diff --git a/src/test/java/com/google/devtools/build/lib/worker/BUILD b/src/test/java/com/google/devtools/build/lib/worker/BUILD index 2f0a1fd..418d0a0 100644 --- a/src/test/java/com/google/devtools/build/lib/worker/BUILD +++ b/src/test/java/com/google/devtools/build/lib/worker/BUILD
@@ -130,6 +130,7 @@ "//src/test/java/com/google/devtools/build/lib/buildtool/util", "//src/test/java/com/google/devtools/build/lib/exec/util", "//src/test/java/com/google/devtools/build/lib/testutil", + "//src/test/java/com/google/devtools/build/lib/testutil:TestThread", "//src/test/java/com/google/devtools/build/lib/testutil:TestUtils", "//src/test/java/com/google/devtools/build/lib/vfs/util", "//third_party:apache_commons_pool2",
diff --git a/src/test/java/com/google/devtools/build/lib/worker/WorkerLifecycleManagerTest.java b/src/test/java/com/google/devtools/build/lib/worker/WorkerLifecycleManagerTest.java index befed83..f50ec75 100644 --- a/src/test/java/com/google/devtools/build/lib/worker/WorkerLifecycleManagerTest.java +++ b/src/test/java/com/google/devtools/build/lib/worker/WorkerLifecycleManagerTest.java
@@ -18,6 +18,7 @@ import static com.google.devtools.build.lib.worker.WorkerTestUtils.createWorkerKey; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.spy; import com.google.common.collect.ImmutableList; import com.google.common.collect.Maps; @@ -26,25 +27,48 @@ import com.google.devtools.build.lib.vfs.FileSystem; import com.google.devtools.build.lib.vfs.inmemoryfs.InMemoryFileSystem; import com.google.devtools.build.lib.worker.WorkerProcessStatus.Status; +import java.io.IOException; import java.time.Instant; +import java.util.Arrays; +import java.util.List; import java.util.Map.Entry; +import java.util.function.Supplier; import org.apache.commons.pool2.PooledObject; import org.apache.commons.pool2.impl.DefaultPooledObject; import org.junit.Before; import org.junit.Rule; import org.junit.Test; import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; +import org.junit.runners.Parameterized; +import org.junit.runners.Parameterized.Parameter; +import org.junit.runners.Parameterized.Parameters; import org.mockito.Mock; import org.mockito.junit.MockitoJUnit; import org.mockito.junit.MockitoRule; -@RunWith(JUnit4.class) +@RunWith(Parameterized.class) public final class WorkerLifecycleManagerTest { @Rule public final MockitoRule mockito = MockitoJUnit.rule(); + + private interface WorkerPoolSupplier { + WorkerPool get( + WorkerFactory factory, + List<Entry<String, Integer>> singlexplexMaxInstances, + List<Entry<String, Integer>> multiplexMaxInstance); + } + + @Parameter(0) + public WorkerPoolSupplier workerPoolSupplier; + + @Parameter(1) + public Supplier<WorkerFactory> workerFactorySupplier; + @Mock WorkerFactory factoryMock; - private FileSystem fileSystem; + + public static final FileSystem fileSystem = + new InMemoryFileSystem(BlazeClock.instance(), DigestHashFunction.SHA256); + private static final WorkerOptions options = new WorkerOptions(); private static final String DUMMY_MNEMONIC = "dummy"; private static final long PROCESS_ID_1 = 1L; private static final long PROCESS_ID_2 = 2L; @@ -54,9 +78,35 @@ private int workerIds = 1; + @Parameters + public static List<Object[]> data() throws IOException { + Supplier<WorkerFactory> workerFactorySupplier = + () -> spy(new WorkerFactory(fileSystem.getPath("/outputbase/bazel-workers"), options)); + + return Arrays.asList( + new Object[][] { + { + (WorkerPoolSupplier) + (factory, singleplexMaxInstances, multiplexMaxInstances) -> + new WorkerPoolImplLegacy( + new WorkerPoolConfig( + factory, singleplexMaxInstances, multiplexMaxInstances)), + workerFactorySupplier, + }, + { + (WorkerPoolSupplier) + (factory, singleplexMaxInstances, multiplexMaxInstances) -> + new WorkerPoolImpl( + new WorkerPoolConfig( + factory, singleplexMaxInstances, multiplexMaxInstances)), + workerFactorySupplier, + } + }); + } + @Before public void setUp() throws Exception { - fileSystem = new InMemoryFileSystem(BlazeClock.instance(), DigestHashFunction.SHA256); + factoryMock = workerFactorySupplier.get(); WorkerOptions options = new WorkerOptions(); doAnswer( args -> { @@ -102,9 +152,8 @@ @Test public void testEvictWorkers_doNothing_lowMemoryUsage() throws Exception { - WorkerPoolImplLegacy workerPool = - new WorkerPoolImplLegacy( - new WorkerPoolConfig(factoryMock, entryList(DUMMY_MNEMONIC, 1), emptyEntryList())); + WorkerPool workerPool = + workerPoolSupplier.get(factoryMock, entryList(DUMMY_MNEMONIC, 1), emptyEntryList()); WorkerKey key = createWorkerKey(DUMMY_MNEMONIC, fileSystem); Worker w1 = workerPool.borrowObject(key); workerPool.returnObject(key, w1); @@ -115,12 +164,12 @@ WorkerLifecycleManager manager = new WorkerLifecycleManager(workerPool, options); - assertThat(workerPool.getNumIdlePerKey(key)).isEqualTo(1); + assertThat(workerPool.getIdleWorkers()).hasSize(1); assertThat(workerPool.getNumActive(key)).isEqualTo(0); manager.evictWorkers(workerMetrics); - assertThat(workerPool.getNumIdlePerKey(key)).isEqualTo(1); + assertThat(workerPool.getIdleWorkers()).hasSize(1); assertThat(workerPool.getNumActive(key)).isEqualTo(0); // It should still have a valid status since it was not killed / marked to be killed. assertThat(w1.getStatus().isValid()).isTrue(); @@ -128,9 +177,8 @@ @Test public void testEvictWorkers_doNothing_zeroThreshold() throws Exception { - WorkerPoolImplLegacy workerPool = - new WorkerPoolImplLegacy( - new WorkerPoolConfig(factoryMock, entryList(DUMMY_MNEMONIC, 1), emptyEntryList())); + WorkerPool workerPool = + workerPoolSupplier.get(factoryMock, entryList(DUMMY_MNEMONIC, 1), emptyEntryList()); WorkerKey key = createWorkerKey(DUMMY_MNEMONIC, fileSystem); Worker w1 = workerPool.borrowObject(key); workerPool.returnObject(key, w1); @@ -142,12 +190,12 @@ WorkerLifecycleManager manager = new WorkerLifecycleManager(workerPool, options); - assertThat(workerPool.getNumIdlePerKey(key)).isEqualTo(1); + assertThat(workerPool.getIdleWorkers()).hasSize(1); assertThat(workerPool.getNumActive(key)).isEqualTo(0); manager.evictWorkers(workerMetrics); - assertThat(workerPool.getNumIdlePerKey(key)).isEqualTo(1); + assertThat(workerPool.getIdleWorkers()).hasSize(1); assertThat(workerPool.getNumActive(key)).isEqualTo(0); // It should still have a valid status since it was not killed / marked to be killed. assertThat(w1.getStatus().isValid()).isTrue(); @@ -155,9 +203,8 @@ @Test public void testEvictWorkers_doNothing_emptyMetrics() throws Exception { - WorkerPoolImplLegacy workerPool = - new WorkerPoolImplLegacy( - new WorkerPoolConfig(factoryMock, entryList(DUMMY_MNEMONIC, 1), emptyEntryList())); + WorkerPool workerPool = + workerPoolSupplier.get(factoryMock, entryList(DUMMY_MNEMONIC, 1), emptyEntryList()); WorkerKey key = createWorkerKey(DUMMY_MNEMONIC, fileSystem); Worker w1 = workerPool.borrowObject(key); workerPool.returnObject(key, w1); @@ -168,12 +215,12 @@ WorkerLifecycleManager manager = new WorkerLifecycleManager(workerPool, options); - assertThat(workerPool.getNumIdlePerKey(key)).isEqualTo(1); + assertThat(workerPool.getIdleWorkers()).hasSize(1); assertThat(workerPool.getNumActive(key)).isEqualTo(0); manager.evictWorkers(workerMetrics); - assertThat(workerPool.getNumIdlePerKey(key)).isEqualTo(1); + assertThat(workerPool.getIdleWorkers()).hasSize(1); assertThat(workerPool.getNumActive(key)).isEqualTo(0); // It should still have a valid status since it was not killed / marked to be killed. assertThat(w1.getStatus().isValid()).isTrue(); @@ -181,9 +228,8 @@ @Test public void testGetEvictionCandidates_selectOnlyWorker() throws Exception { - WorkerPoolImplLegacy workerPool = - new WorkerPoolImplLegacy( - new WorkerPoolConfig(factoryMock, entryList(DUMMY_MNEMONIC, 1), emptyEntryList())); + WorkerPool workerPool = + workerPoolSupplier.get(factoryMock, entryList(DUMMY_MNEMONIC, 1), emptyEntryList()); WorkerKey key = createWorkerKey(DUMMY_MNEMONIC, fileSystem); Worker w1 = workerPool.borrowObject(key); workerPool.returnObject(key, w1); @@ -193,14 +239,14 @@ options.totalWorkerMemoryLimitMb = 1; WorkerLifecycleManager manager = new WorkerLifecycleManager(workerPool, options); - assertThat(workerPool.getNumIdlePerKey(key)).isEqualTo(1); + assertThat(workerPool.getIdleWorkers()).hasSize(1); assertThat(workerPool.getNumActive(key)).isEqualTo(0); // It should still have a valid status since it was not killed / marked to be killed. assertThat(w1.getStatus().isValid()).isTrue(); manager.evictWorkers(workerMetrics); - assertThat(workerPool.getNumIdlePerKey(key)).isEqualTo(0); + assertThat(workerPool.getIdleWorkers()).isEmpty(); assertThat(workerPool.getNumActive(key)).isEqualTo(0); // Directly killed since it is already returned. assertThat(w1.getStatus().get()).isEqualTo(Status.KILLED_DUE_TO_MEMORY_PRESSURE); @@ -208,9 +254,8 @@ @Test public void testGetEvictionCandidates_evictLargestWorkers() throws Exception { - WorkerPoolImplLegacy workerPool = - new WorkerPoolImplLegacy( - new WorkerPoolConfig(factoryMock, entryList(DUMMY_MNEMONIC, 3), emptyEntryList())); + WorkerPool workerPool = + workerPoolSupplier.get(factoryMock, entryList(DUMMY_MNEMONIC, 3), emptyEntryList()); WorkerKey key = createWorkerKey(DUMMY_MNEMONIC, fileSystem); Worker w1 = workerPool.borrowObject(key); Worker w2 = workerPool.borrowObject(key); @@ -229,7 +274,7 @@ options.totalWorkerMemoryLimitMb = 2; WorkerLifecycleManager manager = new WorkerLifecycleManager(workerPool, options); - assertThat(workerPool.getNumIdlePerKey(key)).isEqualTo(3); + assertThat(workerPool.getIdleWorkers()).hasSize(3); assertThat(workerPool.getNumActive(key)).isEqualTo(0); assertThat(w1.getStatus().isValid()).isTrue(); assertThat(w2.getStatus().isValid()).isTrue(); @@ -237,7 +282,7 @@ manager.evictWorkers(workerMetrics); - assertThat(workerPool.getNumIdlePerKey(key)).isEqualTo(1); + assertThat(workerPool.getIdleWorkers()).hasSize(1); assertThat(workerPool.getNumActive(key)).isEqualTo(0); // Only w1 and w3 should have been killed. assertThat(w1.getStatus().get()).isEqualTo(Status.KILLED_DUE_TO_MEMORY_PRESSURE); @@ -248,10 +293,8 @@ @Test public void testGetEvictionCandidates_numberOfWorkersIsMoreThanDefaultNumTests() throws Exception { - // The default number of tests is 3. - WorkerPoolImplLegacy workerPool = - new WorkerPoolImplLegacy( - new WorkerPoolConfig(factoryMock, entryList(DUMMY_MNEMONIC, 4), emptyEntryList())); + WorkerPool workerPool = + workerPoolSupplier.get(factoryMock, entryList(DUMMY_MNEMONIC, 4), emptyEntryList()); WorkerKey key = createWorkerKey(DUMMY_MNEMONIC, fileSystem); Worker w1 = workerPool.borrowObject(key); Worker w2 = workerPool.borrowObject(key); @@ -273,12 +316,12 @@ options.totalWorkerMemoryLimitMb = 1; WorkerLifecycleManager manager = new WorkerLifecycleManager(workerPool, options); - assertThat(workerPool.getNumIdlePerKey(key)).isEqualTo(4); + assertThat(workerPool.getIdleWorkers()).hasSize(4); assertThat(workerPool.getNumActive(key)).isEqualTo(0); manager.evictWorkers(workerMetrics); - assertThat(workerPool.getNumIdlePerKey(key)).isEqualTo(0); + assertThat(workerPool.getIdleWorkers()).isEmpty(); assertThat(workerPool.getNumActive(key)).isEqualTo(0); assertThat(w1.getStatus().get()).isEqualTo(Status.KILLED_DUE_TO_MEMORY_PRESSURE); assertThat(w2.getStatus().get()).isEqualTo(Status.KILLED_DUE_TO_MEMORY_PRESSURE); @@ -289,9 +332,8 @@ @Test public void testGetEvictionCandidates_evictWorkerWithSameMenmonicButDifferentKeys() throws Exception { - WorkerPoolImplLegacy workerPool = - new WorkerPoolImplLegacy( - new WorkerPoolConfig(factoryMock, entryList(DUMMY_MNEMONIC, 3), emptyEntryList())); + WorkerPool workerPool = + workerPoolSupplier.get(factoryMock, entryList(DUMMY_MNEMONIC, 3), emptyEntryList()); WorkerKey key1 = createWorkerKey(DUMMY_MNEMONIC, fileSystem); WorkerKey key2 = createWorkerKey(DUMMY_MNEMONIC, fileSystem, true); @@ -313,20 +355,18 @@ options.workerVerbose = true; WorkerLifecycleManager manager = new WorkerLifecycleManager(workerPool, options); - assertThat(workerPool.getNumIdlePerKey(key1)).isEqualTo(1); - assertThat(workerPool.getNumIdlePerKey(key2)).isEqualTo(2); + assertThat(workerPool.getIdleWorkers()) + .containsExactly(w1.getWorkerId(), w2.getWorkerId(), w3.getWorkerId()); assertThat(w1.getStatus().isValid()).isTrue(); assertThat(w2.getStatus().isValid()).isTrue(); assertThat(w3.getStatus().isValid()).isTrue(); manager.evictWorkers(workerMetrics); - assertThat(workerPool.getNumIdlePerKey(key1)).isEqualTo(0); - assertThat(workerPool.getNumActive(key1)).isEqualTo(0); - - assertThat(workerPool.getNumIdlePerKey(key2)).isEqualTo(1); - assertThat(workerPool.getNumActive(key2)).isEqualTo(0); // Only w3 shouldn't be killed. + assertThat(workerPool.getIdleWorkers()).containsExactly(w3.getWorkerId()); + assertThat(workerPool.getNumActive(key1)).isEqualTo(0); + assertThat(workerPool.getNumActive(key2)).isEqualTo(0); assertThat(w1.getStatus().get()).isEqualTo(Status.KILLED_DUE_TO_MEMORY_PRESSURE); assertThat(w2.getStatus().get()).isEqualTo(Status.KILLED_DUE_TO_MEMORY_PRESSURE); assertThat(w3.getStatus().isValid()).isTrue(); @@ -334,9 +374,9 @@ @Test public void testGetEvictionCandidates_evictOnlyIdleWorkers() throws Exception { - WorkerPoolImplLegacy workerPool = - new WorkerPoolImplLegacy( - new WorkerPoolConfig(factoryMock, entryList(DUMMY_MNEMONIC, 3), emptyEntryList())); + + WorkerPool workerPool = + workerPoolSupplier.get(factoryMock, entryList(DUMMY_MNEMONIC, 3), emptyEntryList()); WorkerKey key = createWorkerKey(DUMMY_MNEMONIC, fileSystem); Worker w1 = workerPool.borrowObject(key); Worker w2 = workerPool.borrowObject(key); @@ -355,7 +395,7 @@ WorkerLifecycleManager manager = new WorkerLifecycleManager(workerPool, options); - assertThat(workerPool.getNumIdlePerKey(key)).isEqualTo(2); + assertThat(workerPool.getIdleWorkers()).hasSize(2); assertThat(workerPool.getNumActive(key)).isEqualTo(1); assertThat(w1.getStatus().isValid()).isTrue(); assertThat(w2.getStatus().isValid()).isTrue(); @@ -363,7 +403,7 @@ manager.evictWorkers(workerMetrics); - assertThat(workerPool.getNumIdlePerKey(key)).isEqualTo(0); + assertThat(workerPool.getIdleWorkers()).isEmpty(); assertThat(workerPool.getNumActive(key)).isEqualTo(1); assertThat(w1.getStatus().get()).isEqualTo(Status.KILLED_DUE_TO_MEMORY_PRESSURE); assertThat(w2.getStatus().get()).isEqualTo(Status.KILLED_DUE_TO_MEMORY_PRESSURE); @@ -373,10 +413,9 @@ @Test public void testGetEvictionCandidates_evictDifferentWorkerKeys() throws Exception { - WorkerPoolImplLegacy workerPool = - new WorkerPoolImplLegacy( - new WorkerPoolConfig( - factoryMock, entryList(DUMMY_MNEMONIC, 2, "smart", 2), emptyEntryList())); + WorkerPool workerPool = + workerPoolSupplier.get( + factoryMock, entryList(DUMMY_MNEMONIC, 2, "smart", 2), emptyEntryList()); WorkerKey key1 = createWorkerKey(DUMMY_MNEMONIC, fileSystem); WorkerKey key2 = createWorkerKey("smart", fileSystem); Worker w1 = workerPool.borrowObject(key1); @@ -400,9 +439,8 @@ WorkerLifecycleManager manager = new WorkerLifecycleManager(workerPool, options); - assertThat(workerPool.getNumIdlePerKey(key1)).isEqualTo(2); + assertThat(workerPool.getIdleWorkers()).hasSize(4); assertThat(workerPool.getNumActive(key1)).isEqualTo(0); - assertThat(workerPool.getNumIdlePerKey(key2)).isEqualTo(2); assertThat(workerPool.getNumActive(key2)).isEqualTo(0); assertThat(w1.getStatus().isValid()).isTrue(); assertThat(w2.getStatus().isValid()).isTrue(); @@ -411,14 +449,12 @@ manager.evictWorkers(workerMetrics); - assertThat(workerPool.getNumIdlePerKey(key1)).isEqualTo(1); + // Only w1 and w4 should be alive. + assertThat(workerPool.getIdleWorkers()).containsExactly(w1.getWorkerId(), w4.getWorkerId()); assertThat(workerPool.getNumActive(key1)).isEqualTo(0); - assertThat(workerPool.borrowObject(key1).getWorkerId()).isEqualTo(w1.getWorkerId()); - - assertThat(workerPool.getNumIdlePerKey(key2)).isEqualTo(1); assertThat(workerPool.getNumActive(key2)).isEqualTo(0); + assertThat(workerPool.borrowObject(key1).getWorkerId()).isEqualTo(w1.getWorkerId()); assertThat(workerPool.borrowObject(key2).getWorkerId()).isEqualTo(w4.getWorkerId()); - assertThat(w1.getStatus().isValid()).isTrue(); assertThat(w2.getStatus().get()).isEqualTo(Status.KILLED_DUE_TO_MEMORY_PRESSURE); assertThat(w3.getStatus().get()).isEqualTo(Status.KILLED_DUE_TO_MEMORY_PRESSURE); @@ -427,9 +463,8 @@ @Test public void testGetEvictionCandidates_testDoomedWorkers() throws Exception { - WorkerPoolImplLegacy workerPool = - new WorkerPoolImplLegacy( - new WorkerPoolConfig(factoryMock, entryList(DUMMY_MNEMONIC, 2), emptyEntryList())); + WorkerPool workerPool = + workerPoolSupplier.get(factoryMock, entryList(DUMMY_MNEMONIC, 2), emptyEntryList()); WorkerKey key = createWorkerKey(DUMMY_MNEMONIC, fileSystem); Worker w1 = workerPool.borrowObject(key); Worker w2 = workerPool.borrowObject(key); @@ -445,7 +480,7 @@ WorkerLifecycleManager manager = new WorkerLifecycleManager(workerPool, options); - assertThat(workerPool.getNumIdlePerKey(key)).isEqualTo(0); + assertThat(workerPool.getIdleWorkers()).isEmpty(); assertThat(workerPool.getNumActive(key)).isEqualTo(2); assertThat(w1.getStatus().isValid()).isTrue(); assertThat(w2.getStatus().isValid()).isTrue(); @@ -459,7 +494,7 @@ workerPool.returnObject(key, w1); // w1 gets destroyed when it is returned, so there are 0 idle workers. - assertThat(workerPool.getNumIdlePerKey(key)).isEqualTo(0); + assertThat(workerPool.getIdleWorkers()).isEmpty(); assertThat(workerPool.getNumActive(key)).isEqualTo(1); // Since w1 is already returned, it is killed on return. assertThat(w1.getStatus().get()).isEqualTo(Status.KILLED_DUE_TO_MEMORY_PRESSURE); @@ -473,9 +508,8 @@ @Test public void testGetEvictionCandidates_testDoomedAndIdleWorkers() throws Exception { - WorkerPoolImplLegacy workerPool = - new WorkerPoolImplLegacy( - new WorkerPoolConfig(factoryMock, entryList(DUMMY_MNEMONIC, 5), emptyEntryList())); + WorkerPool workerPool = + workerPoolSupplier.get(factoryMock, entryList(DUMMY_MNEMONIC, 5), emptyEntryList()); WorkerKey key = createWorkerKey(DUMMY_MNEMONIC, fileSystem); Worker w1 = workerPool.borrowObject(key); Worker w2 = workerPool.borrowObject(key); @@ -499,7 +533,7 @@ WorkerLifecycleManager manager = new WorkerLifecycleManager(workerPool, options); - assertThat(workerPool.getNumIdlePerKey(key)).isEqualTo(2); + assertThat(workerPool.getIdleWorkers()).hasSize(2); assertThat(workerPool.getNumActive(key)).isEqualTo(3); assertThat(w1.getStatus().isValid()).isTrue(); assertThat(w2.getStatus().isValid()).isTrue(); @@ -509,7 +543,7 @@ manager.evictWorkers(workerMetrics); - assertThat(workerPool.getNumIdlePerKey(key)).isEqualTo(0); + assertThat(workerPool.getIdleWorkers()).isEmpty(); assertThat(workerPool.getNumActive(key)).isEqualTo(3); // w1 and w2 are killed immediately. assertThat(w1.getStatus().get()).isEqualTo(Status.KILLED_DUE_TO_MEMORY_PRESSURE); @@ -522,9 +556,8 @@ @Test public void evictWorkers_testMultiplexWorkers() throws Exception { - WorkerPoolImplLegacy workerPool = - new WorkerPoolImplLegacy( - new WorkerPoolConfig(factoryMock, emptyEntryList(), entryList(DUMMY_MNEMONIC, 2))); + WorkerPool workerPool = + workerPoolSupplier.get(factoryMock, emptyEntryList(), entryList(DUMMY_MNEMONIC, 2)); WorkerKey key = createWorkerKey(DUMMY_MNEMONIC, fileSystem, /* multiplex= */ true, /* sandboxed= */ false); Worker w1 = workerPool.borrowObject(key); @@ -545,7 +578,7 @@ manager.evictWorkers(workerMetrics); - assertThat(workerPool.getNumIdlePerKey(key)).isEqualTo(0); + assertThat(workerPool.getIdleWorkers()).isEmpty(); assertThat(workerPool.getNumActive(key)).isEqualTo(0); // Since both w1 and w2 have been returned, it is killed. assertThat(w1.getStatus().get()).isEqualTo(Status.KILLED_DUE_TO_MEMORY_PRESSURE); @@ -553,9 +586,8 @@ @Test public void evictWorkers_doomMultiplexWorker() throws Exception { - WorkerPoolImplLegacy workerPool = - new WorkerPoolImplLegacy( - new WorkerPoolConfig(factoryMock, emptyEntryList(), entryList(DUMMY_MNEMONIC, 2))); + WorkerPool workerPool = + workerPoolSupplier.get(factoryMock, emptyEntryList(), entryList(DUMMY_MNEMONIC, 2)); WorkerKey key = createWorkerKey(DUMMY_MNEMONIC, fileSystem, /* multiplex= */ true, /* sandboxed= */ false); Worker w1 = workerPool.borrowObject(key); @@ -577,13 +609,13 @@ manager.evictWorkers(workerMetrics); // w1 should have been evicted already. - assertThat(workerPool.getNumIdlePerKey(key)).isEqualTo(0); + assertThat(workerPool.getIdleWorkers()).isEmpty(); assertThat(workerPool.getNumActive(key)).isEqualTo(1); // Not yet killed because w2 is still alive (and both share a WorkerProcessStatus). assertThat(w1.getStatus().get()).isEqualTo(Status.PENDING_KILL_DUE_TO_MEMORY_PRESSURE); workerPool.returnObject(key, w2); - assertThat(workerPool.getNumIdlePerKey(key)).isEqualTo(0); + assertThat(workerPool.getIdleWorkers()).isEmpty(); assertThat(workerPool.getNumActive(key)).isEqualTo(0); // Status is only set to killed after the last worker proxy is destroyed. assertThat(w2.getStatus().get()).isEqualTo(Status.KILLED_DUE_TO_MEMORY_PRESSURE);
diff --git a/src/test/java/com/google/devtools/build/lib/worker/WorkerPoolTest.java b/src/test/java/com/google/devtools/build/lib/worker/WorkerPoolTest.java index 3794a3e..18e037e 100644 --- a/src/test/java/com/google/devtools/build/lib/worker/WorkerPoolTest.java +++ b/src/test/java/com/google/devtools/build/lib/worker/WorkerPoolTest.java
@@ -15,40 +15,54 @@ import static com.google.common.truth.Truth.assertThat; import static com.google.devtools.build.lib.worker.WorkerTestUtils.createWorkerKey; +import static org.junit.Assert.assertThrows; +import static org.junit.Assume.assumeTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.spy; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; import com.google.common.collect.Maps; import com.google.devtools.build.lib.clock.BlazeClock; +import com.google.devtools.build.lib.testutil.TestThread; import com.google.devtools.build.lib.vfs.DigestHashFunction; import com.google.devtools.build.lib.vfs.FileSystem; import com.google.devtools.build.lib.vfs.Path; import com.google.devtools.build.lib.vfs.inmemoryfs.InMemoryFileSystem; import com.google.devtools.build.lib.worker.WorkerProcessStatus.Status; -import java.util.Map.Entry; -import org.apache.commons.pool2.PooledObject; -import org.apache.commons.pool2.impl.DefaultPooledObject; +import java.io.IOException; +import java.util.Arrays; +import java.util.List; +import java.util.function.Supplier; import org.junit.Before; -import org.junit.Rule; import org.junit.Test; import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; -import org.mockito.Mock; -import org.mockito.junit.MockitoJUnit; -import org.mockito.junit.MockitoRule; +import org.junit.runners.Parameterized; +import org.junit.runners.Parameterized.Parameter; +import org.junit.runners.Parameterized.Parameters; /** Tests WorkerPool. */ -@RunWith(JUnit4.class) +@RunWith(Parameterized.class) public class WorkerPoolTest { - @Rule public final MockitoRule mockito = MockitoJUnit.rule(); - @Mock WorkerFactory factoryMock; - private FileSystem fileSystem; + + public static final FileSystem fileSystem = + new InMemoryFileSystem(BlazeClock.instance(), DigestHashFunction.SHA256); + private int workerIds = 1; - private final WorkerOptions options = new WorkerOptions(); + @Parameter(0) + public WorkerPoolSupplier workerPoolSupplier; + + @Parameter(1) + public Supplier<WorkerFactory> workerFactorySupplier; + + private WorkerPool workerPool; + private WorkerFactory factoryMock; + + private static final WorkerOptions options = new WorkerOptions(); private static class TestWorker extends SingleplexWorker { TestWorker( @@ -57,68 +71,97 @@ } } + private interface WorkerPoolSupplier { + WorkerPool get(WorkerFactory factory); + } + + @Parameters + public static List<Object[]> data() throws IOException { + Supplier<WorkerFactory> workerFactorySupplier = + () -> spy(new WorkerFactory(fileSystem.getPath("/outputbase/bazel-workers"), options)); + + return Arrays.asList( + new Object[][] { + { + (WorkerPoolSupplier) + (factory) -> + new WorkerPoolImplLegacy( + new WorkerPoolConfig( + factory, + /* workerMaxInstances= */ ImmutableList.of( + Maps.immutableEntry("mnem", 2), Maps.immutableEntry("", 1)), + /* workerMaxMultiplexInstances= */ ImmutableList.of( + Maps.immutableEntry("mnem", 2), Maps.immutableEntry("", 1)))), + workerFactorySupplier, + }, + { + (WorkerPoolSupplier) + (factory) -> + new WorkerPoolImpl( + new WorkerPoolConfig( + factory, + /* workerMaxInstances= */ ImmutableList.of( + Maps.immutableEntry("mnem", 2)), + /* workerMaxMultiplexInstances= */ ImmutableList.of( + Maps.immutableEntry("mnem", 2)))), + workerFactorySupplier, + } + }); + } + @Before public void setUp() throws Exception { - fileSystem = new InMemoryFileSystem(BlazeClock.instance(), DigestHashFunction.SHA256); + factoryMock = workerFactorySupplier.get(); + workerPool = workerPoolSupplier.get(factoryMock); doAnswer( arg -> - new DefaultPooledObject<>( - new TestWorker( - arg.getArgument(0), - workerIds++, - fileSystem.getPath("/workDir"), - fileSystem.getPath("/logDir"), - options))) + new TestWorker( + arg.getArgument(0), + workerIds++, + fileSystem.getPath("/workDir"), + fileSystem.getPath("/logDir"), + options)) .when(factoryMock) - .makeObject(any()); + .create(any()); doAnswer( args -> { - PooledObject<Worker> obj = args.getArgument(1); - return obj.getObject().getStatus().isValid(); + Worker worker = args.getArgument(1); + return worker.getStatus().isValid(); }) .when(factoryMock) - .validateObject(any(), any()); + .validateWorker(any(), any()); doAnswer( args -> { - PooledObject<Worker> obj = args.getArgument(1); - obj.getObject().destroy(); + Worker worker = args.getArgument(1); + worker.destroy(); return null; }) .when(factoryMock) - .destroyObject(any(), any()); + .destroyWorker(any(), any()); } @Test public void testBorrow_createsWhenNeeded() throws Exception { - WorkerPool workerPool = - new WorkerPoolImplLegacy( - new WorkerPoolConfig(factoryMock, entryList("mnem", 2, "", 1), entryList())); WorkerKey workerKey = createWorkerKey(fileSystem, "mnem", false); Worker worker1 = workerPool.borrowObject(workerKey); Worker worker2 = workerPool.borrowObject(workerKey); assertThat(worker1.getWorkerId()).isEqualTo(1); assertThat(worker2.getWorkerId()).isEqualTo(2); - verify(factoryMock, times(2)).makeObject(workerKey); + verify(factoryMock, times(2)).create(workerKey); } @Test public void testBorrow_reusesWhenPossible() throws Exception { - WorkerPool workerPool = - new WorkerPoolImplLegacy( - new WorkerPoolConfig(factoryMock, entryList("mnem", 2, "", 1), entryList())); WorkerKey workerKey = createWorkerKey(fileSystem, "mnem", false); Worker worker1 = workerPool.borrowObject(workerKey); workerPool.returnObject(workerKey, worker1); Worker worker2 = workerPool.borrowObject(workerKey); assertThat(worker1).isSameInstanceAs(worker2); - verify(factoryMock, times(1)).makeObject(workerKey); + verify(factoryMock).create(workerKey); } @Test - public void testBorrow_usesDefault() throws Exception { - WorkerPool workerPool = - new WorkerPoolImplLegacy( - new WorkerPoolConfig(factoryMock, entryList("mnem", 2, "", 1), entryList())); + public void testBorrow_nonSpecifiedKey() throws Exception { WorkerKey workerKey1 = createWorkerKey(fileSystem, "mnem", false); Worker worker1 = workerPool.borrowObject(workerKey1); Worker worker1a = workerPool.borrowObject(workerKey1); @@ -127,15 +170,12 @@ WorkerKey workerKey2 = createWorkerKey(fileSystem, "other", false); Worker worker2 = workerPool.borrowObject(workerKey2); assertThat(worker2.getWorkerId()).isEqualTo(3); - verify(factoryMock, times(2)).makeObject(workerKey1); - verify(factoryMock, times(1)).makeObject(workerKey2); + verify(factoryMock, times(2)).create(workerKey1); + verify(factoryMock).create(workerKey2); } @Test public void testBorrow_pooledByKey() throws Exception { - WorkerPool workerPool = - new WorkerPoolImplLegacy( - new WorkerPoolConfig(factoryMock, entryList("mnem", 2, "", 1), entryList())); WorkerKey workerKey1 = createWorkerKey(fileSystem, "mnem", false); Worker worker1 = workerPool.borrowObject(workerKey1); Worker worker1a = workerPool.borrowObject(workerKey1); @@ -144,16 +184,12 @@ WorkerKey workerKey2 = createWorkerKey(fileSystem, "mnem", false, "arg1"); Worker worker2 = workerPool.borrowObject(workerKey2); assertThat(worker2.getWorkerId()).isEqualTo(3); - verify(factoryMock, times(2)).makeObject(workerKey1); - verify(factoryMock, times(1)).makeObject(workerKey2); + verify(factoryMock, times(2)).create(workerKey1); + verify(factoryMock).create(workerKey2); } @Test public void testBorrow_separateMultiplexWorkers() throws Exception { - WorkerPool workerPool = - new WorkerPoolImplLegacy( - new WorkerPoolConfig( - factoryMock, entryList("mnem", 1, "", 1), entryList("mnem", 2, "", 1))); WorkerKey workerKey = createWorkerKey(fileSystem, "mnem", false); Worker worker1 = workerPool.borrowObject(workerKey); assertThat(worker1.getWorkerId()).isEqualTo(1); @@ -168,15 +204,12 @@ assertThat(multiplexWorker2.getWorkerId()).isEqualTo(3); assertThat(worker1a.getWorkerId()).isEqualTo(1); - verify(factoryMock, times(1)).makeObject(workerKey); - verify(factoryMock, times(2)).makeObject(multiplexKey); + verify(factoryMock).create(workerKey); + verify(factoryMock, times(2)).create(multiplexKey); } @Test public void testBorrow_doomedWorkers() throws Exception { - WorkerPool workerPool = - new WorkerPoolImplLegacy( - new WorkerPoolConfig(factoryMock, entryList("mnem", 2, "", 1), entryList())); WorkerKey workerKey = createWorkerKey(fileSystem, "mnem", false); Worker worker1 = workerPool.borrowObject(workerKey); Worker worker2 = workerPool.borrowObject(workerKey); @@ -192,12 +225,180 @@ assertThat(worker2.getStatus().isKilled()).isFalse(); } - private static ImmutableList<Entry<String, Integer>> entryList() { - return ImmutableList.of(); + @Test + public void testBorrow_blocksWhenUnavailable() throws Exception { + WorkerKey workerKey = createWorkerKey(fileSystem, "mnem", false); + Worker unused1 = workerPool.borrowObject(workerKey); + Worker unused2 = workerPool.borrowObject(workerKey); + TestThread blockedBorrowThread = + new TestThread( + () -> { + Worker unused = workerPool.borrowObject(workerKey); + }); + blockedBorrowThread.start(); + + AssertionError e = + assertThrows(AssertionError.class, () -> blockedBorrowThread.joinAndAssertState(1000)); + assertThat(e).hasCauseThat().hasMessageThat().contains("is still alive"); + assertThat(workerPool.getNumActive(workerKey)).isEqualTo(2); } - private static ImmutableList<Entry<String, Integer>> entryList( - String key1, int value1, String key2, int value2) { - return ImmutableList.of(Maps.immutableEntry(key1, value1), Maps.immutableEntry(key2, value2)); + @Test + public void testBorrow_blockedThread_getsReturnedWorker() throws Exception { + WorkerKey workerKey = createWorkerKey(fileSystem, "mnem", false); + Worker worker1 = workerPool.borrowObject(workerKey); + Worker unused2 = workerPool.borrowObject(workerKey); + TestThread blockedBorrowThread = + new TestThread( + () -> { + // This blocks until worker1 returns its object. + Worker worker = workerPool.borrowObject(workerKey); + assertThat(worker).isSameInstanceAs(worker1); + }); + blockedBorrowThread.start(); + + // We want to 3rd borrow to be blocked for some time. + Thread.sleep(500); + workerPool.returnObject(worker1.getWorkerKey(), worker1); + + blockedBorrowThread.joinAndAssertState(10000); + assertThat(workerPool.getNumActive(workerKey)).isEqualTo(2); + } + + @Test + public void testBorrow_blockedThread_createsWorkerWhenInvalidated() throws Exception { + WorkerKey workerKey = createWorkerKey(fileSystem, "mnem", false); + Worker worker1 = workerPool.borrowObject(workerKey); + Worker unused2 = workerPool.borrowObject(workerKey); + TestThread blockedBorrowThread = + new TestThread( + () -> { + Worker worker = workerPool.borrowObject(workerKey); + // Create a new worker instead. + assertThat(worker.getWorkerId()).isEqualTo(3); + }); + blockedBorrowThread.start(); + + // We want to 3rd borrow to be blocked for some time. + Thread.sleep(500); + workerPool.invalidateObject(worker1.getWorkerKey(), worker1); + + blockedBorrowThread.joinAndAssertState(10000); + assertThat(workerPool.getNumActive(workerKey)).isEqualTo(2); + } + + @Test + public void testBorrow_blockedThread_remainsBlockedWhenInvalidatedAndShrunk() throws Exception { + // This is meant solely for WorkerPoolImpl; WorkerPoolImplLegacy will block only after some + // time because GenericKeyedObjectPool only implements some + assumeTrue(workerPool instanceof WorkerPoolImpl); + WorkerKey workerKey = createWorkerKey(fileSystem, "mnem", false); + Worker worker1 = workerPool.borrowObject(workerKey); + Worker unused2 = workerPool.borrowObject(workerKey); + TestThread blockedBorrowThread = + new TestThread( + () -> { + Worker unused = workerPool.borrowObject(workerKey); + }); + blockedBorrowThread.start(); + + // There's no need to wait here as it doesn't matter whether #invalidateObject gets called + // before or after the 3rd #borrowObject, the pool would not have the quota and borrowing will + // still get blocked. + worker1.getStatus().maybeUpdateStatus(Status.PENDING_KILL_DUE_TO_MEMORY_PRESSURE); + workerPool.invalidateObject(worker1.getWorkerKey(), worker1); + + AssertionError e = + assertThrows(AssertionError.class, () -> blockedBorrowThread.joinAndAssertState(1000)); + assertThat(e).hasCauseThat().hasMessageThat().contains("is still alive"); + assertThat(workerPool.getNumActive(workerKey)).isEqualTo(1); + } + + @Test + public void testEvict_evictsIdleWorkers() throws Exception { + WorkerKey workerKey = createWorkerKey(fileSystem, "mnem", false); + Worker worker1 = workerPool.borrowObject(workerKey); + Worker worker2 = workerPool.borrowObject(workerKey); + workerPool.returnObject(workerKey, worker1); + workerPool.returnObject(workerKey, worker2); + ImmutableSet<Integer> evicted = + workerPool.evictWorkers(ImmutableSet.of(worker1.getWorkerId(), worker2.getWorkerId())); + assertThat(evicted).containsExactly(worker1.getWorkerId(), worker2.getWorkerId()); + assertThat(workerPool.getNumActive(workerKey)).isEqualTo(0); + } + + @Test + public void testEvict_doesNotEvictActiveWorkers() throws Exception { + WorkerKey workerKey = createWorkerKey(fileSystem, "mnem", false); + Worker worker1 = workerPool.borrowObject(workerKey); + Worker worker2 = workerPool.borrowObject(workerKey); + workerPool.returnObject(workerKey, worker1); + ImmutableSet<Integer> evicted = + workerPool.evictWorkers(ImmutableSet.of(worker1.getWorkerId(), worker2.getWorkerId())); + // Worker2 does not get evicted because it is still active. + assertThat(evicted).containsExactly(worker1.getWorkerId()); + assertThat(workerPool.getNumActive(workerKey)).isEqualTo(1); + } + + @Test + public void testGetIdleWorkers() throws Exception { + WorkerKey workerKey = createWorkerKey(fileSystem, "mnem", false); + Worker worker1 = workerPool.borrowObject(workerKey); + Worker worker2 = workerPool.borrowObject(workerKey); + + assertThat(workerPool.getIdleWorkers()).isEmpty(); + workerPool.returnObject(workerKey, worker1); + workerPool.returnObject(workerKey, worker2); + + assertThat(workerPool.getIdleWorkers()) + .containsExactly(worker1.getWorkerId(), worker2.getWorkerId()); + assertThat(workerPool.getNumActive(workerKey)).isEqualTo(0); + } + + @Test + public void testShrinkingPool_doesNotShrinkBelowOneWorker() throws Exception { + WorkerKey workerKey = createWorkerKey(fileSystem, "mnem", false); + assertThat(workerPool.getMaxTotalPerKey(workerKey)).isEqualTo(2); + + Worker worker1 = workerPool.borrowObject(workerKey); + // Shrink the worker pool by 1. + worker1.getStatus().maybeUpdateStatus(Status.PENDING_KILL_DUE_TO_MEMORY_PRESSURE); + workerPool.returnObject(workerKey, worker1); + assertThat(workerPool.getMaxTotalPerKey(workerKey)).isEqualTo(1); + + Worker worker2 = workerPool.borrowObject(workerKey); + // Attempt to shrink the pool again. + worker2.getStatus().maybeUpdateStatus(Status.PENDING_KILL_DUE_TO_MEMORY_PRESSURE); + workerPool.returnObject(workerKey, worker2); + // It should not be shrunk below 1. + assertThat(workerPool.getMaxTotalPerKey(workerKey)).isEqualTo(1); + assertThat(workerPool.getNumActive(workerKey)).isEqualTo(0); + } + + @Test + public void testGetNumActive() throws Exception { + WorkerKey workerKey = createWorkerKey(fileSystem, "mnem", false); + assertThat(workerPool.getNumActive(workerKey)).isEqualTo(0); + Worker worker1 = workerPool.borrowObject(workerKey); + Worker worker2 = workerPool.borrowObject(workerKey); + assertThat(workerPool.getNumActive(workerKey)).isEqualTo(2); + workerPool.returnObject(workerKey, worker1); + workerPool.returnObject(workerKey, worker2); + assertThat(workerPool.getNumActive(workerKey)).isEqualTo(0); + } + + @Test + public void testReset_removesPreviouslyShrunkValues() throws Exception { + WorkerKey workerKey = createWorkerKey(fileSystem, "mnem", false); + assertThat(workerPool.getMaxTotalPerKey(workerKey)).isEqualTo(2); + + Worker worker1 = workerPool.borrowObject(workerKey); + // Shrink the worker pool by 1. + worker1.getStatus().maybeUpdateStatus(Status.PENDING_KILL_DUE_TO_MEMORY_PRESSURE); + workerPool.returnObject(workerKey, worker1); + assertThat(workerPool.getMaxTotalPerKey(workerKey)).isEqualTo(1); + + workerPool.reset(); + assertThat(workerPool.getMaxTotalPerKey(workerKey)).isEqualTo(2); } }