[6.3.0] Wait for outputs downloads before emitting local BEP events that refe…

src/main/java/com/google/devtools/build/lib/buildeventservice/BuildEventServiceModule.java

@@ -38,6 +38,7 @@
 import com.google.devtools.build.lib.buildeventservice.client.BuildEventServiceClient;
 import com.google.devtools.build.lib.buildeventstream.AnnounceBuildEventTransportsEvent;
 import com.google.devtools.build.lib.buildeventstream.BuildEventArtifactUploader;
+import com.google.devtools.build.lib.buildeventstream.BuildEventLocalFileSynchronizer;
 import com.google.devtools.build.lib.buildeventstream.BuildEventProtocolOptions;
 import com.google.devtools.build.lib.buildeventstream.BuildEventStreamProtos.Aborted.AbortReason;
 import com.google.devtools.build.lib.buildeventstream.BuildEventTransport;
@@ -90,6 +91,7 @@
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 
 /**
@@ -753,6 +755,19 @@ private ImmutableSet<BuildEventTransport> createBepTransports(
       CountingArtifactGroupNamer artifactGroupNamer)
       throws IOException {
     ImmutableSet.Builder<BuildEventTransport> bepTransportsBuilder = new ImmutableSet.Builder<>();
+    BuildEventLocalFileSynchronizer synchronizer =
+        localFiles -> {
+          var outputService = cmdEnv.getOutputService();
+          if (outputService == null || localFiles.isEmpty()) {
+            return Futures.immediateVoidFuture();
+          }
+
+          var files =
+              localFiles.stream()
+                  .map(localFile -> localFile.path.asFragment())
+                  .collect(Collectors.toList());
+          return outputService.waitOutputDownloads(files);
+        };
 
     if (!Strings.isNullOrEmpty(besStreamOptions.buildEventTextFile)) {
       try {
@@ -766,7 +781,11 @@ private ImmutableSet<BuildEventTransport> createBepTransports(
                 : new LocalFilesArtifactUploader();
         bepTransportsBuilder.add(
             new TextFormatFileTransport(
-                bepTextOutputStream, bepOptions, localFileUploader, artifactGroupNamer));
+                bepTextOutputStream,
+                bepOptions,
+                localFileUploader,
+                artifactGroupNamer,
+                synchronizer));
       } catch (IOException exception) {
         // TODO(b/125216340): Consider making this a warning instead of an error once the
         //  associated bug has been resolved.
@@ -793,7 +812,11 @@ private ImmutableSet<BuildEventTransport> createBepTransports(
                 : new LocalFilesArtifactUploader();
         bepTransportsBuilder.add(
             new BinaryFormatFileTransport(
-                bepBinaryOutputStream, bepOptions, localFileUploader, artifactGroupNamer));
+                bepBinaryOutputStream,
+                bepOptions,
+                localFileUploader,
+                artifactGroupNamer,
+                synchronizer));
       } catch (IOException exception) {
         // TODO(b/125216340): Consider making this a warning instead of an error once the
         //  associated bug has been resolved.
@@ -819,7 +842,11 @@ private ImmutableSet<BuildEventTransport> createBepTransports(
                 : new LocalFilesArtifactUploader();
         bepTransportsBuilder.add(
             new JsonFormatFileTransport(
-                bepJsonOutputStream, bepOptions, localFileUploader, artifactGroupNamer));
+                bepJsonOutputStream,
+                bepOptions,
+                localFileUploader,
+                artifactGroupNamer,
+                synchronizer));
       } catch (IOException exception) {
         // TODO(b/125216340): Consider making this a warning instead of an error once the
         //  associated bug has been resolved.

src/main/java/com/google/devtools/build/lib/buildeventstream/BuildEventLocalFileSynchronizer.java

@@ -0,0 +1,29 @@
+// 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.buildeventstream;
+
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.devtools.build.lib.buildeventstream.BuildEvent.LocalFile;
+import java.util.Collection;
+
+/**
+ * An interface that is used to wait for downloads of remote outputs before sending out local BEP
+ * events that reference these outputs.
+ */
+public interface BuildEventLocalFileSynchronizer {
+  BuildEventLocalFileSynchronizer NO_OP = localFiles -> Futures.immediateVoidFuture();
+
+  ListenableFuture<Void> waitForLocalFileDownloads(Collection<LocalFile> localFiles);
+}

