Skip to content
Snippets Groups Projects
Commit d124dd52 authored by Georges Racinet's avatar Georges Racinet
Browse files

rhgitaly::streaming::stream_chunks

This Rust version of Python `hgitaly.util.chunked` should be useful
for many gRPC methods. It provide the means to distinguish the first
response, but does not implement the pagination protocol.
parent ce0a1e85
No related branches found
No related tags found
2 merge requests!189Merged stable branch into default and bumped VERSION,!188RHGitaly CommitService.GetTreeEntries implementation
......@@ -65,3 +65,93 @@
});
}
}
/// Wrap an inner iterator over Results to aggregate into chunks.
///
/// As soon as an error is encountered while iterating for a chunk, it is yielded in place
/// of the chunk.
///
/// As of Rust 1.65, the stdlib the [`Iterator`] trait has a `next_chunk` method but
/// - it is unstable
/// - it yields arrays
/// - the remainder still needs to be collected, making it not so practical
/// - it does not have the error handling, of course
struct ChunkedIterator<T, IT, E>
where
IT: Iterator<Item = Result<T, E>>,
{
inner: IT,
chunk_size: usize,
}
impl<T, IT, E> Iterator for ChunkedIterator<T, IT, E>
where
IT: Iterator<Item = Result<T, E>>,
{
type Item = Result<Vec<T>, E>;
fn next(&mut self) -> Option<Result<Vec<T>, E>> {
let mut chunk: Vec<T> = Vec::new();
while chunk.len() < self.chunk_size {
match self.inner.next() {
None => {
if chunk.is_empty() {
return None;
} else {
break;
}
}
Some(Ok(v)) => chunk.push(v),
Some(Err(e)) => {
return Some(Err(e));
}
}
}
Some(Ok(chunk))
}
}
pub const DEFAULT_CHUNK_SIZE: usize = 50;
/// Stream responses by aggregating an inner iterator into chunks
///
/// This is useful when streamed responses have a repeated field: the inner iterator
/// would yield individual elements, and this generic function can be used to
/// collect them in chunks, from which responses are built and sent to the channel.
///
/// The resulting chunks are guaranteed not to be empty.
///
/// In `resp_builder`, the boolean tells whether the response to build is the first one.
/// It is indeed frequent in the Gitaly protocol for streamed responses with a repeated field
/// and additional metadata to put the metadata on the first response only, the subsequent
/// ones being only about chunking the repeated field.
///
/// This parallels Python `hgitaly.util.chunked` except that the chunk size is for now
/// the constant [`DEFAULT_CHUNK_SIZE`].
/// Note that this flexibility was introduced early in the Python implementation, but has
/// not been actually used yet.
/// Gitaly seems to apply a more subtle logic based on actual response size
/// in some cases. It looks like a better avenue for improvement to mimick that rather than
/// making the chunk size customizable.
pub fn stream_chunks<Iter, Resp, Item, E>(
tx: &BlockingResponseSender<Resp>,
iter: Iter,
resp_builder: impl FnOnce(Vec<Item>, bool) -> Resp + Copy,
err_handler: impl FnOnce(E) -> Status + Copy,
) where
Iter: Iterator<Item = Result<Item, E>>,
{
let mut first = true;
for res in (ChunkedIterator {
inner: iter,
chunk_size: DEFAULT_CHUNK_SIZE,
}) {
tx.send(res.map_or_else(
|e| Err(err_handler(e)),
|chunk| Ok(resp_builder(chunk, first)),
));
if first {
first = false;
}
}
}
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment