Add support for fetching repos using a worker thread

Added flag `--experimental_worker_for_repo_fetching`, whose value can be `off` (no change), `platform` (uses a platform/OS thread to do repo fetching), `virtual` (uses a virtual thread feat. Project Loom).

When a worker thread is used, the Skyframe thread does not perform any fetching itself. Instead, it gives its Environment object to the worker thread, and let it do the bulk of the work. When a new Skyframe dependency is registered and is found missing, the worker thread doesn't restart; instead, it signals to the host thread that it should restart, and waits for the host thread to give back a fresh Environment object.

This weird dance can work with both platform threads and virtual threads. Right now, passing `--experimental_worker_for_repo_fetching=virtual` results in a fatal error, but that should be easily fixable when JDK20+ lands.

Work towards https://github.com/bazelbuild/bazel/issues/10515

PiperOrigin-RevId: 538909616
Change-Id: Iacc8b84ca2f90821fba5add171a4778ed2b48ab9
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/BazelRepositoryModule.java b/src/main/java/com/google/devtools/build/lib/bazel/BazelRepositoryModule.java
index f264168..a5c04e9 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/BazelRepositoryModule.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/BazelRepositoryModule.java
@@ -116,6 +116,8 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
@@ -154,6 +156,7 @@
   private LockfileMode bazelLockfileMode = LockfileMode.OFF;
   private List<String> allowedYankedVersions = ImmutableList.of();
   private SingleExtensionEvalFunction singleExtensionEvalFunction;
+  private final ExecutorService repoFetchingWorkerThreadPool = Executors.newFixedThreadPool(100);
 
   @Nullable private CredentialModule credentialModule;
 
