Decouple Skyframe worker thread stuff from repo fetching, part 2
Following up to https://github.com/bazelbuild/bazel/commit/d24f94734266c4331ba8324a470d60ae68c26edc, this CL consolidates the API surface of using worker threads down to one class and one method, `WorkerSkyKeyComputeState.startOrContinueWork`. See Javadoc for that method for more information.
Also took this chance to move these two "Worker"-related classes to jcg/devtools/build/skyframe, as they're now completely decoupled from repo fetching.
Work towards https://github.com/bazelbuild/bazel/issues/22729
PiperOrigin-RevId: 669390115
Change-Id: I9f9a743c9f1a92bc65af8bccb98c53bc22439204
diff --git a/src/main/java/com/google/devtools/build/skyframe/WorkerSkyFunctionEnvironment.java b/src/main/java/com/google/devtools/build/skyframe/WorkerSkyFunctionEnvironment.java
new file mode 100644
index 0000000..1ce7617
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/skyframe/WorkerSkyFunctionEnvironment.java
@@ -0,0 +1,177 @@
+// 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.skyframe;
+
+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.lib.supplier.InterruptibleSupplier;
+import java.util.function.Supplier;
+import javax.annotation.Nullable;
+
+/**
+ * A {@link SkyFunction.Environment} implementation designed to be used in a different thread (the
+ * "worker 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 WorkerSkyFunctionEnvironment
+ implements SkyFunction.Environment, ExtendedEventHandler, SkyframeLookupResult {
+ private SkyFunction.Environment delegate;
+ private final InterruptibleSupplier<SkyFunction.Environment> newDelegateSupplier;
+
+ WorkerSkyFunctionEnvironment(
+ SkyFunction.Environment initialDelegate,
+ InterruptibleSupplier<SkyFunction.Environment> newDelegateSupplier) {
+ this.delegate = initialDelegate;
+ this.newDelegateSupplier = newDelegateSupplier;
+ }
+
+ @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 = newDelegateSupplier.get();
+ 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 = newDelegateSupplier.get();
+ 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 inErrorBubbling() {
+ return delegate.inErrorBubbling();
+ }
+
+ @Override
+ public void dependOnFuture(ListenableFuture<?> future) {
+ delegate.dependOnFuture(future);
+ }
+
+ @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();
+ }
+}
