try harder

Author: fmeum

src/test/java/com/google/devtools/build/lib/remote/disk/DiskCacheClientTest.java

@@ -46,6 +46,7 @@
 import com.google.protobuf.ByteString;
 import com.google.protobuf.Message;
 import java.io.IOException;
+import java.io.OutputStream;
 import java.util.ArrayList;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
@@ -364,7 +365,7 @@ public void concurrentUploadDownload()
       var tasks = new ArrayList<Future<?>>();
       // Use 1 MB blobs to increase the window for concurrent access during write/rename.
       var contentSize = 1024 * 1024;
-      var numConcurrentOps = 5;
+      var numConcurrentOps = 10;
       for (int attempt = 0; attempt < 100; attempt++) {
         var contentArray = new byte[contentSize];
         // Fill with a pattern based on the attempt number.
@@ -373,24 +374,56 @@ public void concurrentUploadDownload()
         }
         var contentBytes = ByteString.copyFrom(contentArray);
         var contentDigest = DIGEST_UTIL.compute(contentArray);
-        // Use a latch to ensure all concurrent tasks start their upload at roughly the same time,
-        // maximizing the chance that multiple threads pass the "file exists" check before any
-        // completes the rename.
+        // Use a latch to ensure all concurrent tasks start at roughly the same time.
         var startLatch = new CountDownLatch(numConcurrentOps);
+        // Half the tasks do uploads, half do downloads with a slow OutputStream to keep the file
+        // open longer. This maximizes the chance of a rename failing because a download has the
+        // file open.
         for (int concurrentOp = 0; concurrentOp < numConcurrentOps; concurrentOp++) {
+          boolean isUploader = concurrentOp % 2 == 0;
           tasks.add(
               executor.submit(
                   () -> {
                     // Signal ready and wait for all tasks to be ready.
                     startLatch.countDown();
                     startLatch.await();
-                    nativeClient.uploadBlob(contentDigest, contentBytes);
-                    try (var out = ByteString.newOutput()) {
-                      getFromFuture(nativeClient.downloadBlob(contentDigest, out));
-                      var downloadedBytes = out.toByteString();
-                      assertThat(downloadedBytes).isEqualTo(contentBytes);
-                    } catch (CacheNotFoundException ignored) {
-                      // This task won the race over the upload task.
+                    if (isUploader) {
+                      getFromFuture(nativeClient.uploadBlob(contentDigest, contentBytes));
+                    } else {
+                      // Use a slow OutputStream that pauses periodically to keep the file open
+                      // longer during download.
+                      var out =
+                          new OutputStream() {
+                            private int bytesWritten = 0;
+
+                            @Override
+                            public void write(int b) throws IOException {
+                              bytesWritten++;
+                              maybeSleep();
+                            }
+
+                            @Override
+                            public void write(byte[] b, int off, int len) throws IOException {
+                              bytesWritten += len;
+                              maybeSleep();
+                            }
+
+                            private void maybeSleep() {
+                              // Sleep every 64KB to slow down the download.
+                              if (bytesWritten % (64 * 1024) < 100) {
+                                try {
+                                  Thread.sleep(1);
+                                } catch (InterruptedException e) {
+                                  Thread.currentThread().interrupt();
+                                }
+                              }
+                            }
+                          };
+                      try {
+                        getFromFuture(nativeClient.downloadBlob(contentDigest, out));
+                      } catch (CacheNotFoundException ignored) {
+                        // File not yet uploaded by another task.
+                      }
                     }
                     return null;
                   }));