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 //! Balloon related control APIs.
6
7 use std::collections::VecDeque;
8
9 use anyhow::bail;
10 use anyhow::Context;
11 use anyhow::Result;
12 pub use balloon_control::BalloonStats;
13 use balloon_control::BalloonTubeCommand;
14 pub use balloon_control::BalloonTubeResult;
15 pub use balloon_control::BalloonWS;
16 pub use balloon_control::WSBucket;
17 pub use balloon_control::VIRTIO_BALLOON_WS_MAX_NUM_BINS;
18 pub use balloon_control::VIRTIO_BALLOON_WS_MIN_NUM_BINS;
19 use base::error;
20 use base::Error as SysError;
21 use base::Tube;
22 use serde::Deserialize;
23 use serde::Serialize;
24
25 use crate::VmResponse;
26
27 // Balloon commands that are sent on the crosvm control socket.
28 #[derive(Serialize, Deserialize, Debug, Clone)]
29 pub enum BalloonControlCommand {
30 /// Set the size of the VM's balloon.
31 Adjust {
32 num_bytes: u64,
33 wait_for_success: bool,
34 },
35 Stats,
36 WorkingSet,
37 WorkingSetConfig {
38 bins: Vec<u32>,
39 refresh_threshold: u32,
40 report_threshold: u32,
41 },
42 }
43
do_send(tube: &Tube, cmd: &BalloonControlCommand) -> Option<VmResponse>44 fn do_send(tube: &Tube, cmd: &BalloonControlCommand) -> Option<VmResponse> {
45 match *cmd {
46 BalloonControlCommand::Adjust {
47 num_bytes,
48 wait_for_success,
49 } => {
50 match tube.send(&BalloonTubeCommand::Adjust {
51 num_bytes,
52 allow_failure: wait_for_success,
53 }) {
54 Ok(_) => {
55 if wait_for_success {
56 None
57 } else {
58 Some(VmResponse::Ok)
59 }
60 }
61 Err(_) => Some(VmResponse::Err(SysError::last())),
62 }
63 }
64 BalloonControlCommand::WorkingSetConfig {
65 ref bins,
66 refresh_threshold,
67 report_threshold,
68 } => {
69 match tube.send(&BalloonTubeCommand::WorkingSetConfig {
70 bins: bins.clone(),
71 refresh_threshold,
72 report_threshold,
73 }) {
74 Ok(_) => Some(VmResponse::Ok),
75 Err(_) => Some(VmResponse::Err(SysError::last())),
76 }
77 }
78 BalloonControlCommand::Stats => match tube.send(&BalloonTubeCommand::Stats) {
79 Ok(_) => None,
80 Err(_) => Some(VmResponse::Err(SysError::last())),
81 },
82 BalloonControlCommand::WorkingSet => match tube.send(&BalloonTubeCommand::WorkingSet) {
83 Ok(_) => None,
84 Err(_) => Some(VmResponse::Err(SysError::last())),
85 },
86 }
87 }
88
89 /// Utility for multiplexing a balloon tube between multiple control tubes. Commands
90 /// are sent and processed serially.
91 pub struct BalloonTube {
92 tube: Tube,
93 pending_queue: VecDeque<(BalloonControlCommand, Option<usize>)>,
94 pending_adjust_with_completion: Option<(u64, usize)>,
95 }
96
97 #[cfg(feature = "balloon")]
98 impl BalloonTube {
new(tube: Tube) -> Self99 pub fn new(tube: Tube) -> Self {
100 BalloonTube {
101 tube,
102 pending_queue: VecDeque::new(),
103 pending_adjust_with_completion: None,
104 }
105 }
106
107 /// Sends or queues the given command to this tube. Associates the
108 /// response with the given key.
send_cmd( &mut self, cmd: BalloonControlCommand, key: Option<usize>, ) -> Option<(VmResponse, usize)>109 pub fn send_cmd(
110 &mut self,
111 cmd: BalloonControlCommand,
112 key: Option<usize>,
113 ) -> Option<(VmResponse, usize)> {
114 match cmd {
115 BalloonControlCommand::Adjust {
116 wait_for_success: true,
117 num_bytes,
118 } => {
119 let Some(key) = key else {
120 error!("Asked for completion without reply key");
121 return None;
122 };
123 let resp = self
124 .pending_adjust_with_completion
125 .take()
126 .map(|(_, key)| (VmResponse::ErrString("Adjust overriden".to_string()), key));
127 if do_send(&self.tube, &cmd).is_some() {
128 unreachable!("Unexpected early reply");
129 }
130 self.pending_adjust_with_completion = Some((num_bytes, key));
131 resp
132 }
133 _ => {
134 if !self.pending_queue.is_empty() {
135 self.pending_queue.push_back((cmd, key));
136 return None;
137 }
138 let resp = do_send(&self.tube, &cmd);
139 if resp.is_none() {
140 self.pending_queue.push_back((cmd, key));
141 }
142 match key {
143 None => None,
144 Some(key) => resp.map(|r| (r, key)),
145 }
146 }
147 }
148 }
149
150 /// Receives responses from the balloon tube, and returns them with
151 /// their assoicated keys.
recv(&mut self) -> Result<Vec<(VmResponse, usize)>>152 pub fn recv(&mut self) -> Result<Vec<(VmResponse, usize)>> {
153 let res = self
154 .tube
155 .recv::<BalloonTubeResult>()
156 .context("failed to read balloon tube")?;
157 if let BalloonTubeResult::Adjusted { num_bytes: actual } = res {
158 let Some((target, key)) = self.pending_adjust_with_completion else {
159 bail!("Unexpected balloon adjust to {}", actual);
160 };
161 if actual != target {
162 return Ok(vec![]);
163 }
164 self.pending_adjust_with_completion.take();
165 return Ok(vec![(VmResponse::Ok, key)]);
166 }
167 let mut responses = vec![];
168 if self.pending_queue.is_empty() {
169 bail!("Unexpected balloon tube result {:?}", res)
170 }
171 let resp = match (
172 &self.pending_queue.front().expect("entry disappeared").0,
173 res,
174 ) {
175 (
176 BalloonControlCommand::Stats,
177 BalloonTubeResult::Stats {
178 stats,
179 balloon_actual,
180 },
181 ) => VmResponse::BalloonStats {
182 stats,
183 balloon_actual,
184 },
185 (
186 BalloonControlCommand::WorkingSet,
187 BalloonTubeResult::WorkingSet { ws, balloon_actual },
188 ) => VmResponse::BalloonWS { ws, balloon_actual },
189 (_, resp) => {
190 bail!("Unexpected balloon tube result {:?}", resp);
191 }
192 };
193 let key = self.pending_queue.pop_front().expect("entry disappeared").1;
194 if let Some(key) = key {
195 responses.push((resp, key))
196 }
197 while let Some((cmd, key)) = self.pending_queue.front() {
198 match do_send(&self.tube, cmd) {
199 Some(resp) => {
200 if let Some(key) = key {
201 responses.push((resp, *key));
202 }
203 self.pending_queue.pop_front();
204 }
205 None => break,
206 }
207 }
208 Ok(responses)
209 }
210 }
211
212 #[cfg(test)]
213 mod tests {
214 use super::*;
215
balloon_device_respond_stats(device: &Tube)216 fn balloon_device_respond_stats(device: &Tube) {
217 let BalloonTubeCommand::Stats = device.recv::<BalloonTubeCommand>().unwrap() else {
218 panic!("unexpected command");
219 };
220
221 device
222 .send(&BalloonTubeResult::Stats {
223 stats: BalloonStats::default(),
224 balloon_actual: 0,
225 })
226 .unwrap();
227 }
228
229 #[test]
test_stat_command()230 fn test_stat_command() {
231 let (host, device) = Tube::pair().unwrap();
232 let mut balloon_tube = BalloonTube::new(host);
233
234 let resp = balloon_tube.send_cmd(BalloonControlCommand::Stats, Some(0xc0ffee));
235 assert!(resp.is_none());
236
237 balloon_device_respond_stats(&device);
238
239 let resp = balloon_tube.recv().unwrap();
240 assert_eq!(resp.len(), 1);
241 assert_eq!(resp[0].1, 0xc0ffee);
242 assert!(matches!(resp[0].0, VmResponse::BalloonStats { .. }));
243 }
244
245 #[test]
test_multiple_stat_command()246 fn test_multiple_stat_command() {
247 let (host, device) = Tube::pair().unwrap();
248 let mut balloon_tube = BalloonTube::new(host);
249
250 let resp = balloon_tube.send_cmd(BalloonControlCommand::Stats, Some(0xc0ffee));
251 assert!(resp.is_none());
252 let resp = balloon_tube.send_cmd(BalloonControlCommand::Stats, Some(0xbadcafe));
253 assert!(resp.is_none());
254
255 balloon_device_respond_stats(&device);
256
257 let resp = balloon_tube.recv().unwrap();
258 assert_eq!(resp.len(), 1);
259 assert_eq!(resp[0].1, 0xc0ffee);
260 assert!(matches!(resp[0].0, VmResponse::BalloonStats { .. }));
261
262 balloon_device_respond_stats(&device);
263
264 let resp = balloon_tube.recv().unwrap();
265 assert_eq!(resp.len(), 1);
266 assert_eq!(resp[0].1, 0xbadcafe);
267 assert!(matches!(resp[0].0, VmResponse::BalloonStats { .. }));
268 }
269
270 #[test]
test_queued_stats_adjust_no_reply()271 fn test_queued_stats_adjust_no_reply() {
272 let (host, device) = Tube::pair().unwrap();
273 let mut balloon_tube = BalloonTube::new(host);
274
275 let resp = balloon_tube.send_cmd(BalloonControlCommand::Stats, Some(0xc0ffee));
276 assert!(resp.is_none());
277 let resp = balloon_tube.send_cmd(
278 BalloonControlCommand::Adjust {
279 num_bytes: 0,
280 wait_for_success: false,
281 },
282 Some(0xbadcafe),
283 );
284 assert!(resp.is_none());
285
286 balloon_device_respond_stats(&device);
287
288 let resp = balloon_tube.recv().unwrap();
289 assert_eq!(resp.len(), 2);
290 assert_eq!(resp[0].1, 0xc0ffee);
291 assert!(matches!(resp[0].0, VmResponse::BalloonStats { .. }));
292 assert_eq!(resp[1].1, 0xbadcafe);
293 assert!(matches!(resp[1].0, VmResponse::Ok));
294
295 let cmd = device.recv::<BalloonTubeCommand>().unwrap();
296 assert!(matches!(cmd, BalloonTubeCommand::Adjust { .. }));
297 }
298
299 #[test]
test_adjust_with_reply()300 fn test_adjust_with_reply() {
301 let (host, device) = Tube::pair().unwrap();
302 let mut balloon_tube = BalloonTube::new(host);
303
304 let resp = balloon_tube.send_cmd(
305 BalloonControlCommand::Adjust {
306 num_bytes: 0xc0ffee,
307 wait_for_success: true,
308 },
309 Some(0xc0ffee),
310 );
311 assert!(resp.is_none());
312 let cmd = device.recv::<BalloonTubeCommand>().unwrap();
313 assert!(matches!(cmd, BalloonTubeCommand::Adjust { .. }));
314
315 let resp = balloon_tube.send_cmd(
316 BalloonControlCommand::Adjust {
317 num_bytes: 0xbadcafe,
318 wait_for_success: true,
319 },
320 Some(0xbadcafe),
321 );
322 assert!(matches!(resp, Some((VmResponse::ErrString(_), 0xc0ffee))));
323 let cmd = device.recv::<BalloonTubeCommand>().unwrap();
324 assert!(matches!(cmd, BalloonTubeCommand::Adjust { .. }));
325
326 device
327 .send(&BalloonTubeResult::Adjusted {
328 num_bytes: 0xc0ffee,
329 })
330 .unwrap();
331 let resp = balloon_tube.recv().unwrap();
332 assert_eq!(resp.len(), 0);
333
334 device
335 .send(&BalloonTubeResult::Adjusted {
336 num_bytes: 0xbadcafe,
337 })
338 .unwrap();
339 let resp = balloon_tube.recv().unwrap();
340 assert_eq!(resp.len(), 1);
341 assert_eq!(resp[0].1, 0xbadcafe);
342 assert!(matches!(resp[0].0, VmResponse::Ok));
343 }
344
345 #[test]
test_stats_and_adjust_with_reply()346 fn test_stats_and_adjust_with_reply() {
347 let (host, device) = Tube::pair().unwrap();
348 let mut balloon_tube = BalloonTube::new(host);
349
350 let resp = balloon_tube.send_cmd(BalloonControlCommand::Stats, Some(0xc0ffee));
351 assert!(resp.is_none());
352
353 let resp = balloon_tube.send_cmd(
354 BalloonControlCommand::Adjust {
355 num_bytes: 0xbadcafe,
356 wait_for_success: true,
357 },
358 Some(0xbadcafe),
359 );
360 assert!(resp.is_none());
361
362 let cmd = device.recv::<BalloonTubeCommand>().unwrap();
363 assert!(matches!(cmd, BalloonTubeCommand::Stats));
364 let cmd = device.recv::<BalloonTubeCommand>().unwrap();
365 assert!(matches!(cmd, BalloonTubeCommand::Adjust { .. }));
366
367 device
368 .send(&BalloonTubeResult::Adjusted {
369 num_bytes: 0xbadcafe,
370 })
371 .unwrap();
372 let resp = balloon_tube.recv().unwrap();
373 assert_eq!(resp.len(), 1);
374 assert_eq!(resp[0].1, 0xbadcafe);
375 assert!(matches!(resp[0].0, VmResponse::Ok));
376
377 device
378 .send(&BalloonTubeResult::Stats {
379 stats: BalloonStats::default(),
380 balloon_actual: 0,
381 })
382 .unwrap();
383 let resp = balloon_tube.recv().unwrap();
384 assert_eq!(resp.len(), 1);
385 assert_eq!(resp[0].1, 0xc0ffee);
386 assert!(matches!(resp[0].0, VmResponse::BalloonStats { .. }));
387 }
388 }
389