From 2fffa0e9359040dbeee05b1eeff7fed1927a34ec Mon Sep 17 00:00:00 2001 From: Fabian Meumertzheim Date: Fri, 18 Sep 2026 11:04:19 +0200 Subject: [PATCH] Close the GrpcRemoteDownloader's channel at the end of a command `RemoteModule` never closed the `GrpcRemoteDownloader` it created for a command, so the reference it held on its gRPC channel was leaked and the channel (a dedicated one for `--remote_downloader`, or the one shared with `--remote_cache`) was never shut down. --- .../build/lib/remote/RemoteModule.java | 20 ++++- .../downloader/GrpcRemoteDownloader.java | 6 +- .../build/lib/remote/RemoteModuleTest.java | 89 ++++++++++++++++++- 3 files changed, 110 insertions(+), 5 deletions(-) diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteModule.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteModule.java index 4f62d47a559541..f273c75a9632b9 100644 --- a/src/main/java/com/google/devtools/build/lib/remote/RemoteModule.java +++ b/src/main/java/com/google/devtools/build/lib/remote/RemoteModule.java @@ -186,7 +186,7 @@ public ManagedChannel newChannel( private final RepositoryRemoteHelpersFactoryDelegate repositoryRemoteHelpersFactoryDelegate = new RepositoryRemoteHelpersFactoryDelegate(); - private Downloader remoteDownloader; + @Nullable private GrpcRemoteDownloader remoteDownloader; private CredentialModule credentialModule; @@ -1130,11 +1130,20 @@ public void afterCommand() { // Some cleanup tasks must wait until every other BlazeModule's afterCommand() has run, as // otherwise we might interfere with asynchronous remote downloads that are in progress. RemoteActionContextProvider actionContextProviderRef = actionContextProvider; + GrpcRemoteDownloader remoteDownloaderRef = remoteDownloader; TempPathGenerator tempPathGeneratorRef = tempPathGenerator; AsynchronousMessageOutputStream rpcLogFileRef = rpcLogFile; - if (actionContextProviderRef != null || tempPathGeneratorRef != null || rpcLogFileRef != null) { + if (actionContextProviderRef != null + || remoteDownloaderRef != null + || tempPathGeneratorRef != null + || rpcLogFileRef != null) { blockWaitingModule.submit( - () -> afterCommandTask(actionContextProviderRef, tempPathGeneratorRef, rpcLogFileRef)); + () -> + afterCommandTask( + actionContextProviderRef, + remoteDownloaderRef, + tempPathGeneratorRef, + rpcLogFileRef)); } lastRemoteOutputChecker = remoteOutputChecker; @@ -1159,6 +1168,7 @@ public void afterCommand() { private static void afterCommandTask( @Nullable RemoteActionContextProvider actionContextProvider, + @Nullable GrpcRemoteDownloader remoteDownloader, @Nullable TempPathGenerator tempPathGenerator, @Nullable AsynchronousMessageOutputStream rpcLogFile) throws AbruptExitException { @@ -1166,6 +1176,10 @@ private static void afterCommandTask( actionContextProvider.afterCommand(); } + if (remoteDownloader != null) { + remoteDownloader.close(); + } + if (tempPathGenerator != null) { Path tempDir = tempPathGenerator.getTempDir(); try { diff --git a/src/main/java/com/google/devtools/build/lib/remote/downloader/GrpcRemoteDownloader.java b/src/main/java/com/google/devtools/build/lib/remote/downloader/GrpcRemoteDownloader.java index a66fc191f02f40..647a0234ed9a8c 100644 --- a/src/main/java/com/google/devtools/build/lib/remote/downloader/GrpcRemoteDownloader.java +++ b/src/main/java/com/google/devtools/build/lib/remote/downloader/GrpcRemoteDownloader.java @@ -125,12 +125,16 @@ public GrpcRemoteDownloader( this.remoteDownloaderLocalFallback = remoteDownloaderLocalFallback; } + /** + * Releases the reference to the channel held by this downloader. + * + *

