From a1db1ed6397ec682504de08a2a187ce0661a0e4c Mon Sep 17 00:00:00 2001
From: felly <felly@google.com>
Date: Wed, 24 Jul 2019 19:05:52 -0700
Subject: [PATCH] Fix bug where we were cancelling in-flight async uploads in
 sync/async mixed modes.

RELNOTES: None
PiperOrigin-RevId: 259865252
---
 .../BuildEventServiceModule.java                | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

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 3bc25730b26cfe..290b886b894411 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
@@ -33,6 +33,7 @@
 import com.google.common.util.concurrent.Uninterruptibles;
 import com.google.devtools.build.lib.analysis.test.TestConfiguration.TestOptions;
 import com.google.devtools.build.lib.authandtls.AuthAndTLSOptions;
+import com.google.devtools.build.lib.buildeventservice.BuildEventServiceOptions.BesUploadMode;
 import com.google.devtools.build.lib.buildeventservice.client.BuildEventServiceClient;
 import com.google.devtools.build.lib.buildeventstream.AnnounceBuildEventTransportsEvent;
 import com.google.devtools.build.lib.buildeventstream.BuildEventArtifactUploader;
@@ -425,7 +426,8 @@ public void blazeShutdown() {
   }
 
   private void waitForBuildEventTransportsToClose(
-      Map<BuildEventTransport, ListenableFuture<Void>> transportFutures)
+      Map<BuildEventTransport, ListenableFuture<Void>> transportFutures,
+      boolean besUploadModeIsSynchronous)
       throws AbruptExitException {
     final ScheduledExecutorService executor =
         Executors.newSingleThreadScheduledExecutor(
@@ -461,7 +463,9 @@ private void waitForBuildEventTransportsToClose(
               "Unexpected Exception '%s' when closing BEP transports, this is a bug.",
               e.getCause().getMessage()));
     } finally {
-      cancelAndResetPendingUploads();
+      if (besUploadModeIsSynchronous) {
+        cancelAndResetPendingUploads();
+      }
       if (waitMessageFuture != null) {
         waitMessageFuture.cancel(/* mayInterruptIfRunning= */ true);
       }
@@ -514,14 +518,13 @@ private void closeBepTransports() throws AbruptExitException {
         constructCloseFuturesMapWithTimeouts(streamer.getCloseFuturesMap());
     halfCloseFuturesWithTimeoutsMap =
         constructCloseFuturesMapWithTimeouts(streamer.getHalfClosedMap());
-
+    boolean besUploadModeIsSynchronous =
+        besOptions.besUploadMode == BesUploadMode.WAIT_FOR_UPLOAD_COMPLETE;
     Map<BuildEventTransport, ListenableFuture<Void>> blockingTransportFutures = new HashMap<>();
     for (Map.Entry<BuildEventTransport, ListenableFuture<Void>> entry :
         closeFuturesWithTimeoutsMap.entrySet()) {
       BuildEventTransport bepTransport = entry.getKey();
-      if (!bepTransport.mayBeSlow()
-          || besOptions.besUploadMode
-              == BuildEventServiceOptions.BesUploadMode.WAIT_FOR_UPLOAD_COMPLETE) {
+      if (!bepTransport.mayBeSlow() || besUploadModeIsSynchronous) {
         blockingTransportFutures.put(bepTransport, entry.getValue());
       } else {
         // When running asynchronously notify the UI immediately since we won't wait for the
@@ -530,7 +533,7 @@ private void closeBepTransports() throws AbruptExitException {
       }
     }
     if (!blockingTransportFutures.isEmpty()) {
-      waitForBuildEventTransportsToClose(blockingTransportFutures);
+      waitForBuildEventTransportsToClose(blockingTransportFutures, besUploadModeIsSynchronous);
     }
   }