@@ -312,6 +315,19 @@
 
     RepositoryOptions repoOptions = env.getOptions().getOptions(RepositoryOptions.class);
     if (repoOptions != null) {
+      switch (repoOptions.workerForRepoFetching) {
+        case OFF:
+          starlarkRepositoryFunction.setWorkerExecutorService(null);
+          break;
+        case PLATFORM:
+          starlarkRepositoryFunction.setWorkerExecutorService(repoFetchingWorkerThreadPool);
+          break;
+        case VIRTUAL:
+          throw new AbruptExitException(
+              detailedExitCode(
+                  "using a virtual worker thread for repo fetching is not yet supported",
+                  Code.BAD_DOWNLOADER_CONFIG));
+      }
       downloadManager.setDisableDownload(repoOptions.disableDownload);
       if (repoOptions.repositoryDownloaderRetries >= 0) {
         downloadManager.setRetries(repoOptions.repositoryDownloaderRetries);
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/repository/RepositoryOptions.java b/src/main/java/com/google/devtools/build/lib/bazel/repository/RepositoryOptions.java
index 4f28fac..7be04e8 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/repository/RepositoryOptions.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/repository/RepositoryOptions.java
@@ -236,6 +236,31 @@
               + "give, and in this case multiple URLs will be returned.")
   public String downloaderConfig;
 
+  /** See {@link #workerForRepoFetching}. */
+  public enum WorkerForRepoFetching {
+    OFF,
+    PLATFORM,
+    VIRTUAL;
+
+    static class Converter extends EnumConverter<WorkerForRepoFetching> {
+      public Converter() {
+        super(WorkerForRepoFetching.class, "worker for repo fetching");
+      }
+    }
+  }
+
+  @Option(
+      name = "experimental_worker_for_repo_fetching",
+      defaultValue = "off",
+      converter = WorkerForRepoFetching.Converter.class,
+      documentationCategory = OptionDocumentationCategory.REMOTE,
+      effectTags = {OptionEffectTag.UNKNOWN},
+      help =
+          "The threading mode to use for repo fetching. If set to 'off', no worker thread is used,"
+              + " and the repo fetching is subject to restarts. Otherwise, uses a platform thread"
+              + " (i.e. OS thread) if set to 'platform' or a virtual thread if set to 'virtual'.")
+  public WorkerForRepoFetching workerForRepoFetching;
+
   @Option(
       name = "ignore_dev_dependency",
       defaultValue = "false",
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/BUILD b/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/BUILD
index 51564c2..62ed8b7 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/BUILD
@@ -55,6 +55,7 @@
         "//src/main/java/net/starlark/java/annot",
         "//src/main/java/net/starlark/java/eval",
         "//src/main/java/net/starlark/java/syntax",
+        "//third_party:auto_value",
         "//third_party:error_prone_annotations",
         "//third_party:guava",
         "//third_party:java-diff-utils",
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/RepoFetchingSkyKeyComputeState.java b/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/RepoFetchingSkyKeyComputeState.java
new file mode 100644
index 0000000..55473dd
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/RepoFetchingSkyKeyComputeState.java
@@ -0,0 +1,82 @@
+// Copyright 2023 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.bazel.repository.starlark;
+
+import com.google.devtools.build.lib.rules.repository.RepositoryDirectoryValue;
+import com.google.devtools.build.skyframe.SkyFunction;
+import com.google.devtools.build.skyframe.SkyFunction.Environment.SkyKeyComputeState;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.Future;
+import java.util.concurrent.SynchronousQueue;
+import javax.annotation.Nullable;
+
+/**
+ * Captures state that persists across different invocations of {@link
+ * com.google.devtools.build.lib.rules.repository.RepositoryDelegatorFunction}, specifically {@link
+ * StarlarkRepositoryFunction}.
+ *
+ * <p>This class is used to hold on to a worker thread (in reality just a {@link Future} object)
+ * when fetching repos using a worker thread is enabled. The worker thread uses a {@link
+ * SkyFunction.Environment} object acquired from the host thread, and can signal the host thread to
+ * restart to get a fresh environment object.
+ */
+class RepoFetchingSkyKeyComputeState implements SkyKeyComputeState {
+
+  /** A signal that the worker thread can send to the host Skyframe thread. */
+  enum Signal {
+    /**
+     * Indicates that the host thread should return {@code null}, causing a Skyframe restart. After
+     * sending this signal, the client will immediately block on {@code delegateEnvQueue}, waiting
+     * for the host thread to send a fresh {@link SkyFunction.Environment} over.
+     */
+    RESTART,
+    /**
+     * Indicates that the worker thread has finished running, either yielding a result or an
+     * exception.
+     */
+    DONE
+  }
+
+  /** The channel for the worker thread to send a signal to the host Skyframe thread. */
+  final BlockingQueue<Signal> signalQueue = new SynchronousQueue<>();
+  /**
+   * The channel for the host Skyframe thread to send fresh {@link SkyFunction.Environment} objects
+   * back to the worker thread.
+   */
+  final BlockingQueue<SkyFunction.Environment> delegateEnvQueue = new SynchronousQueue<>();
+  /**
+   * This future holds on to the worker thread in order to cancel it when necessary; it also serves
+   * to tell whether a worker thread is already running.
+   */
+  // This is volatile since we set it to null to indicate the worker thread isn't running, and this
+  // could happen on multiple threads. Canceling a future multiple times is safe, though, so we
+  // only need to worry about nullness. Using a mutex/synchronization is an alternative but it means
+  // we might block in `close()`, which is potentially bad (see its javadoc).
+  @Nullable volatile Future<RepositoryDirectoryValue.Builder> workerFuture = null;
+
+  SkyFunction.Environment signalForFreshEnv() throws InterruptedException {
+    signalQueue.put(Signal.RESTART);
+    return delegateEnvQueue.take();
+  }
+
+  @Override
+  public void close() {
+    var myWorkerFuture = workerFuture;
+    workerFuture = null;
+    if (myWorkerFuture != null) {
+      myWorkerFuture.cancel(true);
+    }
+  }
+}
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/RepoFetchingWorkerSkyFunctionEnvironment.java b/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/RepoFetchingWorkerSkyFunctionEnvironment.java
new file mode 100644
index 0000000..e9c38fa
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/RepoFetchingWorkerSkyFunctionEnvironment.java
@@ -0,0 +1,185 @@
+// Copyright 2023 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.bazel.repository.starlark;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.devtools.build.lib.events.Event;
+import com.google.devtools.build.lib.events.ExtendedEventHandler;
+import com.google.devtools.build.skyframe.SkyFunction;
+import com.google.devtools.build.skyframe.SkyKey;
+import com.google.devtools.build.skyframe.SkyValue;
+import com.google.devtools.build.skyframe.SkyframeLookupResult;
+import com.google.devtools.build.skyframe.Version;
+import java.util.function.Supplier;
+import javax.annotation.Nullable;
+
+/**
+ * A {@link SkyFunction.Environment} implementation designed to be used in a different thread than
+ * the corresponding SkyFunction runs in. It relies on a delegate Environment object to do
+ * underlying work. Its {@link #getValue} and {@link #getValueOrThrow} methods do not return {@code
+ * null} when the {@link SkyValue} in question is not available. Instead, it blocks and waits for
+ * the host Skyframe thread to restart, and replaces the delegate Environment with a fresh one from
+ * the restarted SkyFunction before continuing. (Note that those methods <em>do</em> return {@code
+ * null} if the SkyValue was evaluated but found to be in error.)
+ *
+ * <p>Crucially, the delegate Environment object must not be used by multiple threads at the same
+ * time. In effect, this is guaranteed by only one of the worker thread and host thread being active
+ * at any given time.
+ */
+class RepoFetchingWorkerSkyFunctionEnvironment
+    implements SkyFunction.Environment, ExtendedEventHandler, SkyframeLookupResult {
+  private final RepoFetchingSkyKeyComputeState state;
+  private SkyFunction.Environment delegate;
+
+  RepoFetchingWorkerSkyFunctionEnvironment(
+      RepoFetchingSkyKeyComputeState state, SkyFunction.Environment delegate) {
+    this.state = state;
+    this.delegate = delegate;
+  }
+
+  @Override
+  public boolean valuesMissing() {
+    return delegate.valuesMissing();
+  }
+
+  @Override
+  public SkyframeLookupResult getValuesAndExceptions(Iterable<? extends SkyKey> depKeys)
+      throws InterruptedException {
+    delegate.getValuesAndExceptions(depKeys);
+    if (!delegate.valuesMissing()) {
+      // Do NOT just return the return value of `delegate.getValuesAndExceptions` here! That would
+      // cause anyone holding onto the returned // result object to potentially use a stale version
+      // of it after a skyfunction restart.
+      return this;
+    }
+    // We null out `delegate` before blocking for the fresh env so that the old one becomes
+    // eligible for GC.
+    delegate = null;
+    delegate = state.signalForFreshEnv();
+    delegate.getValuesAndExceptions(depKeys);
+    return this;
+  }
+
+  @Nullable
+  @Override
+  public <E1 extends Exception, E2 extends Exception, E3 extends Exception> SkyValue getOrThrow(
+      SkyKey skyKey, Class<E1> e1, Class<E2> e2, Class<E3> e3) throws E1, E2, E3 {
+    return delegate.getLookupHandleForPreviouslyRequestedDeps().getOrThrow(skyKey, e1, e2, e3);
+  }
+
+  @Override
+  public boolean queryDep(SkyKey key, QueryDepCallback resultCallback) {
+    return delegate.getLookupHandleForPreviouslyRequestedDeps().queryDep(key, resultCallback);
+  }
+
+  @Nullable
+  @Override
+  public SkyValue getValue(SkyKey depKey) throws InterruptedException {
+    return getValuesAndExceptions(ImmutableList.of(depKey)).get(depKey);
+  }
+
+  @Nullable
+  @Override
+  public <E1 extends Exception> SkyValue getValueOrThrow(SkyKey depKey, Class<E1> e1)
+      throws E1, InterruptedException {
+    return getValuesAndExceptions(ImmutableList.of(depKey)).getOrThrow(depKey, e1);
+  }
+
+  @Nullable
+  @Override
+  public <E1 extends Exception, E2 extends Exception> SkyValue getValueOrThrow(
+      SkyKey depKey, Class<E1> e1, Class<E2> e2) throws E1, E2, InterruptedException {
+    return getValuesAndExceptions(ImmutableList.of(depKey)).getOrThrow(depKey, e1, e2);
+  }
+
+  @Nullable
+  @Override
+  public <E1 extends Exception, E2 extends Exception, E3 extends Exception>
+      SkyValue getValueOrThrow(SkyKey depKey, Class<E1> e1, Class<E2> e2, Class<E3> e3)
+          throws E1, E2, E3, InterruptedException {
+    return getValuesAndExceptions(ImmutableList.of(depKey)).getOrThrow(depKey, e1, e2, e3);
+  }
+
+  @Nullable
+  @Override
+  public <E1 extends Exception, E2 extends Exception, E3 extends Exception, E4 extends Exception>
+      SkyValue getValueOrThrow(
+          SkyKey depKey, Class<E1> e1, Class<E2> e2, Class<E3> e3, Class<E4> e4)
+          throws E1, E2, E3, E4, InterruptedException {
+    SkyValue value = delegate.getValueOrThrow(depKey, e1, e2, e3, e4);
+    if (value != null) {
+      return value;
+    }
+    // We null out `delegate` before blocking for the fresh env so that the old one becomes
+    // eligible for GC.
+    delegate = null;
+    delegate = state.signalForFreshEnv();
+    return delegate.getValueOrThrow(depKey, e1, e2, e3, e4);
+  }
+
+  @Override
+  public ExtendedEventHandler getListener() {
+    // Do NOT just return `delegate.getListener()` here! That would cause anyone holding onto the
+    // returned listener to potentially post events to a stale listener.
+    return this;
+  }
+
+  @Override
+  public void post(Postable obj) {
+    delegate.getListener().post(obj);
+  }
+
+  @Override
+  public void handle(Event event) {
+    delegate.getListener().handle(event);
+  }
+
+  @Override
+  public void registerDependencies(Iterable<SkyKey> keys) {
+    delegate.registerDependencies(keys);
+  }
+
+  @Override
+  public boolean inErrorBubblingForSkyFunctionsThatCanFullyRecoverFromErrors() {
+    return delegate.inErrorBubblingForSkyFunctionsThatCanFullyRecoverFromErrors();
+  }
+
+  @Override
+  public void dependOnFuture(ListenableFuture<?> future) {
+    delegate.dependOnFuture(future);
+  }
+
+  @Override
+  public boolean restartPermitted() {
+    return delegate.restartPermitted();
+  }
+
+  @Override
+  public SkyframeLookupResult getLookupHandleForPreviouslyRequestedDeps() {
+    return delegate.getLookupHandleForPreviouslyRequestedDeps();
+  }
+
+  @Override
+  public <T extends SkyKeyComputeState> T getState(Supplier<T> stateSupplier) {
+    return delegate.getState(stateSupplier);
+  }
+
+  @Nullable
+  @Override
+  public Version getMaxTransitiveSourceVersionSoFar() {
+    return delegate.getMaxTransitiveSourceVersionSoFar();
+  }
+}
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkRepositoryFunction.java b/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkRepositoryFunction.java
index 1406974..fc3a7e1 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkRepositoryFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkRepositoryFunction.java
@@ -16,6 +16,7 @@
 
 import static com.google.common.base.Preconditions.checkNotNull;
 
+import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
@@ -23,7 +24,9 @@
 import com.google.devtools.build.lib.analysis.RuleDefinition;
 import com.google.devtools.build.lib.bazel.repository.RepositoryResolvedEvent;
 import com.google.devtools.build.lib.bazel.repository.downloader.DownloadManager;
+import com.google.devtools.build.lib.bazel.repository.starlark.RepoFetchingSkyKeyComputeState.Signal;
 import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.packages.BazelStarlarkContext;
 import com.google.devtools.build.lib.packages.Rule;
@@ -52,6 +55,8 @@
 import java.io.IOException;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
 import javax.annotation.Nullable;
 import net.starlark.java.eval.EvalException;
 import net.starlark.java.eval.Mutability;
@@ -61,11 +66,12 @@
 import net.starlark.java.eval.StarlarkThread;
 
 /** A repository function to delegate work done by Starlark remote repositories. */
-public class StarlarkRepositoryFunction extends RepositoryFunction {
+public final class StarlarkRepositoryFunction extends RepositoryFunction {
   static final String SEMANTICS = "STARLARK_SEMANTICS";
 
   private final DownloadManager downloadManager;
   private double timeoutScaling = 1.0;
+  @Nullable private ExecutorService workerExecutorService = null;
   @Nullable private ProcessWrapper processWrapper = null;
   @Nullable private RepositoryRemoteExecutor repositoryRemoteExecutor;
   @Nullable private SyscallCache syscallCache;
@@ -86,6 +92,10 @@
     this.syscallCache = checkNotNull(syscallCache);
   }
 
+  public void setWorkerExecutorService(@Nullable ExecutorService workerExecutorService) {
+    this.workerExecutorService = workerExecutorService;
+  }
+
   static String describeSemantics(StarlarkSemantics semantics) {
     // Here we use the hash code provided by AutoValue. This is unique, as long
     // as the number of bits in the StarlarkSemantics is small enough. We will have to
@@ -105,6 +115,24 @@
     return describeSemantics(starlarkSemantics).equals(markerData.get(SEMANTICS));
   }
 
+  @Override
+  protected void setupRepoRootBeforeFetching(Path repoRoot) throws RepositoryFunctionException {
+    // DON'T delete the repo root here if we're using a worker thread, since when this SkyFunction
+    // restarts, fetching is still happening inside the worker thread.
+    if (workerExecutorService == null) {
+      setupRepoRoot(repoRoot);
+    }
+  }
+
+  @Override
+  public void reportSkyframeRestart(Environment env, RepositoryName repoName) {
+    // DON'T report a "restarting." event if we're using a worker thread, since the actual fetch
+    // function run by the worker thread never restarts.
+    if (workerExecutorService == null) {
+      super.reportSkyframeRestart(env, repoName);
+    }
+  }
+
   @Nullable
   @Override
   public RepositoryDirectoryValue.Builder fetch(
@@ -115,6 +143,70 @@
       Map<String, String> markerData,
       SkyKey key)
       throws RepositoryFunctionException, InterruptedException {
+    if (workerExecutorService == null) {
+      return fetchInternal(rule, outputDirectory, directories, env, markerData, key);
+    }
+    var state = env.getState(RepoFetchingSkyKeyComputeState::new);
+    var workerFuture = state.workerFuture;
+    if (workerFuture == null) {
+      // No worker is running yet, which means we're just starting to fetch this repo. Start with a
+      // clean slate, and create the worker.
+      setupRepoRoot(outputDirectory);
+      Environment workerEnv = new RepoFetchingWorkerSkyFunctionEnvironment(state, env);
+      workerFuture =
+          workerExecutorService.submit(
+              () -> {
+                try {
+                  return fetchInternal(
+                      rule, outputDirectory, directories, workerEnv, markerData, key);
+                } finally {
+                  state.signalQueue.put(Signal.DONE);
+                }
+              });
+      state.workerFuture = workerFuture;
+    } else {
+      // A worker is already running. This can only mean one thing -- we just had a Skyframe
+      // restart, and need to send over a fresh Environment.
+      state.delegateEnvQueue.put(env);
+    }
+    switch (state.signalQueue.take()) {
+      case RESTART:
+        return null;
+      case DONE:
+        try {
+          return workerFuture.get();
+        } catch (ExecutionException e) {
+          Throwables.throwIfInstanceOf(e.getCause(), RepositoryFunctionException.class);
+          Throwables.throwIfUnchecked(e.getCause());
+          throw new IllegalStateException(
+              "unexpected exception type: " + e.getClass(), e.getCause());
+        } finally {
+          // Make sure we interrupt the worker thread if work on the Skyframe thread were cut short
+          // for any reason.
+          state.close();
+          try {
+            // Synchronously wait for the worker thread to finish any remaining work.
+            workerFuture.get();
+          } catch (ExecutionException e) {
+            // When this happens, we either already dealt with the exception (see `catch` clause
+            // above), or we're in the middle of propagating an InterruptedException in which case
+            // we don't care about the result of execution anyway.
+          }
+        }
+    }
+    // TODO(wyv): use a switch expression above instead and remove this.
+    throw new IllegalStateException();
+  }
+
+  @Nullable
+  private RepositoryDirectoryValue.Builder fetchInternal(
+      Rule rule,
+      Path outputDirectory,
+      BlazeDirectories directories,
+      Environment env,
+      Map<String, String> markerData,
+      SkyKey key)
+      throws RepositoryFunctionException, InterruptedException {
 
     String defInfo = RepositoryResolvedEvent.getRuleDefinitionInformation(rule);
     env.getListener().post(new StarlarkRepositoryDefinitionLocationEvent(rule.getName(), defInfo));
diff --git a/src/main/java/com/google/devtools/build/lib/rules/repository/RepositoryDelegatorFunction.java b/src/main/java/com/google/devtools/build/lib/rules/repository/RepositoryDelegatorFunction.java
index b523730..aa92bfd1a 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/repository/RepositoryDelegatorFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/repository/RepositoryDelegatorFunction.java
@@ -223,15 +223,6 @@
     return value;
   }
 
-  private void setupRepositoryRoot(Path repoRoot) throws RepositoryFunctionException {
-    try {
-      repoRoot.deleteTree();
-      Preconditions.checkNotNull(repoRoot.getParentDirectory()).createDirectoryAndParents();
-    } catch (IOException e) {
-      throw new RepositoryFunctionException(e, Transience.TRANSIENT);
-    }
-  }
-
   @Nullable
   @Override
   public SkyValue compute(SkyKey skyKey, Environment env)
@@ -401,7 +392,7 @@
       Rule rule)
       throws InterruptedException, RepositoryFunctionException {
 
-    setupRepositoryRoot(repoRoot);
+    handler.setupRepoRootBeforeFetching(repoRoot);
 
     RepositoryName repoName = (RepositoryName) skyKey.argument();
     env.getListener().post(RepositoryFetchProgress.ongoing(repoName, "starting"));
@@ -425,7 +416,7 @@
     }
 
     if (env.valuesMissing()) {
-      env.getListener().post(RepositoryFetchProgress.ongoing(repoName, "Restarting."));
+      handler.reportSkyframeRestart(env, repoName);
       return null;
     }
     env.getListener().post(RepositoryFetchProgress.finished(repoName));
@@ -463,7 +454,7 @@
   private RepositoryDirectoryValue setupOverride(
       PathFragment sourcePath, Environment env, Path repoRoot, String pathAttr)
       throws RepositoryFunctionException, InterruptedException {
-    setupRepositoryRoot(repoRoot);
+    RepositoryFunction.setupRepoRoot(repoRoot);
     RepositoryDirectoryValue.Builder directoryValue =
         symlinkRepoRoot(
             directories,
diff --git a/src/main/java/com/google/devtools/build/lib/rules/repository/RepositoryFunction.java b/src/main/java/com/google/devtools/build/lib/rules/repository/RepositoryFunction.java
index 55084b9..1854bc1 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/repository/RepositoryFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/repository/RepositoryFunction.java
@@ -34,6 +34,7 @@
 import com.google.devtools.build.lib.packages.Rule;
 import com.google.devtools.build.lib.packages.Type;
 import com.google.devtools.build.lib.repository.ExternalPackageException;
+import com.google.devtools.build.lib.repository.RepositoryFetchProgress;
 import com.google.devtools.build.lib.skyframe.ActionEnvironmentFunction;
 import com.google.devtools.build.lib.skyframe.AlreadyReportedException;
 import com.google.devtools.build.lib.skyframe.PackageLookupFunction;
@@ -133,6 +134,23 @@
     }
   }
 
+  public static void setupRepoRoot(Path repoRoot) throws RepositoryFunctionException {
+    try {
+      repoRoot.deleteTree();
+      Preconditions.checkNotNull(repoRoot.getParentDirectory()).createDirectoryAndParents();
+    } catch (IOException e) {
+      throw new RepositoryFunctionException(e, Transience.TRANSIENT);
+    }
+  }
+
+  protected void setupRepoRootBeforeFetching(Path repoRoot) throws RepositoryFunctionException {
+    setupRepoRoot(repoRoot);
+  }
+
+  public void reportSkyframeRestart(Environment env, RepositoryName repoName) {
+    env.getListener().post(RepositoryFetchProgress.ongoing(repoName, "Restarting."));
+  }
+
   /**
    * Fetch the remote repository represented by the given rule.
    *
diff --git a/src/main/java/com/google/devtools/build/skyframe/SkyFunction.java b/src/main/java/com/google/devtools/build/skyframe/SkyFunction.java
index cdd262f..d39bc3a 100644
--- a/src/main/java/com/google/devtools/build/skyframe/SkyFunction.java
+++ b/src/main/java/com/google/devtools/build/skyframe/SkyFunction.java
@@ -396,6 +396,8 @@
        * being held on to is approaching starvation, we currently don't do anything to alleviate
        * that pressure. So think *hard* before you start doing that!
        *
+       * <p>Implementations <strong>MUST</strong> be idempotent.
+       *
        * <p>Note also that this method should not perform any heavy work (especially blocking
        * operations).
        */
diff --git a/src/main/java/com/google/devtools/build/skyframe/SkyFunctionException.java b/src/main/java/com/google/devtools/build/skyframe/SkyFunctionException.java
index ea5298a..2110503 100644
--- a/src/main/java/com/google/devtools/build/skyframe/SkyFunctionException.java
+++ b/src/main/java/com/google/devtools/build/skyframe/SkyFunctionException.java
@@ -102,7 +102,8 @@
     }
   }
 