The {@link RemoteCacheClient} is owned by the caller and thus not closed here. + */ @Override public void close() { if (closed.getAndSet(true)) { return; } - cacheClient.close(); channel.release(); } diff --git a/src/test/java/com/google/devtools/build/lib/remote/RemoteModuleTest.java b/src/test/java/com/google/devtools/build/lib/remote/RemoteModuleTest.java index ce16e388136861..7b3cbe7a5d193a 100644 --- a/src/test/java/com/google/devtools/build/lib/remote/RemoteModuleTest.java +++ b/src/test/java/com/google/devtools/build/lib/remote/RemoteModuleTest.java @@ -80,6 +80,7 @@ import com.google.devtools.common.options.OptionsParsingResult; import com.google.errorprone.annotations.CanIgnoreReturnValue; import io.grpc.BindableService; +import io.grpc.ManagedChannel; import io.grpc.Server; import io.grpc.ServerInterceptors; import io.grpc.inprocess.InProcessChannelBuilder; @@ -90,6 +91,7 @@ import java.net.URI; import java.time.Duration; import java.util.ArrayList; +import java.util.List; import java.util.HashMap; import java.util.Map; import java.util.Optional; @@ -104,6 +106,7 @@ public final class RemoteModuleTest { private static final String EXECUTION_SERVER_NAME = "execution-server"; private static final String CACHE_SERVER_NAME = "cache-server"; + private static final String DOWNLOADER_SERVER_NAME = "downloader-server"; private static final String OUTPUT_SERVICE_SERVER_NAME = "output-service"; private static final ServerCapabilities CACHE_ONLY_CAPS = ServerCapabilities.newBuilder() @@ -261,15 +264,20 @@ private static RemoteOptions parseRemoteOptions(String... args) throws Exception private RemoteModule remoteModule; private RemoteOptions remoteOptions; private Map> serviceConfigsByTarget; + private List createdChannels; @Before public void initialize() { serviceConfigsByTarget = new HashMap<>(); + createdChannels = new ArrayList<>(); remoteModule = new RemoteModule(); remoteModule.setChannelFactory( (target, proxy, options, interceptors, serviceConfig) -> { serviceConfigsByTarget.put(target, serviceConfig); - return InProcessChannelBuilder.forName(target).directExecutor().build(); + ManagedChannel channel = + InProcessChannelBuilder.forName(target).directExecutor().build(); + createdChannels.add(channel); + return channel; }); remoteOptions = Options.getDefaults(RemoteOptions.class); } @@ -738,11 +746,90 @@ private CommandEnvironment beforeCommand(WorkspaceInitializer workspaceInitializ throws IOException, AbruptExitException { CommandEnvironment env = createTestCommandEnvironment(remoteModule, remoteOptions, workspaceInitializer); + env.getRuntime().getBlazeModule(BlockWaitingModule.class).beforeCommand(env); remoteModule.beforeCommand(env); env.throwPendingException(); return env; } + /** Runs the remote module's after-command cleanup and waits for it to complete. */ + private void afterCommand(CommandEnvironment env) throws AbruptExitException { + remoteModule.afterCommand(); + env.getRuntime().getBlazeModule(BlockWaitingModule.class).afterCommand(); + } + + /** Waits for the eagerly created channels of the current command to be connected. */ + private void awaitChannelsConnected() throws Exception { + var combinedCache = remoteModule.getActionContextProvider().getCombinedCache(); + if (combinedCache != null) { + var unused = combinedCache.getRemoteCacheCapabilities(); + } + if (remoteModule.getRemoteDownloader() instanceof GrpcRemoteDownloader downloader) { + var unused = downloader.getChannel().withChannelBlocking(ch -> new Object()); + } + } + + @Test + public void remoteDownloader_separateEndpoint_channelsAreClosedAfterCommand() throws Exception { + Server cacheServer = createFakeServer(CACHE_SERVER_NAME, new CapabilitiesImpl(CACHE_ONLY_CAPS)); + cacheServer.start(); + Server downloaderServer = createFakeServer(DOWNLOADER_SERVER_NAME); + downloaderServer.start(); + + try { + remoteOptions.setRemoteCache(CACHE_SERVER_NAME); + remoteOptions.setRemoteDownloader(DOWNLOADER_SERVER_NAME); + + var env = beforeCommand(); + awaitChannelsConnected(); + assertThat(createdChannels).hasSize(2); + + afterCommand(env); + + for (ManagedChannel channel : createdChannels) { + assertThat(channel.isTerminated()).isTrue(); + } + } finally { + cacheServer.shutdownNow(); + downloaderServer.shutdownNow(); + cacheServer.awaitTermination(); + downloaderServer.awaitTermination(); + } + } + + @Test + public void remoteDownloader_sharedEndpoint_channelIsClosedOnceCacheIsReleased() + throws Exception { + Server cacheServer = createFakeServer(CACHE_SERVER_NAME, new CapabilitiesImpl(CACHE_ONLY_CAPS)); + cacheServer.start(); + + try { + remoteOptions.setRemoteCache(CACHE_SERVER_NAME); + remoteOptions.setRemoteDownloader(CACHE_SERVER_NAME); + + var env = beforeCommand(); + awaitChannelsConnected(); + assertThat(createdChannels).hasSize(1); + ManagedChannel channel = createdChannels.get(0); + + // Retain the cache beyond the end of the command, as e.g. the artifact uploader of an + // asynchronous BES upload does. + var combinedCache = remoteModule.getActionContextProvider().getCombinedCache(); + combinedCache.retain(); + afterCommand(env); + + // Closing the downloader must not shut down the channel it shares with the cache. + assertThat(channel.isShutdown()).isFalse(); + + combinedCache.release(); + + assertThat(channel.isTerminated()).isTrue(); + } finally { + cacheServer.shutdownNow(); + cacheServer.awaitTermination(); + } + } + @Test public void diskCache_defaultLocation_resolvesToOutputUserRoot() throws Exception { remoteOptions.setDiskCache(PathFragment.EMPTY_FRAGMENT);