From be9fb708e84cc9a8648f6adc47172ac25f3e6cc6 Mon Sep 17 00:00:00 2001
From: Jordan Mele <mele@canva.com>
Date: Tue, 5 Dec 2023 09:04:45 -0800
Subject: [PATCH] Restart at most once when prepopulating repository rule
 environment

When a repository rule is fetch attributes are iterated over in `enforceLabelAttributes` to prepopulate the environment, restarting the fetch each time a new dependency is discovered.

This is faster than calling `repository_ctx.path(...)` early in the repository rule implementation function but still has considerable overhead.

This PR defers throwing `NeedsSkyframeRestartException` till after every attribute has been processed, greatly reducing the number of restarts and iterations.

In an internal project these optimisations are particularly noticeable.
Before: 2min 8s, 2min 7s
After: 1min 35s, 1min 27s

Closes #20434.

PiperOrigin-RevId: 588090528
Change-Id: I7917b137d6e60b6d6a73189cf396418a85b3ec28
---
 .../starlark/StarlarkRepositoryContext.java   | 27 ++++++++++++++-----
 1 file changed, 20 insertions(+), 7 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkRepositoryContext.java b/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkRepositoryContext.java
index dd8ad14c9b9fc8..246228875f8788 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkRepositoryContext.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkRepositoryContext.java
@@ -516,43 +516,56 @@ public String toString() {
    * potentially more expensive operations.
    */
   // TODO(wyv): somehow migrate this to the base context too.
-  public void enforceLabelAttributes() throws EvalException, InterruptedException {
+  public void enforceLabelAttributes()
+      throws EvalException, InterruptedException, NeedsSkyframeRestartException {
     // TODO: If a labels fails to resolve to an existing regular file, we do not add a dependency on
     //  that fact - if the file is created later or changes its type, it will not trigger a rerun of
     //  the repository function.
     StructImpl attr = getAttr();
+    // Batch restarts as they are expensive
+    boolean needsRestart = false;
     for (String name : attr.getFieldNames()) {
       Object value = attr.getValue(name);
       if (value instanceof Label) {
-        dependOnLabelIgnoringErrors((Label) value);
+        if (dependOnLabelIgnoringErrors((Label) value)) {
+          needsRestart = true;
+        }
       }
       if (value instanceof Sequence) {
         for (Object entry : (Sequence) value) {
           if (entry instanceof Label) {
-            dependOnLabelIgnoringErrors((Label) entry);
+            if (dependOnLabelIgnoringErrors((Label) entry)) {
+              needsRestart = true;
+            }
           }
         }
       }
       if (value instanceof Dict) {
         for (Object entry : ((Dict) value).keySet()) {
           if (entry instanceof Label) {
-            dependOnLabelIgnoringErrors((Label) entry);
+            if (dependOnLabelIgnoringErrors((Label) entry)) {
+              needsRestart = true;
+            }
           }
         }
       }
     }
+
+    if (needsRestart) {
+      throw new NeedsSkyframeRestartException();
+    }
   }
 
-  private void dependOnLabelIgnoringErrors(Label label)
-      throws InterruptedException, NeedsSkyframeRestartException {
+  private boolean dependOnLabelIgnoringErrors(Label label) throws InterruptedException {
     try {
       getPathFromLabel(label);
     } catch (NeedsSkyframeRestartException e) {
-      throw e;
+      return true;
     } catch (EvalException e) {
       // EvalExceptions indicate labels not referring to existing files. This is fine,
       // as long as they are never resolved to files in the execution of the rule; we allow
       // non-strict rules.
     }
+    return false;
   }
 }