From 621649d5f927d5087bfc7f8a9b82f929723fc9cd Mon Sep 17 00:00:00 2001 From: Chi Wang Date: Wed, 26 Jan 2022 02:14:29 -0800 Subject: [PATCH] Remote: Postpone the block waiting in `afterCommand` to `BlockWaitingModule` When implementing async upload, we introduced a block waiting behaviour in `RemoteModule#afterCommand` so that uploads happened in the background can be waited before the whole build complete. However, there are other block waiting code in other module's `afterCommand` method (e.g. BES module). Block waiting in remote module will prevent other modules' `afterCommand` from executing until it's completed. This causes issues like #14576. This PR adds a new module `BlockWaitingModule` whose sole purpose is to accept tasks submitted by other modules in `afterCommand` and block waiting all the tasks in its own `afterCommand` method. So those tasks can be executed in parallel. This PR only updates RemoteModule's `afterCommand` method to submit block waiting task. Other modules should be updated to use `BlockWaitingModule` as well but that's beyond the scope this this PR. This PR along with 73a76a8af802f721f9de0f01026c50c4e72b2d35 fix #14576. Closes #14618. PiperOrigin-RevId: 424295121 --- .../devtools/build/lib/bazel/Bazel.java | 5 +- .../remote/RemoteActionContextProvider.java | 5 +- .../build/lib/runtime/BlockWaitingModule.java | 59 +++++++++++++++++++ 3 files changed, 67 insertions(+), 2 deletions(-) create mode 100644 src/main/java/com/google/devtools/build/lib/runtime/BlockWaitingModule.java diff --git a/src/main/java/com/google/devtools/build/lib/bazel/Bazel.java b/src/main/java/com/google/devtools/build/lib/bazel/Bazel.java index 54180cf4f5a6f5..68d0fbee7b942e 100644 --- a/src/main/java/com/google/devtools/build/lib/bazel/Bazel.java +++ b/src/main/java/com/google/devtools/build/lib/bazel/Bazel.java @@ -80,7 +80,10 @@ public final class Bazel { com.google.devtools.build.lib.packages.metrics.PackageMetricsModule.class, com.google.devtools.build.lib.metrics.MetricsModule.class, BazelBuiltinCommandModule.class, - com.google.devtools.build.lib.includescanning.IncludeScanningModule.class); + com.google.devtools.build.lib.includescanning.IncludeScanningModule.class, + // This module needs to be registered after any module submitting tasks with its {@code + // submit} method. + com.google.devtools.build.lib.runtime.BlockWaitingModule.class); public static void main(String[] args) { BlazeVersionInfo.setBuildInfo(tryGetBuildInfo()); diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteActionContextProvider.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteActionContextProvider.java index ba37cbba0c390e..d382364cd1fdc9 100644 --- a/src/main/java/com/google/devtools/build/lib/remote/RemoteActionContextProvider.java +++ b/src/main/java/com/google/devtools/build/lib/remote/RemoteActionContextProvider.java @@ -31,6 +31,7 @@ import com.google.devtools.build.lib.remote.common.RemotePathResolver.SiblingRepositoryLayoutResolver; import com.google.devtools.build.lib.remote.options.RemoteOptions; import com.google.devtools.build.lib.remote.util.DigestUtil; +import com.google.devtools.build.lib.runtime.BlockWaitingModule; import com.google.devtools.build.lib.runtime.CommandEnvironment; import com.google.devtools.build.lib.vfs.Path; import java.util.concurrent.Executor; @@ -211,7 +212,9 @@ void setFilesToDownload(ImmutableSet topLevelOutputs) { public void afterCommand() { if (remoteExecutionService != null) { - remoteExecutionService.shutdown(); + BlockWaitingModule blockWaitingModule = + checkNotNull(env.getRuntime().getBlazeModule(BlockWaitingModule.class)); + blockWaitingModule.submit(() -> remoteExecutionService.shutdown()); } else { if (remoteCache != null) { remoteCache.release(); diff --git a/src/main/java/com/google/devtools/build/lib/runtime/BlockWaitingModule.java b/src/main/java/com/google/devtools/build/lib/runtime/BlockWaitingModule.java new file mode 100644 index 00000000000000..6d5b74a8f51e88 --- /dev/null +++ b/src/main/java/com/google/devtools/build/lib/runtime/BlockWaitingModule.java @@ -0,0 +1,59 @@ +// Copyright 2022 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.runtime; + +import static com.google.common.base.Preconditions.checkNotNull; +import static com.google.common.base.Preconditions.checkState; +import static java.util.concurrent.TimeUnit.SECONDS; + +import com.google.common.util.concurrent.ThreadFactoryBuilder; +import com.google.devtools.build.lib.util.AbruptExitException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import javax.annotation.Nullable; + +/** A {@link BlazeModule} that waits for submitted tasks to terminate after every command. */ +public class BlockWaitingModule extends BlazeModule { + @Nullable private ExecutorService executorService; + + @Override + public void beforeCommand(CommandEnvironment env) throws AbruptExitException { + checkState(executorService == null, "executorService must be null"); + + executorService = + Executors.newCachedThreadPool( + new ThreadFactoryBuilder().setNameFormat("block-waiting-%d").build()); + } + + @SuppressWarnings("FutureReturnValueIgnored") + public void submit(Runnable task) { + checkNotNull(executorService, "executorService must not be null"); + + executorService.submit(task); + } + + @Override + public void afterCommand() throws AbruptExitException { + checkNotNull(executorService, "executorService must not be null"); + + executorService.shutdown(); + try { + executorService.awaitTermination(Long.MAX_VALUE, SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + + executorService = null; + } +}