From f6344ffcacdea6c4a61e112d0f60beda8068eac5 Mon Sep 17 00:00:00 2001
From: Googler <jhorvitz@google.com>
Date: Fri, 21 Jul 2023 14:51:05 -0700
Subject: [PATCH] Run idle server GC immediately following a
 `--nokeep_state_after_build_command` instead of waiting 10 seconds.

This helps to ensure that weakly reachable objects are not resurrected. Using the idle server task avoids blocking command completion on a full GC, which can take multiple seconds.

PiperOrigin-RevId: 550056326
Change-Id: Ia55a2360a196ee0f0dcbb0ac6b6baa11e39e244d
---
 .../build/lib/runtime/BlazeCommandResult.java | 24 ++++++++++---
 .../build/lib/runtime/BlazeRuntime.java       |  2 +-
 .../build/lib/server/CommandManager.java      | 19 ++++++----
 .../build/lib/server/GrpcServerImpl.java      |  7 ++++
 .../build/lib/server/IdleServerTasks.java     | 36 ++++++++++++-------
 .../build/lib/server/GrpcServerTest.java      |  3 +-
 6 files changed, 66 insertions(+), 25 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/runtime/BlazeCommandResult.java b/src/main/java/com/google/devtools/build/lib/runtime/BlazeCommandResult.java
