fix TryChunksStreamError usage (#24256)

using `try_chunks` directly loses all context on errors.
i tried to create a new trait that has a better `try_chunks` implementation, but the types were annoying so i decided to put `map_err` everywhere (for now, to fix the bleeding).

added a regression test.

GitOrigin-RevId: e1723a714d34b7687b594e476bd47af545d54948

Author: ldanilek

crates/common/src/persistence_helpers.rs

@@ -49,7 +49,7 @@ pub async fn stream_revision_pairs<'a>(
     documents: impl Stream<Item = RevisionStreamEntry> + 'a,
     reader: &'a RepeatablePersistence,
 ) {
-    let documents = documents.try_chunks(*DOCUMENTS_IN_MEMORY);
+    let documents = documents.try_chunks(*DOCUMENTS_IN_MEMORY).map_err(|e| e.1);
     futures::pin_mut!(documents);
 
     while let Some(read_chunk) = documents.try_next().await? {

crates/database/src/retention.rs

@@ -520,6 +520,7 @@ impl<RT: Runtime> LeaderRetentionManager<RT> {
         let mut index_entry_chunks = reader
             .load_documents(TimestampRange::new(cursor..min_snapshot_ts)?, Order::Asc)
             .try_chunks(*RETENTION_READ_CHUNK)
+            .map_err(|e| e.1)
             .map(move |chunk| async move {
                 let chunk = chunk?.to_vec();
                 let mut entries_to_delete = vec![];
@@ -638,7 +639,8 @@ impl<RT: Runtime> LeaderRetentionManager<RT> {
             all_indexes,
             persistence_version,
         )
-        .try_chunks(*RETENTION_DELETE_CHUNK);
+        .try_chunks(*RETENTION_DELETE_CHUNK)
+        .map_err(|e| e.1);
         pin_mut!(expired_chunks);
         while let Some(delete_chunk) = expired_chunks.try_next().await? {
             tracing::trace!(
@@ -694,6 +696,7 @@ impl<RT: Runtime> LeaderRetentionManager<RT> {
                 Arc::new(NoopRetentionValidator),
             )
             .try_chunks(*RETENTION_READ_CHUNK)
+            .map_err(|e| e.1)
             .map(move |chunk| async move {
                 let chunk = chunk?.to_vec();
                 let mut entries_to_delete: Vec<(Timestamp, InternalDocumentId)> = vec![];
@@ -795,7 +798,8 @@ impl<RT: Runtime> LeaderRetentionManager<RT> {
 
         tracing::trace!("delete_documents: about to grab chunks");
         let expired_chunks = Self::expired_documents(rt, reader, cursor, min_snapshot_ts)
-            .try_chunks(*RETENTION_DELETE_CHUNK);
+            .try_chunks(*RETENTION_DELETE_CHUNK)
+            .map_err(|e| e.1);
         pin_mut!(expired_chunks);
         while let Some(delete_chunk) = expired_chunks.try_next().await? {
             tracing::trace!(
@@ -1521,13 +1525,45 @@ mod tests {
             TableName,
         },
     };
-    use futures::TryStreamExt;
+    use errors::ErrorMetadataAnyhowExt;
+    use futures::{
+        pin_mut,
+        stream,
+        TryStreamExt,
+    };
     use maplit::{
         btreemap,
         btreeset,
     };
 
     use super::LeaderRetentionManager;
+    use crate::retention::{
+        snapshot_invalid_error,
+        RetentionType,
+    };
+
+    #[convex_macro::test_runtime]
+    async fn test_chunks_is_out_of_retention(_rt: TestRuntime) -> anyhow::Result<()> {
+        let throws = || -> anyhow::Result<()> {
+            anyhow::bail!(snapshot_invalid_error(
+                Timestamp::must(1),
+                Timestamp::must(30),
+                RetentionType::Document
+            ));
+        };
+        let stream_throws = stream::once(async move { throws() });
+        // IMPORTANT: the map_err is required here and whenever we use try_chunks.
+        // Otherwise the error gets re-wrapped and loses context.
+        let chunks = stream_throws.try_chunks(1).map_err(|e| e.1);
+        let chunk_throws = async move || -> anyhow::Result<()> {
+            pin_mut!(chunks);
+            chunks.try_next().await?;
+            anyhow::Ok(())
+        };
+        let err = chunk_throws().await.unwrap_err();
+        assert!(err.is_out_of_retention());
+        Ok(())
+    }
 
     #[convex_macro::test_runtime]
     async fn test_expired_index_entries(_rt: TestRuntime) -> anyhow::Result<()> {

crates/database/src/table_iteration.rs

@@ -318,7 +318,8 @@ impl<RT: Runtime> TableIterator<RT> {
             RepeatablePersistence::new(reader, end_ts, self.retention_validator.clone());
         let documents = repeatable_persistence
             .load_documents(TimestampRange::new(start_ts.succ()?..=*end_ts)?, Order::Asc)
-            .try_chunks(self.page_size);
+            .try_chunks(self.page_size)
+            .map_err(|e| e.1);
         pin_mut!(documents);
         while let Some(chunk) = documents.try_next().await? {
             while let Err(not_until) = rate_limiter.check() {
@@ -408,7 +409,7 @@ impl<RT: Runtime> TableIterator<RT> {
 
         // Note even though `previous_revisions` can paginate internally, we don't want
         // to hold the entire result set in memory, because documents can be large.
-        let id_chunks = ids.try_chunks(self.page_size);
+        let id_chunks = ids.try_chunks(self.page_size).map_err(|e| e.1);
         pin_mut!(id_chunks);
 
         while let Some(chunk) = id_chunks.try_next().await? {