Skip to content

Commit

Permalink
Handle OOM when writing
Browse files Browse the repository at this point in the history
  • Loading branch information
kornelski committed Apr 17, 2024
1 parent 80e00b3 commit 06414f5
Showing 1 changed file with 26 additions and 4 deletions.
30 changes: 26 additions & 4 deletions rmp-serde/src/encode.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1219,9 +1219,9 @@ pub fn to_vec<T>(val: &T) -> Result<Vec<u8>, Error>
where
T: Serialize + ?Sized,
{
let mut wr = Vec::with_capacity(128);
let mut wr = FallibleWriter(Vec::new());
write(&mut wr, val)?;
Ok(wr)
Ok(wr.0)
}

/// Serializes data structure into byte vector as a map
Expand All @@ -1235,7 +1235,29 @@ pub fn to_vec_named<T>(val: &T) -> Result<Vec<u8>, Error>
where
T: Serialize + ?Sized,
{
let mut wr = Vec::with_capacity(128);
let mut wr = FallibleWriter(Vec::new());
write_named(&mut wr, val)?;
Ok(wr)
Ok(wr.0)
}

#[repr(transparent)]
struct FallibleWriter(Vec<u8>);

impl Write for FallibleWriter {
#[inline(always)]
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
self.write_all(buf)?;
Ok(buf.len())
}

#[inline]
fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
self.0.try_reserve(buf.len()).map_err(|_| std::io::ErrorKind::OutOfMemory)?;
self.0.extend_from_slice(buf);
Ok(())
}

fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}

0 comments on commit 06414f5

Please sign in to comment.