From 4935f2799b16f86ca0d2127b203281ece7c74039 Mon Sep 17 00:00:00 2001
From: Yannic Bonenberger <yannic@engflow.com>
Date: Wed, 8 Mar 2023 18:49:52 +0100
Subject: [PATCH] [remote] upload: treat `ALREADY_EXISTS` as success

If the service returns an `ALREADY_EXISTS` error, then we assume that
the proper file is present remotely.

Prior art: https://github.com/bazelbuild/bazel/pull/12112

Fixes https://github.com/bazelbuild/bazel/issues/12111
---
 .../build/lib/remote/ByteStreamUploader.java  | 20 ++++++++++++++++++-
 .../build/lib/remote/GrpcCacheClientTest.java | 14 ++++++++++---
 2 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/remote/ByteStreamUploader.java b/src/main/java/com/google/devtools/build/lib/remote/ByteStreamUploader.java
index b7774ecf942a19..106207d6f10e69 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/ByteStreamUploader.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/ByteStreamUploader.java
@@ -30,6 +30,7 @@
 import com.google.bytestream.ByteStreamProto.WriteResponse;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Ascii;
+import com.google.common.base.Preconditions;
 import com.google.common.base.Strings;
 import com.google.common.util.concurrent.AsyncCallable;
 import com.google.common.util.concurrent.Futures;
@@ -391,7 +392,24 @@ private ListenableFuture<Long> upload(long pos) {
           channel -> {
             SettableFuture<Long> uploadResult = SettableFuture.create();
             bsAsyncStub(channel).write(new Writer(resourceName, chunker, pos, uploadResult));
-            return uploadResult;
+            return Futures.catchingAsync(
+                uploadResult,
+                Throwable.class,
+                throwable -> {
+                  Preconditions.checkNotNull(throwable);
+
+                  Status status = Status.fromThrowable(throwable);
+                  switch (status.getCode()) {
+                    case ALREADY_EXISTS:
+                      // Server indicated the blob already exists, so we translate the error to a
+                      // successful upload.
+                      return Futures.immediateFuture(chunker.getSize());
+
+                    default:
+                      return Futures.immediateFailedFuture(throwable);
+                  }
+                },
+                MoreExecutors.directExecutor());
           });
     }
   }
diff --git a/src/test/java/com/google/devtools/build/lib/remote/GrpcCacheClientTest.java b/src/test/java/com/google/devtools/build/lib/remote/GrpcCacheClientTest.java
index 60dac85b559779..6fb66f92ce00a3 100644
--- a/src/test/java/com/google/devtools/build/lib/remote/GrpcCacheClientTest.java
+++ b/src/test/java/com/google/devtools/build/lib/remote/GrpcCacheClientTest.java
@@ -749,9 +749,12 @@ public void testUploadCacheMissesWithRetries() throws Exception {
         fakeFileCache.createScratchInput(ActionInputHelper.fromPath("bar"), "x");
     final Digest bazDigest =
         fakeFileCache.createScratchInput(ActionInputHelper.fromPath("baz"), "z");
+    final Digest foobarDigest =
+        fakeFileCache.createScratchInput(ActionInputHelper.fromPath("foobar"), "foobar");
     final Path fooFile = execRoot.getRelative("a/foo");
     final Path barFile = execRoot.getRelative("bar");
     final Path bazFile = execRoot.getRelative("baz");
+    final Path foobarFile = execRoot.getRelative("foobar");
     ActionKey actionKey = DIGEST_UTIL.asActionKey(fooDigest); // Could be any key.
     barFile.setExecutable(true);
     serviceRegistry.addService(
@@ -769,6 +772,7 @@ public void findMissingBlobs(
                       .addMissingBlobDigests(fooDigest)
                       .addMissingBlobDigests(barDigest)
                       .addMissingBlobDigests(bazDigest)
+                      .addMissingBlobDigests(foobarDigest)
                       .build());
               responseObserver.onCompleted();
             } else {
@@ -782,6 +786,7 @@ public void findMissingBlobs(
     rb.addOutputFilesBuilder().setPath("a/foo").setDigest(fooDigest).setIsExecutable(true);
     rb.addOutputFilesBuilder().setPath("bar").setDigest(barDigest).setIsExecutable(true);
     rb.addOutputFilesBuilder().setPath("baz").setDigest(bazDigest).setIsExecutable(true);
+    rb.addOutputFilesBuilder().setPath("foobar").setDigest(foobarDigest).setIsExecutable(true);
     ActionResult result = rb.build();
     serviceRegistry.addService(
         new ActionCacheImplBase() {
@@ -837,6 +842,9 @@ public void onNext(WriteRequest request) {
                     } else if (resourceName.contains(bazDigest.getHash())) {
                       assertThat(dataStr).isEqualTo("z");
                       size = 1;
+                    } else if (resourceName.contains(foobarDigest.getHash())) {
+                      responseObserver.onError(Status.ALREADY_EXISTS.asRuntimeException());
+                      return;
                     } else {
                       fail("Unexpected resource name in upload: " + resourceName);
                     }
@@ -874,9 +882,9 @@ public void onError(Throwable t) {
         actionKey,
         Action.getDefaultInstance(),
         Command.getDefaultInstance(),
-        ImmutableList.<Path>of(fooFile, barFile, bazFile));
-    // 4 times for the errors, 3 times for the successful uploads.
-    Mockito.verify(mockByteStreamImpl, Mockito.times(7))
+        ImmutableList.<Path>of(fooFile, barFile, bazFile, foobarFile));
+    // 4 times for the errors, 4 times for the successful uploads.
+    Mockito.verify(mockByteStreamImpl, Mockito.times(8))
         .write(ArgumentMatchers.<StreamObserver<WriteResponse>>any());
   }