src/tools.rs

   1 //! Tools and utilities
   2 //!
   3 //! This is a collection of small and useful tools.
   4
   5 use failure::*;
   6 use nix::unistd;
   7 use nix::sys::stat;
   8
   9 use std::fs::{File, OpenOptions};
  10 use std::io::Write;
  11 use std::path::Path;
  12 use std::io::Read;
  13 use std::io::ErrorKind;
  14 use std::time::Duration;
  15
  16 use std::os::unix::io::AsRawFd;
  17
  18 pub mod timer;
  19
  20 /// The `BufferedReader` trait provides a single function
  21 /// `buffered_read`. It returns a reference to an internal buffer. The
  22 /// purpose of this traid is to avoid unnecessary data copies.
  23 pub trait BufferedReader {
  24     /// This functions tries to fill the internal buffers, then
  25     /// returns a reference to the available data. It returns an empty
  26     /// buffer if `offset` points to the end of the file.
  27     fn buffered_read(&mut self, offset: u64) -> Result<&[u8], Error>;
  28 }
  29
  30 /// Directly map a type into a binary buffer. This is mostly useful
  31 /// for reading structured data from a byte stream (file). You need to
  32 /// make sure that the buffer location does not change, so please
  33 /// avoid vec resize while you use such map.
  34 ///
  35 /// This function panics if the buffer is not large enough.
  36 pub fn map_struct<T>(buffer: &[u8]) -> Result<&T, Error> {
  37     if buffer.len() < ::std::mem::size_of::<T>() {
  38         bail!("unable to map struct - buffer too small");
  39     }
  40     Ok(unsafe { & * (buffer.as_ptr() as *const T) })
  41 }
  42
  43 /// Directly map a type into a mutable binary buffer. This is mostly
  44 /// useful for writing structured data into a byte stream (file). You
  45 /// need to make sure that the buffer location does not change, so
  46 /// please avoid vec resize while you use such map.
  47 ///
  48 /// This function panics if the buffer is not large enough.
  49 pub fn map_struct_mut<T>(buffer: &mut [u8]) -> Result<&mut T, Error> {
  50     if buffer.len() < ::std::mem::size_of::<T>() {
  51         bail!("unable to map struct - buffer too small");
  52     }
  53     Ok(unsafe { &mut * (buffer.as_ptr() as *mut T) })
  54 }
  55
  56 /// Atomically write a file. We first create a temporary file, which
  57 /// is then renamed.
  58 pub fn file_set_contents<P: AsRef<Path>>(
  59     path: P,
  60     data: &[u8],
  61     perm: Option<stat::Mode>,
  62 ) -> Result<(), Error> {
  63
  64     let path = path.as_ref();
  65
  66     // Note: we use mkstemp heŕe, because this worka with different
  67     // processes, threads, and even tokio tasks.
  68     let mut template = path.to_owned();
  69     template.set_extension("tmp_XXXXXX");
  70     let (fd, tmp_path) = match unistd::mkstemp(&template) {
  71         Ok((fd, path)) => (fd, path),
  72         Err(err) => bail!("mkstemp {:?} failed: {}", template, err),
  73     };
  74
  75     let tmp_path = tmp_path.as_path();
  76
  77     let mode : stat::Mode = perm.unwrap_or(stat::Mode::from(
  78         stat::Mode::S_IRUSR | stat::Mode::S_IWUSR |
  79         stat::Mode::S_IRGRP | stat::Mode::S_IROTH
  80     ));
  81
  82     if let Err(err) = stat::fchmod(fd, mode) {
  83         let _ = unistd::unlink(tmp_path);
  84         bail!("fchmod {:?} failed: {}", tmp_path, err);
  85     }
  86
  87     use std::os::unix::io::FromRawFd;
  88     let mut file = unsafe { File::from_raw_fd(fd) };
  89
  90     if let Err(err) = file.write_all(data) {
  91         let _ = unistd::unlink(tmp_path);
  92         bail!("write failed: {}", err);
  93     }
  94
  95     if let Err(err) = std::fs::rename(tmp_path, path) {
  96         let _ = unistd::unlink(tmp_path);
  97         bail!("Atomic rename failed for file {:?} - {}", path, err);
  98     }
  99
 100     Ok(())
 101 }
 102
 103 /// Create a file lock using fntl. This function allows you to specify
 104 /// a timeout if you want to avoid infinite blocking.
 105 pub fn lock_file<F: AsRawFd>(
 106     file: &mut F,
 107     exclusive: bool,
 108     timeout: Option<Duration>,
 109     ) -> Result<(), Error>
 110 {
 111     let lockarg =
 112         if exclusive {
 113             nix::fcntl::FlockArg::LockExclusive
 114         } else {
 115             nix::fcntl::FlockArg::LockShared
 116         };
 117
 118     let timeout = match timeout {
 119         None => {
 120             nix::fcntl::flock(file.as_raw_fd(), lockarg)?;
 121             return Ok(());
 122         }
 123         Some(t) => t,
 124     };
 125
 126     // unblock the timeout signal temporarily
 127     let _sigblock_guard = timer::unblock_timeout_signal();
 128
 129     // setup a timeout timer
 130     let mut timer = timer::Timer::create(
 131         timer::Clock::Realtime,
 132         timer::TimerEvent::ThisThreadSignal(timer::SIGTIMEOUT))?;
 133
 134     timer.arm(timer::TimerSpec::new()
 135         .value(Some(timeout))
 136         .interval(Some(Duration::from_millis(10))))?;
 137
 138     nix::fcntl::flock(file.as_raw_fd(), lockarg)?;
 139     Ok(())
 140 }
 141
 142 /// Open or create a lock file (append mode). Then try to
 143 /// aquire a lock using `lock_file()`.
 144 pub fn open_file_locked<P: AsRef<Path>>(path: P, timeout: Duration)
 145     -> Result<File, Error>
 146 {
 147     let path = path.as_ref();
 148     let mut file =
 149         match OpenOptions::new()
 150             .create(true)
 151             .append(true)
 152             .open(path)
 153         {
 154             Ok(file) => file,
 155             Err(err) => bail!("Unable to open lock {:?} - {}",
 156                               path, err),
 157         };
 158     match lock_file(&mut file, true, Some(timeout)) {
 159         Ok(_) => Ok(file),
 160         Err(err) => bail!("Unable to aquire lock {:?} - {}",
 161                           path, err),
 162     }
 163 }
 164
 165 /// Split a file into equal sized chunks. The last chunk may be
 166 /// smaller. Note: We cannot implement an `Iterator`, because iterators
 167 /// cannot return a borrowed buffer ref (we want zero-copy)
 168 pub fn file_chunker<C, R>(
 169     mut file: R,
 170     chunk_size: usize,
 171     mut chunk_cb: C
 172 ) -> Result<(), Error>
 173     where C: FnMut(usize, &[u8]) -> Result<bool, Error>,
 174           R: Read,
 175 {
 176
 177     const READ_BUFFER_SIZE: usize = 4*1024*1024; // 4M
 178
 179     if chunk_size > READ_BUFFER_SIZE { bail!("chunk size too large!"); }
 180
 181     let mut buf = vec![0u8; READ_BUFFER_SIZE];
 182
 183     let mut pos = 0;
 184     let mut file_pos = 0;
 185     loop {
 186         let mut eof = false;
 187         let mut tmp = &mut buf[..];
 188        // try to read large portions, at least chunk_size
 189         while pos < chunk_size {
 190             match file.read(tmp) {
 191                 Ok(0) => { eof = true; break; },
 192                 Ok(n) => {
 193                     pos += n;
 194                     if pos > chunk_size { break; }
 195                     tmp = &mut tmp[n..];
 196                 }
 197                 Err(ref e) if e.kind() == ErrorKind::Interrupted => { /* try again */ }
 198                 Err(e) => bail!("read chunk failed - {}", e.to_string()),
 199             }
 200         }
 201         let mut start = 0;
 202         while start + chunk_size <= pos {
 203             if !(chunk_cb)(file_pos, &buf[start..start+chunk_size])? { break; }
 204             file_pos += chunk_size;
 205             start += chunk_size;
 206         }
 207         if eof {
 208             if start < pos {
 209                 (chunk_cb)(file_pos, &buf[start..pos])?;
 210                 //file_pos += pos - start;
 211             }
 212             break;
 213         } else {
 214             let rest = pos - start;
 215             if rest > 0 {
 216                 let ptr = buf.as_mut_ptr();
 217                 unsafe { std::ptr::copy_nonoverlapping(ptr.add(start), ptr, rest); }
 218                 pos = rest;
 219             } else {
 220                 pos = 0;
 221             }
 222         }
 223     }
 224
 225     Ok(())
 226 }