src/main/java/com/google/devtools/build/lib/buildeventstream/transports/BinaryFormatFileTransport.java

@@ -17,6 +17,7 @@
 import com.google.devtools.build.lib.buildeventstream.ArtifactGroupNamer;
 import com.google.devtools.build.lib.buildeventstream.BuildEvent;
 import com.google.devtools.build.lib.buildeventstream.BuildEventArtifactUploader;
+import com.google.devtools.build.lib.buildeventstream.BuildEventLocalFileSynchronizer;
 import com.google.devtools.build.lib.buildeventstream.BuildEventProtocolOptions;
 import com.google.devtools.build.lib.buildeventstream.BuildEventStreamProtos;
 import com.google.devtools.build.lib.buildeventstream.BuildEventTransport;
@@ -34,8 +35,9 @@ public BinaryFormatFileTransport(
       BufferedOutputStream outputStream,
       BuildEventProtocolOptions options,
       BuildEventArtifactUploader uploader,
-      ArtifactGroupNamer namer) {
-    super(outputStream, options, uploader, namer);
+      ArtifactGroupNamer namer,
+      BuildEventLocalFileSynchronizer synchronizer) {
+    super(outputStream, options, uploader, namer, synchronizer);
   }
 
   @Override

src/main/java/com/google/devtools/build/lib/buildeventstream/transports/FileTransport.java

@@ -28,6 +28,7 @@
 import com.google.devtools.build.lib.buildeventstream.BuildEvent;
 import com.google.devtools.build.lib.buildeventstream.BuildEventArtifactUploader;
 import com.google.devtools.build.lib.buildeventstream.BuildEventContext;
+import com.google.devtools.build.lib.buildeventstream.BuildEventLocalFileSynchronizer;
 import com.google.devtools.build.lib.buildeventstream.BuildEventProtocolOptions;
 import com.google.devtools.build.lib.buildeventstream.BuildEventStreamProtos;
 import com.google.devtools.build.lib.buildeventstream.BuildEventTransport;
