puffin/puffin_manager/fs_puffin_manager/
reader.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
15#[cfg(test)]
16mod tests;
17
18use std::io;
19use std::ops::Range;
20use std::sync::Arc;
21
22use async_compression::futures::bufread::ZstdDecoder;
23use async_trait::async_trait;
24use bytes::{BufMut, Bytes};
25use common_base::range_read::{AsyncReadAdapter, Metadata, RangeReader, SizeAwareRangeReader};
26use futures::io::BufReader;
27use futures::{AsyncRead, AsyncWrite};
28use snafu::{OptionExt, ResultExt, ensure};
29use tokio::sync::OnceCell;
30
31use crate::blob_metadata::{BlobMetadata, CompressionCodec};
32use crate::error::{
33    BlobIndexOutOfBoundSnafu, BlobNotFoundSnafu, DeserializeJsonSnafu, FileKeyNotMatchSnafu,
34    MetadataSnafu, ReadSnafu, Result, UnsupportedDecompressionSnafu, WriteSnafu,
35};
36use crate::file_format::reader::{AsyncReader, PuffinFileReader};
37use crate::file_metadata::FileMetadata;
38use crate::partial_reader::PartialReader;
39use crate::puffin_manager::file_accessor::PuffinFileAccessor;
40use crate::puffin_manager::file_key::PuffinFileKey;
41use crate::puffin_manager::fs_puffin_manager::PuffinMetadataCacheRef;
42use crate::puffin_manager::fs_puffin_manager::dir_meta::DirMetadata;
43use crate::puffin_manager::stager::{BoxWriter, DirWriterProviderRef, Stager};
44use crate::puffin_manager::{BlobGuard, DirMetrics, GuardWithMetadata, PuffinReader};
45
46/// `FsPuffinReader` is a `PuffinReader` that provides fs readers for puffin files.
47/// Metadata is retained for this immutable file identity until the reader is dropped.
48/// Replacing file contents requires a new reader and invalidating shared caches.
49pub struct FsPuffinReader<S, F>
50where
51    S: Stager + 'static,
52    F: PuffinFileAccessor + Clone,
53{
54    /// The handle of the puffin file.
55    handle: F::FileHandle,
56
57    /// The file size hint.
58    file_size_hint: Option<u64>,
59
60    /// The stager.
61    stager: S,
62
63    /// The puffin file accessor.
64    puffin_file_accessor: F,
65
66    /// The puffin file metadata cache.
67    puffin_file_metadata_cache: Option<PuffinMetadataCacheRef>,
68
69    /// Successful metadata retained for this immutable file handle.
70    /// Failures leave the cell empty so a later operation can retry.
71    metadata: OnceCell<Arc<FileMetadata>>,
72}
73
74impl<S, F> FsPuffinReader<S, F>
75where
76    S: Stager + 'static,
77    F: PuffinFileAccessor + Clone,
78{
79    /// Returns this reader's storage-scoped file identity without opening the file.
80    pub fn file_key(&self) -> PuffinFileKey<F::FileHandle> {
81        self.puffin_file_accessor.file_key(&self.handle)
82    }
83
84    pub(crate) fn new(
85        handle: F::FileHandle,
86        stager: S,
87        puffin_file_accessor: F,
88        puffin_file_metadata_cache: Option<PuffinMetadataCacheRef>,
89    ) -> Self {
90        Self {
91            handle,
92            file_size_hint: None,
93            stager,
94            puffin_file_accessor,
95            puffin_file_metadata_cache,
96            metadata: OnceCell::new(),
97        }
98    }
99}
100
101#[async_trait]
102impl<S, F> PuffinReader for FsPuffinReader<S, F>
103where
104    F: PuffinFileAccessor + Clone,
105    S: Stager<FileHandle = F::FileHandle> + 'static,
106{
107    type Blob = Either<RandomReadBlob<F>, S::Blob>;
108    type Dir = S::Dir;
109
110    fn with_file_size_hint(mut self, file_size_hint: Option<u64>) -> Self {
111        self.file_size_hint = file_size_hint;
112        self
113    }
114
115    async fn metadata(&self) -> Result<Arc<FileMetadata>> {
116        if let Some(metadata) = self.metadata.get() {
117            return Ok(metadata.clone());
118        }
119        let mut file = self.puffin_reader().await?;
120        self.get_puffin_file_metadata(&mut file).await
121    }
122
123    async fn blob(&self, key: &str) -> Result<GuardWithMetadata<Self::Blob>> {
124        let metadata = self.metadata().await?;
125        let blob_metadata = get_blob_metadata(&metadata, key)?;
126        let blob = if blob_metadata.compression_codec.is_none() {
127            // If the blob is not compressed, we can directly read it from the puffin file.
128            Either::L(RandomReadBlob {
129                handle: self.handle.clone(),
130                accessor: self.puffin_file_accessor.clone(),
131                blob_metadata: blob_metadata.clone(),
132            })
133        } else {
134            // If the blob is compressed, we need to decompress it into staging space before reading.
135            let file = self.puffin_reader().await?;
136            let blob_metadata = blob_metadata.clone();
137            let staged_blob = self
138                .stager
139                .get_blob(
140                    &self.handle,
141                    key,
142                    Box::new(|writer| {
143                        Box::pin(Self::init_blob_to_stager(file, blob_metadata, writer))
144                    }),
145                )
146                .await?;
147
148            Either::R(staged_blob)
149        };
150
151        Ok(GuardWithMetadata::new(blob, blob_metadata))
152    }
153
154    async fn dir(&self, key: &str) -> Result<(GuardWithMetadata<Self::Dir>, DirMetrics)> {
155        let mut file = self.puffin_reader().await?;
156        let metadata = self.get_puffin_file_metadata(&mut file).await?;
157        let blob_metadata = get_blob_metadata(&metadata, key)?;
158        let (dir, metrics) = self
159            .stager
160            .get_dir(
161                &self.handle,
162                key,
163                Box::new(|writer_provider| {
164                    let accessor = self.puffin_file_accessor.clone();
165                    let handle = self.handle.clone();
166                    let blob_metadata = blob_metadata.clone();
167                    Box::pin(Self::init_dir_to_stager(
168                        file,
169                        metadata,
170                        blob_metadata,
171                        handle,
172                        writer_provider,
173                        accessor,
174                    ))
175                }),
176            )
177            .await?;
178
179        Ok((GuardWithMetadata::new(dir, blob_metadata), metrics))
180    }
181}
182
183impl<S, F> FsPuffinReader<S, F>
184where
185    S: Stager,
186    F: PuffinFileAccessor + Clone,
187{
188    async fn get_puffin_file_metadata(
189        &self,
190        reader: &mut PuffinFileReader<F::Reader>,
191    ) -> Result<Arc<FileMetadata>> {
192        let cached = self
193            .metadata
194            .get_or_try_init(|| async {
195                let metadata = if let Some(cache) = &self.puffin_file_metadata_cache {
196                    let key = self
197                        .puffin_file_accessor
198                        .file_key(&self.handle)
199                        .map_handle(|handle| handle.to_string());
200                    cache
201                        .try_get_with(&key, async { reader.metadata().await.map(Arc::new) })
202                        .await?
203                } else {
204                    Arc::new(reader.metadata().await?)
205                };
206                Ok::<_, crate::error::Error>(metadata)
207            })
208            .await?;
209        Ok(cached.clone())
210    }
211
212    async fn puffin_reader(&self) -> Result<PuffinFileReader<F::Reader>> {
213        let mut reader = self.puffin_file_accessor.reader(&self.handle).await?;
214        if let Some(file_size_hint) = self.file_size_hint {
215            reader.with_file_size_hint(file_size_hint);
216        }
217        Ok(PuffinFileReader::new(reader))
218    }
219
220    async fn init_blob_to_stager(
221        reader: PuffinFileReader<F::Reader>,
222        blob_metadata: BlobMetadata,
223        mut writer: BoxWriter,
224    ) -> Result<u64> {
225        let reader = reader.into_blob_reader(&blob_metadata);
226        let reader = AsyncReadAdapter::new(reader).await.context(MetadataSnafu)?;
227        let compression = blob_metadata.compression_codec;
228        let size = Self::handle_decompress(reader, &mut writer, compression).await?;
229        Ok(size)
230    }
231
232    async fn init_dir_to_stager(
233        mut file: PuffinFileReader<F::Reader>,
234        puffin_metadata: Arc<FileMetadata>,
235        blob_metadata: BlobMetadata,
236        handle: F::FileHandle,
237        writer_provider: DirWriterProviderRef,
238        accessor: F,
239    ) -> Result<u64> {
240        let reader = file.blob_reader(&blob_metadata)?;
241        let meta = reader.metadata().await.context(MetadataSnafu)?;
242        let buf = reader
243            .read(0..meta.content_length)
244            .await
245            .context(ReadSnafu)?;
246        let dir_meta: DirMetadata = serde_json::from_slice(&buf).context(DeserializeJsonSnafu)?;
247
248        let mut tasks = vec![];
249        for file_meta in dir_meta.files {
250            let blob_meta = puffin_metadata
251                .blobs
252                .get(file_meta.blob_index)
253                .context(BlobIndexOutOfBoundSnafu {
254                    index: file_meta.blob_index,
255                    max_index: puffin_metadata.blobs.len(),
256                })?
257                .clone();
258            ensure!(
259                blob_meta.blob_type == file_meta.key,
260                FileKeyNotMatchSnafu {
261                    expected: file_meta.key,
262                    actual: &blob_meta.blob_type,
263                }
264            );
265
266            let reader = accessor.reader(&handle).await?;
267            let writer = writer_provider.writer(&file_meta.relative_path).await?;
268            let task = common_runtime::spawn_global(async move {
269                let reader = PuffinFileReader::new(reader).into_blob_reader(&blob_meta);
270                let reader = AsyncReadAdapter::new(reader).await.context(MetadataSnafu)?;
271                let compression = blob_meta.compression_codec;
272                let size = Self::handle_decompress(reader, writer, compression).await?;
273                Ok(size)
274            });
275            tasks.push(task);
276        }
277
278        let size = futures::future::try_join_all(tasks)
279            .await
280            .into_iter()
281            .flatten()
282            .sum::<Result<_>>()?;
283
284        Ok(size)
285    }
286
287    /// Handles the decompression of the reader and writes the decompressed data to the writer.
288    /// Returns the number of bytes written.
289    async fn handle_decompress(
290        reader: impl AsyncRead,
291        mut writer: impl AsyncWrite + Unpin,
292        compression: Option<CompressionCodec>,
293    ) -> Result<u64> {
294        match compression {
295            Some(CompressionCodec::Lz4) => UnsupportedDecompressionSnafu {
296                decompression: "lz4",
297            }
298            .fail(),
299            Some(CompressionCodec::Zstd) => {
300                let reader = ZstdDecoder::new(BufReader::new(reader));
301                futures::io::copy(reader, &mut writer)
302                    .await
303                    .context(WriteSnafu)
304            }
305            None => futures::io::copy(reader, &mut writer)
306                .await
307                .context(WriteSnafu),
308        }
309    }
310}
311
312/// `RandomReadBlob` is a `BlobGuard` that directly reads the blob from the puffin file.
313pub struct RandomReadBlob<F: PuffinFileAccessor> {
314    handle: F::FileHandle,
315    accessor: F,
316    blob_metadata: BlobMetadata,
317}
318
319#[async_trait]
320impl<F: PuffinFileAccessor + Clone> BlobGuard for RandomReadBlob<F> {
321    type Reader = PartialReader<F::Reader>;
322
323    async fn reader(&self) -> Result<Self::Reader> {
324        ensure!(
325            self.blob_metadata.compression_codec.is_none(),
326            UnsupportedDecompressionSnafu {
327                decompression: self.blob_metadata.compression_codec.unwrap().to_string()
328            }
329        );
330
331        let reader = self.accessor.reader(&self.handle).await?;
332        let blob_reader = PuffinFileReader::new(reader).into_blob_reader(&self.blob_metadata);
333        Ok(blob_reader)
334    }
335}
336
337/// `Either` is a type that represents either `A` or `B`.
338///
339/// Used to:
340/// impl `RangeReader` for `Either<A: RangeReader, B: RangeReader>`,
341/// impl `BlobGuard` for `Either<A: BlobGuard, B: BlobGuard>`.
342pub enum Either<A, B> {
343    L(A),
344    R(B),
345}
346
347#[async_trait]
348impl<A, B> RangeReader for Either<A, B>
349where
350    A: RangeReader,
351    B: RangeReader,
352{
353    async fn metadata(&self) -> io::Result<Metadata> {
354        match self {
355            Either::L(a) => a.metadata().await,
356            Either::R(b) => b.metadata().await,
357        }
358    }
359    async fn read(&self, range: Range<u64>) -> io::Result<Bytes> {
360        match self {
361            Either::L(a) => a.read(range).await,
362            Either::R(b) => b.read(range).await,
363        }
364    }
365    async fn read_into(&self, range: Range<u64>, buf: &mut (impl BufMut + Send)) -> io::Result<()> {
366        match self {
367            Either::L(a) => a.read_into(range, buf).await,
368            Either::R(b) => b.read_into(range, buf).await,
369        }
370    }
371    async fn read_vec(&self, ranges: &[Range<u64>]) -> io::Result<Vec<Bytes>> {
372        match self {
373            Either::L(a) => a.read_vec(ranges).await,
374            Either::R(b) => b.read_vec(ranges).await,
375        }
376    }
377}
378
379#[async_trait]
380impl<A, B> BlobGuard for Either<A, B>
381where
382    A: BlobGuard + Sync,
383    B: BlobGuard + Sync,
384{
385    type Reader = Either<A::Reader, B::Reader>;
386    async fn reader(&self) -> Result<Self::Reader> {
387        match self {
388            Either::L(a) => Ok(Either::L(a.reader().await?)),
389            Either::R(b) => Ok(Either::R(b.reader().await?)),
390        }
391    }
392}
393
394/// Looks up the first matching blob without allocating a per-reader key index.
395fn get_blob_metadata(metadata: &FileMetadata, key: &str) -> Result<BlobMetadata> {
396    metadata
397        .blobs
398        .iter()
399        .find(|blob| blob.blob_type == key)
400        .cloned()
401        .context(BlobNotFoundSnafu { blob: key })
402}
puffin/puffin_manager/fs_puffin_manager/reader.rs

puffin/puffin_manager/fs_puffin_manager/
reader.rs
