fix(ecstore): paginate ListMultipartUploads across pools (#4522)

2026-07-26 16:28:15 +00:00 · 2026-07-09 01:07:46 +08:00
parent 7b87d4d13b
commit 2055044cb4
2 changed files with 223 additions and 21 deletions
@@ -128,6 +128,25 @@ fn reduce_quorum_part_numbers(object_parts: Vec<Vec<String>>, read_quorum: usize
    part_numbers
 }

+/// Applies the `max_uploads` page cap to an already marker-offset, sorted slice
+/// of multipart uploads.
+///
+/// At most `max_uploads` entries are returned. When more uploads remain beyond
+/// the page, the first overflow element acts purely as a truncation probe: it is
+/// never returned, but flips `is_truncated` to `true` and yields a
+/// `next_upload_id_marker` pointing at the last returned upload so the caller can
+/// resume paging.
+fn paginate_upload_page(remaining: &[MultipartInfo], max_uploads: usize) -> (Vec<MultipartInfo>, bool, Option<String>) {
+    let is_truncated = remaining.len() > max_uploads;
+    let page: Vec<MultipartInfo> = remaining.iter().take(max_uploads).cloned().collect();
+    let next_upload_id_marker = if is_truncated {
+        page.last().map(|upload| upload.upload_id.clone())
+    } else {
+        None
+    };
+    (page, is_truncated, next_upload_id_marker)
+}
+
 impl SetDisks {
    async fn acquire_multipart_upload_read_lock(
        &self,
@@ -804,27 +823,10 @@ impl crate::storage_api_contracts::multipart::MultipartOperations for SetDisks {
            }
        }

-        let mut ret_uploads = Vec::new();
-        let mut next_upload_id_marker = None;
-        while upload_idx < uploads.len() {
-            if ret_uploads.len() >= max_uploads {
-                break;
-            }
-
-            ret_uploads.push(uploads[upload_idx].clone());
-            next_upload_id_marker = Some(uploads[upload_idx].upload_id.clone());
-            upload_idx += 1;
-        }
-
-        // Truncated iff there are still unconsumed uploads left after this page.
-        // `upload_idx` tracks the position in the full (post-marker) list, so it
-        // stays correct even when a marker skipped entries before the page start,
-        // unlike comparing the page length against the total.
-        let is_truncated = upload_idx < uploads.len();
-
-        if !is_truncated {
-            next_upload_id_marker = None;
-        }
+        // `upload_idx` is the post-marker offset, so `uploads[upload_idx..]` is the
+        // page candidate. The helper applies the `max_uploads` cap, using the first
+        // overflow element only as a truncation probe (never returned).
+        let (ret_uploads, is_truncated, next_upload_id_marker) = paginate_upload_page(&uploads[upload_idx..], max_uploads);

        Ok(ListMultipartsInfo {
            key_marker: key_marker.to_owned(),
@@ -1920,4 +1922,66 @@ mod tests {
        assert!(matches!(err, StorageError::InvalidUploadID(bucket, object, upload_id)
            if bucket == "bucket" && object == "object" && upload_id == "upload-id"));
    }
+
+    fn test_multipart_upload(object: &str, upload_id: &str) -> MultipartInfo {
+        MultipartInfo {
+            bucket: "bucket".to_string(),
+            object: object.to_string(),
+            upload_id: upload_id.to_string(),
+            initiated: None,
+            ..Default::default()
+        }
+    }
+
+    #[test]
+    fn paginate_upload_page_caps_at_max_uploads_and_probes_truncation() {
+        // max_uploads = N (3) with N + 1 (4) uploads must return exactly N,
+        // report truncation, and expose the last returned id as the next marker.
+        let uploads: Vec<MultipartInfo> = (0..4).map(|i| test_multipart_upload("obj", &format!("u{i}"))).collect();
+
+        let (page, is_truncated, next_upload_id_marker) = paginate_upload_page(&uploads, 3);
+
+        assert_eq!(page.len(), 3);
+        assert!(is_truncated);
+        assert_eq!(next_upload_id_marker.as_deref(), Some("u2"));
+        // The overflow element `u3` is only a truncation probe and must not leak
+        // into the returned page.
+        assert!(page.iter().all(|upload| upload.upload_id != "u3"));
+    }
+
+    #[test]
+    fn paginate_upload_page_reports_complete_when_within_cap() {
+        let uploads: Vec<MultipartInfo> = (0..3).map(|i| test_multipart_upload("obj", &format!("u{i}"))).collect();
+
+        let (page, is_truncated, next_upload_id_marker) = paginate_upload_page(&uploads, 3);
+
+        assert_eq!(page.len(), 3);
+        assert!(!is_truncated);
+        assert!(next_upload_id_marker.is_none());
+    }
+
+    #[test]
+    fn paginate_upload_page_marker_resumes_without_loss() {
+        let uploads: Vec<MultipartInfo> = (0..4).map(|i| test_multipart_upload("obj", &format!("u{i}"))).collect();
+
+        // First page stops at the cap and hands back `u2` as the resume marker.
+        let (first, first_truncated, first_marker) = paginate_upload_page(&uploads, 3);
+        assert!(first_truncated);
+        assert_eq!(first_marker.as_deref(), Some("u2"));
+
+        // The caller applies the marker offset (drops through `u2`) and pages the
+        // remainder; the final upload must appear exactly once.
+        let resume_offset = uploads
+            .iter()
+            .position(|u| u.upload_id == "u2")
+            .expect("marker must be present")
+            + 1;
+        let (second, second_truncated, _) = paginate_upload_page(&uploads[resume_offset..], 3);
+        assert!(!second_truncated);
+        assert_eq!(second.len(), 1);
+        assert_eq!(second[0].upload_id, "u3");
+
+        let paged: Vec<&str> = first.iter().chain(second.iter()).map(|u| u.upload_id.as_str()).collect();
+        assert_eq!(paged, vec!["u0", "u1", "u2", "u3"]);
+    }
 }
@@ -135,10 +135,19 @@ impl ECStore {
            uploads.extend(res.uploads);
        }

+        // Each pool caps its own page at `max_uploads`, so the concatenation is
+        // unordered across pools and may exceed the global cap. Re-sort, re-cap,
+        // and derive the truncation markers so a bucket whose uploads span pools
+        // pages correctly instead of being silently reported complete.
+        let (uploads, is_truncated, next_key_marker, next_upload_id_marker) = merge_multipart_upload_pages(uploads, max_uploads);
+
        Ok(ListMultipartsInfo {
            key_marker,
            upload_id_marker,
+            next_key_marker,
+            next_upload_id_marker,
            max_uploads,
+            is_truncated,
            uploads,
            prefix: prefix.to_owned(),
            delimiter: delimiter.to_owned(),
@@ -379,6 +388,36 @@ impl ECStore {
    }
 }

+/// Merges per-pool `ListMultipartUploads` pages into a single globally paginated
+/// page.
+///
+/// Each pool independently applies the `max_uploads` cap, so the concatenated
+/// input can hold up to `pools * max_uploads` entries and is unordered across
+/// pools. This re-sorts the union by `(key, upload_id)` — the order S3 clients
+/// page through — caps it to `max_uploads`, and derives the truncation markers
+/// from the first overflow element (used only as a probe, never returned) so a
+/// bucket whose uploads span pools can be paged without loss or duplication.
+fn merge_multipart_upload_pages(
+    mut uploads: Vec<MultipartInfo>,
+    max_uploads: usize,
+) -> (Vec<MultipartInfo>, bool, Option<String>, Option<String>) {
+    uploads.sort_by(|a, b| a.object.cmp(&b.object).then_with(|| a.upload_id.cmp(&b.upload_id)));
+
+    let is_truncated = uploads.len() > max_uploads;
+    uploads.truncate(max_uploads);
+
+    let (next_key_marker, next_upload_id_marker) = if is_truncated {
+        match uploads.last() {
+            Some(last) => (Some(last.object.clone()), Some(last.upload_id.clone())),
+            None => (None, None),
+        }
+    } else {
+        (None, None)
+    };
+
+    (uploads, is_truncated, next_key_marker, next_upload_id_marker)
+}
+
 #[cfg(test)]
 mod tests {
    use super::*;
@@ -388,6 +427,105 @@ mod tests {
    };
    use std::time::Duration;

+    fn mp(object: &str, upload_id: &str) -> MultipartInfo {
+        MultipartInfo {
+            bucket: "bucket".to_string(),
+            object: object.to_string(),
+            upload_id: upload_id.to_string(),
+            initiated: None,
+            ..Default::default()
+        }
+    }
+
+    /// Models a single pool's `list_multipart_uploads`: returns uploads strictly
+    /// after the `(key, upload_id)` marker in `(key, upload_id)` order, capped at
+    /// `max_uploads` (mirroring the per-pool page cap).
+    fn pool_query(
+        pool: &[MultipartInfo],
+        key_marker: Option<&str>,
+        upload_id_marker: Option<&str>,
+        max_uploads: usize,
+    ) -> Vec<MultipartInfo> {
+        pool.iter()
+            .filter(|u| match (key_marker, upload_id_marker) {
+                (Some(k), Some(uid)) => (u.object.as_str(), u.upload_id.as_str()) > (k, uid),
+                (Some(k), None) => u.object.as_str() > k,
+                _ => true,
+            })
+            .take(max_uploads)
+            .cloned()
+            .collect()
+    }
+
+    #[test]
+    fn merge_multipart_upload_pages_sorts_and_caps_across_pools() {
+        // Union of two pools, unordered and exceeding the global cap.
+        let uploads = vec![mp("b", "u1"), mp("a", "u2"), mp("a", "u1"), mp("c", "u1"), mp("b", "u2")];
+
+        let (page, is_truncated, next_key_marker, next_upload_id_marker) = merge_multipart_upload_pages(uploads, 3);
+
+        assert!(is_truncated);
+        assert_eq!(page.len(), 3);
+        let ordered: Vec<(&str, &str)> = page.iter().map(|u| (u.object.as_str(), u.upload_id.as_str())).collect();
+        assert_eq!(ordered, vec![("a", "u1"), ("a", "u2"), ("b", "u1")]);
+        assert_eq!(next_key_marker.as_deref(), Some("b"));
+        assert_eq!(next_upload_id_marker.as_deref(), Some("u1"));
+    }
+
+    #[test]
+    fn merge_multipart_upload_pages_reports_complete_within_cap() {
+        let uploads = vec![mp("b", "u1"), mp("a", "u1")];
+
+        let (page, is_truncated, next_key_marker, next_upload_id_marker) = merge_multipart_upload_pages(uploads, 3);
+
+        assert_eq!(page.len(), 2);
+        assert!(!is_truncated);
+        assert!(next_key_marker.is_none());
+        assert!(next_upload_id_marker.is_none());
+    }
+
+    #[test]
+    fn merge_multipart_upload_pages_paginates_across_pools_without_loss() {
+        // Uploads for the same bucket spread across two pools, together exceeding
+        // the cap, so pagination must span multiple pages.
+        let pool0 = vec![mp("a", "u1"), mp("a", "u3"), mp("c", "u1"), mp("e", "u1")];
+        let pool1 = vec![mp("a", "u2"), mp("b", "u1"), mp("d", "u1"), mp("f", "u1")];
+
+        let mut expected: Vec<(String, String)> = pool0
+            .iter()
+            .chain(pool1.iter())
+            .map(|u| (u.object.clone(), u.upload_id.clone()))
+            .collect();
+        expected.sort();
+
+        let max_uploads = 3;
+        let mut key_marker: Option<String> = None;
+        let mut upload_id_marker: Option<String> = None;
+        let mut collected: Vec<(String, String)> = Vec::new();
+
+        for _ in 0..16 {
+            let mut merged = Vec::new();
+            for pool in [&pool0, &pool1] {
+                merged.extend(pool_query(pool, key_marker.as_deref(), upload_id_marker.as_deref(), max_uploads));
+            }
+
+            let (page, is_truncated, next_key_marker, next_upload_id_marker) = merge_multipart_upload_pages(merged, max_uploads);
+            assert!(page.len() <= max_uploads);
+            collected.extend(page.iter().map(|u| (u.object.clone(), u.upload_id.clone())));
+
+            if !is_truncated {
+                break;
+            }
+            key_marker = next_key_marker;
+            upload_id_marker = next_upload_id_marker;
+        }
+
+        assert_eq!(collected, expected, "pagination must return every upload exactly once, in sorted order");
+        let mut deduped = collected.clone();
+        deduped.dedup();
+        assert_eq!(deduped.len(), collected.len(), "pagination must not duplicate uploads");
+    }
+
    async fn new_multipart_lock_test_store() -> ECStore {
        let format = FormatV3::new(1, 2);
        let endpoints = vec![