[7.2.1] Enforce and await cleanup in StarlarkBaseExternalContext (#22814)

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionContext.java

@@ -71,6 +71,7 @@ protected ModuleExtensionContext(
         timeoutScaling,
         processWrapper,
         starlarkSemantics,
+        ModuleExtensionEvaluationProgress.moduleExtensionEvaluationContextString(extensionId),
         remoteExecutor,
         /* allowWatchingPathsOutsideWorkspace= */ false);
     this.extensionId = extensionId;
@@ -83,8 +84,10 @@ public Path getWorkingDirectory() {
   }
 
   @Override
-  protected String getIdentifyingStringForLogging() {
-    return ModuleExtensionEvaluationProgress.moduleExtensionEvaluationContextString(extensionId);
+  protected boolean shouldDeleteWorkingDirectoryOnClose(boolean successful) {
+    // The contents of the working directory are purely ephemeral, only the repos instantiated by
+    // the extension are considered its results.
+    return true;
   }
 
   @Override

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/SingleExtensionEvalFunction.java

@@ -891,15 +891,15 @@ public RunModuleExtensionResult run(
               BazelModuleContext.of(bzlLoadValue.getModule()).repoMapping(),
               directories,
               env.getListener());
-      ModuleExtensionContext moduleContext;
       Optional<ModuleExtensionMetadata> moduleExtensionMetadata;
       var repoMappingRecorder = new Label.RepoMappingRecorder();
       repoMappingRecorder.mergeEntries(bzlLoadValue.getRecordedRepoMappings());
       try (Mutability mu =
-          Mutability.create("module extension", usagesValue.getExtensionUniqueName())) {
+              Mutability.create("module extension", usagesValue.getExtensionUniqueName());
+          ModuleExtensionContext moduleContext =
+              createContext(env, usagesValue, starlarkSemantics, extensionId)) {
         StarlarkThread thread = new StarlarkThread(mu, starlarkSemantics);
         thread.setPrintHandler(Event.makeDebugPrintHandler(env.getListener()));
-        moduleContext = createContext(env, usagesValue, starlarkSemantics, extensionId);
         threadContext.storeInThread(thread);
         new BazelStarlarkContext(
                 Phase.WORKSPACE,
@@ -938,39 +938,34 @@ public RunModuleExtensionResult run(
             moduleExtensionMetadata = Optional.empty();
           }
         } catch (NeedsSkyframeRestartException e) {
-          // Clean up and restart by returning null.
-          try {
-            if (moduleContext.getWorkingDirectory().exists()) {
-              moduleContext.getWorkingDirectory().deleteTree();
-            }
-          } catch (IOException e1) {
-            ExternalDepsException externalDepsException =
-                ExternalDepsException.withCauseAndMessage(
-                    ExternalDeps.Code.EXTERNAL_DEPS_UNKNOWN,
-                    e1,
-                    "Failed to clean up module context directory");
-            throw new SingleExtensionEvalFunctionException(
-                externalDepsException, Transience.TRANSIENT);
-          }
+          // Restart by returning null.
           return null;
-        } catch (EvalException e) {
-          env.getListener().handle(Event.error(e.getMessageWithStack()));
-          throw new SingleExtensionEvalFunctionException(
-              ExternalDepsException.withMessage(
-                  ExternalDeps.Code.BAD_MODULE,
-                  "error evaluating module extension %s in %s",
-                  extensionId.getExtensionName(),
-                  extensionId.getBzlFileLabel()),
-              Transience.TRANSIENT);
         }
+        moduleContext.markSuccessful();
+        return RunModuleExtensionResult.create(
+            moduleContext.getRecordedFileInputs(),
+            moduleContext.getRecordedDirentsInputs(),
+            moduleContext.getRecordedEnvVarInputs(),
+            threadContext.getGeneratedRepoSpecs(),
+            moduleExtensionMetadata,
+            repoMappingRecorder.recordedEntries());
+      } catch (EvalException e) {
+        env.getListener().handle(Event.error(e.getMessageWithStack()));
+        throw new SingleExtensionEvalFunctionException(
+            ExternalDepsException.withMessage(
+                ExternalDeps.Code.BAD_MODULE,
+                "error evaluating module extension %s in %s",
+                extensionId.getExtensionName(),
+                extensionId.getBzlFileLabel()),
+            Transience.TRANSIENT);
+      } catch (IOException e) {
+        ExternalDepsException externalDepsException =
+            ExternalDepsException.withCauseAndMessage(
+                ExternalDeps.Code.EXTERNAL_DEPS_UNKNOWN,
+                e,
+                "Failed to clean up module context directory");
+        throw new SingleExtensionEvalFunctionException(externalDepsException, Transience.TRANSIENT);
       }
-      return RunModuleExtensionResult.create(
-          moduleContext.getRecordedFileInputs(),
-          moduleContext.getRecordedDirentsInputs(),
-          moduleContext.getRecordedEnvVarInputs(),
-          threadContext.getGeneratedRepoSpecs(),
-          moduleExtensionMetadata,
-          repoMappingRecorder.recordedEntries());
     }
 
     private ModuleExtensionContext createContext(

src/main/java/com/google/devtools/build/lib/bazel/repository/downloader/DownloadManager.java

@@ -24,7 +24,6 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
-import com.google.common.util.concurrent.ThreadFactoryBuilder;
 import com.google.devtools.build.lib.authandtls.StaticCredentials;
 import com.google.devtools.build.lib.bazel.repository.cache.RepositoryCache;
 import com.google.devtools.build.lib.bazel.repository.cache.RepositoryCache.KeyType;
@@ -48,7 +47,6 @@
 import java.util.Optional;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import javax.annotation.Nullable;
 
@@ -59,17 +57,6 @@
  * to disk.
  */
 public class DownloadManager {
-  private static final ExecutorService DOWNLOAD_EXECUTOR =
-      Executors.newFixedThreadPool(
-          // There is also GrpcRemoteDownloader so if we set the thread pool to the same size as
-          // the allowed number of HTTP downloads, it might unnecessarily block. No, this is not a
-          // very
-          // principled approach; ideally, we'd grow the thread pool as needed with some generous
-          // upper
-          // limit.
-          2 * HttpDownloader.MAX_PARALLEL_DOWNLOADS,
-          new ThreadFactoryBuilder().setNameFormat("download-manager-%d").build());
-
   private final RepositoryCache repositoryCache;
   private List<Path> distdir = ImmutableList.of();
   private UrlRewriter rewriter;
@@ -115,6 +102,7 @@ public void setCredentialFactory(CredentialFactory credentialFactory) {
   }
 
   public Future<Path> startDownload(
+      ExecutorService executorService,
       List<URL> originalUrls,
       Map<String, List<String>> headers,
       Map<URI, Map<String, List<String>>> authHeaders,
@@ -125,7 +113,7 @@ public Future<Path> startDownload(
       ExtendedEventHandler eventHandler,
       Map<String, String> clientEnv,
       String context) {
-    return DOWNLOAD_EXECUTOR.submit(
+    return executorService.submit(
         () -> {
           try (SilentCloseable c = Profiler.instance().profile("fetching: " + context)) {
             return downloadInExecutor(
@@ -154,33 +142,6 @@ public Path finalizeDownload(Future<Path> download) throws IOException, Interrup
     }
   }
 
-  public Path download(
-      List<URL> originalUrls,
-      Map<String, List<String>> headers,
-      Map<URI, Map<String, List<String>>> authHeaders,
-      Optional<Checksum> checksum,
-      String canonicalId,
-      Optional<String> type,
-      Path output,
-      ExtendedEventHandler eventHandler,
-      Map<String, String> clientEnv,
-      String context)
-      throws IOException, InterruptedException {
-    Future<Path> future =
-        startDownload(
-            originalUrls,
-            headers,
-            authHeaders,
-            checksum,
-            canonicalId,
-            type,
-            output,
-            eventHandler,
-            clientEnv,
-            context);
-    return finalizeDownload(future);
-  }
-
   /**
    * Downloads file to disk and returns path.
    *