Improve performance of sise stats

This introduces two limits to the concurrency:

* max memory keeping for in flight manifests
* max number of concurrent manifest fetches

With this, we don't need to limit the concurrency to a specific number
of manifests, but now we can fetch more manifests if they are small.

Combining these two limits, we get an important performance optimization,
particularly for repos with a large number of smallish manifests.

I used this script to test and observe the performance of this:


```python
import arraylake
import time
import icechunk

icechunk.set_logs_filter("icechunk::stream_utils=trace")

token = '...'

client = arraylake.Client(token=token)
repo = client.get_repo("earthmover-public/era5-surface-aws")

start_time = time.time()
repo.total_chunks_storage(None, None)
print(" %s seconds" % (time.time() - start_time))
```

At the time of running this, ERA5 took around 4 seconds to compute.

Testing with real world repos with small manifests shows
performance improvements of the order of 30x.

Author: paraseba

icechunk-python/python/icechunk/_icechunk_python.pyi

@@ -1357,7 +1357,9 @@ class PyRepository:
         self, message: str, *, branch: str, metadata: dict[str, Any] | None = None
     ) -> str: ...
     def total_chunks_storage(
-        self, process_manifests_concurrently: int | None = None
+        self,
+        max_manifest_mem_bytes: int | None = None,
+        max_concurrent_manifest_fetches: int | None = None
     ) -> int: ...
 
 class PySession:

icechunk-python/python/icechunk/repository.py

@@ -729,7 +729,9 @@ def rewrite_manifests(
         )
 
     def total_chunks_storage(
-        self, process_manifests_concurrently: int | None = None
+        self,
+        max_manifest_mem_bytes: int | None = None,
+        max_concurrent_manifest_fetches: int | None = None
     ) -> int:
         """Calculate the total storage used for chunks, in bytes .
 
@@ -743,8 +745,10 @@ def total_chunks_storage(
 
         Parameters
         ----------
-        process_manifests_concurrently : int | None
-            Process this many manifests concurrently. Defaults to 10.
+        max_manifest_mem_bytes : int | None
+            Don't use more than this memory to store in-flight manifests. Defaults to 512 MB.
+        max_concurrent_manifest_fetches : int | None
+            Don't run more than this many concurrent manifest fetches. Defaults to 500.
         """
 
-        return self._repository.total_chunks_storage(process_manifests_concurrently)
+        return self._repository.total_chunks_storage(max_manifest_mem_bytes, max_concurrent_manifest_fetches)

icechunk-python/src/repository.rs

@@ -1,7 +1,7 @@
 use std::{
     borrow::Cow,
     collections::{BTreeMap, BTreeSet, HashMap, HashSet},
-    num::NonZeroU16,
+    num::{NonZeroU16, NonZeroUsize},
     sync::Arc,
 };
 
