Skip to main content

polars_io/cloud/
polars_object_store.rs

1use std::fmt::Display;
2use std::ops::Range;
3use std::sync::Arc;
4use std::time::Instant;
5
6use futures::{Stream, StreamExt as _, TryStreamExt as _};
7use hashbrown::hash_map::RawEntryMut;
8use object_store::path::Path;
9use object_store::{ObjectMeta, ObjectStore, ObjectStoreExt};
10use polars_buffer::Buffer;
11use polars_core::prelude::{InitHashMaps, PlHashMap};
12use polars_error::{PolarsError, PolarsResult};
13use polars_utils::pl_path::PlRefPath;
14use tokio::io::AsyncWriteExt;
15
16use super::concurrency::{ConcurrencyController, IoSample};
17use super::concurrency_config::{ConcurrencyStrategy, FetchConfig, get_download_chunk_size};
18use crate::pl_async::{
19    self, MAX_BUDGET_PER_REQUEST, get_concurrency_limit, tune_with_concurrency_budget,
20    with_concurrency_budget,
21};
22
23#[derive(Debug)]
24pub struct PolarsObjectStoreError {
25    pub base_url: PlRefPath,
26    pub source: object_store::Error,
27}
28
29impl PolarsObjectStoreError {
30    pub fn from_url(base_url: &PlRefPath) -> impl FnOnce(object_store::Error) -> Self {
31        |error| Self {
32            base_url: base_url.clone(),
33            source: error,
34        }
35    }
36}
37
38impl Display for PolarsObjectStoreError {
39    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
40        write!(
41            f,
42            "object-store error: {} (path: {})",
43            self.source, self.base_url
44        )
45    }
46}
47
48impl std::error::Error for PolarsObjectStoreError {
49    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
50        Some(&self.source)
51    }
52}
53
54impl From<PolarsObjectStoreError> for std::io::Error {
55    fn from(value: PolarsObjectStoreError) -> Self {
56        std::io::Error::other(value)
57    }
58}
59
60impl From<PolarsObjectStoreError> for PolarsError {
61    fn from(value: PolarsObjectStoreError) -> Self {
62        PolarsError::IO {
63            error: Arc::new(value.into()),
64            msg: None,
65        }
66    }
67}
68
69mod inner {
70
71    use std::borrow::Cow;
72    use std::future::Future;
73    use std::sync::Arc;
74
75    use object_store::ObjectStore;
76    use polars_core::config;
77    use polars_error::{PolarsError, PolarsResult};
78    use polars_utils::relaxed_cell::RelaxedCell;
79
80    use crate::cloud::concurrency::{ConcurrencyController, ControllerConfig};
81    use crate::cloud::{ObjectStoreErrorContext, PolarsObjectStoreBuilder};
82    use crate::metrics::{IOMetrics, OptIOMetrics};
83
84    #[derive(Debug)]
85    struct Inner {
86        store: tokio::sync::RwLock<Arc<dyn ObjectStore>>,
87        builder: PolarsObjectStoreBuilder,
88        rebuilt: RelaxedCell<bool>,
89    }
90
91    /// Polars wrapper around [`ObjectStore`] functionality. This struct is cheaply cloneable.
92    #[derive(Clone, Debug)]
93    pub struct PolarsObjectStore {
94        inner: Arc<Inner>,
95        /// Avoid contending the Mutex `lock()` until the first re-build.
96        initial_store: std::sync::Arc<dyn ObjectStore>,
97        io_metrics: OptIOMetrics,
98        /// In-flight concurrency control using the (new) BDP model.
99        concurrency: Arc<std::sync::OnceLock<Arc<ConcurrencyController>>>,
100    }
101
102    impl PolarsObjectStore {
103        pub(crate) fn new_from_inner(
104            store: Arc<dyn ObjectStore>,
105            builder: PolarsObjectStoreBuilder,
106        ) -> Self {
107            let initial_store = store.clone();
108            Self {
109                inner: Arc::new(Inner {
110                    store: tokio::sync::RwLock::new(store),
111                    builder,
112                    rebuilt: RelaxedCell::from(false),
113                }),
114                initial_store,
115                io_metrics: OptIOMetrics(None),
116                concurrency: Arc::new(std::sync::OnceLock::new()),
117            }
118        }
119
120        pub fn set_io_metrics(&mut self, io_metrics: Option<Arc<IOMetrics>>) -> &mut Self {
121            self.io_metrics = OptIOMetrics(io_metrics);
122            self
123        }
124
125        pub fn io_metrics(&self) -> &OptIOMetrics {
126            &self.io_metrics
127        }
128
129        pub fn get_or_init_concurrency(&self) -> &Arc<ConcurrencyController> {
130            self.concurrency.get_or_init(|| {
131                // Valid across rebuilds: the signal wraps the stable cell.
132                let signal = self.inner.builder.rate_limit_signal();
133                Arc::new(ConcurrencyController::new(
134                    ControllerConfig::default(),
135                    signal,
136                ))
137            })
138        }
139
140        /// Gets the underlying [`ObjectStore`] implementation.
141        pub async fn to_dyn_object_store(&self) -> Cow<'_, Arc<dyn ObjectStore>> {
142            if !self.inner.rebuilt.load() {
143                Cow::Borrowed(&self.initial_store)
144            } else {
145                Cow::Owned(self.inner.store.read().await.clone())
146            }
147        }
148
149        pub async fn rebuild_inner(
150            &self,
151            from_version: &Arc<dyn ObjectStore>,
152        ) -> PolarsResult<Arc<dyn ObjectStore>> {
153            let mut current_store = self.inner.store.write().await;
154
155            // If this does not eq, then `inner` was already re-built by another thread.
156            if Arc::ptr_eq(&*current_store, from_version) {
157                *current_store =
158                    self.inner
159                        .builder
160                        .clone()
161                        .build_impl(true)
162                        .await
163                        .map_err(|e| {
164                            e.wrap_msg(|e| format!("attempt to rebuild object store failed: {e}"))
165                        })?;
166            }
167
168            self.inner.rebuilt.store(true);
169
170            Ok((*current_store).clone())
171        }
172
173        pub async fn exec_with_rebuild_retry_on_err<'s, 'f, Fn, Fut, O>(
174            &'s self,
175            mut func: Fn,
176        ) -> PolarsResult<O>
177        where
178            Fn: FnMut(Cow<'s, Arc<dyn ObjectStore>>) -> Fut + 'f,
179            Fut: Future<Output = object_store::Result<O>>,
180        {
181            let store = self.to_dyn_object_store().await;
182
183            let out = func(store.clone()).await;
184
185            let orig_err = match out {
186                Ok(v) => return Ok(v),
187                Err(e) => e,
188            };
189
190            if config::verbose() {
191                eprintln!(
192                    "[PolarsObjectStore]: got error: {}, will rebuild store and retry",
193                    orig_err
194                );
195            }
196
197            let store = self
198                .rebuild_inner(&store)
199                .await
200                .map_err(|e| e.wrap_msg(|e| format!("{e}; original error: {orig_err}")))?;
201
202            func(Cow::Owned(store)).await.map_err(|e| {
203                let e: PolarsError = self.error_context().attach_err_info(e).into();
204
205                if self.inner.builder.is_azure()
206                    && std::env::var("POLARS_AUTO_USE_AZURE_STORAGE_ACCOUNT_KEY").as_deref()
207                        != Ok("1")
208                {
209                    // Note: This error is intended for Python audiences. The logic for retrieving
210                    // these keys exist only on the Python side.
211                    e.wrap_msg(|e| {
212                        format!(
213                            "{e}; note: if you are using Python, consider setting \
214POLARS_AUTO_USE_AZURE_STORAGE_ACCOUNT_KEY=1 if you would like polars to try to retrieve \
215and use the storage account keys from Azure CLI to authenticate"
216                        )
217                    })
218                } else {
219                    e
220                }
221            })
222        }
223
224        pub fn error_context(&self) -> ObjectStoreErrorContext {
225            ObjectStoreErrorContext::new(self.inner.builder.path().clone())
226        }
227    }
228}
229
230#[derive(Clone)]
231pub struct ObjectStoreErrorContext {
232    path: PlRefPath,
233}
234
235impl ObjectStoreErrorContext {
236    pub fn new(path: PlRefPath) -> Self {
237        Self { path }
238    }
239
240    pub fn attach_err_info(self, err: object_store::Error) -> PolarsObjectStoreError {
241        let ObjectStoreErrorContext { path } = self;
242
243        PolarsObjectStoreError {
244            base_url: path,
245            source: err,
246        }
247    }
248}
249
250pub use inner::PolarsObjectStore;
251
252pub type ObjectStorePath = object_store::path::Path;
253
254impl PolarsObjectStore {
255    pub fn build_buffered_ranges_stream<'a, T: Iterator<Item = Range<usize>>>(
256        &'a self,
257        path: &'a Path,
258        ranges: T,
259        strategy: ConcurrencyStrategy,
260    ) -> impl Stream<Item = PolarsResult<Buffer<u8>>> + use<'a, T> {
261        let controller = match strategy {
262            ConcurrencyStrategy::BytesBased => Some(self.get_or_init_concurrency().clone()),
263            ConcurrencyStrategy::Unbounded | ConcurrencyStrategy::Legacy => None,
264        };
265
266        let n_buffered = match strategy {
267            // In case of bytes-based concurrency, the concurrency is controlled by the
268            // admission semaphore in the pipeline.
269            // The buffered size is set to a large constant as a backstop. Once all
270            // callsites are verified to pass finite, metadata-derived ranges, this can be
271            // set to usize::MAX.
272            ConcurrencyStrategy::BytesBased => 4096,
273            ConcurrencyStrategy::Unbounded | ConcurrencyStrategy::Legacy => {
274                get_concurrency_limit() as usize
275            },
276        };
277
278        futures::stream::iter(ranges.map(move |range| {
279            let controller = controller.clone();
280            async move {
281                if range.is_empty() {
282                    return Ok(Buffer::new());
283                }
284                let bytes_req = range.len() as u64;
285
286                // Held until end of block to bound in-flight bytes.
287                let _permit = match &controller {
288                    Some(controller) => Some(controller.acquire(bytes_req).await),
289                    None => None,
290                };
291
292                let (out, ttfb) = self
293                    .io_metrics()
294                    .record_io_read(
295                        bytes_req,
296                        self.exec_with_rebuild_retry_on_err(|s| async move {
297                            let t0 = Instant::now();
298                            let response = s
299                                .get_opts(
300                                    path,
301                                    object_store::GetOptions {
302                                        range: Some((range.start as u64..range.end as u64).into()),
303                                        ..Default::default()
304                                    },
305                                )
306                                .await?;
307                            let ttfb = t0.elapsed();
308                            let out = response.bytes().await?;
309
310                            Ok((out, ttfb))
311                        }),
312                    )
313                    .await?;
314
315                if let Some(controller) = &controller {
316                    controller.record_io(IoSample {
317                        n_bytes: out.len() as u64,
318                        ttfb,
319                        completion_time: Instant::now(),
320                    });
321                }
322
323                Ok(Buffer::from_owner(out))
324            }
325        }))
326        .buffered(n_buffered)
327    }
328
329    pub async fn get_range(
330        &self,
331        path: &Path,
332        range: Range<usize>,
333        config: FetchConfig,
334    ) -> PolarsResult<Buffer<u8>> {
335        if range.is_empty() {
336            return Ok(Buffer::new());
337        }
338
339        let parts = split_range(range.clone(), Some(config.chunk_size));
340
341        match config.strategy {
342            ConcurrencyStrategy::Legacy => self.get_range_legacy(path, range).await,
343            ConcurrencyStrategy::Unbounded | ConcurrencyStrategy::BytesBased => self
344                .build_buffered_ranges_stream(path, parts, config.strategy)
345                .try_collect::<Vec<_>>()
346                .await
347                .map(|parts| {
348                    if parts.len() == 1 {
349                        return parts.into_iter().next().unwrap();
350                    }
351                    let mut combined = Vec::with_capacity(range.len());
352                    for part in parts {
353                        combined.extend_from_slice(&part);
354                    }
355                    assert_eq!(combined.len(), range.len());
356                    Buffer::from_vec(combined)
357                }),
358        }
359    }
360
361    async fn get_range_legacy(&self, path: &Path, range: Range<usize>) -> PolarsResult<Buffer<u8>> {
362        if range.is_empty() {
363            return Ok(Buffer::new());
364        }
365
366        let parts = split_range(range.clone(), None);
367
368        if parts.len() == 1 {
369            let out = tune_with_concurrency_budget(1, move || async move {
370                let bytes = self
371                    .io_metrics()
372                    .record_io_read(
373                        range.len() as u64,
374                        self.exec_with_rebuild_retry_on_err(|s| async move {
375                            s.get_range(path, range.start as u64..range.end as u64)
376                                .await
377                        }),
378                    )
379                    .await?;
380
381                PolarsResult::Ok(Buffer::from_owner(bytes))
382            })
383            .await?;
384
385            Ok(out)
386        } else {
387            let parts = tune_with_concurrency_budget(
388                parts.len().clamp(0, MAX_BUDGET_PER_REQUEST) as u32,
389                || {
390                    self.build_buffered_ranges_stream(path, parts, ConcurrencyStrategy::Legacy)
391                        .try_collect::<Vec<Buffer<u8>>>()
392                },
393            )
394            .await?;
395
396            let mut combined = Vec::with_capacity(range.len());
397
398            for part in parts {
399                combined.extend_from_slice(&part)
400            }
401
402            assert_eq!(combined.len(), range.len());
403
404            PolarsResult::Ok(Buffer::from_vec(combined))
405        }
406    }
407
408    pub async fn get_ranges_sort(
409        &self,
410        path: &Path,
411        ranges: &mut [Range<usize>],
412        config: FetchConfig,
413    ) -> PolarsResult<PlHashMap<usize, Buffer<u8>>> {
414        if ranges.is_empty() {
415            return Ok(Default::default());
416        }
417
418        ranges.sort_unstable_by_key(|x| x.start);
419
420        let ranges_len = ranges.len();
421        let (merged_ranges, merged_ends): (Vec<_>, Vec<_>) =
422            merge_ranges(ranges, Some(config.chunk_size)).unzip();
423
424        let mut out = PlHashMap::with_capacity(ranges_len);
425
426        // Build an inflight admission-aware stream over the merged ranges.
427        let mut stream =
428            self.build_buffered_ranges_stream(path, merged_ranges.iter().cloned(), config.strategy);
429
430        let mut current_offset = 0;
431        let mut ends_iter = merged_ends.iter();
432        let mut splitted_parts: Vec<Buffer<u8>> = vec![];
433
434        while let Some(bytes) = stream.try_next().await? {
435            let end = *ends_iter.next().unwrap();
436
437            if end == 0 {
438                splitted_parts.push(bytes);
439                continue;
440            }
441
442            let full_range = ranges[current_offset..end]
443                .iter()
444                .cloned()
445                .reduce(|l, r| l.start.min(r.start)..l.end.max(r.end))
446                .unwrap();
447
448            let bytes = if splitted_parts.is_empty() {
449                bytes
450            } else {
451                let mut out = Vec::with_capacity(full_range.len());
452                for x in splitted_parts.drain(..) {
453                    out.extend_from_slice(&x);
454                }
455                out.extend_from_slice(&bytes);
456                Buffer::from(out)
457            };
458
459            assert_eq!(bytes.len(), full_range.len());
460
461            for range in &ranges[current_offset..end] {
462                let slice = bytes
463                    .clone()
464                    .sliced(range.start - full_range.start..range.end - full_range.start);
465
466                match out.raw_entry_mut().from_key(&range.start) {
467                    RawEntryMut::Vacant(slot) => {
468                        slot.insert(range.start, slice);
469                    },
470                    RawEntryMut::Occupied(mut slot) => {
471                        if slot.get_mut().len() < slice.len() {
472                            *slot.get_mut() = slice;
473                        }
474                    },
475                }
476            }
477
478            current_offset = end;
479        }
480
481        assert!(splitted_parts.is_empty());
482
483        Ok(out)
484    }
485
486    // TODO: Refactor for updated concurrency strategy.
487    pub async fn download(&self, path: &Path, file: &mut tokio::fs::File) -> PolarsResult<()> {
488        let size = self.head(path, ConcurrencyStrategy::Unbounded).await?.size;
489        let parts = split_range(0..size as usize, None);
490
491        // TODO: Replace the legacy concurrency_budget call and switch to BytesBased inflight
492        // admission control.
493        tune_with_concurrency_budget(
494            parts.len().clamp(0, MAX_BUDGET_PER_REQUEST) as u32,
495            || async {
496                let mut stream =
497                    self.build_buffered_ranges_stream(path, parts, ConcurrencyStrategy::Unbounded);
498                let mut len = 0;
499                while let Some(bytes) = stream.try_next().await? {
500                    len += bytes.len();
501                    file.write_all(&bytes).await?;
502                }
503
504                assert_eq!(len, size as usize);
505
506                PolarsResult::Ok(pl_async::Size::from(len as u64))
507            },
508        )
509        .await?;
510
511        // Dropping is delayed for tokio async files so we need to explicitly
512        // flush here (https://github.com/tokio-rs/tokio/issues/2307#issuecomment-596336451).
513        file.sync_all().await.map_err(PolarsError::from)?;
514
515        Ok(())
516    }
517
518    /// Fetch the metadata of the parquet file, do not memoize it.
519    pub async fn head(
520        &self,
521        path: &Path,
522        strategy: ConcurrencyStrategy,
523    ) -> PolarsResult<ObjectMeta> {
524        match strategy {
525            ConcurrencyStrategy::BytesBased => {
526                let controller = self.get_or_init_concurrency();
527                let _permit = controller.acquire(0).await;
528                self.head_inner(path, Some(&**controller)).await
529            },
530            ConcurrencyStrategy::Legacy => {
531                with_concurrency_budget(1, || self.head_inner(path, None)).await
532            },
533            ConcurrencyStrategy::Unbounded => self.head_inner(path, None).await,
534        }
535    }
536
537    async fn head_inner(
538        &self,
539        path: &Path,
540        controller: Option<&ConcurrencyController>,
541    ) -> PolarsResult<ObjectMeta> {
542        self.exec_with_rebuild_retry_on_err(|s| {
543            async move {
544                let t0 = Instant::now();
545                let head_result = self.io_metrics().record_io_read(0, s.head(path)).await;
546
547                if head_result.is_err() {
548                    let t0 = Instant::now();
549
550                    // Pre-signed URLs forbid the HEAD method, but we can still retrieve the header
551                    // information with a range 0-1 request.
552                    let get_range_0_1_result = self
553                        .io_metrics()
554                        .record_io_read(
555                            0,
556                            s.get_opts(
557                                path,
558                                object_store::GetOptions {
559                                    range: Some((0..1).into()),
560                                    ..Default::default()
561                                },
562                            ),
563                        )
564                        .await;
565
566                    if let Ok(v) = get_range_0_1_result {
567                        if let Some(controller) = controller {
568                            controller.record_head_rtt(t0.elapsed());
569                        }
570                        return Ok(v.meta);
571                    }
572                }
573
574                let out = head_result?;
575
576                if let Some(controller) = controller {
577                    controller.record_head_rtt(t0.elapsed());
578                }
579
580                Ok(out)
581            }
582        })
583        .await
584    }
585}
586
587/// Splits a single range into multiple smaller ranges, which can be downloaded concurrently for
588/// much higher throughput.
589fn split_range(
590    range: Range<usize>,
591    chunk_size: Option<usize>,
592) -> impl ExactSizeIterator<Item = Range<usize>> {
593    let chunk_size = chunk_size.unwrap_or_else(get_download_chunk_size);
594
595    // Calculate n_parts such that we are as close as possible to the `chunk_size`.
596    let n_parts = [
597        (range.len().div_ceil(chunk_size)).max(1),
598        (range.len() / chunk_size).max(1),
599    ]
600    .into_iter()
601    .min_by_key(|x| (range.len() / *x).abs_diff(chunk_size))
602    .unwrap();
603
604    let chunk_size = (range.len() / n_parts).max(1);
605
606    assert_eq!(n_parts, (range.len() / chunk_size).max(1));
607    let bytes_rem = range.len() % chunk_size;
608
609    (0..n_parts).map(move |part_no| {
610        let (start, end) = if part_no == 0 {
611            // Download remainder length in the first chunk since it starts downloading first.
612            let end = range.start + chunk_size + bytes_rem;
613            let end = if end > range.end { range.end } else { end };
614            (range.start, end)
615        } else {
616            let start = bytes_rem + range.start + part_no * chunk_size;
617            (start, start + chunk_size)
618        };
619
620        start..end
621    })
622}
623
624/// Note: For optimal performance, `ranges` should be sorted. More generally,
625/// ranges placed next to each other should also be close in range value.
626///
627/// # Returns
628/// `[(range1, end1), (range2, end2)]`, where:
629/// * `range1` contains bytes for the ranges from `ranges[0..end1]`
630/// * `range2` contains bytes for the ranges from `ranges[end1..end2]`
631/// * etc..
632///
633/// Note that if an end value is 0, it means the range is a splitted part and should be combined.
634fn merge_ranges(
635    ranges: &[Range<usize>],
636    chunk_size: Option<usize>,
637) -> impl Iterator<Item = (Range<usize>, usize)> + '_ {
638    let chunk_size = chunk_size.unwrap_or_else(get_download_chunk_size);
639
640    let mut current_merged_range = ranges.first().map_or(0..0, Clone::clone);
641    // Number of fetched bytes excluding excess.
642    let mut current_n_bytes = current_merged_range.len();
643
644    (0..ranges.len())
645        .filter_map(move |current_idx| {
646            let current_idx = 1 + current_idx;
647
648            if current_idx == ranges.len() {
649                // No more items - flush current state.
650                Some((current_merged_range.clone(), current_idx))
651            } else {
652                let range = ranges[current_idx].clone();
653
654                let new_merged = current_merged_range.start.min(range.start)
655                    ..current_merged_range.end.max(range.end);
656
657                // E.g.:
658                // |--------|
659                //  oo        // range1
660                //       oo   // range2
661                //    ^^^     // distance = 3, is_overlapping = false
662                // E.g.:
663                // |--------|
664                //  ooooo     // range1
665                //     ooooo  // range2
666                //     ^^     // distance = 2, is_overlapping = true
667                let (distance, is_overlapping) = {
668                    let l = current_merged_range.end.min(range.end);
669                    let r = current_merged_range.start.max(range.start);
670
671                    (r.abs_diff(l), r < l)
672                };
673
674                let should_merge = is_overlapping || {
675                    let leq_current_len_dist_to_chunk_size = new_merged.len().abs_diff(chunk_size)
676                        <= current_merged_range.len().abs_diff(chunk_size);
677                    let gap_tolerance =
678                        (current_n_bytes.max(range.len()) / 8).clamp(1024 * 1024, 8 * 1024 * 1024);
679
680                    leq_current_len_dist_to_chunk_size && distance <= gap_tolerance
681                };
682
683                if should_merge {
684                    // Merge to existing range
685                    current_merged_range = new_merged;
686                    current_n_bytes += if is_overlapping {
687                        range.len() - distance
688                    } else {
689                        range.len()
690                    };
691                    None
692                } else {
693                    let out = (current_merged_range.clone(), current_idx);
694                    current_merged_range = range;
695                    current_n_bytes = current_merged_range.len();
696                    Some(out)
697                }
698            }
699        })
700        .flat_map(move |x| {
701            // Split large individual ranges within the list of ranges.
702            let (range, end) = x;
703            let split = split_range(range, Some(chunk_size));
704            let len = split.len();
705
706            split
707                .enumerate()
708                .map(move |(i, range)| (range, if 1 + i == len { end } else { 0 }))
709        })
710}
711
712#[cfg(test)]
713mod tests {
714
715    #[test]
716    fn test_split_range() {
717        use super::{get_download_chunk_size, split_range};
718
719        let chunk_size = get_download_chunk_size();
720
721        assert_eq!(chunk_size, 64 * 1024 * 1024);
722
723        #[allow(clippy::single_range_in_vec_init)]
724        {
725            // Round-trip empty ranges.
726            assert_eq!(split_range(0..0, None).collect::<Vec<_>>(), [0..0]);
727            assert_eq!(split_range(3..3, None).collect::<Vec<_>>(), [3..3]);
728        }
729
730        // Threshold to start splitting to 2 ranges
731        //
732        // n - chunk_size == chunk_size - n / 2
733        // n + n / 2 == 2 * chunk_size
734        // 3 * n == 4 * chunk_size
735        // n = 4 * chunk_size / 3
736        let n = 4 * chunk_size / 3;
737
738        #[allow(clippy::single_range_in_vec_init)]
739        {
740            assert_eq!(split_range(0..n, None).collect::<Vec<_>>(), [0..89478485]);
741        }
742
743        assert_eq!(
744            split_range(0..n + 1, None).collect::<Vec<_>>(),
745            [0..44739243, 44739243..89478486]
746        );
747
748        // Threshold to start splitting to 3 ranges
749        //
750        // n / 2 - chunk_size == chunk_size - n / 3
751        // n / 2 + n / 3 == 2 * chunk_size
752        // 5 * n == 12 * chunk_size
753        // n == 12 * chunk_size / 5
754        let n = 12 * chunk_size / 5;
755
756        assert_eq!(
757            split_range(0..n, None).collect::<Vec<_>>(),
758            [0..80530637, 80530637..161061273]
759        );
760
761        assert_eq!(
762            split_range(0..n + 1, None).collect::<Vec<_>>(),
763            [0..53687092, 53687092..107374183, 107374183..161061274]
764        );
765    }
766
767    #[test]
768    fn test_merge_ranges() {
769        use super::{get_download_chunk_size, merge_ranges};
770
771        let chunk_size = get_download_chunk_size();
772
773        assert_eq!(chunk_size, 64 * 1024 * 1024);
774
775        // Round-trip empty slice
776        assert_eq!(merge_ranges(&[], None).collect::<Vec<_>>(), []);
777
778        // We have 1 tiny request followed by 1 huge request. They are combined as it reduces the
779        // `abs_diff()` to the `chunk_size`, but afterwards they are split to 2 evenly sized
780        // requests.
781        assert_eq!(
782            merge_ranges(&[0..1, 1..127 * 1024 * 1024], None).collect::<Vec<_>>(),
783            [(0..66584576, 0), (66584576..133169152, 2)]
784        );
785
786        // <= 1MiB gap, merge
787        assert_eq!(
788            merge_ranges(&[0..1, 1024 * 1024 + 1..1024 * 1024 + 2], None).collect::<Vec<_>>(),
789            [(0..1048578, 2)]
790        );
791
792        // > 1MiB gap, do not merge
793        assert_eq!(
794            merge_ranges(&[0..1, 1024 * 1024 + 2..1024 * 1024 + 3], None).collect::<Vec<_>>(),
795            [(0..1, 1), (1048578..1048579, 2)]
796        );
797
798        // <= 12.5% gap, merge
799        assert_eq!(
800            merge_ranges(&[0..8, 10..11], None).collect::<Vec<_>>(),
801            [(0..11, 2)]
802        );
803
804        // <= 12.5% gap relative to RHS, merge
805        assert_eq!(
806            merge_ranges(&[0..1, 3..11], None).collect::<Vec<_>>(),
807            [(0..11, 2)]
808        );
809
810        // Overlapping range, merge
811        assert_eq!(
812            merge_ranges(
813                &[0..80 * 1024 * 1024, 10 * 1024 * 1024..70 * 1024 * 1024],
814                None
815            )
816            .collect::<Vec<_>>(),
817            [(0..80 * 1024 * 1024, 2)]
818        );
819    }
820}