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
19fn 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
38fn 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
58enum Step<T> {
62 Group(Vec<Result<T>>, usize),
63 Big,
64}
65
66#[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 pub fn new(offset: u64, data: Vec<u8>) -> WriteItem {
77 WriteItem { offset, data }
78 }
79}
80
81fn 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
130fn 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
185pub 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 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 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 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 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 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 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 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}