-
Notifications
You must be signed in to change notification settings - Fork 9
/
metadata.rs
174 lines (147 loc) · 6.16 KB
/
metadata.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
use std::io::{Read, Seek, SeekFrom};
use bytes::Bytes;
use prost::Message;
use snafu::{OptionExt, ResultExt};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeek, AsyncSeekExt};
use crate::error::{self, Result};
use crate::proto::{Footer, Metadata, PostScript, StripeFooter};
use crate::reader::decompress::Decompressor;
use super::decompress::Compression;
const DEFAULT_FOOTER_SIZE: u64 = 16 * 1024;
/// The file's metadata.
#[derive(Debug)]
pub struct FileMetadata {
pub postscript: PostScript,
pub footer: Footer,
pub metadata: Metadata,
pub stripe_footers: Vec<StripeFooter>,
}
macro_rules! impl_read_metadata {
($reader:ident $($_await:tt)*) => {
{
let file_len = {
let old_pos = $reader.stream_position()$($_await)*.context(error::SeekSnafu)?;
let len = $reader.seek(SeekFrom::End(0))$($_await)*.context(error::SeekSnafu)?;
// Avoid seeking a third time when we were already at the end of the
// stream. The branch is usually way cheaper than a seek operation.
if old_pos != len {
$reader.seek(SeekFrom::Start(old_pos))$($_await)*
.context(error::SeekSnafu)?;
}
len
};
// initial read of the footer
let assume_footer_len = if file_len < DEFAULT_FOOTER_SIZE {
file_len
} else {
DEFAULT_FOOTER_SIZE
};
$reader
.seek(SeekFrom::End(-(assume_footer_len as i64)))$($_await)*
.context(error::SeekSnafu)?;
let mut tail_bytes = Vec::with_capacity(assume_footer_len as usize);
$reader
.take(assume_footer_len)
.read_to_end(&mut tail_bytes)$($_await)*
.context(error::IoSnafu)?;
// The final byte of the file contains the serialized length of the Postscript,
// which must be less than 256 bytes.
let postscript_len = tail_bytes[tail_bytes.len() - 1] as usize;
tail_bytes.truncate(tail_bytes.len() - 1);
// next is the postscript
let postscript = PostScript::decode(&tail_bytes[tail_bytes.len() - postscript_len..])
.context(error::DecodeProtoSnafu)?;
let compression = Compression::from_proto(
postscript.compression(),
postscript.compression_block_size
);
tail_bytes.truncate(tail_bytes.len() - postscript_len);
// next is the footer
let footer_length = postscript.footer_length.context(error::OutOfSpecSnafu {
msg: "Footer length is empty",
})? as usize; // todo: throw error
let footer_offset = file_len - footer_length as u64 - postscript_len as u64 - 1;
$reader
.seek(SeekFrom::Start(footer_offset))$($_await)*
.context(error::SeekSnafu)?;
let mut footer = vec![0; footer_length];
$reader
.read_exact(&mut footer)$($_await)*
.context(error::SeekSnafu)?;
let footer = deserialize_footer(&footer, compression)?;
// finally the metadata
let metadata_length = postscript.metadata_length.context(error::OutOfSpecSnafu {
msg: "Metadata length is empty",
})? as usize;
let metadata_offset =
file_len - metadata_length as u64 - footer_length as u64 - postscript_len as u64 - 1;
$reader
.seek(SeekFrom::Start(metadata_offset))$($_await)*
.context(error::SeekSnafu)?;
let mut metadata = vec![0; metadata_length];
$reader.read_exact(&mut metadata)$($_await)*.context(error::IoSnafu)?;
let metadata = deserialize_footer_metadata(&metadata, compression)?;
let mut stripe_footers = Vec::with_capacity(footer.stripes.len());
let mut scratch = Vec::<u8>::new();
for stripe in &footer.stripes {
let start = stripe.offset() + stripe.index_length() + stripe.data_length();
let len = stripe.footer_length();
$reader
.seek(SeekFrom::Start(start))$($_await)*
.context(error::SeekSnafu)?;
scratch.clear();
scratch.reserve(len as usize);
$reader
.take(len)
.read_to_end(&mut scratch)$($_await)*
.context(error::IoSnafu)?;
stripe_footers.push(deserialize_stripe_footer(
&scratch,
compression,
)?);
}
Ok(FileMetadata {
postscript,
footer,
metadata,
stripe_footers,
})
}
};
}
pub fn read_metadata<R>(reader: &mut R) -> Result<FileMetadata>
where
R: Read + Seek,
{
impl_read_metadata!(reader)
}
pub async fn read_metadata_async<R>(reader: &mut R) -> Result<FileMetadata>
where
R: AsyncRead + AsyncSeek + Unpin + Send,
{
impl_read_metadata!(reader.await)
}
fn deserialize_footer(bytes: &[u8], compression: Option<Compression>) -> Result<Footer> {
let mut buffer = vec![];
Decompressor::new(Bytes::copy_from_slice(bytes), compression, vec![])
.read_to_end(&mut buffer)
.context(error::IoSnafu)?;
Footer::decode(&*buffer).context(error::DecodeProtoSnafu)
}
fn deserialize_footer_metadata(bytes: &[u8], compression: Option<Compression>) -> Result<Metadata> {
let mut buffer = vec![];
Decompressor::new(Bytes::copy_from_slice(bytes), compression, vec![])
.read_to_end(&mut buffer)
.context(error::IoSnafu)?;
Metadata::decode(&*buffer).context(error::DecodeProtoSnafu)
}
fn deserialize_stripe_footer(
bytes: &[u8],
compression: Option<Compression>,
) -> Result<StripeFooter> {
let mut buffer = vec![];
Decompressor::new(Bytes::copy_from_slice(bytes), compression, vec![])
.read_to_end(&mut buffer)
.context(error::IoSnafu)?;
StripeFooter::decode(&*buffer).context(error::DecodeProtoSnafu)
}