From ce8836dd05de94b92a494be12a47a56bf9819107 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 9a984df82213a5..ad32126995154e 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
@@ -676,7 +676,11 @@ public void afterCommand(AfterCommandEvent event) {
   public void downloadProgress(FetchProgress event) {
     maybeAddDate();
     stateTracker.downloadProgress(event);
-    refresh();
+    if (!event.isFinished()) {
+      refresh();
+    } else {
+      checkActivities();
+    }
   }
 
   @Subscribe