index 77191d71b5ca01..4d8b3b901ab86c 100644
--- a/src/main/java/com/google/devtools/build/lib/runtime/BlazeCommandResult.java
+++ b/src/main/java/com/google/devtools/build/lib/runtime/BlazeCommandResult.java
@@ -38,21 +38,24 @@ public final class BlazeCommandResult {
   @Nullable private final ExecRequest execDescription;
   private final ImmutableList<Any> responseExtensions;
   private final boolean shutdown;
+  private final boolean stateKeptAfterBuild;
 
   private BlazeCommandResult(
       DetailedExitCode detailedExitCode,
       @Nullable ExecRequest execDescription,
       boolean shutdown,
-      ImmutableList<Any> responseExtensions) {
+      ImmutableList<Any> responseExtensions,
+      boolean stateKeptAfterBuild) {
     this.detailedExitCode = Preconditions.checkNotNull(detailedExitCode);
     this.execDescription = execDescription;
     this.shutdown = shutdown;
     this.responseExtensions = responseExtensions;
+    this.stateKeptAfterBuild = stateKeptAfterBuild;
   }
 
   private BlazeCommandResult(
       DetailedExitCode detailedExitCode, @Nullable ExecRequest execDescription, boolean shutdown) {
-    this(detailedExitCode, execDescription, shutdown, ImmutableList.of());
+    this(detailedExitCode, execDescription, shutdown, ImmutableList.of(), true);
   }
 
   public ExitCode getExitCode() {
@@ -85,6 +88,11 @@ public ImmutableList<Any> getResponseExtensions() {
     return responseExtensions;
   }
 
+  /** Reflects the value of {@link CommonCommandOptions#keepStateAfterBuild}. */
+  public boolean stateKeptAfterBuild() {
+    return stateKeptAfterBuild;
+  }
+
   @Override
   public String toString() {
     return MoreObjects.toStringHelper(this)
@@ -92,6 +100,8 @@ public String toString() {
         .add("failureDetail", getFailureDetail())
         .add("execDescription", execDescription)
         .add("shutdown", shutdown)
+        .add("responseExtensions", responseExtensions)
+        .add("stateKeptAfterBuild", stateKeptAfterBuild)
         .toString();
   }
 
@@ -116,9 +126,15 @@ public static BlazeCommandResult detailedExitCode(DetailedExitCode detailedExitC
   }
 
   public static BlazeCommandResult withResponseExtensions(
-      BlazeCommandResult result, ImmutableList<Any> responseExtensions) {
+      BlazeCommandResult result,
+      ImmutableList<Any> responseExtensions,
+      boolean stateKeptAfterBuild) {
     return new BlazeCommandResult(
-        result.detailedExitCode, result.execDescription, result.shutdown, responseExtensions);
+        result.detailedExitCode,
+        result.execDescription,
+        result.shutdown,
+        responseExtensions,
+        stateKeptAfterBuild);
   }
 
   public static BlazeCommandResult execute(ExecRequest execDescription) {
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/BlazeRuntime.java b/src/main/java/com/google/devtools/build/lib/runtime/BlazeRuntime.java
index 1ec96abad74336..2e00120e1fc66c 100644
--- a/src/main/java/com/google/devtools/build/lib/runtime/BlazeRuntime.java
+++ b/src/main/java/com/google/devtools/build/lib/runtime/BlazeRuntime.java
@@ -680,7 +680,7 @@ public BlazeCommandResult afterCommand(CommandEnvironment env, BlazeCommandResul
     DebugLoggerConfigurator.flushServerLog();
     storedExitCode.set(null);
     return BlazeCommandResult.withResponseExtensions(
-        finalCommandResult, env.getResponseExtensions());
+        finalCommandResult, env.getResponseExtensions(), commonOptions.keepStateAfterBuild);
   }
 
   /**
diff --git a/src/main/java/com/google/devtools/build/lib/server/CommandManager.java b/src/main/java/com/google/devtools/build/lib/server/CommandManager.java
index 617ec047c523be..60e65513aac9c7 100644
--- a/src/main/java/com/google/devtools/build/lib/server/CommandManager.java
+++ b/src/main/java/com/google/devtools/build/lib/server/CommandManager.java
@@ -18,6 +18,7 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.common.flogger.GoogleLogger;
 import com.google.devtools.build.lib.server.CommandProtos.CancelRequest;
+import com.google.devtools.build.lib.server.IdleServerTasks.IdleServerCleanupStrategy;
 import com.google.devtools.build.lib.util.ThreadUtils;
 import java.util.Collections;
 import java.util.HashMap;
@@ -43,7 +44,7 @@ class CommandManager {
   CommandManager(boolean doIdleServerTasks, @Nullable String slowInterruptMessageSuffix) {
     this.doIdleServerTasks = doIdleServerTasks;
     this.slowInterruptMessageSuffix = slowInterruptMessageSuffix;
-    idle();
+    idle(IdleServerCleanupStrategy.DELAYED);
   }
 
   void preemptEligibleCommands() {
@@ -132,11 +133,11 @@ private void registerCommand(RunningCommand command) {
     logger.atInfo().log("Starting command %s on thread %s", command.id, command.thread.getName());
   }
 
-  private void idle() {
+  private void idle(IdleServerCleanupStrategy cleanupStrategy) {
     Preconditions.checkState(idleServerTasks == null);
     if (doIdleServerTasks) {
       idleServerTasks = new IdleServerTasks();
-      idleServerTasks.idle();
+      idleServerTasks.idle(cleanupStrategy);
     }
   }
 
@@ -176,10 +177,11 @@ private void startSlowInterruptWatcher(final ImmutableSet<String> commandIds) {
     interruptWatcherThread.start();
   }
 
-  class RunningCommand implements AutoCloseable {
+  final class RunningCommand implements AutoCloseable {
     private final Thread thread;
     private final String id;
     private final boolean preemptible;
+    private IdleServerCleanupStrategy cleanupStrategy = IdleServerCleanupStrategy.DELAYED;
 
     private RunningCommand(boolean preemptible) {
       thread = Thread.currentThread();
@@ -192,7 +194,7 @@ public void close() {
       synchronized (runningCommandsMap) {
         runningCommandsMap.remove(id);
         if (runningCommandsMap.isEmpty()) {
-          idle();
+          idle(cleanupStrategy);
         }
         runningCommandsMap.notify();
       }
@@ -205,7 +207,12 @@ String getId() {
     }
 
     boolean isPreemptible() {
-      return this.preemptible;
+      return preemptible;
+    }
+
+    /** Requests a manual GC as soon as the server becomes idle. */
+    void requestEagerIdleServerCleanup() {
+      cleanupStrategy = IdleServerCleanupStrategy.EAGER;
     }
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/server/GrpcServerImpl.java b/src/main/java/com/google/devtools/build/lib/server/GrpcServerImpl.java
index 9f73f5a3db5e85..8efb5da16d3473 100644
--- a/src/main/java/com/google/devtools/build/lib/server/GrpcServerImpl.java
+++ b/src/main/java/com/google/devtools/build/lib/server/GrpcServerImpl.java
@@ -568,6 +568,13 @@ private void executeCommand(RunRequest request, BlockingStreamObserver<RunRespon
                                 .setCode(Command.Code.INVOCATION_POLICY_PARSE_FAILURE))
                         .build()));
       }
