1 // Copyright 2023 The ChromiumOS Authors
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4
5 #![deny(missing_docs)]
6
7 use std::fs::File;
8 use std::io::Seek;
9 use std::io::SeekFrom;
10 use std::time::Duration;
11
12 use anyhow::Context;
13 use anyhow::Result;
14 use base::Descriptor;
15 use base::Event;
16 use base::EventToken;
17 use base::Timer;
18 use base::TimerTrait;
19 use base::WaitContext;
20 use base::WorkerThread;
21
22 /// Truncates a file to length 0, in the background when possible.
23 ///
24 /// Truncating a large file can result in a significant amount of IO when
25 /// updating filesystem metadata. When possible, [FileTruncator] truncates a
26 /// given file gradually over time to avoid competing with higher prioirty IO.
27 pub struct FileTruncator {
28 worker: Option<WorkerThread<Result<File>>>,
29 }
30
31 // The particular values here are relatively arbitrary values that
32 // result in a "slow-enough" background truncation.
33 const TRUNCATE_STEP_BYTES: u64 = 64 * 1024 * 1024; // 64 MiB
34 const TRUNCATE_INITIAL_WAIT: Duration = Duration::from_secs(30);
35 const TRUNCATE_INTERVAL: Duration = Duration::from_secs(5);
36
truncate_worker( mut timer: Box<dyn TimerTrait>, mut file: File, kill_evt: Event, ) -> Result<File>37 fn truncate_worker(
38 mut timer: Box<dyn TimerTrait>,
39 mut file: File,
40 kill_evt: Event,
41 ) -> Result<File> {
42 #[derive(EventToken)]
43 enum Token {
44 Alarm,
45 Kill,
46 }
47
48 let mut len = file
49 .seek(SeekFrom::End(0))
50 .context("Failed to determine size")?;
51
52 let descriptor = Descriptor(timer.as_raw_descriptor());
53 let wait_ctx: WaitContext<Token> =
54 WaitContext::build_with(&[(&descriptor, Token::Alarm), (&kill_evt, Token::Kill)])
55 .context("worker context failed")?;
56
57 while len > 0 {
58 let events = wait_ctx.wait().context("wait failed")?;
59 for event in events.iter().filter(|e| e.is_readable) {
60 match event.token {
61 Token::Alarm => {
62 let _ = timer.mark_waited().context("failed to reset timer")?;
63 len = len.saturating_sub(TRUNCATE_STEP_BYTES);
64 file.set_len(len).context("failed to truncate file")?;
65 }
66 Token::Kill => {
67 file.set_len(0).context("failed to clear file")?;
68 return Ok(file);
69 }
70 }
71 }
72 }
73 Ok(file)
74 }
75
76 impl FileTruncator {
77 /// Creates an new [FileTruncator] to truncate the given file.
78 ///
79 /// # Arguments
80 ///
81 /// * `file` - The file to truncate.
new(file: File) -> Result<Self>82 pub fn new(file: File) -> Result<Self> {
83 let timer = Timer::new().context("failed to create truncate timer")?;
84 Self::new_inner(Box::new(timer), file)
85 }
86
new_inner(mut timer: Box<dyn TimerTrait>, file: File) -> Result<Self>87 fn new_inner(mut timer: Box<dyn TimerTrait>, file: File) -> Result<Self> {
88 timer
89 .reset(TRUNCATE_INITIAL_WAIT, Some(TRUNCATE_INTERVAL))
90 .context("failed to arm timer")?;
91 Ok(Self {
92 worker: Some(WorkerThread::start(
93 "truncate_worker",
94 move |kill_evt| -> Result<File> { truncate_worker(timer, file, kill_evt) },
95 )),
96 })
97 }
98
99 /// Retrieves the underlying file, which is guaranteed to be truncated.
100 ///
101 /// If this function is called while the background worker thread has not
102 /// finished, it may block briefly while stopping the background worker.
take_file(mut self) -> Result<File>103 pub fn take_file(mut self) -> Result<File> {
104 let file = self
105 .worker
106 .take()
107 .context("missing worker")?
108 .stop()
109 .context("worker failure")?;
110 Ok(file)
111 }
112 }
113
114 impl Drop for FileTruncator {
drop(&mut self)115 fn drop(&mut self) {
116 if let Some(worker) = self.worker.take() {
117 let _ = worker.stop();
118 }
119 }
120 }
121
122 #[cfg(test)]
123 mod tests {
124 use std::sync::Arc;
125
126 use base::FakeClock;
127 use base::FakeTimer;
128 use sync::Mutex;
129
130 use super::*;
131
wait_for_target_length(file: &mut File, len: u64)132 fn wait_for_target_length(file: &mut File, len: u64) {
133 let mut count = 0;
134 while file.seek(SeekFrom::End(0)).unwrap() != len && count < 100 {
135 std::thread::sleep(Duration::from_millis(1));
136 count += 1;
137 }
138 assert_eq!(file.seek(SeekFrom::End(0)).unwrap(), len);
139 }
140
141 #[test]
test_full_truncate()142 fn test_full_truncate() {
143 let mut file = tempfile::tempfile().unwrap();
144 let clock = Arc::new(Mutex::new(FakeClock::new()));
145 let timer = Box::new(FakeTimer::new(clock.clone()));
146
147 file.set_len(2 * TRUNCATE_STEP_BYTES).unwrap();
148
149 let worker = FileTruncator::new_inner(timer, file.try_clone().unwrap()).unwrap();
150 clock.lock().add_ns(TRUNCATE_INITIAL_WAIT.as_nanos() as u64);
151 wait_for_target_length(&mut file, TRUNCATE_STEP_BYTES);
152 clock.lock().add_ns(TRUNCATE_INTERVAL.as_nanos() as u64);
153 wait_for_target_length(&mut file, 0);
154
155 let _ = worker.take_file().unwrap();
156 }
157
158 #[test]
test_early_exit()159 fn test_early_exit() {
160 let mut file = tempfile::tempfile().unwrap();
161 let clock = Arc::new(Mutex::new(FakeClock::new()));
162 let timer = Box::new(FakeTimer::new(clock));
163
164 file.set_len(2 * TRUNCATE_STEP_BYTES).unwrap();
165
166 let worker = FileTruncator::new_inner(timer, file.try_clone().unwrap()).unwrap();
167
168 let _ = worker.take_file().unwrap();
169 assert_eq!(file.seek(SeekFrom::End(0)).unwrap(), 0);
170 }
171 }
172