tempest_rt/fs/
write_exact.rs

1//! Async exact-length write operation.
2
3use std::{
4    io,
5    mem::replace,
6    pin::Pin,
7    task::{Context, Poll},
8};
9
10use tempest_io::{Io, IoBuf, Slice};
11
12use super::{WriteAt, write_at};
13
14enum WriteExactState<B: IoBuf, I: Io> {
15    Idle {
16        buf: B,
17        done: usize,
18    },
19    Writing {
20        fut: WriteAt<Slice<B>, I>,
21        done: usize,
22    },
23    Done,
24}
25
26/// Future that writes exactly `buf.bytes_init()` bytes to a file, retrying on partial writes.
27#[must_use = "futures do nothing unless awaited"]
28pub struct WriteExact<B: IoBuf, I: Io> {
29    state: WriteExactState<B, I>,
30    fd: I::Fd,
31    offset: u64,
32    total: usize,
33}
34
35/// Writes exactly `buf.bytes_init()` bytes from `buf` to `fd` at `offset`, retrying on partials.
36/// Returns `(Ok(()), buf)` on success, or `(Err(e), buf)` on I/O error.
37// TODO: we could optimize this by inlining the WriteAt and reusing the OpHandle...
38// but that buys us only a really little advantage at most
39pub fn write_exact<B: IoBuf, I: Io>(fd: I::Fd, buf: B, offset: u64) -> WriteExact<B, I> {
40    let total = buf.bytes_init();
41    WriteExact {
42        state: WriteExactState::Idle { buf, done: 0 },
43        fd,
44        offset,
45        total,
46    }
47}
48
49impl<B: IoBuf, I: Io> Unpin for WriteExact<B, I> {}
50
51impl<B: IoBuf, I: Io> Future for WriteExact<B, I> {
52    type Output = (io::Result<()>, B);
53
54    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
55        let this = self.get_mut();
56
57        loop {
58            match replace(&mut this.state, WriteExactState::Done) {
59                WriteExactState::Idle { buf, done } => {
60                    let fut = write_at::<_, I>(
61                        this.fd,
62                        buf.slice(done..this.total),
63                        this.offset + done as u64,
64                    );
65                    this.state = WriteExactState::Writing { fut, done };
66                    // fall through to poll the WriteAt immediately
67                }
68                WriteExactState::Writing { mut fut, done } => {
69                    match Pin::new(&mut fut).poll(cx) {
70                        Poll::Pending => {
71                            this.state = WriteExactState::Writing { fut, done };
72                            return Poll::Pending;
73                        }
74                        Poll::Ready((Ok(n), slice)) => {
75                            let buf = slice.into_inner();
76                            let new_done = done + n;
77                            if new_done >= this.total {
78                                return Poll::Ready((Ok(()), buf));
79                            }
80                            this.state = WriteExactState::Idle {
81                                buf,
82                                done: new_done,
83                            };
84                            // loop to submit the next partial write
85                        }
86                        Poll::Ready((Err(e), slice)) => {
87                            return Poll::Ready((Err(e), slice.into_inner()));
88                        }
89                    }
90                }
91                WriteExactState::Done => panic!("WriteExact polled after completion"),
92            }
93        }
94    }
95}
tempest_rt/fs/write_exact.rs

tempest_rt/fs/
write_exact.rs
