Copilot commented on code in PR #3671:
URL: https://github.com/apache/celeborn/pull/3671#discussion_r3214317627


##########
common/src/main/java/org/apache/celeborn/common/write/DataBatches.java:
##########
@@ -57,18 +58,20 @@ public ArrayList<DataBatch> requireBatches() {
   public ArrayList<DataBatch> requireBatches(int requestSize) {
     if (requestSize >= totalSize) {
       totalSize = 0;
-      return batches;
+      ArrayList<DataBatch> allBatches = batches;
+      batches = new ArrayList<>();
+      return allBatches;
     }
-    // TODO: ArrayList.remove(0) in a loop is O(n^2) due to element shifting.
-    // Investigate subList(0, count).clear() or use LinkedList/ArrayDeque.
-    ArrayList<DataBatch> retBatches = new ArrayList<>();
+    int count = 0;
     int currentSize = 0;
     while (currentSize < requestSize) {
-      DataBatch elem = batches.remove(0);
-      retBatches.add(elem);
-      currentSize += elem.body.length;
-      totalSize -= elem.body.length;
+      currentSize += batches.get(count).body.length;
+      count++;
     }
+    List<DataBatch> head = batches.subList(0, count);
+    ArrayList<DataBatch> retBatches = new ArrayList<>(head);
+    head.clear();
+    totalSize -= currentSize;

Review Comment:
   `requireBatches(int)` mutates `batches` and `totalSize` without any 
synchronization, while `addDataBatch(...)` is `synchronized` and `DataBatches` 
instances are shared via a `ConcurrentHashMap` (e.g., `PushState.batchesMap`). 
This can lead to races where one thread is adding while another is draining 
(e.g., `batches.get(count)` / `subList(...).clear()` can throw or corrupt 
`totalSize`). Consider making `requireBatches(int)` `synchronized` (and keep 
the entire drain logic under the same monitor as `addDataBatch`) so add/drain 
are atomic and consistent.



-- 
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: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to