shimkit/sandbox/shim/
instance_data.rs

1use chrono::{DateTime, Utc};
2use tokio::sync::{OnceCell, RwLock};
3
4use crate::sandbox::shim::task_state::TaskState;
5use crate::sandbox::{Instance, InstanceConfig, Result};
6
7pub(super) struct InstanceData<T: Instance> {
8    pub instance: T,
9    pub config: InstanceConfig,
10    pid: OnceCell<u32>,
11    state: RwLock<TaskState>,
12}
13
14impl<T: Instance> InstanceData<T> {
15    #[cfg_attr(feature = "tracing", tracing::instrument(level = "Debug"))]
16    pub async fn new(
17        id: impl AsRef<str> + std::fmt::Debug,
18        config: InstanceConfig,
19    ) -> Result<Self> {
20        let id = id.as_ref().to_string();
21        let instance = T::new(id, &config).await?;
22        Ok(Self {
23            instance,
24            config,
25            pid: OnceCell::default(),
26            state: RwLock::new(TaskState::Created),
27        })
28    }
29
30    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self), level = "Debug"))]
31    pub fn pid(&self) -> Option<u32> {
32        self.pid.get().copied()
33    }
34
35    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self), level = "Debug"))]
36    pub async fn start(&self) -> Result<u32> {
37        let mut s = self.state.write().await;
38        s.start()?;
39
40        let res = self.instance.start().await;
41
42        // These state transitions are always `Ok(())` because
43        // we hold the lock since `s.start()`
44        let _ = match res {
45            Ok(pid) => {
46                let _ = self.pid.set(pid);
47                s.started()
48            }
49            Err(_) => s.stop(),
50        };
51
52        res
53    }
54
55    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self), level = "Debug"))]
56    pub async fn kill(&self, signal: u32) -> Result<()> {
57        let mut s = self.state.write().await;
58        s.kill()?;
59
60        self.instance.kill(signal).await
61    }
62
63    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self), level = "Debug"))]
64    pub async fn delete(&self) -> Result<()> {
65        let mut s = self.state.write().await;
66        s.delete()?;
67
68        let res = self.instance.delete().await;
69
70        if res.is_err() {
71            // Always `Ok(())` because we hold the lock since `s.delete()`
72            let _ = s.stop();
73        }
74
75        res
76    }
77
78    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self), level = "Debug"))]
79    pub async fn wait(&self) -> (u32, DateTime<Utc>) {
80        let res = self.instance.wait().await;
81        let mut s = self.state.write().await;
82        *s = TaskState::Exited;
83        res
84    }
85}