Stop remote blob upload if upload is complete. (bazelbuild#14467)

benjaminp · web-flow · commit 90965b072eb4 · 2021-12-23T04:40:01.000+11:00
If a ByteStream/Write RPC fails, but ByteStream/QueryWriteStatus reveals the upload is in fact complete, avoid a NullPointerException. This CL is the dual fix of 78b89a0 for uploads. On bazel-6.0.0-pre.20211117.1, I observed: ``` java.lang.NullPointerException at com.google.devtools.build.lib.remote.Chunker.seek(Chunker.java:156) at com.google.devtools.build.lib.remote.ByteStreamUploader$AsyncUpload.lambda$start$0(ByteStreamUploader.java:416) at com.google.devtools.build.lib.remote.Retrier.executeAsync(Retrier.java:277) at com.google.devtools.build.lib.remote.Retrier.lambda$onExecuteAsyncFailure$1(Retrier.java:293) at com.google.common.util.concurrent.TrustedListenableFutureTask$TrustedFutureInterruptibleAsyncTask.runInterruptibly(TrustedListenableFutureTask.java:160) at com.google.common.util.concurrent.TrustedListenableFutureTask$TrustedFutureInterruptibleAsyncTask.runInterruptibly(TrustedListenableFutureTask.java:143) at com.google.common.util.concurrent.InterruptibleTask.run(InterruptibleTask.java:69) at com.google.common.util.concurrent.TrustedListenableFutureTask.run(TrustedListenableFutureTask.java:78) at java.base/java.util.concurrent.Executors$RunnableAdapter.call(Unknown Source) at com.google.common.util.concurrent.TrustedListenableFutureTask$TrustedFutureInterruptibleTask.runInterruptibly(TrustedListenableFutureTask.java:125) at com.google.common.util.concurrent.InterruptibleTask.run(InterruptibleTask.java:69) at com.google.common.util.concurrent.TrustedListenableFutureTask.run(TrustedListenableFutureTask.java:78) at java.base/java.util.concurrent.Executors$RunnableAdapter.call(Unknown Source) at java.base/java.util.concurrent.FutureTask.run(Unknown Source) at java.base/java.util.concurrent.ScheduledThreadPoolExecutor$ScheduledFutureTask.run(Unknown Source) ``` Closes bazelbuild#14464. PiperOrigin-RevId: 417795715
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
@@ -248,7 +248,7 @@ public ListenableFuture<Void> uploadBlobAsync(
       checkState(!isShutdown, "Must not call uploadBlobs after shutdown.");
 
       if (!forceUpload && uploadedBlobs.contains(HashCode.fromString(digest.getHash()))) {
-        return Futures.immediateFuture(null);
+        return immediateVoidFuture();
       }
 
       ListenableFuture<Void> inProgress = uploadsInProgress.get(digest);
@@ -410,7 +410,7 @@ ListenableFuture<Void> start() {
               () ->
                   retrier.executeAsync(
                       () -> {
-                        if (chunker.getSize() == 0) {
+                        if (chunker.getSize() == committedOffset.get()) {
                           return immediateVoidFuture();
                         }
                         try {
@@ -426,7 +426,7 @@ ListenableFuture<Void> start() {
                         if (chunker.hasNext()) {
                           return callAndQueryOnFailure(committedOffset, progressiveBackoff);
                         }
-                        return Futures.immediateFuture(null);
+                        return immediateVoidFuture();
                       },
                       progressiveBackoff),
               callCredentialsProvider);
@@ -448,7 +448,7 @@ ListenableFuture<Void> start() {
                 return Futures.immediateFailedFuture(new IOException(message));
               }
             }
-            return Futures.immediateFuture(null);
+            return immediateVoidFuture();
           },
           MoreExecutors.directExecutor());
     }
@@ -536,7 +536,7 @@ private ListenableFuture<Void> query(
               progressiveBackoff.reset();
             }
             committedOffset.set(committedSize);
-            return Futures.immediateFuture(null);
+            return immediateVoidFuture();
           },
           MoreExecutors.directExecutor());
     }
diff --git a/src/test/java/com/google/devtools/build/lib/remote/ByteStreamUploaderTest.java b/src/test/java/com/google/devtools/build/lib/remote/ByteStreamUploaderTest.java
@@ -1517,6 +1517,68 @@ public void onCompleted() {
     blockUntilInternalStateConsistent(uploader);
   }
 
