snix_store/nar/renderer/seekable/
mod.rs

1use std::{
2    collections::HashMap,
3    io::{self, SeekFrom},
4    pin::Pin,
5    task::{Context, Poll},
6};
7
8use futures::ready;
9use pin_project::pin_project;
10use segments::Segments;
11use snix_castore::Node;
12use snix_castore::directoryservice::DirectoryService;
13use snix_castore::{blobservice::BlobService, directoryservice::DirectoryGraphBuilder};
14use tokio::io::{AsyncBufRead, AsyncRead, AsyncSeek, AsyncWrite};
15use tokio_stream::StreamExt;
16use tracing::{instrument, warn};
17
18use crate::nar::RenderError;
19
20mod segments;
21
22#[cfg(test)]
23mod test;
24
25/// The number of segments to poll data from concurrently.
26const SEGMENT_CONCURRENCY: usize = 24;
27
28pub async fn write_nar<W, BS, DS>(
29    mut w: W,
30    root_node: &Node,
31    blob_service: &BS,
32    directory_service: &DS,
33) -> Result<(), RenderError>
34where
35    W: AsyncWrite + Unpin + Send,
36    BS: BlobService,
37    DS: DirectoryService,
38{
39    let mut reader = Reader::new(root_node, blob_service, directory_service).await?;
40    tokio::io::copy_buf(&mut reader, &mut w)
41        .await
42        // FUTUREWORK: RenderError makes no sense
43        .map_err(RenderError::BlobService)?;
44
45    Ok(())
46}
47
48#[pin_project]
49pub struct Reader<'bs, BS: BlobService + 'bs> {
50    segments: Segments,
51    pos: u64,
52    blob_service: BS,
53    #[pin]
54    rd: Box<dyn AsyncBufRead + Send + Unpin + 'bs>,
55}
56
57impl<'bs, BS: BlobService + Clone + 'bs> Reader<'bs, BS> {
58    /// Creates a new seekable NAR renderer for the given castore root node.
59    ///
60    /// This function pre-fetches the directory closure using `get_recursive()` and assembles the
61    /// NAR structure, except the file contents which are stored as 'holes' with references to a blob
62    /// of a specific BLAKE3 digest and known size.
63    /// The AsyncRead implementation will then switch between serving the
64    /// precomputed literal segments, and the appropriate blob for the file
65    /// contents.
66    #[instrument(skip(blob_service, directory_service), err)]
67    pub async fn new(
68        root_node: &Node,
69        blob_service: BS,
70        directory_service: impl DirectoryService,
71        // FUTUREWORK: add concurrency arg
72    ) -> Result<Self, RenderError> {
73        // If this is a directory, resolve all subdirectories
74        let directories = if let Node::Directory { digest, .. } = root_node {
75            let mut builder = DirectoryGraphBuilder::new_root_to_leaves(digest.to_owned());
76            let mut directories = directory_service.get_recursive(digest);
77            while let Some(directory) = directories
78                .try_next()
79                .await
80                .map_err(RenderError::DirectoryService)?
81            {
82                builder
83                    .try_insert(directory)
84                    .map_err(RenderError::OrderingError)?;
85            }
86
87            let directory_graph = builder.build().map_err(|err| {
88                if err == snix_castore::directoryservice::OrderingError::EmptySet {
89                    // The graph should at least contain the root, if there's no child directories.
90                    // The only way we could run into this is by the
91                    // DirectoryService not having the root directory we asked
92                    // for, which hints to misconfiguration, so explicitly warn!.
93                    let err = RenderError::DirectoryNotFound(*digest, "root".into());
94                    warn!(%err, "tried to render NAR, but DirectoryService didn't contain the root directory");
95                    err
96                } else {
97                    RenderError::OrderingError(err)
98                }
99            })?;
100
101            HashMap::from_iter(
102                directory_graph
103                    // drain order doesn't really matter
104                    .drain_leaves_to_root()
105                    .map(|d| (d.digest(), d)),
106            )
107        } else {
108            // If the top-level node is a file or a symlink, there is no directory graph.
109            Default::default()
110        };
111
112        let segments = Segments::from_root_node_and_directories(root_node, &directories);
113        let rd = segments.reader_for_offset(0, SEGMENT_CONCURRENCY, blob_service.clone());
114
115        Ok(Self {
116            segments,
117            pos: 0,
118            blob_service,
119            rd,
120        })
121    }
122
123    pub fn nar_size(&self) -> u64 {
124        self.segments.total_len()
125    }
126}
127
128impl<'bs, BS: BlobService> AsyncRead for Reader<'bs, BS> {
129    fn poll_read(
130        self: Pin<&mut Self>,
131        cx: &mut Context,
132        buf: &mut tokio::io::ReadBuf,
133    ) -> Poll<io::Result<()>> {
134        let this = self.project();
135
136        let bytes_read = {
137            let filled = buf.filled().len();
138            ready!(this.rd.poll_read(cx, buf))?;
139            buf.filled().len() - filled
140        };
141        *this.pos = this
142            .pos
143            .checked_add(bytes_read as u64)
144            .ok_or(std::io::Error::new(
145                std::io::ErrorKind::OutOfMemory,
146                "position > u64::MAX bytes",
147            ))?;
148
149        Poll::Ready(Ok(()))
150    }
151}
152
153impl<'bs, BS: BlobService> AsyncBufRead for Reader<'bs, BS> {
154    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
155        let this = self.project();
156        this.rd.poll_fill_buf(cx)
157    }
158
159    fn consume(self: Pin<&mut Self>, amt: usize) {
160        let this = self.project();
161
162        this.rd.consume(amt);
163        *this.pos = this
164            .pos
165            .checked_add(amt as u64)
166            .expect("consume would increase pos > u64::MAX bytes");
167    }
168}
169
170impl<'bs, BS: BlobService + Clone + 'bs> AsyncSeek for Reader<'bs, BS> {
171    fn start_seek(self: Pin<&mut Self>, pos: io::SeekFrom) -> io::Result<()> {
172        let nar_size = self.nar_size();
173        let new_pos = calc_pos(self.pos, nar_size, pos)?;
174
175        if new_pos != self.pos {
176            // FUTUREWORK: seek forward small amounts by skipping?
177            let mut this = self.project();
178
179            *this.rd = this.segments.reader_for_offset(
180                new_pos,
181                SEGMENT_CONCURRENCY,
182                this.blob_service.clone(),
183            );
184            *this.pos = new_pos;
185        }
186
187        Ok(())
188    }
189    fn poll_complete(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<io::Result<u64>> {
190        Poll::Ready(Ok(self.pos))
191    }
192}
193
194/// For a given nar_size and current position, returns the position that seek_from would seek to.
195fn calc_pos(cur_pos: u64, nar_size: u64, seek_from: SeekFrom) -> std::io::Result<u64> {
196    let new_pos = match seek_from {
197        SeekFrom::Start(p) => p,
198        SeekFrom::End(p) => nar_size.checked_sub_signed(p).ok_or(std::io::Error::new(
199            std::io::ErrorKind::InvalidInput,
200            "tried to seek before beginning of NAR",
201        ))?,
202        SeekFrom::Current(p) => cur_pos.checked_add_signed(p).ok_or(std::io::Error::new(
203            std::io::ErrorKind::UnexpectedEof,
204            "tried to seek way past end of NAR",
205        ))?,
206    };
207
208    if new_pos > nar_size {
209        Err(std::io::Error::new(
210            std::io::ErrorKind::UnexpectedEof,
211            "tried to seek past end of NAR",
212        ))
213    } else {
214        Ok(new_pos)
215    }
216}
snix_store/nar/renderer/seekable/mod.rs

snix_store/nar/renderer/seekable/
mod.rs
