mx_tester/
cleanup.rs

1use crate::Config;
2use log::warn;
3use std::sync::Arc;
4
5/// Cleanup any Docker images at the end of a block,
6/// even in case of panic.
7///
8/// This reduces the chances that Docker will assume
9/// that the images must be restarted on your next computer
10/// startup and decide to bring them up as `root`.
11///
12/// As a side-effect, all tests MUST be prefixed with
13/// `#[tokio::test(flavor = "multi_thread")]`
14pub struct Cleanup {
15    /// If `true`, cleanup is still needed.
16    is_armed: bool,
17
18    /// The container name used during `build`.
19    setup_container_name: Arc<str>,
20
21    /// The container name used during `up` and `run`.
22    run_container_name: Arc<str>,
23
24    /// The network to which this container is attached.
25    network_name: Arc<str>,
26
27    /// If `true`, during cleanup, also take down the network.
28    /// `false` by default.
29    cleanup_network: bool,
30}
31impl Cleanup {
32    pub fn new(config: &Config) -> Self {
33        Cleanup {
34            is_armed: true,
35            setup_container_name: config.setup_container_name().into(),
36            run_container_name: config.run_container_name().into(),
37            network_name: config.network().into(),
38            cleanup_network: false,
39        }
40    }
41
42    /// Enable or disable network cleanup.
43    ///
44    /// `false` by default.
45    ///
46    /// Note that `disarm()` prevents *all* cleanup, regardless of `cleanup_network()`.
47    pub fn cleanup_network(&mut self, value: bool) {
48        self.cleanup_network = value;
49    }
50
51    /// Disarm this guard.
52    ///
53    /// Once disarmed, it will not cause cleanup anymore when it leaves scope.
54    pub fn disarm(mut self) {
55        self.is_armed = false;
56    }
57}
58impl Drop for Cleanup {
59    fn drop(&mut self) {
60        if !self.is_armed {
61            return;
62        }
63        let docker = bollard::Docker::connect_with_local_defaults()
64            .expect("Failed to connect to Docker daemon");
65        let setup_container_name = self.setup_container_name.clone();
66        let run_container_name = self.run_container_name.clone();
67        let network_name = self.network_name.clone();
68        let cleanup_network = self.cleanup_network;
69        tokio::task::block_in_place(move || {
70            tokio::runtime::Handle::current().block_on(async move {
71                warn!("Auto-cleanup...");
72                let _ = docker.stop_container(&setup_container_name, None).await;
73                let _ = docker.remove_container(&setup_container_name, None).await;
74                let _ = docker.stop_container(&run_container_name, None).await;
75                let _ = docker.remove_container(&run_container_name, None).await;
76                if cleanup_network {
77                    let _ = docker.remove_network(&network_name).await;
78                }
79                warn!("Auto-cleanup... DONE");
80            });
81        });
82    }
83}
84
85/// A utility trait used to call `foo.disarm()` on a container
86/// holding an instance of `Cleanup`.
87///
88/// The main utility at the time of this writing is to be able
89/// to disarm a `Option<Cleanup>`.
90pub trait Disarm {
91    fn disarm(self);
92}
93impl<T> Disarm for T
94where
95    T: IntoIterator<Item = Cleanup>,
96{
97    fn disarm(self) {
98        for item in self {
99            // In case of panic during a call to `disarm()`,
100            // the remaining items will be auto-cleaned up.
101            item.disarm();
102        }
103    }
104}
mx_tester/cleanup.rs

mx_tester/
cleanup.rs