+  @Test
+  public void failureAfterUploadCompletes() throws Exception {
+    AtomicInteger numUploads = new AtomicInteger();
+    RemoteRetrier retrier =
+        TestUtils.newRemoteRetrier(
+            () -> mockBackoff, e -> e instanceof StatusRuntimeException, retryService);
+    ByteStreamUploader uploader =
+        new ByteStreamUploader(
+            INSTANCE_NAME,
+            new ReferenceCountedChannel(channelConnectionFactory),
+            CallCredentialsProvider.NO_CREDENTIALS,
+            /* callTimeoutSecs= */ 60,
+            retrier);
+
+    byte[] blob = new byte[CHUNK_SIZE - 1];
+    new Random().nextBytes(blob);
+
+    serviceRegistry.addService(
+        new ByteStreamImplBase() {
+          @Override
+          public StreamObserver<WriteRequest> write(StreamObserver<WriteResponse> streamObserver) {
+            numUploads.incrementAndGet();
+            return new StreamObserver<WriteRequest>() {
+              @Override
+              public void onNext(WriteRequest writeRequest) {}
+
+              @Override
+              public void onError(Throwable throwable) {
+                fail("onError should never be called.");
+              }
+
+              @Override
+              public void onCompleted() {
+                streamObserver.onNext(
+                    WriteResponse.newBuilder().setCommittedSize(blob.length).build());
+                streamObserver.onError(Status.UNAVAILABLE.asException());
+              }
+            };
+          }
+
+          @Override
+          public void queryWriteStatus(
+              QueryWriteStatusRequest request, StreamObserver<QueryWriteStatusResponse> response) {
+            response.onNext(
+                QueryWriteStatusResponse.newBuilder()
+                    .setCommittedSize(blob.length)
+                    .setComplete(true)
+                    .build());
+            response.onCompleted();
+          }
+        });
+
+    Chunker chunker = Chunker.builder().setInput(blob).setChunkSize(CHUNK_SIZE).build();
+    HashCode hash = HashCode.fromString(DIGEST_UTIL.compute(blob).getHash());
+
+    uploader.uploadBlob(context, hash, chunker, true);
+
+    blockUntilInternalStateConsistent(uploader);
+
+    assertThat(numUploads.get()).isEqualTo(1);
+  }
+
   @Test
   public void testCompressedUploads() throws Exception {
     RemoteRetrier retrier =

Original file line number	Diff line number	Diff line change
`@@ -248,7 +248,7 @@ public ListenableFuture<Void> uploadBlobAsync(`
`248`	`248`	`checkState(!isShutdown, "Must not call uploadBlobs after shutdown.");`
`249`	`249`
`250`	`250`	`if (!forceUpload && uploadedBlobs.contains(HashCode.fromString(digest.getHash()))) {`
`251`		`- return Futures.immediateFuture(null);`
	`251`	`+ return immediateVoidFuture();`
`252`	`252`	`}`
`253`	`253`
`254`	`254`	`ListenableFuture<Void> inProgress = uploadsInProgress.get(digest);`
`@@ -410,7 +410,7 @@ ListenableFuture<Void> start() {`
`410`	`410`	`() ->`
`411`	`411`	`retrier.executeAsync(`
`412`	`412`	`() -> {`
`413`		`- if (chunker.getSize() == 0) {`
	`413`	`+ if (chunker.getSize() == committedOffset.get()) {`
`414`	`414`	`return immediateVoidFuture();`
`415`	`415`	`}`
`416`	`416`	`try {`
`@@ -426,7 +426,7 @@ ListenableFuture<Void> start() {`
`426`	`426`	`if (chunker.hasNext()) {`
`427`	`427`	`return callAndQueryOnFailure(committedOffset, progressiveBackoff);`
`428`	`428`	`}`
`429`		`- return Futures.immediateFuture(null);`
	`429`	`+ return immediateVoidFuture();`
`430`	`430`	`},`
`431`	`431`	`progressiveBackoff),`
`432`	`432`	`callCredentialsProvider);`
`@@ -448,7 +448,7 @@ ListenableFuture<Void> start() {`
`448`	`448`	`return Futures.immediateFailedFuture(new IOException(message));`
`449`	`449`	`}`
`450`	`450`	`}`
`451`		`- return Futures.immediateFuture(null);`
	`451`	`+ return immediateVoidFuture();`
`452`	`452`	`},`
`453`	`453`	`MoreExecutors.directExecutor());`
`454`	`454`	`}`
`@@ -536,7 +536,7 @@ private ListenableFuture<Void> query(`
`536`	`536`	`progressiveBackoff.reset();`
`537`	`537`	`}`
`538`	`538`	`committedOffset.set(committedSize);`
`539`		`- return Futures.immediateFuture(null);`
	`539`	`+ return immediateVoidFuture();`
`540`	`540`	`},`
`541`	`541`	`MoreExecutors.directExecutor());`
`542`	`542`	`}`