vm_control/
balloon_tube.rs

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

vm_control/
balloon_tube.rs
