Skip to main content

mountos_native/
file.rs

1use crate::bulk::{batch_errors, item, MAX_BATCH_ITEMS};
2use crate::client::{Slot, Sock};
3use crate::error::{Error, Result};
4use crate::frame::MAX_PAYLOAD_BYTES;
5use crate::wire::*;
6use std::io::{self, Read, Seek, SeekFrom, Write};
7use std::sync::{Arc, Mutex, MutexGuard};
8
9pub(crate) fn to_i64(v: u64) -> Result<i64> {
10  i64::try_from(v).map_err(|_| Error::server("EINVAL", "value out of range"))
11}
12
13pub(crate) struct State {
14  pub(crate) handle: u64,
15  pos: u64,
16  closed: bool,
17}
18
19/// One ReadAt request of at most one payload from `base`. Returns the data and the length asked.
20fn read_chunk(c: &Sock, handle: u64, base: i64, len: usize) -> Result<(Vec<u8>, usize)> {
21  let want = len.min(c.max_payload());
22  let req = ReadAtRequest {
23    handle,
24    offset: base,
25    length: want as u32,
26  };
27  let (r, data) = c.call_msg::<ReadAtResponse>(method::READ_AT, &req, &[])?;
28  if data.len() != r.count as usize || data.len() > want {
29    return Err(Error::Protocol(format!(
30      "read reply count {} payload {} want {want}",
31      r.count,
32      data.len()
33    )));
34  }
35  Ok((data, want))
36}
37
38/// One WriteAt request of at most one payload from the start of `rest`. Returns the bytes stored.
39fn write_chunk(c: &Sock, handle: u64, base: i64, rest: &[u8]) -> Result<usize> {
40  let chunk = &rest[..rest.len().min(c.max_payload())];
41  let (r, _) = c.call_msg::<WriteAtResponse>(method::WRITE_AT, &WriteAtRequest { handle, offset: base }, chunk)?;
42  let n = r.count as usize;
43  if n == 0 || n > chunk.len() {
44    return Err(Error::Protocol(format!(
45      "write reply count {n} for {} bytes",
46      chunk.len()
47    )));
48  }
49  Ok(n)
50}
51
52fn offset_at(base: i64, done: usize) -> Result<i64> {
53  base
54    .checked_add(done as i64)
55    .ok_or_else(|| Error::server("EINVAL", "offset overflow"))
56}
57
58/// What one pass of a batch call produced. A `Group` holds the results of the leading items and
59/// how many items they cover. `Big` means the first item needs more than one payload and goes
60/// through the chunked path.
61enum Step<T> {
62  Group(Vec<Result<T>>, usize),
63  Big,
64}
65
66/// One write in `File::write_batch`. Build it with `WriteItem::new`.
67#[derive(Clone, Debug, Default, PartialEq, Eq)]
68#[non_exhaustive]
69pub struct WriteItem {
70  pub offset: u64,
71  pub data: Vec<u8>,
72}
73
74impl WriteItem {
75  /// Writes `data` at `offset`.
76  pub fn new(offset: u64, data: Vec<u8>) -> WriteItem {
77    WriteItem { offset, data }
78  }
79}
80
81/// Reads the leading ranges of `ranges` that fit in one ReadBatch request.
82fn read_group(c: &Sock, handle: u64, ranges: &[Range]) -> Result<Step<Vec<u8>>> {
83  if ranges[0].length as usize > c.max_payload() {
84    return Ok(Step::Big);
85  }
86  let (mut end, mut bytes) = (0, 0usize);
87  while end < ranges.len() && end < MAX_BATCH_ITEMS && bytes + ranges[end].length as usize <= c.max_payload() {
88    bytes += ranges[end].length as usize;
89    end += 1;
90  }
91  let part = &ranges[..end];
92  let req = ReadBatchRequest {
93    handle,
94    ranges: part.to_vec(),
95  };
96  let (resp, data) = c.call_msg::<ReadBatchResponse>(method::READ_BATCH, &req, &[])?;
97  if resp.counts.len() != part.len() {
98    return Err(Error::Protocol(format!(
99      "read batch returned {} counts for {} ranges",
100      resp.counts.len(),
101      part.len()
102    )));
103  }
104  let mut errs = batch_errors(part.len(), resp.errors)?;
105  let mut out = Vec::with_capacity(part.len());
106  let mut pos = 0usize;
107  for (i, r) in part.iter().enumerate() {
108    if let Some(e) = errs[i].take() {
109      out.push(Err(e));
110      continue;
111    }
112    let n = resp.counts[i] as usize;
113    if n > r.length as usize || pos + n > data.len() {
114      return Err(Error::Protocol(format!(
115        "read batch count {n} for range {i} is out of bounds"
116      )));
117    }
118    out.push(Ok(data[pos..pos + n].to_vec()));
119    pos += n;
120  }
121  if pos != data.len() {
122    return Err(Error::Protocol(format!(
123      "read batch payload {} bytes, counts sum to {pos}",
124      data.len()
125    )));
126  }
127  Ok(Step::Group(out, end))
128}
129
130/// Writes the leading items of `items` that fit in one WriteBatch request. An item whose offset
131/// cannot be sent fails alone without a request.
132fn write_group(c: &Sock, handle: u64, items: &[WriteItem]) -> Result<Step<usize>> {
133  if let Err(e) = to_i64(items[0].offset) {
134    return Ok(Step::Group(vec![Err(e)], 1));
135  }
136  if items[0].data.len() > c.max_payload() {
137    return Ok(Step::Big);
138  }
139  let (mut end, mut bytes) = (0, 0usize);
140  while end < items.len()
141    && end < MAX_BATCH_ITEMS
142    && items[end].offset <= i64::MAX as u64
143    && bytes + items[end].data.len() <= c.max_payload()
144  {
145    bytes += items[end].data.len();
146    end += 1;
147  }
148  let part = &items[..end];
149  let mut ranges = Vec::with_capacity(part.len());
150  let mut payload = Vec::with_capacity(bytes);
151  for it in part {
152    ranges.push(Range {
153      offset: it.offset as i64,
154      length: it.data.len() as u32,
155    });
156    payload.extend_from_slice(&it.data);
157  }
158  let (resp, _) =
159    c.call_msg::<WriteBatchResponse>(method::WRITE_BATCH, &WriteBatchRequest { handle, ranges }, &payload)?;
160  if resp.counts.len() != part.len() {
161    return Err(Error::Protocol(format!(
162      "write batch returned {} counts for {} items",
163      resp.counts.len(),
164      part.len()
165    )));
166  }
167  let mut errs = batch_errors(part.len(), resp.errors)?;
168  let mut out = Vec::with_capacity(part.len());
169  for (i, it) in part.iter().enumerate() {
170    if let Some(e) = errs[i].take() {
171      out.push(Err(e));
172      continue;
173    }
174    let n = resp.counts[i] as usize;
175    if n > it.data.len() {
176      return Err(Error::Protocol(format!(
177        "write batch count {n} for item {i} exceeds its length"
178      )));
179    }
180    out.push(Ok(n));
181  }
182  Ok(Step::Group(out, end))
183}
184
185/// An open remote file. Calls on one file are serialized by a mutex. Different files
186/// on different pool connections run in parallel.
187pub struct File {
188  slot: Arc<Slot>,
189  epoch: u64,
190  st: Mutex<State>,
191}
192
193impl File {
194  pub(crate) fn new(slot: Arc<Slot>, handle: u64, epoch: u64) -> File {
195    File {
196      slot,
197      epoch,
198      st: Mutex::new(State {
199        handle,
200        pos: 0,
201        closed: false,
202      }),
203    }
204  }
205
206  pub(crate) fn state(&self) -> Result<MutexGuard<'_, State>> {
207    let st = self.st.lock().unwrap_or_else(|e| e.into_inner());
208    if st.closed {
209      return Err(Error::server("EBADF", "file is closed"));
210    }
211    Ok(st)
212  }
213
214  pub(crate) fn unary<R: prost::Message + Default>(&self, m: u16, req: &impl prost::Message) -> Result<R> {
215    Ok(self.slot.run(Some(self.epoch), |c| c.call_msg::<R>(m, req, &[]))?.0 .0)
216  }
217
218  /// Reads `len` bytes from `base` one request at a time, so the connection is free between
219  /// requests. Stops early at end of file.
220  fn read_from(&self, handle: u64, base: i64, len: usize) -> Result<Vec<u8>> {
221    let mut out = Vec::with_capacity(len.min(MAX_PAYLOAD_BYTES as usize));
222    while out.len() < len {
223      let off = offset_at(base, out.len())?;
224      let ((data, want), _) = self
225        .slot
226        .run(Some(self.epoch), |c| read_chunk(c, handle, off, len - out.len()))?;
227      let short = data.len() < want;
228      out.extend_from_slice(&data);
229      if short {
230        break;
231      }
232    }
233    Ok(out)
234  }
235
236  /// Writes all of `buf` at `base` one request at a time. On failure after some bytes were
237  /// stored, the error carries the count.
238  fn write_to(&self, handle: u64, base: i64, buf: &[u8]) -> std::result::Result<usize, (usize, Error)> {
239    let mut done = 0;
240    while done < buf.len() {
241      let off = offset_at(base, done).map_err(|e| (done, e))?;
242      let (n, _) = self
243        .slot
244        .run(Some(self.epoch), |c| write_chunk(c, handle, off, &buf[done..]))
245        .map_err(|e| (done, e))?;
246      done += n;
247    }
248    Ok(done)
249  }
250
251  fn do_read(&self, handle: u64, offset: u64, len: usize) -> Result<Vec<u8>> {
252    self.read_from(handle, to_i64(offset)?, len)
253  }
254
255  fn do_write(&self, handle: u64, offset: u64, buf: &[u8]) -> Result<usize> {
256    self.write_to(handle, to_i64(offset)?, buf).map_err(|(_, e)| e)
257  }
258
259  fn do_fstat(&self, handle: u64) -> Result<FileInfo> {
260    let r: StatResponse = self.unary(method::FSTAT, &FstatRequest { handle })?;
261    r.info.ok_or_else(|| Error::Protocol("fstat reply has no info".into()))
262  }
263
264  /// Reads up to `len` bytes at `offset`. The result is shorter than `len` only at end of file.
265  pub fn read_at(&self, offset: u64, len: usize) -> Result<Vec<u8>> {
266    let st = self.state()?;
267    self.do_read(st.handle, offset, len)
268  }
269
270  /// Writes all of `buf` at `offset`, split into frames within the payload limit.
271  pub fn write_at(&self, offset: u64, buf: &[u8]) -> Result<usize> {
272    let st = self.state()?;
273    self.do_write(st.handle, offset, buf)
274  }
275
276  /// Reads many ranges in as few requests as the limits allow. The outer error means the call
277  /// failed. Each inner result is one range: a range past end of file returns fewer bytes,
278  /// a failed range returns its own error. The connection is free between requests.
279  pub fn read_batch(&self, ranges: &[Range]) -> Result<Vec<Result<Vec<u8>>>> {
280    let st = self.state()?;
281    let handle = st.handle;
282    let mut out = Vec::with_capacity(ranges.len());
283    let mut start = 0;
284    while start < ranges.len() {
285      let (step, _) = self
286        .slot
287        .run(Some(self.epoch), |c| read_group(c, handle, &ranges[start..]))?;
288      match step {
289        Step::Group(items, n) => {
290          out.extend(items);
291          start += n;
292        }
293        Step::Big => {
294          out.push(item(self.read_from(
295            handle,
296            ranges[start].offset,
297            ranges[start].length as usize,
298          ))?);
299          start += 1;
300        }
301      }
302    }
303    Ok(out)
304  }
305
306  /// Writes many items in as few requests as the limits allow. Each inner result is the
307  /// byte count written for that item or its own error. An item larger than one payload that
308  /// fails after some bytes were stored returns `Error::Partial` with the count.
309  pub fn write_batch(&self, items: &[WriteItem]) -> Result<Vec<Result<usize>>> {
310    let st = self.state()?;
311    let handle = st.handle;
312    let mut out = Vec::with_capacity(items.len());
313    let mut start = 0;
314    while start < items.len() {
315      let (step, _) = self
316        .slot
317        .run(Some(self.epoch), |c| write_group(c, handle, &items[start..]))?;
318      match step {
319        Step::Group(res, n) => {
320          out.extend(res);
321          start += n;
322        }
323        Step::Big => {
324          let it = &items[start];
325          let r = match self.write_to(handle, it.offset as i64, &it.data) {
326            Ok(n) => Ok(Ok(n)),
327            Err((0, e)) => item(Err(e)),
328            Err((written, e @ Error::Server { .. })) => Ok(Err(Error::Partial {
329              written,
330              source: Box::new(e),
331            })),
332            Err((_, e)) => Err(e),
333          };
334          out.push(r?);
335          start += 1;
336        }
337      }
338    }
339    Ok(out)
340  }
341
342  pub fn fsync(&self, data_only: bool) -> Result<()> {
343    let st = self.state()?;
344    self
345      .unary::<Empty>(
346        method::FSYNC,
347        &FsyncRequest {
348          handle: st.handle,
349          data_only,
350        },
351      )
352      .map(drop)
353  }
354
355  /// Closes the handle. A second close returns EBADF.
356  pub fn close(&self) -> Result<()> {
357    let mut st = self.state()?;
358    st.closed = true;
359    self
360      .unary::<Empty>(method::CLOSE, &CloseRequest { handle: st.handle })
361      .map(drop)
362  }
363}
364
365impl Drop for File {
366  fn drop(&mut self) {
367    let (closed, handle) = {
368      let st = self.st.get_mut().unwrap_or_else(|e| e.into_inner());
369      (st.closed, st.handle)
370    };
371    if !closed {
372      let _ = self.unary::<Empty>(method::CLOSE, &CloseRequest { handle });
373    }
374  }
375}
376
377impl Read for &File {
378  fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
379    let mut st = self.state()?;
380    let data = self.do_read(st.handle, st.pos, buf.len())?;
381    buf[..data.len()].copy_from_slice(&data);
382    st.pos += data.len() as u64;
383    Ok(data.len())
384  }
385}
386
387impl Write for &File {
388  fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
389    let mut st = self.state()?;
390    let n = self.do_write(st.handle, st.pos, buf)?;
391    st.pos += n as u64;
392    Ok(n)
393  }
394
395  fn flush(&mut self) -> io::Result<()> {
396    File::flush(self).map_err(Into::into)
397  }
398}
399
400impl Seek for &File {
401  fn seek(&mut self, from: SeekFrom) -> io::Result<u64> {
402    let mut st = self.state()?;
403    let (base, delta) = match from {
404      SeekFrom::Start(p) => (0i128, p as i128),
405      SeekFrom::Current(d) => (st.pos as i128, d as i128),
406      SeekFrom::End(d) => (self.do_fstat(st.handle)?.size as i128, d as i128),
407    };
408    let pos = base + delta;
409    if pos < 0 || pos > i64::MAX as i128 {
410      return Err(io::Error::new(io::ErrorKind::InvalidInput, "seek out of range"));
411    }
412    st.pos = pos as u64;
413    Ok(st.pos)
414  }
415}
416
417impl Read for File {
418  fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
419    (&*self).read(buf)
420  }
421}
422
423impl Write for File {
424  fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
425    (&*self).write(buf)
426  }
427
428  fn flush(&mut self) -> io::Result<()> {
429    File::flush(self).map_err(Into::into)
430  }
431}
432
433impl Seek for File {
434  fn seek(&mut self, from: SeekFrom) -> io::Result<u64> {
435    (&*self).seek(from)
436  }
437}