forbid bare TryStreamExt::try_chunks (#24261)

avoid issues with losing error context by forbidding bare `try_chunks` and requiring `TryChunksExt::try_chunks2` instead.

better name? i don't like using `try_chunks` because then you have to qualify it with `TryChunksExt::try_chunks(stream, cap)` whenever it's used in a file that also imports `TryStreamExt`.

GitOrigin-RevId: f1e758e21301c6affdb0a27868254cabe1de7a37

Author: ldanilek

crates/common/src/lib.rs

@@ -70,6 +70,7 @@ pub mod types;
 pub mod utils;
 pub use value;
 pub mod bounded_thread_pool;
+pub mod try_chunks;
 pub mod version;
 pub mod ws;
 

crates/common/src/persistence_helpers.rs

@@ -10,6 +10,7 @@ use crate::{
     document::ResolvedDocument,
     knobs::DOCUMENTS_IN_MEMORY,
     persistence::RepeatablePersistence,
+    try_chunks::TryChunksExt,
     types::Timestamp,
 };
 
@@ -49,7 +50,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).map_err(|e| e.1);
+    let documents = documents.try_chunks2(*DOCUMENTS_IN_MEMORY);
     futures::pin_mut!(documents);
 
     while let Some(read_chunk) = documents.try_next().await? {

crates/common/src/try_chunks.rs

@@ -0,0 +1,17 @@
+use futures::{
+    Stream,
+    TryStreamExt,
+};
+
+/// Use this instead of `try_chunks` to pass through errors directly.
+/// TryStreamExt::try_chunks wraps errors which can lose stacktrace and context.
+pub trait TryChunksExt<T, E> {
+    fn try_chunks2(self, cap: usize) -> impl Stream<Item = Result<Vec<T>, E>>;
+}
+
+impl<T, E, S: Stream<Item = Result<T, E>>> TryChunksExt<T, E> for S {
+    fn try_chunks2(self, cap: usize) -> impl Stream<Item = Result<Vec<T>, E>> {
+        #[allow(clippy::disallowed_methods)]
+        TryStreamExt::try_chunks(self, cap).map_err(|e| e.1)
+    }
+}

crates/database/src/retention.rs

@@ -82,6 +82,7 @@ use common::{
         Reader,
         Writer,
     },
+    try_chunks::TryChunksExt,
     types::{
         GenericIndexName,
         IndexId,
@@ -533,8 +534,7 @@ impl<RT: Runtime> LeaderRetentionManager<RT> {
         let reader_ = &reader;
         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)
+            .try_chunks2(*RETENTION_READ_CHUNK)
             .map(move |chunk| async move {
                 let chunk = chunk?.to_vec();
                 let mut entries_to_delete = vec![];
@@ -653,8 +653,7 @@ impl<RT: Runtime> LeaderRetentionManager<RT> {
             all_indexes,
             persistence_version,
         )
-        .try_chunks(*RETENTION_DELETE_CHUNK)
-        .map_err(|e| e.1);
+        .try_chunks2(*RETENTION_DELETE_CHUNK);
         pin_mut!(expired_chunks);
         while let Some(delete_chunk) = expired_chunks.try_next().await? {
             tracing::trace!(
@@ -709,8 +708,7 @@ impl<RT: Runtime> LeaderRetentionManager<RT> {
                 Order::Asc,
                 Arc::new(NoopRetentionValidator),
             )
-            .try_chunks(*RETENTION_READ_CHUNK)
-            .map_err(|e| e.1)
+            .try_chunks2(*RETENTION_READ_CHUNK)
             .map(move |chunk| async move {
                 let chunk = chunk?.to_vec();
                 let mut entries_to_delete: Vec<(Timestamp, InternalDocumentId)> = vec![];
@@ -812,8 +810,7 @@ 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)
-            .map_err(|e| e.1);
+            .try_chunks2(*RETENTION_DELETE_CHUNK);
         pin_mut!(expired_chunks);
         while let Some(delete_chunk) = expired_chunks.try_next().await? {
             tracing::trace!(
@@ -1530,6 +1527,7 @@ mod tests {
             TestIdGenerator,
             TestPersistence,
         },
+        try_chunks::TryChunksExt,
         types::{
             unchecked_repeatable_ts,
             DatabaseIndexUpdate,
@@ -1570,9 +1568,8 @@ mod tests {
             ));
         };
         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);
+        // IMPORTANT: try_chunks fails here. try_chunks2 is necessary.
+        let chunks = stream_throws.try_chunks2(1);
         let chunk_throws = async move || -> anyhow::Result<()> {
             pin_mut!(chunks);
             chunks.try_next().await?;

crates/database/src/table_iteration.rs

@@ -35,6 +35,7 @@ use common::{
         Runtime,
         RuntimeInstant,
     },
+    try_chunks::TryChunksExt,
     types::{
         IndexId,
         RepeatableTimestamp,
@@ -318,8 +319,7 @@ 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)
-            .map_err(|e| e.1);
+            .try_chunks2(self.page_size);
         pin_mut!(documents);
         while let Some(chunk) = documents.try_next().await? {
             while let Err(not_until) = rate_limiter.check() {
@@ -409,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).map_err(|e| e.1);
+        let id_chunks = ids.try_chunks2(self.page_size);
         pin_mut!(id_chunks);
 
         while let Some(chunk) = id_chunks.try_next().await? {