puffin/file_format/writer/
footer.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::HashMap;
16use std::io::Write;
17use std::mem;
18
19use snafu::ResultExt;
20
21use crate::blob_metadata::BlobMetadata;
22use crate::error::{Lz4CompressionSnafu, Result, SerializeJsonSnafu};
23use crate::file_format::{Flags, MAGIC, MIN_FOOTER_SIZE};
24use crate::file_metadata::FileMetadataBuilder;
25
26/// Writer for the footer of a Puffin file.
27///
28/// ```text
29/// Footer layout: HeadMagic Payload PayloadSize Flags FootMagic
30///                [4]       [?]     [4]         [4]   [4]
31/// ```
32pub struct FooterWriter {
33    blob_metadata: Vec<BlobMetadata>,
34    file_properties: HashMap<String, String>,
35    lz4_compressed: bool,
36}
37
38impl FooterWriter {
39    pub fn new(
40        blob_metadata: Vec<BlobMetadata>,
41        file_properties: HashMap<String, String>,
42        lz4_compressed: bool,
43    ) -> Self {
44        Self {
45            blob_metadata,
46            file_properties,
47            lz4_compressed,
48        }
49    }
50
51    /// Serializes the footer to bytes
52    pub fn into_footer_bytes(mut self) -> Result<Vec<u8>> {
53        let payload = self.footer_payload()?;
54        let payload_size = payload.len();
55
56        let capacity = MIN_FOOTER_SIZE as usize + payload_size;
57        let mut buf = Vec::with_capacity(capacity);
58
59        self.write_magic(&mut buf); // HeadMagic
60        self.write_payload(&mut buf, &payload); // Payload
61        self.write_footer_payload_size(payload_size as _, &mut buf); // PayloadSize
62        self.write_flags(&mut buf); // Flags
63        self.write_magic(&mut buf); // FootMagic
64        Ok(buf)
65    }
66
67    fn write_magic(&self, buf: &mut Vec<u8>) {
68        buf.extend_from_slice(&MAGIC);
69    }
70
71    fn write_payload(&self, buf: &mut Vec<u8>, payload: &[u8]) {
72        buf.extend_from_slice(payload);
73    }
74
75    fn write_footer_payload_size(&self, payload_size: i32, buf: &mut Vec<u8>) {
76        buf.extend_from_slice(&payload_size.to_le_bytes());
77    }
78
79    /// Appends reserved flags (currently zero-initialized) to the given buffer.
80    fn write_flags(&self, buf: &mut Vec<u8>) {
81        let mut flags = Flags::DEFAULT;
82        if self.lz4_compressed {
83            flags |= Flags::FOOTER_PAYLOAD_COMPRESSED_LZ4;
84        } else {
85            flags &= !Flags::FOOTER_PAYLOAD_COMPRESSED_LZ4;
86        }
87
88        buf.extend_from_slice(&flags.bits().to_le_bytes());
89    }
90
91    fn footer_payload(&mut self) -> Result<Vec<u8>> {
92        let file_metadata = FileMetadataBuilder::default()
93            .blobs(mem::take(&mut self.blob_metadata))
94            .properties(mem::take(&mut self.file_properties))
95            .build()
96            .expect("Required fields are not set");
97
98        if self.lz4_compressed {
99            let mut buf = vec![];
100            let mut encoder = lz4_flex::frame::FrameEncoder::new(&mut buf);
101            serde_json::to_writer(&mut encoder, &file_metadata).context(SerializeJsonSnafu)?;
102            encoder.flush().context(Lz4CompressionSnafu)?;
103            Ok(buf)
104        } else {
105            serde_json::to_vec(&file_metadata).context(SerializeJsonSnafu)
106        }
107    }
108}
puffin/file_format/writer/footer.rs

puffin/file_format/writer/
footer.rs
