diff --git a/src/main/java/com/google/devtools/build/lib/buildeventservice/BuildEventServiceModule.java b/src/main/java/com/google/devtools/build/lib/buildeventservice/BuildEventServiceModule.java index 92b215e3e36281..06c113dfc73ea9 100644 --- a/src/main/java/com/google/devtools/build/lib/buildeventservice/BuildEventServiceModule.java +++ b/src/main/java/com/google/devtools/build/lib/buildeventservice/BuildEventServiceModule.java @@ -23,6 +23,7 @@ import com.google.common.collect.ImmutableSet; import com.google.common.collect.Iterables; import com.google.common.flogger.GoogleLogger; +import com.google.common.util.concurrent.ForwardingListenableFuture.SimpleForwardingListenableFuture; import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.MoreExecutors; @@ -533,14 +534,26 @@ private void waitForBuildEventTransportsToClose( final ListenableFuture enclosingFuture = Futures.nonCancellationPropagating(closeFuture); - closeFutureWithTimeout = + ListenableFuture timeoutFuture = Futures.withTimeout( enclosingFuture, bepTransport.getTimeout().toMillis(), TimeUnit.MILLISECONDS, timeoutExecutor); - closeFutureWithTimeout.addListener( - timeoutExecutor::shutdown, MoreExecutors.directExecutor()); + timeoutFuture.addListener(timeoutExecutor::shutdown, MoreExecutors.directExecutor()); + + // Cancellation is not propagated to the `closeFuture` for the reasons above. But in + // order to cancel the returned future by our explicit mechanism elsewhere in this + // class, we need to delegate the `cancel` to `closeFuture` so that cancellation + // from Futures.withTimeout is ignored and cancellation from our mechanism is properly + // handled. + closeFutureWithTimeout = + new SimpleForwardingListenableFuture<>(timeoutFuture) { + @Override + public boolean cancel(boolean mayInterruptIfRunning) { + return closeFuture.cancel(mayInterruptIfRunning); + } + }; } builder.put(bepTransport, closeFutureWithTimeout); }); diff --git a/src/test/java/com/google/devtools/build/lib/buildeventservice/BazelBuildEventServiceModuleTest.java b/src/test/java/com/google/devtools/build/lib/buildeventservice/BazelBuildEventServiceModuleTest.java index ebe29d8c7dbe08..7249671516423d 100644 --- a/src/test/java/com/google/devtools/build/lib/buildeventservice/BazelBuildEventServiceModuleTest.java +++ b/src/test/java/com/google/devtools/build/lib/buildeventservice/BazelBuildEventServiceModuleTest.java @@ -24,6 +24,7 @@ import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; import com.google.common.collect.Iterables; import com.google.common.collect.Sets; import com.google.common.util.concurrent.Uninterruptibles; @@ -317,8 +318,13 @@ public void testAfterCommand_waitForUploadComplete_slowFullCloseError() throws E "--bes_backend=inprocess", "--bes_upload_mode=WAIT_FOR_UPLOAD_COMPLETE", "--bes_timeout=5s"); + ImmutableSet bepTransports = besModule.getBepTransports(); + assertThat(bepTransports).hasSize(1); afterBuildCommand(); assertContainsError("The Build Event Protocol upload timed out"); + for (BuildEventTransport bepTransport : bepTransports) { + assertThat(bepTransport.close().isDone()).isTrue(); + } } @Test