From 59748fbf7bef4424dcbac1b15fd22215b8767c6e Mon Sep 17 00:00:00 2001
From: Chi Wang <chiwang@google.com>
Date: Tue, 13 Jun 2023 04:38:43 -0700
Subject: [PATCH] Fix potential memory leak in UI

We report download progress to UI when downloading outputs from remote cache. UI thread keeps track of active downloads. There are two cases the UI thread could leak memory:

1. If we failed to close the output stream, the `reporter.finished()` will never be called, prevent UI thread from releasing the active download. This is fixed by calling `reporter.finished()` in `finally` block.
2. Normally, UI thread stops after `BuildCompleted` event. However, if we have background download after build is completed, UI thread is not stopped to continue printing out download progress. But after all downloads are done, we forgot to stop the UI thread, resulting all referenced objects leaked. This is fixed by calling `checkActivities()` for every download progress.

Fixes #18145.

Closes #18593.

PiperOrigin-RevId: 539923685
Change-Id: I7e2887035e540b39e382ab5fcbc06bad03b10427
---
 .../com/google/devtools/build/lib/remote/RemoteCache.java   | 3 ++-
 .../google/devtools/build/lib/runtime/UiEventHandler.java   | 6 +++++-
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteCache.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteCache.java
index 1fe79e16906f0b..2ad56b0239896c 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/RemoteCache.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/RemoteCache.java
@@ -379,11 +379,12 @@ public ListenableFuture<Void> downloadFile(
         () -> {
           try {
             out.close();
-            reporter.finished();
           } catch (IOException e) {
             logger.atWarning().withCause(e).log(
                 "Unexpected exception closing output stream after downloading %s/%d to %s",
                 digest.getHash(), digest.getSizeBytes(), path);
+          } finally {
+            reporter.finished();
           }
         },
         directExecutor());
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/UiEventHandler.java b/src/main/java/com/google/devtools/build/lib/runtime/UiEventHandler.java
index bea9c3854ae3a3..9b0f54f7b8ec51 100644
--- a/src/main/java/com/google/devtools/build/lib/runtime/UiEventHandler.java
+++ b/src/main/java/com/google/devtools/build/lib/runtime/UiEventHandler.java
@@ -670,7 +670,11 @@ public void afterCommand(AfterCommandEvent event) {
   public void downloadProgress(FetchProgress event) {
     maybeAddDate();
     stateTracker.downloadProgress(event);
-    refresh();
+    if (!event.isFinished()) {
+      refresh();
+    } else {
+      checkActivities();
+    }
   }
 
   @Subscribe