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 e5ec69136256e5..bb44b7d9108e4e 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 @@ -76,7 +76,10 @@ public final class Bazel { com.google.devtools.build.lib.metrics.PostGCMemoryUseRecorder.GcAfterBuildModule.class, com.google.devtools.build.lib.packages.metrics.PackageMetricsModule.class, com.google.devtools.build.lib.metrics.MetricsModule.class, - BazelBuiltinCommandModule.class); + BazelBuiltinCommandModule.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; + } +}