-  static <E1 extends Exception, E2 extends Exception, E3 extends Exception, E4 extends Exception>
+  public static <
+          E1 extends Exception, E2 extends Exception, E3 extends Exception, E4 extends Exception>
       void throwIfInstanceOf(
           @Nullable Exception e,
           @Nullable Class<E1> exceptionClass1,
diff --git a/src/test/shell/bazel/starlark_repository_test.sh b/src/test/shell/bazel/starlark_repository_test.sh
index cbf8f18..b2f4503 100755
--- a/src/test/shell/bazel/starlark_repository_test.sh
+++ b/src/test/shell/bazel/starlark_repository_test.sh
@@ -2358,4 +2358,32 @@
   bazel build --repository_disable_download //:it || fail "Failed to build"
 }
 
+function test_no_restarts_fetching_with_worker_thread() {
+  setup_starlark_repository
+
+  echo foo > file1
+  echo bar > file2
+
+  cat >test.bzl <<EOF
+def _impl(rctx):
+  print("hello world!")
+  print(rctx.read(Label("//:file1")))
+  print(rctx.read(Label("//:file2")))
+  rctx.file("BUILD", "filegroup(name='bar')")
+
+repo = repository_rule(implementation=_impl, local=True)
+EOF
+
+  # no worker thread, restarts twice
+  bazel build @foo//:bar --experimental_worker_for_repo_fetching=off >& $TEST_log \
+    || fail "Expected build to succeed"
+  expect_log_n "hello world!" 3
+
+  # platform worker thread, never restarts
+  bazel shutdown
+  bazel build @foo//:bar --experimental_worker_for_repo_fetching=platform >& $TEST_log \
+    || fail "Expected build to succeed"
+  expect_log_n "hello world!" 1
+}
+
 run_suite "local repository tests"