]>
Commit | Line | Data |
---|---|---|
51b499db DM |
1 | //! Tools and utilities |
2 | //! | |
3 | //! This is a collection of small and useful tools. | |
4 | ||
f12f8ff1 DM |
5 | use failure::*; |
6 | use nix::unistd; | |
7 | use nix::sys::stat; | |
8 | ||
365bb90f | 9 | use std::fs::{File, OpenOptions}; |
f12f8ff1 DM |
10 | use std::io::Write; |
11 | use std::path::Path; | |
43eeef28 DM |
12 | use std::io::Read; |
13 | use std::io::ErrorKind; | |
1628a4c7 | 14 | use std::time::Duration; |
f12f8ff1 | 15 | |
eae8aa3a | 16 | use std::os::unix::io::AsRawFd; |
365bb90f | 17 | |
8cf6e764 WB |
18 | pub mod timer; |
19 | ||
0a72e267 DM |
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 { | |
318564ac DM |
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. | |
0a72e267 DM |
27 | fn buffered_read(&mut self, offset: u64) -> Result<&[u8], Error>; |
28 | } | |
29 | ||
51b499db DM |
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. | |
50ea4396 | 36 | pub fn map_struct<T>(buffer: &[u8]) -> Result<&T, Error> { |
dc3de618 DM |
37 | if buffer.len() < ::std::mem::size_of::<T>() { |
38 | bail!("unable to map struct - buffer too small"); | |
39 | } | |
95bd5dfe | 40 | Ok(unsafe { & * (buffer.as_ptr() as *const T) }) |
dc3de618 DM |
41 | } |
42 | ||
51b499db DM |
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. | |
50ea4396 | 49 | pub fn map_struct_mut<T>(buffer: &mut [u8]) -> Result<&mut T, Error> { |
dc3de618 DM |
50 | if buffer.len() < ::std::mem::size_of::<T>() { |
51 | bail!("unable to map struct - buffer too small"); | |
52 | } | |
95bd5dfe | 53 | Ok(unsafe { &mut * (buffer.as_ptr() as *mut T) }) |
dc3de618 DM |
54 | } |
55 | ||
51b499db DM |
56 | /// Atomically write a file. We first create a temporary file, which |
57 | /// is then renamed. | |
f12f8ff1 DM |
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 | ||
d64d80d2 DM |
66 | // Note: we use mkstemp heŕe, because this worka with different |
67 | // processes, threads, and even tokio tasks. | |
f12f8ff1 DM |
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 | ||
1a7bc3dd | 77 | let mode : stat::Mode = perm.unwrap_or(stat::Mode::from( |
f12f8ff1 DM |
78 | stat::Mode::S_IRUSR | stat::Mode::S_IWUSR | |
79 | stat::Mode::S_IRGRP | stat::Mode::S_IROTH | |
1a7bc3dd | 80 | )); |
f12f8ff1 DM |
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 | } | |
43eeef28 | 102 | |
51b499db DM |
103 | /// Create a file lock using fntl. This function allows you to specify |
104 | /// a timeout if you want to avoid infinite blocking. | |
1628a4c7 WB |
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 | } | |
365bb90f | 141 | |
51b499db DM |
142 | /// Open or create a lock file (append mode). Then try to |
143 | /// aquire a lock using `lock_file()`. | |
1628a4c7 WB |
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 | { | |
365bb90f DM |
154 | Ok(file) => file, |
155 | Err(err) => bail!("Unable to open lock {:?} - {}", | |
156 | path, err), | |
157 | }; | |
28b96b56 DM |
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 | } | |
365bb90f DM |
163 | } |
164 | ||
51b499db DM |
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) | |
43eeef28 DM |
168 | pub fn file_chunker<C, R>( |
169 | mut file: R, | |
170 | chunk_size: usize, | |
606ce64b | 171 | mut chunk_cb: C |
43eeef28 | 172 | ) -> Result<(), Error> |
606ce64b | 173 | where C: FnMut(usize, &[u8]) -> Result<bool, Error>, |
43eeef28 DM |
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 */ } | |
5f0c2d56 | 198 | Err(e) => bail!("read chunk failed - {}", e.to_string()), |
43eeef28 DM |
199 | } |
200 | } | |
43eeef28 DM |
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(()) | |
43eeef28 | 226 | } |