Fix flakey test testOverflowDisabledAsynchronous (#18891)

Signed-off-by: guojialiang <guojialiang.2012@bytedance.com>

Author: guojialiang92

server/src/main/java/org/opensearch/index/store/remote/utils/TransferManager.java

@@ -113,7 +113,7 @@ public IndexInput fetchBlob(BlobFetchRequest blobFetchRequest) throws IOExceptio
     }
 
     @ExperimentalApi
-    public void fetchBlobAsync(BlobFetchRequest blobFetchRequest) throws IOException {
+    public CompletableFuture<IndexInput> fetchBlobAsync(BlobFetchRequest blobFetchRequest) throws IOException {
         final Path key = blobFetchRequest.getFilePath();
         logger.trace("Asynchronous fetchBlob called for {}", key.toString());
         try {
@@ -133,7 +133,7 @@ public void fetchBlobAsync(BlobFetchRequest blobFetchRequest) throws IOException
             // decrement this reference _after_ creating the clone to be returned.
             // Making sure remote recovery thread-pool take care of background download
             try {
-                cacheEntry.asyncLoadIndexInput(threadPool.executor(ThreadPool.Names.REMOTE_RECOVERY));
+                return cacheEntry.asyncLoadIndexInput(threadPool.executor(ThreadPool.Names.REMOTE_RECOVERY));
             } catch (Exception exception) {
                 fileCache.decRef(key);
                 throw exception;

server/src/test/java/org/opensearch/index/store/remote/utils/TransferManagerTestCase.java

@@ -28,6 +28,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -313,14 +314,14 @@ public void testAsyncFetchesToDifferentBlobsDoNotBlockOnEachOther() throws Excep
         executor.shutdownNow();
     }
 
-    public void testRefCount() throws IOException, InterruptedException {
+    public void testRefCount() throws Exception {
         List<BlobFetchRequest.BlobPart> blobParts = new ArrayList<>();
         String blobname = "test-blob";
         blobParts.add(new BlobFetchRequest.BlobPart("blob", 0, EIGHT_MB));
         BlobFetchRequest blobFetchRequest = BlobFetchRequest.builder().fileName(blobname).directory(directory).blobParts(blobParts).build();
         // It will trigger async load
-        transferManager.fetchBlobAsync(blobFetchRequest);
-        Thread.sleep(2000);
+        CompletableFuture<IndexInput> future = transferManager.fetchBlobAsync(blobFetchRequest);
+        future.get();
         assertEquals(Optional.of(0), Optional.of(fileCache.getRef(blobFetchRequest.getFilePath())));
         assertNotNull(fileCache.get(blobFetchRequest.getFilePath()));
         fileCache.decRef(blobFetchRequest.getFilePath());
@@ -347,8 +348,6 @@ private IndexInput asyncFetchBlobWithName(String blobname) throws IOException {
         BlobFetchRequest blobFetchRequest = BlobFetchRequest.builder().fileName(blobname).directory(directory).blobParts(blobParts).build();
         // It will trigger async load
         transferManager.fetchBlobAsync(blobFetchRequest);
-        assertNotNull(fileCache.get(blobFetchRequest.getFilePath()));
-        fileCache.decRef(blobFetchRequest.getFilePath());
         // Making the read call to fetch from file cache
         return transferManager.fetchBlob(blobFetchRequest);
     }