Fix that bepTransports are not properly closed if --bes_timeout is set.

Related bazelbuild#14576.

PiperOrigin-RevId: 424026771

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<Void> enclosingFuture =
  Futures.nonCancellationPropagating(closeFuture);
 
- closeFutureWithTimeout =
+ ListenableFuture<Void> 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);
  });

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<BuildEventTransport> 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