+      if (!result.stateKeptAfterBuild()) {
+        // If state was not kept, GC as soon as the server becomes idle. This ensures that weakly
+        // reachable objects are not "resurrected" on a subsequent command. See b/291641466. Without
+        // this call, a manual GC will only be triggered if the server remains idle for at least 10
+        // seconds before the next command starts.
+        command.requestEagerIdleServerCleanup();
+      }
     } catch (InterruptedException e) {
       result =
           BlazeCommandResult.detailedExitCode(
diff --git a/src/main/java/com/google/devtools/build/lib/server/IdleServerTasks.java b/src/main/java/com/google/devtools/build/lib/server/IdleServerTasks.java
index 4d286e336269d9..5186db8296cf32 100644
--- a/src/main/java/com/google/devtools/build/lib/server/IdleServerTasks.java
+++ b/src/main/java/com/google/devtools/build/lib/server/IdleServerTasks.java
@@ -28,25 +28,35 @@
 import java.util.concurrent.TimeUnit;
 
 /**
- * Run cleanup-related tasks during idle periods in the server.
- * idle() and busy() must be called in that order, and only once.
+ * Run cleanup-related tasks during idle periods in the server. idle() and busy() must be called in
+ * that order, and only once.
  */
-class IdleServerTasks {
-  private final ScheduledThreadPoolExecutor executor;
+final class IdleServerTasks {
+
   private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();
 
+  enum IdleServerCleanupStrategy {
+    EAGER(0),
+    DELAYED(10);
+
+    private final long delaySeconds;
+
+    IdleServerCleanupStrategy(long delaySeconds) {
+      this.delaySeconds = delaySeconds;
+    }
+  }
+
+  private final ScheduledThreadPoolExecutor executor;
+
   /** Must be called from the main thread. */
   public IdleServerTasks() {
-    this.executor = new ScheduledThreadPoolExecutor(
-        1,
-        new ThreadFactoryBuilder().setNameFormat("idle-server-tasks-%d").build());
+    this.executor =
+        new ScheduledThreadPoolExecutor(
+            1, new ThreadFactoryBuilder().setNameFormat("idle-server-tasks-%d").build());
   }
 
-  /**
-   * Called when the server becomes idle. Should not block, but may invoke
-   * new threads.
-   */
-  public void idle() {
+  /** Called when the server becomes idle. Should not block, but may invoke new threads. */
+  public void idle(IdleServerCleanupStrategy cleanupStrategy) {
     Preconditions.checkState(!executor.isShutdown());
 
     @SuppressWarnings("unused")
@@ -66,7 +76,7 @@ public void idle() {
                   StringUtilities.prettyPrintBytes(before.getCommitted()),
                   StringUtilities.prettyPrintBytes(after.getCommitted()));
             },
-            10,
+            cleanupStrategy.delaySeconds,
             TimeUnit.SECONDS);
   }
 
diff --git a/src/test/java/com/google/devtools/build/lib/server/GrpcServerTest.java b/src/test/java/com/google/devtools/build/lib/server/GrpcServerTest.java
index e0593a93f5fa63..c3fdfb00450a01 100644
--- a/src/test/java/com/google/devtools/build/lib/server/GrpcServerTest.java
+++ b/src/test/java/com/google/devtools/build/lib/server/GrpcServerTest.java
@@ -241,7 +241,8 @@ public BlazeCommandResult exec(
                 BlazeCommandResult.success(),
                 ImmutableList.of(
                     Any.pack(StringValue.of("foo")),
-                    Any.pack(BytesValue.of(ByteString.copyFromUtf8("bar")))));
+                    Any.pack(BytesValue.of(ByteString.copyFromUtf8("bar")))),
+                true);
           }
         };
     createServer(dispatcher);