From 3ed9d3681c3e130aafcf3c405ff1795c601bdf95 Mon Sep 17 00:00:00 2001
From: felly <felly@google.com>
Date: Mon, 3 Jun 2019 13:22:36 -0700
Subject: [PATCH] Properly wire up BES half-close. The lack thereof was a
 simple oversight.

With this change, --bes_upload_mode=fully_async works properly.

RELNOTES: None
PiperOrigin-RevId: 251295660
---
 .../buildeventservice/BuildEventServiceTransport.java    | 5 +++++
 .../lib/buildeventservice/BuildEventServiceUploader.java | 9 +++++++++
 2 files changed, 14 insertions(+)

diff --git a/src/main/java/com/google/devtools/build/lib/buildeventservice/BuildEventServiceTransport.java b/src/main/java/com/google/devtools/build/lib/buildeventservice/BuildEventServiceTransport.java
index 88797dd237819c..2c576f95e6ee02 100644
--- a/src/main/java/com/google/devtools/build/lib/buildeventservice/BuildEventServiceTransport.java
+++ b/src/main/java/com/google/devtools/build/lib/buildeventservice/BuildEventServiceTransport.java
@@ -67,6 +67,11 @@ public ListenableFuture<Void> close() {
     return besUploader.close();
   }
 
+  @Override
+  public ListenableFuture<Void> getHalfCloseFuture() {
+    return besUploader.getHalfCloseFuture();
+  }
+
   @Override
   public BuildEventArtifactUploader getUploader() {
     return besUploader.getLocalFileUploader();
diff --git a/src/main/java/com/google/devtools/build/lib/buildeventservice/BuildEventServiceUploader.java b/src/main/java/com/google/devtools/build/lib/buildeventservice/BuildEventServiceUploader.java
index 66ce615e2588e9..3f4d8e8bc3080e 100644
--- a/src/main/java/com/google/devtools/build/lib/buildeventservice/BuildEventServiceUploader.java
+++ b/src/main/java/com/google/devtools/build/lib/buildeventservice/BuildEventServiceUploader.java
@@ -154,6 +154,10 @@ private BuildEventServiceUploader(
     this.clock = clock;
     this.namer = namer;
     this.eventBus = eventBus;
+    // Ensure the half-close future is closed once the upload is complete. This is usually a no-op,
+    // but makes sure we half-close in case of error / interrupt.
+    closeFuture.addListener(
+        () -> halfCloseFuture.setFuture(closeFuture), MoreExecutors.directExecutor());
   }
 
   BuildEventArtifactUploader getLocalFileUploader() {
@@ -255,6 +259,10 @@ private void closeNow() {
     }
   }
 
+  ListenableFuture<Void> getHalfCloseFuture() {
+    return halfCloseFuture;
+  }
+
   private void logAndExitAbruptly(String message, ExitCode exitCode, Throwable cause) {
     checkState(!exitCode.equals(ExitCode.SUCCESS));
     logger.severe(message);
@@ -422,6 +430,7 @@ private void publishBuildEvents()
               streamContext.sendOverStream(request);
               streamContext.halfCloseStream();
               halfCloseFuture.set(null);
+              logger.info("BES uploader is half-closed");
             }
             break;