@@ -1053,7 +1053,8 @@ impl PyRepository {
     pub fn total_chunks_storage(
         &self,
         py: Python<'_>,
-        process_manifests_concurrently: Option<NonZeroU16>,
+        max_manifest_mem_bytes: Option<NonZeroUsize>,
+        max_concurrent_manifest_fetches: Option<NonZeroU16>,
     ) -> PyResult<u64> {
         // This function calls block_on, so we need to allow other thread python to make progress
         py.allow_threads(move || {
@@ -1071,8 +1072,12 @@ impl PyRepository {
                         storage.as_ref(),
                         &storage_settings,
                         asset_manager,
-                        process_manifests_concurrently.unwrap_or(
-                            NonZeroU16::try_from(10).unwrap_or(NonZeroU16::MIN),
+                        max_manifest_mem_bytes.unwrap_or(
+                            NonZeroUsize::try_from(512 * 1020 * 1024)
+                                .unwrap_or(NonZeroUsize::MIN),
+                        ),
+                        max_concurrent_manifest_fetches.unwrap_or(
+                            NonZeroU16::try_from(500).unwrap_or(NonZeroU16::MIN),
                         ),
                     )
                     .await

icechunk/src/lib.rs

@@ -32,6 +32,7 @@ pub mod storage;
 pub mod store;
 #[cfg(test)]
 pub mod strategies;
+mod stream_utils;
 pub mod virtual_chunks;
 
 pub use config::{ObjectStoreConfig, RepositoryConfig};

icechunk/src/ops/gc.rs

@@ -12,7 +12,7 @@ use crate::{
         ChunkId, IcechunkFormatError, IcechunkFormatErrorKind, ManifestId, SnapshotId,
         manifest::ChunkPayload,
     },
-    ops::pointed_snapshot_ids,
+    ops::pointed_snapshots,
     refs::{Ref, RefError, delete_branch, delete_tag, list_refs},
     repository::{RepositoryError, RepositoryErrorKind},
     storage::{self, DeleteObjectsResult, ListInfo},
@@ -174,23 +174,22 @@ pub async fn garbage_collect(
     }
 
     tracing::info!("Finding GC roots");
-    let all_snaps = pointed_snapshot_ids(
+    let all_snaps = pointed_snapshots(
         storage,
         storage_settings,
         Arc::clone(&asset_manager),
         &config.extra_roots,
     )
     .await?;
 
-    // FIXME: add attribute files
     let mut keep_chunks = HashSet::new();
     let mut keep_manifests = HashSet::new();
     let mut keep_snapshots = HashSet::new();
 
     tracing::info!("Calculating retained objects");
     pin!(all_snaps);
-    while let Some(snap_id) = all_snaps.try_next().await? {
-        let snap = asset_manager.fetch_snapshot(&snap_id).await?;
+    while let Some(snap) = all_snaps.try_next().await? {
+        let snap_id = snap.id();
         if config.deletes_snapshots() && keep_snapshots.insert(snap_id.clone()) {
             tracing::trace!("Adding snapshot to keep list: {}", &snap_id);
         }
@@ -240,6 +239,8 @@ pub async fn garbage_collect(
 
     let mut summary = GCSummary::default();
 
+    tracing::info!("Starting deletes");
+
     if config.deletes_snapshots() {
         let res = gc_snapshots(
             asset_manager.as_ref(),

icechunk/src/ops/stats.rs

@@ -1,26 +1,30 @@
-use futures::{TryStream, TryStreamExt as _, future::ready, stream};
+use futures::{TryStreamExt, future::ready, stream};
 use std::{
     collections::HashSet,
-    num::NonZeroU16,
+    num::{NonZeroU16, NonZeroUsize},
     sync::{Arc, Mutex},
 };
+use tracing::trace;
 
 use crate::{
     Storage,
     asset_manager::AssetManager,
-    format::{ChunkId, ManifestId, manifest::ChunkPayload},
+    format::{
+        ChunkId,
+        manifest::{ChunkPayload, Manifest},
+        snapshot::ManifestFileInfo,
+    },
     ops::pointed_snapshots,
-    repository::{RepositoryErrorKind, RepositoryResult},
+    repository::{RepositoryError, RepositoryErrorKind, RepositoryResult},
     storage,
+    stream_utils::{StreamLimiter, try_unique_stream},
 };
 
-async fn manifest_chunks_storage(
-    manifest_id: ManifestId,
-    manifest_size: u64,
-    asset_manager: Arc<AssetManager>,
+fn calculate_manifest_storage(
+    manifest: Arc<Manifest>,
     seen_chunks: Arc<Mutex<HashSet<ChunkId>>>,
 ) -> RepositoryResult<u64> {
-    let manifest = asset_manager.fetch_manifest(&manifest_id, manifest_size).await?;
+    trace!(manifest_id = %manifest.id(), "Processing manifest");
     let mut size = 0;
     for payload in manifest.chunk_payloads() {
         match payload {
@@ -47,36 +51,18 @@ async fn manifest_chunks_storage(
             }
         }
     }
+    trace!(manifest_id = %manifest.id(), "Manifest done");
     Ok(size)
 }
 
-pub fn try_unique_stream<S, T, E, F, V>(
-    f: F,
-    stream: S,
-) -> impl TryStream<Ok = T, Error = E>
-where
-    F: Fn(&S::Ok) -> V,
-    S: TryStream<Ok = T, Error = E>,
-    V: Eq + std::hash::Hash,
-{
-    let mut seen = HashSet::new();
-    stream.try_filter(move |item| {
-        let v = f(item);
-        if seen.insert(v) {
-            futures::future::ready(true)
-        } else {
-            futures::future::ready(false)
-        }
-    })
-}
-
 /// Compute the total size in bytes of all committed repo chunks.
 /// It doesn't include inline or virtual chunks.
 pub async fn repo_chunks_storage(
     storage: &(dyn Storage + Send + Sync),
     storage_settings: &storage::Settings,
     asset_manager: Arc<AssetManager>,
-    process_manifests_concurrently: NonZeroU16,
+    max_manifest_mem_bytes: NonZeroUsize,
+    max_concurrent_manifest_fetches: NonZeroU16,
 ) -> RepositoryResult<u64> {
     let extra_roots = Default::default();
     let all_snaps = pointed_snapshots(
@@ -90,27 +76,59 @@ pub async fn repo_chunks_storage(
     let all_manifest_infos = all_snaps
         // this could be slightly optimized by not collecting all manifest info records into a vec
         // but we don't expect too many, and they are small anyway
-        .map_ok(|snap| stream::iter(snap.manifest_files().map(Ok).collect::<Vec<_>>()))
+        .map_ok(|snap| {
+            stream::iter(
+                snap.manifest_files().map(Ok::<_, RepositoryError>).collect::<Vec<_>>(),
+            )
+        })
         .try_flatten();
+
+    // we don't want to check manifests more than once, so we unique them by their id
     let unique_manifest_infos = try_unique_stream(|mi| mi.id.clone(), all_manifest_infos);
 
+    // we want to fetch many manifests in parallel, but not more than memory allows
+    // for this we use the StreamLimiter using the manifest size in bytes for usage
+    let limiter = &Arc::new(StreamLimiter::new(
+        max_manifest_mem_bytes.get(),
+        |m: &ManifestFileInfo| m.size_bytes as usize,
+    ));
+
+    // The StreamLimiter works by calling limit on every element before they are processed
+    let rate_limited_manifests = unique_manifest_infos
+        .and_then(|m| async move { Ok(limiter.clone().limit(m).await) });
+
     let seen_chunks = &Arc::new(Mutex::new(HashSet::new()));
     let asset_manager = &asset_manager;
 
-    let res = unique_manifest_infos
-        .map_ok(|manifest_info| async move {
-            let manifest_size = manifest_info.size_bytes;
-            manifest_chunks_storage(
-                manifest_info.id,
-                manifest_size,
-                Arc::clone(asset_manager),
-                Arc::clone(seen_chunks),
-            )
-            .await
+    let (_, res) = rate_limited_manifests
+        .map_ok(|m| async move {
+            let manifest =
+                Arc::clone(asset_manager).fetch_manifest(&m.id, m.size_bytes).await?;
+            Ok((manifest, m))
+        })
+        // Now we can buffer a bunch of fetch_manifest operations. Because we are using
+        // StreamLimiter we know memory is not going to blow up
+        .try_buffer_unordered(max_concurrent_manifest_fetches.get() as usize)
+        .map_ok(|(manifest, minfo)| async move {
+            let size = calculate_manifest_storage(manifest, Arc::clone(seen_chunks))?;
+            Ok((size, minfo))
+        })
+        // We do some more buffering to get some concurrency on the processing of the manifest file
+        // TODO: this should actually happen in a CPU bounded worker pool
+        .try_buffer_unordered(4)
+        // Now StreamLimiter requires us to call free, this will make room for more manifests to be
+        // fetched into the previous buffer
+        .and_then(|(size, minfo)| async move {
+            limiter.clone().free(minfo).await;
+            Ok(size)
+        })
+        .try_fold((0u64, 0), |(processed, total_size), partial| {
+            //info!("Processed {processed} manifests");
+            ready(Ok((processed + 1, total_size + partial)))
         })
-        .try_buffered(process_manifests_concurrently.get() as usize)
-        .try_fold(0, |total, partial| ready(Ok(total + partial)))
         .await?;
 
+    debug_assert_eq!(limiter.current_usage().await, (0, 0));
+
     Ok(res)
 }

icechunk/src/stream_utils.rs

@@ -0,0 +1,123 @@
+use std::{collections::HashSet, marker::PhantomData};
+
+use futures::{TryStream, TryStreamExt};
+use tokio::sync::{Mutex, Notify};
+use tracing::trace;
+
+/// Return the unique elements of the input stream
+/// Unique is computed according to `f`, and only the first
+/// matching element is yielded.
+pub(crate) fn try_unique_stream<S, T, E, F, V>(
+    f: F,
+    stream: S,
+) -> impl TryStream<Ok = T, Error = E>
+where
+    F: Fn(&S::Ok) -> V,
+    S: TryStream<Ok = T, Error = E>,
+    V: Eq + std::hash::Hash,
+{
+    let mut seen = HashSet::new();
+    stream.try_filter(move |item| {
+        let v = f(item);
+        if seen.insert(v) {
+            futures::future::ready(true)
+        } else {
+            futures::future::ready(false)
+        }
+    })
+}
+
+#[derive(Debug, Default)]
+struct StreamLimiterState {
+    // TODO: we may want to switch to atomic values instead
+    current_usage: usize,
+    in_flight: usize,
+}
+
+/// A component to limit the rate at which a stream produces items
+/// Call `limit` on each element of the stream before they are processed.
+/// Call `free` on each element of the stream once they are done.
+/// No more than `max_usage` usage will be generated in total by the in-flight
+/// yielded elements. Usage is measured according to the function passed to `new`.
+/// Usually, this is used in combination with `TryStreamExt::try_buffer`.
+///
+/// Note: this could have a much better interface with a bit of extra work:
+/// - We could simple wrap the stream and it would take care of mapping every element
+/// - Elements could be freed automatically using the Drop instance
+#[derive(Debug)]
+pub(crate) struct StreamLimiter<T, F: Fn(&'_ T) -> usize> {
+    get_usage: F,
+    max_usage: usize,
+    state: Mutex<StreamLimiterState>,
+    waiter: Notify,
+    _phantom: PhantomData<T>,
+}
+
+impl<T, F: Fn(&'_ T) -> usize> StreamLimiter<T, F> {
+    /// Create a new StreamLimiter
+    /// Don't yield more than max_usage in-flight usage.
+    /// Use `get_usage` to calculate the usage of an element in the stream.
+    pub(crate) fn new(max_usage: usize, get_usage: F) -> Self {
+        Self {
+            get_usage,
+            max_usage,
+            state: Mutex::new(Default::default()),
+            waiter: Notify::new(),
+            _phantom: Default::default(),
+        }
+    }
+
+    /// Expend some usage. This should be called on every element
+    /// of the input stream, before they are processed.
+    pub(crate) async fn limit(&self, item: T) -> T {
+        let this_item_usage = (self.get_usage)(&item);
+        loop {
+            let mut state = self.state.lock().await;
+            if state.current_usage == 0
+                || state.current_usage + this_item_usage <= self.max_usage
+            {
+                trace!(
+                    item_usage = this_item_usage,
+                    current_usage = state.current_usage,
+                    in_flight = state.in_flight,
+                    "Stats rate limiter: Allowing item"
+                );
+                state.current_usage += this_item_usage;
+                state.in_flight += 1;
+                return item;
+            } else {
+                trace!(
+                    item_usage = this_item_usage,
+                    current_usage = state.current_usage,
+                    in_flight = state.in_flight,
+                    "Stats rate limiter: Blocking item"
+                );
+                drop(state);
+                self.waiter.notified().await;
+            }
+        }
+    }
+
+    /// Free some usage. This should be called on every element
+    /// of the output stream, after they are processed.
+    pub(crate) async fn free(&self, item: T) -> T {
+        let this_item_usage = (self.get_usage)(&item);
+        let mut state = self.state.lock().await;
+        trace!(
+            item_usage = this_item_usage,
+            current_usage = state.current_usage,
+            in_flight = state.in_flight,
+            "Stats rate limiter: Freeing item"
+        );
+        state.current_usage -= this_item_usage;
+        state.in_flight -= 1;
+        self.waiter.notify_waiters();
+        item
+    }
+
+    /// Get the current usage and number of in-flight items.
+    pub(crate) async fn current_usage(&self) -> (usize, usize) {
+        let state = self.state.lock().await;
+        (state.current_usage, state.in_flight)
+    }
+}