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

mountos_native/
file.rs