@@ -66,6 +67,7 @@ abstract class FileTransport implements BuildEventTransport {
   private final BuildEventArtifactUploader uploader;
   private final SequentialWriter writer;
   private final ArtifactGroupNamer namer;
+  private final BuildEventLocalFileSynchronizer synchronizer;
 
   private final ScheduledExecutorService timeoutExecutor =
       MoreExecutors.listeningDecorator(
@@ -76,12 +78,14 @@ abstract class FileTransport implements BuildEventTransport {
       BufferedOutputStream outputStream,
       BuildEventProtocolOptions options,
       BuildEventArtifactUploader uploader,
-      ArtifactGroupNamer namer) {
+      ArtifactGroupNamer namer,
+      BuildEventLocalFileSynchronizer synchronizer) {
     this.uploader = uploader;
     this.options = options;
     this.writer =
         new SequentialWriter(outputStream, this::serializeEvent, uploader, timeoutExecutor);
     this.namer = namer;
+    this.synchronizer = synchronizer;
   }
 
   @ThreadSafe
@@ -280,12 +284,14 @@ private ListenableFuture<BuildEventStreamProtos.BuildEvent> asStreamProto(
       BuildEvent event, ArtifactGroupNamer namer) {
     checkNotNull(event);
 
+    var localFiles = event.referencedLocalFiles();
+    ListenableFuture<?> localFileDownloads = synchronizer.waitForLocalFileDownloads(localFiles);
     ListenableFuture<PathConverter> converterFuture =
-        uploader.uploadReferencedLocalFiles(event.referencedLocalFiles());
+        uploader.uploadReferencedLocalFiles(localFiles);
     ListenableFuture<?> remoteUploads =
         uploader.waitForRemoteUploads(event.remoteUploads(), timeoutExecutor);
     return Futures.transform(
-        Futures.allAsList(converterFuture, remoteUploads),
+        Futures.allAsList(localFileDownloads, converterFuture, remoteUploads),
         results -> {
           BuildEventContext context =
               new BuildEventContext() {

src/main/java/com/google/devtools/build/lib/buildeventstream/transports/JsonFormatFileTransport.java

@@ -18,6 +18,7 @@
 
 import com.google.devtools.build.lib.buildeventstream.ArtifactGroupNamer;
 import com.google.devtools.build.lib.buildeventstream.BuildEventArtifactUploader;
+import com.google.devtools.build.lib.buildeventstream.BuildEventLocalFileSynchronizer;
 import com.google.devtools.build.lib.buildeventstream.BuildEventProtocolOptions;
 import com.google.devtools.build.lib.buildeventstream.BuildEventStreamProtos;
 import com.google.devtools.build.lib.buildeventstream.BuildEventTransport;
@@ -34,8 +35,9 @@ public JsonFormatFileTransport(
       BufferedOutputStream outputStream,
       BuildEventProtocolOptions options,
       BuildEventArtifactUploader uploader,
-      ArtifactGroupNamer namer) {
-    super(outputStream, options, uploader, namer);
+      ArtifactGroupNamer namer,
+      BuildEventLocalFileSynchronizer synchronizer) {
+    super(outputStream, options, uploader, namer, synchronizer);
   }
 
   @Override

src/main/java/com/google/devtools/build/lib/buildeventstream/transports/TextFormatFileTransport.java

@@ -18,6 +18,7 @@
 
 import com.google.devtools.build.lib.buildeventstream.ArtifactGroupNamer;
 import com.google.devtools.build.lib.buildeventstream.BuildEventArtifactUploader;
+import com.google.devtools.build.lib.buildeventstream.BuildEventLocalFileSynchronizer;
 import com.google.devtools.build.lib.buildeventstream.BuildEventProtocolOptions;
 import com.google.devtools.build.lib.buildeventstream.BuildEventStreamProtos;
 import com.google.devtools.build.lib.buildeventstream.BuildEventTransport;
@@ -35,8 +36,9 @@ public TextFormatFileTransport(
       BufferedOutputStream outputStream,
       BuildEventProtocolOptions options,
       BuildEventArtifactUploader uploader,
-      ArtifactGroupNamer namer) {
-    super(outputStream, options, uploader, namer);
+      ArtifactGroupNamer namer,
+      BuildEventLocalFileSynchronizer synchronizer) {
+    super(outputStream, options, uploader, namer, synchronizer);
   }
 
   @Override

src/main/java/com/google/devtools/build/lib/remote/AbstractActionInputPrefetcher.java

@@ -48,6 +48,7 @@
 import com.google.devtools.build.lib.events.Reporter;
 import com.google.devtools.build.lib.remote.common.CacheNotFoundException;
 import com.google.devtools.build.lib.remote.util.AsyncTaskCache;
+import com.google.devtools.build.lib.remote.util.RxFutures;
 import com.google.devtools.build.lib.remote.util.RxUtils.TransferResult;
 import com.google.devtools.build.lib.remote.util.TempPathGenerator;
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
@@ -60,13 +61,15 @@
 import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.Deque;
 import java.util.List;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 
 /**
@@ -760,6 +763,12 @@ public void flushOutputTree() throws InterruptedException {
     downloadCache.awaitInProgressTasks();
   }
 
+  public ListenableFuture<Void> waitDownloads(Collection<PathFragment> files) {
+    var convertedFiles = files.stream().map(file -> execRoot.getFileSystem().getPath(file)).collect(
+        Collectors.toList());
+    return RxFutures.toListenableFuture(downloadCache.waitInProgressTasks(convertedFiles));
+  }
+
   public ImmutableSet<ActionInput> getMissingActionInputs() {
     return ImmutableSet.copyOf(missingActionInputs);
   }

src/main/java/com/google/devtools/build/lib/remote/RemoteOutputService.java

@@ -19,6 +19,8 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.eventbus.Subscribe;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
 import com.google.devtools.build.lib.actions.Action;
 import com.google.devtools.build.lib.actions.ActionExecutionMetadata;
 import com.google.devtools.build.lib.actions.ActionInputMap;
@@ -34,10 +36,12 @@
 import com.google.devtools.build.lib.vfs.FileSystem;
 import com.google.devtools.build.lib.vfs.ModifiedFileSet;
 import com.google.devtools.build.lib.vfs.OutputService;
+import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.devtools.build.lib.vfs.Root;
 import com.google.devtools.build.skyframe.SkyFunction.Environment;
 import java.io.IOException;
+import java.util.Collection;
 import java.util.Map;
 import java.util.UUID;
 import javax.annotation.Nullable;
@@ -111,6 +115,13 @@ public void flushOutputTree() throws InterruptedException {
     }
   }
 
+  public ListenableFuture<Void> waitOutputDownloads(Collection<PathFragment> files) {
+    if (actionInputFetcher != null) {
+      return actionInputFetcher.waitDownloads(files);
+    }
+    return Futures.immediateVoidFuture();
+  }
+
   @Override
   public void finalizeBuild(boolean buildSuccessful) {
     // Intentionally left empty.

src/main/java/com/google/devtools/build/lib/remote/util/AsyncTaskCache.java

@@ -28,6 +28,7 @@
 import io.reactivex.rxjava3.functions.Action;
 import io.reactivex.rxjava3.subjects.AsyncSubject;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -357,6 +358,32 @@ public void shutdown() {
     }
   }
 
+  /**
+   * Returns a {@link Completable} which will completes once the in-progress tasks identified by
+   * {@code keys} are completed. Tasks submitted after the call are not waited.
+   */
+  public Completable waitInProgressTasks(Collection<KeyT> keys) {
+    return Completable.defer(
+        () -> {
+          List<Execution> executions = new ArrayList<>();
+          synchronized (lock) {
+            for (var key : keys) {
+              var execution = inProgress.get(key);
+              if (execution != null) {
+                executions.add(execution);
+              }
+            }
+          }
+          if (executions.isEmpty()) {
+            return Completable.complete();
+          }
+
+          return Completable.fromPublisher(
+              Flowable.fromIterable(executions)
+                  .flatMapSingle(e -> Single.fromObservable(e.completion)));
+        });
+  }
+
   /**
    * Waits for the in-progress tasks to finish. Any tasks that are submitted after the call are not
    * waited.
@@ -535,6 +562,14 @@ public void shutdown() {
       cache.shutdown();
     }
 
+    /**
+     * Returns a {@link Completable} which will completes once the in-progress tasks identified by
+     * {@code keys} are completed. Tasks submitted after the call are not waited.
+     */
+    public Completable waitInProgressTasks(Collection<KeyT> keys) {
+      return cache.waitInProgressTasks(keys);
+    }
+
     /**
      * Waits for the in-progress tasks to finish. Any tasks that are submitted after the call are
      * not waited.

src/main/java/com/google/devtools/build/lib/vfs/OutputService.java

@@ -17,6 +17,8 @@
 import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
 import com.google.devtools.build.lib.actions.Action;
 import com.google.devtools.build.lib.actions.ActionExecutionMetadata;
 import com.google.devtools.build.lib.actions.ActionInputMap;
@@ -33,6 +35,7 @@
 import com.google.devtools.build.lib.util.AbruptExitException;
 import com.google.devtools.build.skyframe.SkyFunction.Environment;
 import java.io.IOException;
+import java.util.Collection;
 import java.util.Map;
 import java.util.UUID;
 import javax.annotation.Nullable;
@@ -121,6 +124,10 @@ ModifiedFileSet startBuild(EventHandler eventHandler, UUID buildId, boolean fina
   /** Flush and wait for in-progress downloads. */
   default void flushOutputTree() throws InterruptedException {}
 
+  default ListenableFuture<Void> waitOutputDownloads(Collection<PathFragment> files) {
+    return Futures.immediateVoidFuture();
+  }
+
   /**
    * Finish the build.
    *