reclog 0.1.3

Command-line tool to capture command output to a file.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

use crate::error::SysError;
use crate::shim::{self, SelectFd};
use rustix::io::{Errno, retry_on_intr};
use rustix::pipe;
use std::io::{Error, Write};
use std::os::fd::{AsFd, OwnedFd};
use std::sync::{Arc, Mutex};

#[derive(PartialEq)]
enum WriterMode {
    Open,
    Closed,
}

/// Allows to write from fd in one thwrite and interrupt write or change
/// write timeout from another thwrite.
pub struct InterruptibleWriter<Fd: AsFd> {
    mode: Mutex<WriterMode>,
    fd: Fd,
    pipe_rd: OwnedFd,
    pipe_wr: OwnedFd,
}

impl<Fd: AsFd> InterruptibleWriter<Fd> {
    /// Construct new writer.
    /// Gains ownership of the fd.
    pub fn open(fd: Fd) -> Result<Self, SysError> {
        let (pipe_rd, pipe_wr) = match retry_on_intr(|| pipe::pipe()) {
            Ok(fds) => fds,
            Err(err) => return Err(SysError("pipe()", err)),
        };

        shim::fcntl_nonblock(&fd, true).map_err(|err| SysError("fcntl(fd)", err))?;
        shim::fcntl_nonblock(&pipe_rd, true).map_err(|err| SysError("fcntl(pipe)", err))?;
        shim::fcntl_nonblock(&pipe_wr, true).map_err(|err| SysError("fcntl(pipe)", err))?;

        Ok(InterruptibleWriter {
            mode: Mutex::new(WriterMode::Open),
            fd,
            pipe_rd,
            pipe_wr,
        })
    }

    /// Close writer.
    /// Will wake up and abort ongoing writes.
    pub fn close(&self) -> Result<(), SysError> {
        {
            // update mode
            let mut locked_mode = self.mode.lock().unwrap();
            if *locked_mode == WriterMode::Closed {
                return Ok(());
            }
            *locked_mode = WriterMode::Closed;
        }

        // wake up and abort blocked write
        if let Err(err) = shim::write(&self.pipe_wr, &[0u8]) {
            if err != Errno::AGAIN {
                return Err(SysError("write(pipe)", err));
            }
        }

        Ok(())
    }

    /// Construct blocking writer.
    /// Waits until descriptor is writable, or writer is closed.
    pub fn blocking_writer(self: &Arc<Self>) -> ArcTimeoutWriter<Fd> {
        ArcTimeoutWriter(Arc::clone(self))
    }

    /// Invoked by ArcTimeoutWriter::write().
    fn write_imp(&self, buf: &[u8]) -> Result<usize, Error> {
        loop {
            // re-read mode
            {
                let locked_mode = self.mode.lock().unwrap();
                if *locked_mode == WriterMode::Closed {
                    // closed, silently discard all bytes
                    return Ok(buf.len());
                }
            };

            // wait until descriptor is ready
            let mut pipe_fd = SelectFd {
                fd: self.pipe_rd.as_fd(),
                mask: SelectFd::READABLE,
            };
            let mut data_fd = SelectFd {
                fd: self.fd.as_fd(),
                mask: SelectFd::WRITEABLE,
            };
            shim::select(&mut [&mut pipe_fd, &mut data_fd], None)?;

            if pipe_fd.mask != 0 {
                // wake up from close()
                // drain bytes from pipe
                _ = shim::read(&self.pipe_rd, &mut [0u8; 128]);
            }
            if data_fd.mask != 0 {
                // file is writeable
                break;
            }
        }

        // if we're here, file is writeable
        match shim::write(&self.fd, buf) {
            Ok(n) => Ok(n),
            Err(err) => Err(Error::new(err.kind(), err)),
        }
    }
}

/// Wrapper for Arc<TimeoutWriter> that implements Write trait.
pub struct ArcTimeoutWriter<Fd: AsFd>(Arc<InterruptibleWriter<Fd>>);

impl<Fd: AsFd> Write for ArcTimeoutWriter<Fd> {
    fn write(&mut self, buf: &[u8]) -> Result<usize, Error> {
        self.0.write_imp(buf)
    }

    fn flush(&mut self) -> Result<(), Error> {
        Ok(())
    }
}