From e1ad7af10ce3f432934cce6b8d3c5f2f0ce981e5 Mon Sep 17 00:00:00 2001 From: Zhenchi Date: Wed, 3 Jan 2024 14:55:09 +0800 Subject: [PATCH] feat(puffin): finish return written bytes (#3082) Signed-off-by: Zhenchi --- src/puffin/src/file_format/writer.rs | 8 ++--- src/puffin/src/file_format/writer/file.rs | 40 ++++++++++++++--------- src/puffin/src/tests.rs | 12 ++++--- 3 files changed, 36 insertions(+), 24 deletions(-) diff --git a/src/puffin/src/file_format/writer.rs b/src/puffin/src/file_format/writer.rs index 95760df0fe..24d1fb0e9f 100644 --- a/src/puffin/src/file_format/writer.rs +++ b/src/puffin/src/file_format/writer.rs @@ -44,8 +44,8 @@ pub trait PuffinSyncWriter { /// Add a blob to the Puffin file fn add_blob(&mut self, blob: Blob) -> Result<()>; - /// Finish writing the Puffin file - fn finish(&mut self) -> Result<()>; + /// Finish writing the Puffin file, returns the number of bytes written + fn finish(&mut self) -> Result; } /// The trait for writing Puffin files asynchronously @@ -57,6 +57,6 @@ pub trait PuffinAsyncWriter { /// Add a blob to the Puffin file async fn add_blob(&mut self, blob: Blob) -> Result<()>; - /// Finish writing the Puffin file - async fn finish(&mut self) -> Result<()>; + /// Finish writing the Puffin file, returns the number of bytes written + async fn finish(&mut self) -> Result; } diff --git a/src/puffin/src/file_format/writer/file.rs b/src/puffin/src/file_format/writer/file.rs index 3f65b9c89d..7337ee2d26 100644 --- a/src/puffin/src/file_format/writer/file.rs +++ b/src/puffin/src/file_format/writer/file.rs @@ -36,8 +36,8 @@ pub struct PuffinFileWriter { /// The metadata of the blobs blob_metadata: Vec, - /// The offset of the next blob - next_blob_offset: u64, + /// The number of bytes written + written_bytes: u64, } impl PuffinFileWriter { @@ -46,7 +46,7 @@ impl PuffinFileWriter { writer, properties: HashMap::new(), blob_metadata: Vec::new(), - next_blob_offset: 0, + written_bytes: 0, } } @@ -59,7 +59,7 @@ impl PuffinFileWriter { BlobMetadataBuilder::default() .blob_type(typ) .properties(properties) - .offset(self.next_blob_offset as _) + .offset(self.written_bytes as _) .length(size as _) .build() .expect("Required fields are not set") @@ -79,14 +79,16 @@ impl PuffinSyncWriter for PuffinFileWriter { let blob_metadata = self.create_blob_metadata(blob.blob_type, blob.properties, size); self.blob_metadata.push(blob_metadata); - self.next_blob_offset += size; + self.written_bytes += size; Ok(()) } - fn finish(&mut self) -> Result<()> { + fn finish(&mut self) -> Result { self.write_header_if_needed_sync()?; self.write_footer_sync()?; - self.writer.flush().context(FlushSnafu) + self.writer.flush().context(FlushSnafu)?; + + Ok(self.written_bytes as usize) } } @@ -106,23 +108,25 @@ impl PuffinAsyncWriter for PuffinFileWriter { let blob_metadata = self.create_blob_metadata(blob.blob_type, blob.properties, size); self.blob_metadata.push(blob_metadata); - self.next_blob_offset += size; + self.written_bytes += size; Ok(()) } - async fn finish(&mut self) -> Result<()> { + async fn finish(&mut self) -> Result { self.write_header_if_needed_async().await?; self.write_footer_async().await?; self.writer.flush().await.context(FlushSnafu)?; - self.writer.close().await.context(CloseSnafu) + self.writer.close().await.context(CloseSnafu)?; + + Ok(self.written_bytes as usize) } } impl PuffinFileWriter { fn write_header_if_needed_sync(&mut self) -> Result<()> { - if self.next_blob_offset == 0 { + if self.written_bytes == 0 { self.writer.write_all(&MAGIC).context(WriteSnafu)?; - self.next_blob_offset += MAGIC.len() as u64; + self.written_bytes += MAGIC.len() as u64; } Ok(()) } @@ -134,15 +138,17 @@ impl PuffinFileWriter { ) .into_footer_bytes()?; - self.writer.write_all(&bytes).context(WriteSnafu) + self.writer.write_all(&bytes).context(WriteSnafu)?; + self.written_bytes += bytes.len() as u64; + Ok(()) } } impl PuffinFileWriter { async fn write_header_if_needed_async(&mut self) -> Result<()> { - if self.next_blob_offset == 0 { + if self.written_bytes == 0 { self.writer.write_all(&MAGIC).await.context(WriteSnafu)?; - self.next_blob_offset += MAGIC.len() as u64; + self.written_bytes += MAGIC.len() as u64; } Ok(()) } @@ -154,6 +160,8 @@ impl PuffinFileWriter { ) .into_footer_bytes()?; - self.writer.write_all(&bytes).await.context(WriteSnafu) + self.writer.write_all(&bytes).await.context(WriteSnafu)?; + self.written_bytes += bytes.len() as u64; + Ok(()) } } diff --git a/src/puffin/src/tests.rs b/src/puffin/src/tests.rs index b05da1aeb5..09792085c4 100644 --- a/src/puffin/src/tests.rs +++ b/src/puffin/src/tests.rs @@ -125,7 +125,8 @@ fn test_writer_reader_with_empty_sync() { "Test 1234".to_string(), )])); - writer.finish().unwrap(); + let written_bytes = writer.finish().unwrap(); + assert!(written_bytes > 0); let mut buf = Cursor::new(buf.into_inner()); let mut reader = PuffinFileReader::new(&mut buf); @@ -150,7 +151,8 @@ async fn test_writer_reader_empty_async() { "Test 1234".to_string(), )])); - writer.finish().await.unwrap(); + let written_bytes = writer.finish().await.unwrap(); + assert!(written_bytes > 0); let mut buf = AsyncCursor::new(buf.into_inner()); let mut reader = PuffinFileReader::new(&mut buf); @@ -194,7 +196,8 @@ fn test_writer_reader_sync() { "Test 1234".to_string(), )])); - writer.finish().unwrap(); + let written_bytes = writer.finish().unwrap(); + assert!(written_bytes > 0); let mut buf = Cursor::new(buf.into_inner()); let mut reader = PuffinFileReader::new(&mut buf); @@ -257,7 +260,8 @@ async fn test_writer_reader_async() { "Test 1234".to_string(), )])); - writer.finish().await.unwrap(); + let written_bytes = writer.finish().await.unwrap(); + assert!(written_bytes > 0); let mut buf = AsyncCursor::new(buf.into_inner()); let mut reader = PuffinFileReader::new(&mut buf);