feat(core): Allow retry in concurrent write operations (#3958)

* Fix concurrent retry Signed-off-by: Xuanwo <[email protected]> * FIx retry for range write Signed-off-by: Xuanwo <[email protected]> * Fix typo Signed-off-by: Xuanwo <[email protected]> * Add test for retry in concurrent Signed-off-by: Xuanwo <[email protected]> * FIx Signed-off-by: Xuanwo <[email protected]> * Fix build Signed-off-by: Xuanwo <[email protected]> --------- Signed-off-by: Xuanwo <[email protected]>
apache · Jan 10, 2024 · 8315b7e · 8315b7e
1 parent 8714ad7
commit 8315b7e
Show file tree

Hide file tree

Showing 5 changed files with 440 additions and 109 deletions.
diff --git a/core/src/raw/futures_util.rs b/core/src/raw/futures_util.rs
@@ -134,8 +134,8 @@ where
         self.remaining() > 0
     }
 
-    /// Push new future into the queue.
-    pub fn push(&mut self, f: F) {
+    /// Push new future into the end of queue.
+    pub fn push_back(&mut self, f: F) {
         debug_assert!(
             self.has_remaining(),
             "concurrent futures must have remaining space"
@@ -149,6 +149,22 @@ where
             Tasks::Large(v) => v.push_back(f),
         }
     }
+
+    /// Push new future into the start of queue, this task will be exactly the next to poll.
+    pub fn push_front(&mut self, f: F) {
+        debug_assert!(
+            self.has_remaining(),
+            "concurrent futures must have remaining space"
+        );
+
+        match &mut self.tasks {
+            Tasks::Once(fut) => {
+                *fut = Some(f);
+            }
+            Tasks::Small(v) => v.push_front(TaskResult::Polling(f)),
+            Tasks::Large(v) => v.push_front(f),
+        }
+    }
 }
 
 impl<F> futures::Stream for ConcurrentFutures<F>
@@ -230,7 +246,7 @@ mod tests {
                     idx
                 };
                 self.idx += 1;
-                self.tasks.push(Box::pin(fut));
+                self.tasks.push_back(Box::pin(fut));
             }
 
             if let Some(v) = ready!(self.tasks.poll_next_unpin(cx)) {

diff --git a/core/src/raw/oio/write/block_write.rs b/core/src/raw/oio/write/block_write.rs
@@ -164,10 +164,11 @@ where
                         self.block_ids.push(block_id.clone());
                         let w = self.w.clone();
                         let size = cache.len();
-                        self.futures.push(WriteBlockFuture(Box::pin(async move {
-                            w.write_block(size as u64, block_id, AsyncBody::ChunkedBytes(cache))
-                                .await
-                        })));
+                        self.futures
+                            .push_back(WriteBlockFuture(Box::pin(async move {
+                                w.write_block(size as u64, block_id, AsyncBody::ChunkedBytes(cache))
+                                    .await
+                            })));
                         let size = self.fill_cache(bs);
                         return Poll::Ready(Ok(size));
                     } else if let Some(res) = ready!(self.futures.poll_next_unpin(cx)) {
@@ -221,14 +222,15 @@ where
                                 self.block_ids.push(block_id.clone());
                                 let size = cache.len();
                                 let w = self.w.clone();
-                                self.futures.push(WriteBlockFuture(Box::pin(async move {
-                                    w.write_block(
-                                        size as u64,
-                                        block_id,
-                                        AsyncBody::ChunkedBytes(cache),
-                                    )
-                                    .await
-                                })));
+                                self.futures
+                                    .push_back(WriteBlockFuture(Box::pin(async move {
+                                        w.write_block(
+                                            size as u64,
+                                            block_id,
+                                            AsyncBody::ChunkedBytes(cache),
+                                        )
+                                        .await
+                                    })));
                             }
                         }
                         while let Some(res) = ready!(self.futures.poll_next_unpin(cx)) {