[GitHub] [incubator-uniffle] advancedxy commented on a diff in pull request #603: [#291] feat: Support async read

["advancedxy (via GitHub)" <gi...@apache.org>]

advancedxy commented on code in PR #603:
URL: https://github.com/apache/incubator-uniffle/pull/603#discussion_r1118704402


##########
client/src/main/java/org/apache/uniffle/client/request/CreateShuffleReadClientRequest.java:
##########
@@ -45,6 +45,44 @@ public class CreateShuffleReadClientRequest {
   private IdHelper idHelper;
   private ShuffleDataDistributionType shuffleDataDistributionType = ShuffleDataDistributionType.NORMAL;
   private boolean expectedTaskIdsBitmapFilterEnable = false;
+  private int prefetchBatchNum;
+
+  public CreateShuffleReadClientRequest(
+      String appId,
+      int shuffleId,
+      int partitionId,
+      String storageType,
+      String basePath,
+      int indexReadLimit,
+      int readBufferSize,
+      int partitionNumPerRange,
+      int partitionNum,
+      Roaring64NavigableMap blockIdBitmap,
+      Roaring64NavigableMap taskIdBitmap,
+      List<ShuffleServerInfo> shuffleServerInfoList,
+      Configuration hadoopConf,
+      IdHelper idHelper,
+      boolean expectedTaskIdsBitmapFilterEnable,
+      ShuffleDataDistributionType shuffleDataDistributionType,
+      int prefetchSize) {

Review Comment:
   -> prefetchBatchNum?



##########
client/src/main/java/org/apache/uniffle/client/util/AsyncDataLoader.java:
##########
@@ -0,0 +1,147 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.uniffle.client.util;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.Iterator;
+import java.util.List;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Supplier;
+
+import com.google.common.collect.Queues;
+import org.apache.commons.collections.CollectionUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.uniffle.common.exception.RssException;
+
+public class AsyncDataLoader<T> implements Closeable {
+  private static final Logger LOGGER = LoggerFactory.getLogger(AsyncDataLoader.class);
+
+  private final BlockingQueue<BatchDataHolder> dataQueue;
+  private BatchDataHolder<T> currentHolder;
+  private volatile boolean noNewlyData = false;
+  private final Semaphore semaphore;
+
+  // Internal sync vars
+  private volatile int finishedTaskCnt = 0;
+  private volatile boolean dataFinished = false;
+
+  private ExecutorService executorService;
+
+  public AsyncDataLoader(int prefetchBatch, Supplier<List<T>> dataLoadFunc) {
+    /**
+     * Currently, only 1 thread will fetch data, which is limited by no thread-safe
+     * {@link org.apache.uniffle.storage.handler.api.ClientReadHandler}.
+     *
+     * TODO: Support multiple threads to fetch data.
+     */
+    this(prefetchBatch, 1, dataLoadFunc);
+  }
+
+  protected AsyncDataLoader(int prefetchBatch, int concurrency, Supplier<List<T>> dataLoadFunc) {
+    this.dataQueue = Queues.newLinkedBlockingQueue(prefetchBatch);
+    this.semaphore = new Semaphore(prefetchBatch);
+
+    this.executorService = Executors.newFixedThreadPool(concurrency);
+    for (int i = 0; i < concurrency; i++) {
+      executorService.submit(() -> {
+        try {
+          while (true) {
+            if (!semaphore.tryAcquire(1, TimeUnit.MILLISECONDS)) {
+              if (dataFinished) {
+                break;
+              }
+              continue;
+            }
+            List<T> result = dataLoadFunc.get();
+            if (CollectionUtils.isEmpty(result)) {
+              this.dataFinished = true;
+              break;
+            }
+            dataQueue.put(new BatchDataHolder(result));
+          }
+        } catch (Exception e) {
+          LOGGER.error("Errors on fetching data.", e);

Review Comment:
   If any exception occurred, it should be propagated to `get` call?



##########
client/src/main/java/org/apache/uniffle/client/util/AsyncDataLoader.java:
##########
@@ -0,0 +1,147 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.uniffle.client.util;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.Iterator;
+import java.util.List;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Supplier;
+
+import com.google.common.collect.Queues;
+import org.apache.commons.collections.CollectionUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.uniffle.common.exception.RssException;
+
+public class AsyncDataLoader<T> implements Closeable {
+  private static final Logger LOGGER = LoggerFactory.getLogger(AsyncDataLoader.class);
+
+  private final BlockingQueue<BatchDataHolder> dataQueue;
+  private BatchDataHolder<T> currentHolder;
+  private volatile boolean noNewlyData = false;
+  private final Semaphore semaphore;
+
+  // Internal sync vars
+  private volatile int finishedTaskCnt = 0;
+  private volatile boolean dataFinished = false;
+
+  private ExecutorService executorService;
+
+  public AsyncDataLoader(int prefetchBatch, Supplier<List<T>> dataLoadFunc) {
+    /**
+     * Currently, only 1 thread will fetch data, which is limited by no thread-safe
+     * {@link org.apache.uniffle.storage.handler.api.ClientReadHandler}.
+     *
+     * TODO: Support multiple threads to fetch data.
+     */
+    this(prefetchBatch, 1, dataLoadFunc);
+  }
+
+  protected AsyncDataLoader(int prefetchBatch, int concurrency, Supplier<List<T>> dataLoadFunc) {
+    this.dataQueue = Queues.newLinkedBlockingQueue(prefetchBatch);
+    this.semaphore = new Semaphore(prefetchBatch);
+
+    this.executorService = Executors.newFixedThreadPool(concurrency);
+    for (int i = 0; i < concurrency; i++) {
+      executorService.submit(() -> {
+        try {
+          while (true) {
+            if (!semaphore.tryAcquire(1, TimeUnit.MILLISECONDS)) {
+              if (dataFinished) {
+                break;
+              }
+              continue;
+            }
+            List<T> result = dataLoadFunc.get();
+            if (CollectionUtils.isEmpty(result)) {
+              this.dataFinished = true;
+              break;
+            }
+            dataQueue.put(new BatchDataHolder(result));
+          }
+        } catch (Exception e) {
+          LOGGER.error("Errors on fetching data.", e);
+        }
+        finishedTaskCnt++;
+        if (finishedTaskCnt == concurrency) {
+          LOGGER.info("All threads of async fetching data are finished.");
+          this.noNewlyData = true;
+        }
+      });
+    }
+  }
+
+  public T get() {
+    try {
+      return doGet();
+    } catch (Exception e) {
+      throw new RssException("Errors on getting data.", e);
+    }
+  }
+
+  public T doGet() throws Exception {
+    while (currentHolder == null) {
+      final BatchDataHolder<T> holder = dataQueue.poll(1, TimeUnit.MILLISECONDS);
+      if (holder == null && noNewlyData) {
+        return null;
+      }
+      if (holder == null) {
+        continue;
+      }
+      currentHolder = holder;
+    }
+    T val = currentHolder.next();
+    if (!currentHolder.hasNext()) {
+      currentHolder = null;
+      semaphore.release();
+    }
+    return val;
+  }
+
+  class BatchDataHolder<T> implements Iterator<T> {

Review Comment:
   I didn't see why this is valuable? Why couldn't just pass Iterator<T> to where `BatchDataHolder` was passed.



##########
client/src/main/java/org/apache/uniffle/client/util/AsyncDataLoader.java:
##########
@@ -0,0 +1,147 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.uniffle.client.util;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.Iterator;
+import java.util.List;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Supplier;
+
+import com.google.common.collect.Queues;
+import org.apache.commons.collections.CollectionUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.uniffle.common.exception.RssException;
+
+public class AsyncDataLoader<T> implements Closeable {
+  private static final Logger LOGGER = LoggerFactory.getLogger(AsyncDataLoader.class);
+
+  private final BlockingQueue<BatchDataHolder> dataQueue;
+  private BatchDataHolder<T> currentHolder;
+  private volatile boolean noNewlyData = false;
+  private final Semaphore semaphore;
+
+  // Internal sync vars
+  private volatile int finishedTaskCnt = 0;
+  private volatile boolean dataFinished = false;
+
+  private ExecutorService executorService;
+
+  public AsyncDataLoader(int prefetchBatch, Supplier<List<T>> dataLoadFunc) {
+    /**
+     * Currently, only 1 thread will fetch data, which is limited by no thread-safe
+     * {@link org.apache.uniffle.storage.handler.api.ClientReadHandler}.
+     *
+     * TODO: Support multiple threads to fetch data.
+     */
+    this(prefetchBatch, 1, dataLoadFunc);
+  }
+
+  protected AsyncDataLoader(int prefetchBatch, int concurrency, Supplier<List<T>> dataLoadFunc) {
+    this.dataQueue = Queues.newLinkedBlockingQueue(prefetchBatch);
+    this.semaphore = new Semaphore(prefetchBatch);
+
+    this.executorService = Executors.newFixedThreadPool(concurrency);
+    for (int i = 0; i < concurrency; i++) {
+      executorService.submit(() -> {
+        try {
+          while (true) {
+            if (!semaphore.tryAcquire(1, TimeUnit.MILLISECONDS)) {
+              if (dataFinished) {
+                break;
+              }
+              continue;
+            }
+            List<T> result = dataLoadFunc.get();
+            if (CollectionUtils.isEmpty(result)) {
+              this.dataFinished = true;

Review Comment:
   so, if any data load func returns null, it means we have loaded all the data?



##########
client/src/main/java/org/apache/uniffle/client/util/AsyncDataLoader.java:
##########
@@ -0,0 +1,147 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.uniffle.client.util;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.Iterator;
+import java.util.List;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Supplier;
+
+import com.google.common.collect.Queues;
+import org.apache.commons.collections.CollectionUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.uniffle.common.exception.RssException;
+
+public class AsyncDataLoader<T> implements Closeable {
+  private static final Logger LOGGER = LoggerFactory.getLogger(AsyncDataLoader.class);
+
+  private final BlockingQueue<BatchDataHolder> dataQueue;
+  private BatchDataHolder<T> currentHolder;
+  private volatile boolean noNewlyData = false;
+  private final Semaphore semaphore;
+
+  // Internal sync vars
+  private volatile int finishedTaskCnt = 0;
+  private volatile boolean dataFinished = false;
+
+  private ExecutorService executorService;
+
+  public AsyncDataLoader(int prefetchBatch, Supplier<List<T>> dataLoadFunc) {
+    /**
+     * Currently, only 1 thread will fetch data, which is limited by no thread-safe
+     * {@link org.apache.uniffle.storage.handler.api.ClientReadHandler}.
+     *
+     * TODO: Support multiple threads to fetch data.
+     */
+    this(prefetchBatch, 1, dataLoadFunc);
+  }
+
+  protected AsyncDataLoader(int prefetchBatch, int concurrency, Supplier<List<T>> dataLoadFunc) {
+    this.dataQueue = Queues.newLinkedBlockingQueue(prefetchBatch);
+    this.semaphore = new Semaphore(prefetchBatch);
+
+    this.executorService = Executors.newFixedThreadPool(concurrency);
+    for (int i = 0; i < concurrency; i++) {
+      executorService.submit(() -> {
+        try {
+          while (true) {
+            if (!semaphore.tryAcquire(1, TimeUnit.MILLISECONDS)) {
+              if (dataFinished) {
+                break;
+              }
+              continue;
+            }
+            List<T> result = dataLoadFunc.get();
+            if (CollectionUtils.isEmpty(result)) {
+              this.dataFinished = true;

Review Comment:
   so, if any data load func returns null? it means we have loaded all the data?



##########
client/src/main/java/org/apache/uniffle/client/impl/ShuffleReadClientImpl.java:
##########
@@ -78,7 +98,8 @@ public ShuffleReadClientImpl(
       Configuration hadoopConf,
       IdHelper idHelper,
       ShuffleDataDistributionType dataDistributionType,
-      boolean expectedTaskIdsBitmapFilterEnable) {
+      boolean expectedTaskIdsBitmapFilterEnable,
+      int prefetchSize) {

Review Comment:
   Let's be consistent. And this should be called as prefetchBatchNum.



##########
client/src/main/java/org/apache/uniffle/client/request/CreateShuffleReadClientRequest.java:
##########
@@ -61,11 +99,11 @@ public CreateShuffleReadClientRequest(
       List<ShuffleServerInfo> shuffleServerInfoList,
       Configuration hadoopConf,
       ShuffleDataDistributionType dataDistributionType,
-      boolean expectedTaskIdsBitmapFilterEnable) {
+      boolean expectedTaskIdsBitmapFilterEnable,
+      int prefetchSize) {

Review Comment:
   ditto



##########
client/src/test/java/org/apache/uniffle/client/AsyncDataLoaderTest.java:
##########
@@ -0,0 +1,127 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.uniffle.client;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Supplier;
+import java.util.stream.IntStream;
+import java.util.stream.Stream;
+
+import com.google.common.collect.Sets;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.MethodSource;
+
+import org.apache.uniffle.client.util.AsyncDataLoader;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
+
+public class AsyncDataLoaderTest {

Review Comment:
   could you add:
   1. more UTs in ShuffleReadClientImpl test with more 2,4 prefetch batch num
   2. integration test with 2,4 prefetch batch num



##########
client/src/main/java/org/apache/uniffle/client/impl/ShuffleReadClientImpl.java:
##########
@@ -142,7 +182,7 @@ public ShuffleReadClientImpl(
     this(storageType, appId, shuffleId, partitionId, indexReadLimit,
         partitionNumPerRange, partitionNum, readBufferSize, storageBasePath,
         blockIdBitmap, taskIdBitmap, shuffleServerInfoList, hadoopConf,
-        idHelper, ShuffleDataDistributionType.NORMAL, false);
+        idHelper, ShuffleDataDistributionType.NORMAL, false, 1);
   }
 
   @Override

Review Comment:
   could we add more comment here for this method? 



##########
common/src/main/java/org/apache/uniffle/common/config/RssClientConf.java:
##########
@@ -43,4 +43,10 @@ public class RssClientConf {
       .defaultValue(ShuffleDataDistributionType.NORMAL)
       .withDescription("The type of partition shuffle data distribution, including normal and local_order. "
           + "The default value is normal. This config is only valid in Spark3.x");
+
+  public static final ConfigOption<Integer> ASYNC_DATA_PREFETCH_BATCH_NUM = ConfigOptions
+      .key("rss.client.data.prefetch.batch.num")
+      .intType()
+      .defaultValue(2)

Review Comment:
   Are we confident enough to turn aync data prefetch by default? 
   
   How about make this default to 1, and then turn it on in later releases, such as 0.9 or 1.0?



##########
client/src/main/java/org/apache/uniffle/client/impl/ShuffleReadClientImpl.java:
##########
@@ -45,22 +44,43 @@
 import org.apache.uniffle.storage.request.CreateShuffleReadHandlerRequest;
 
 public class ShuffleReadClientImpl implements ShuffleReadClient {
-
   private static final Logger LOG = LoggerFactory.getLogger(ShuffleReadClientImpl.class);
+
   private final List<ShuffleServerInfo> shuffleServerInfoList;
   private int shuffleId;
   private int partitionId;
-  private byte[] readBuffer;
   private Roaring64NavigableMap blockIdBitmap;
   private Roaring64NavigableMap taskIdBitmap;
   private Roaring64NavigableMap pendingBlockIds;
   private Roaring64NavigableMap processedBlockIds = Roaring64NavigableMap.bitmapOf();
-  private Queue<BufferSegment> bufferSegmentQueue = Queues.newLinkedBlockingQueue();
   private AtomicLong readDataTime = new AtomicLong(0);
   private AtomicLong copyTime = new AtomicLong(0);
   private AtomicLong crcCheckTime = new AtomicLong(0);
   private ClientReadHandler clientReadHandler;
   private final IdHelper idHelper;
+  private final AsyncDataLoader<BufferSegmentElement> dataLoader;
+
+  public Roaring64NavigableMap getProcessedBlockIds() {
+    return processedBlockIds;
+  }
+
+  class BufferSegmentElement {

Review Comment:
   a new class file or `private static`.
   
   When declaring inner class, it's always better to add static modifier.



##########
client/src/main/java/org/apache/uniffle/client/util/AsyncDataLoader.java:
##########
@@ -0,0 +1,147 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.uniffle.client.util;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.Iterator;
+import java.util.List;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Supplier;
+
+import com.google.common.collect.Queues;
+import org.apache.commons.collections.CollectionUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.uniffle.common.exception.RssException;
+
+public class AsyncDataLoader<T> implements Closeable {
+  private static final Logger LOGGER = LoggerFactory.getLogger(AsyncDataLoader.class);
+
+  private final BlockingQueue<BatchDataHolder> dataQueue;
+  private BatchDataHolder<T> currentHolder;
+  private volatile boolean noNewlyData = false;
+  private final Semaphore semaphore;
+
+  // Internal sync vars
+  private volatile int finishedTaskCnt = 0;
+  private volatile boolean dataFinished = false;
+
+  private ExecutorService executorService;
+
+  public AsyncDataLoader(int prefetchBatch, Supplier<List<T>> dataLoadFunc) {
+    /**
+     * Currently, only 1 thread will fetch data, which is limited by no thread-safe
+     * {@link org.apache.uniffle.storage.handler.api.ClientReadHandler}.
+     *
+     * TODO: Support multiple threads to fetch data.
+     */
+    this(prefetchBatch, 1, dataLoadFunc);
+  }
+
+  protected AsyncDataLoader(int prefetchBatch, int concurrency, Supplier<List<T>> dataLoadFunc) {
+    this.dataQueue = Queues.newLinkedBlockingQueue(prefetchBatch);
+    this.semaphore = new Semaphore(prefetchBatch);
+
+    this.executorService = Executors.newFixedThreadPool(concurrency);
+    for (int i = 0; i < concurrency; i++) {
+      executorService.submit(() -> {
+        try {
+          while (true) {
+            if (!semaphore.tryAcquire(1, TimeUnit.MILLISECONDS)) {
+              if (dataFinished) {
+                break;
+              }
+              continue;
+            }
+            List<T> result = dataLoadFunc.get();
+            if (CollectionUtils.isEmpty(result)) {
+              this.dataFinished = true;
+              break;
+            }
+            dataQueue.put(new BatchDataHolder(result));
+          }
+        } catch (Exception e) {
+          LOGGER.error("Errors on fetching data.", e);
+        }
+        finishedTaskCnt++;
+        if (finishedTaskCnt == concurrency) {
+          LOGGER.info("All threads of async fetching data are finished.");
+          this.noNewlyData = true;
+        }
+      });
+    }
+  }
+
+  public T get() {
+    try {
+      return doGet();
+    } catch (Exception e) {
+      throw new RssException("Errors on getting data.", e);
+    }
+  }
+
+  public T doGet() throws Exception {
+    while (currentHolder == null) {
+      final BatchDataHolder<T> holder = dataQueue.poll(1, TimeUnit.MILLISECONDS);
+      if (holder == null && noNewlyData) {
+        return null;
+      }
+      if (holder == null) {
+        continue;
+      }
+      currentHolder = holder;
+    }
+    T val = currentHolder.next();
+    if (!currentHolder.hasNext()) {
+      currentHolder = null;
+      semaphore.release();
+    }
+    return val;
+  }
+
+  class BatchDataHolder<T> implements Iterator<T> {
+    private Iterator<T> iterator;
+
+    BatchDataHolder(List<T> elements) {
+      this.iterator = elements.listIterator();
+    }
+
+    @Override
+    public boolean hasNext() {
+      return iterator.hasNext();
+    }
+
+    @Override
+    public T next() {
+      return iterator.next();
+    }
+  }
+
+  @Override
+  public void close() throws IOException {
+    if (executorService != null) {
+      executorService.shutdown();

Review Comment:
   shutdownNow? Then waitForCompletion?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@uniffle.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@uniffle.apache.org
For additional commands, e-mail: issues-help@uniffle.apache.org