Skip to main content

vmm_core/
vmbus_unit.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! State unit definitions for vmbus components.
5
6#![warn(missing_docs)]
7
8use inspect::Inspect;
9use pal_async::task::Spawn;
10use state_unit::NameInUse;
11use state_unit::SpawnedUnit;
12use state_unit::StateUnit;
13use state_unit::StateUnits;
14use state_unit::UnitBuilder;
15use state_unit::UnitHandle;
16use state_unit::run_async_unit;
17use std::sync::Arc;
18use vm_resource::Resource;
19use vm_resource::ResourceResolver;
20use vm_resource::kind::VmbusDeviceHandleKind;
21use vmbus_channel::channel::ChannelHandle;
22use vmbus_channel::channel::VmbusDevice;
23use vmbus_channel::channel::offer_channel;
24use vmbus_channel::channel::offer_generic_channel;
25use vmbus_channel::resources::ResolveVmbusDeviceHandleParams;
26use vmbus_channel::simple::InitialDeviceState;
27use vmbus_channel::simple::SimpleDeviceHandle;
28use vmbus_channel::simple::SimpleVmbusDevice;
29use vmbus_channel::simple::offer_simple_device;
30use vmbus_server::VmbusServer;
31use vmbus_server::VmbusServerControl;
32use vmcore::save_restore::RestoreError;
33use vmcore::save_restore::SaveError;
34use vmcore::save_restore::SavedStateBlob;
35use vmcore::vm_task::VmTaskDriverSource;
36
37/// A handle to a vmbus server that is registered as a state unit.
38///
39/// FUTURE: incorporate the state unit handling directly into `VmbusServer`.
40pub struct VmbusServerHandle {
41    unit: SpawnedUnit<VmbusServerUnit>,
42    control: Arc<VmbusServerControl>,
43}
44
45impl VmbusServerHandle {
46    /// Makes a new handle, registering the server via `builder`.
47    pub fn new(
48        spawner: &impl Spawn,
49        builder: UnitBuilder<'_>,
50        server: VmbusServer,
51    ) -> Result<Self, NameInUse> {
52        let control = server.control();
53        let unit = builder.spawn(spawner, |recv| {
54            run_async_unit(VmbusServerUnit(server), recv)
55        })?;
56        Ok(Self { unit, control })
57    }
58
59    /// Gets the vmbus control interface.
60    pub fn control(&self) -> &Arc<VmbusServerControl> {
61        &self.control
62    }
63
64    /// Gets the vmbus unit handle.
65    pub fn unit_handle(&self) -> &UnitHandle {
66        self.unit.handle()
67    }
68
69    /// Removes the server.
70    pub async fn remove(self) -> VmbusServer {
71        self.unit.remove().await.0
72    }
73}
74
75/// A newtype over `VmbusServer` implementing [`StateUnit`].
76#[derive(Inspect)]
77#[inspect(transparent)]
78struct VmbusServerUnit(VmbusServer);
79
80impl StateUnit for &'_ VmbusServerUnit {
81    async fn start(&mut self) {
82        self.0.start();
83    }
84
85    async fn stop(&mut self) {
86        self.0.stop().await;
87    }
88
89    async fn reset(&mut self) -> anyhow::Result<()> {
90        self.0.reset().await;
91        Ok(())
92    }
93
94    async fn save(&mut self) -> Result<Option<SavedStateBlob>, SaveError> {
95        Ok(Some(SavedStateBlob::new(self.0.save().await)))
96    }
97
98    async fn restore(&mut self, buffer: SavedStateBlob) -> Result<(), RestoreError> {
99        self.0
100            .restore(buffer.parse()?)
101            .await
102            .map_err(|err| RestoreError::Other(err.into()))
103    }
104}
105
106/// A type wrapping a [`ChannelHandle`] and implementing [`StateUnit`].
107#[must_use]
108#[derive(Debug, Inspect)]
109#[inspect(transparent)]
110pub struct ChannelUnit<T: ?Sized>(ChannelHandle<T>);
111
112/// Offers a channel, creates a unit for it, and adds it to `state_units`.
113pub async fn offer_channel_unit<T: 'static + VmbusDevice>(
114    driver: &impl Spawn,
115    state_units: &StateUnits,
116    vmbus: &VmbusServerHandle,
117    channel: T,
118) -> anyhow::Result<SpawnedUnit<ChannelUnit<T>>> {
119    let offer = channel.offer();
120    let name = format!("{}:{}", offer.interface_name, offer.instance_id);
121    let handle = offer_channel(driver, vmbus.control.as_ref(), channel).await?;
122    let unit = state_units
123        .add(name)
124        .depends_on(vmbus.unit.handle())
125        .spawn(driver, |recv| run_async_unit(ChannelUnit(handle), recv))?;
126    Ok(unit)
127}
128
129impl<T: 'static + VmbusDevice> ChannelUnit<T> {
130    /// Revokes a channel.
131    pub async fn revoke(self) -> T {
132        self.0.revoke().await.unwrap()
133    }
134}
135
136impl<T: 'static + VmbusDevice + ?Sized> StateUnit for &'_ ChannelUnit<T> {
137    async fn start(&mut self) {
138        self.0.start();
139    }
140
141    async fn stop(&mut self) {
142        self.0.stop().await;
143    }
144
145    async fn reset(&mut self) -> anyhow::Result<()> {
146        self.0.reset().await;
147        Ok(())
148    }
149
150    async fn save(&mut self) -> Result<Option<SavedStateBlob>, SaveError> {
151        let state = self.0.save().await.map_err(SaveError::Other)?;
152        Ok(state)
153    }
154
155    async fn restore(&mut self, state: SavedStateBlob) -> Result<(), RestoreError> {
156        self.0.restore(state).await.map_err(RestoreError::Other)
157    }
158}
159
160/// A type wrapping a [`ChannelHandle`] and implementing [`StateUnit`].
161#[must_use]
162#[derive(Debug)]
163pub struct SimpleChannelUnit<T: SimpleVmbusDevice>(SimpleDeviceHandle<T>);
164
165/// Offers a simple vmbus device, creates a unit for it, and adds it to `state_units`.
166pub async fn offer_simple_device_unit<T: SimpleVmbusDevice>(
167    driver_source: &VmTaskDriverSource,
168    state_units: &StateUnits,
169    vmbus: &VmbusServerHandle,
170    device: T,
171) -> anyhow::Result<SpawnedUnit<SimpleChannelUnit<T>>> {
172    anyhow::ensure!(
173        !state_units.is_running(),
174        "cannot offer a simple VMBus device while state units are running"
175    );
176    let offer = device.offer();
177    let name = format!("{}:{}", offer.interface_name, offer.instance_id);
178    let handle = offer_simple_device(
179        driver_source,
180        vmbus.control.as_ref(),
181        device,
182        InitialDeviceState::Stopped,
183    )
184    .await?;
185    let unit = state_units
186        .add(name)
187        .depends_on(vmbus.unit.handle())
188        .spawn(driver_source.simple(), |recv| {
189            run_async_unit(SimpleChannelUnit(handle), recv)
190        })?;
191    Ok(unit)
192}
193
194impl<T: SimpleVmbusDevice> SimpleChannelUnit<T> {
195    /// Revokes the channel and returns it.
196    pub async fn revoke(self) -> T {
197        self.0.revoke().await.unwrap()
198    }
199}
200
201impl<T: SimpleVmbusDevice> Inspect for SimpleChannelUnit<T> {
202    fn inspect(&self, req: inspect::Request<'_>) {
203        self.0.inspect(req);
204    }
205}
206
207impl<T: SimpleVmbusDevice> StateUnit for &'_ SimpleChannelUnit<T> {
208    async fn start(&mut self) {
209        self.0.start();
210    }
211
212    async fn stop(&mut self) {
213        self.0.stop().await;
214    }
215
216    async fn reset(&mut self) -> anyhow::Result<()> {
217        self.0.reset().await;
218        Ok(())
219    }
220
221    async fn save(&mut self) -> Result<Option<SavedStateBlob>, SaveError> {
222        let state = self.0.save().await.map_err(SaveError::Other)?;
223        Ok(state)
224    }
225
226    async fn restore(&mut self, state: SavedStateBlob) -> Result<(), RestoreError> {
227        self.0.restore(state).await.map_err(RestoreError::Other)
228    }
229}
230
231/// Offers a channel, creates a unit for it, and adds it to `state_units`.
232pub async fn offer_vmbus_device_handle_unit(
233    driver_source: &VmTaskDriverSource,
234    state_units: &StateUnits,
235    vmbus: &VmbusServerHandle,
236    resolver: &ResourceResolver,
237    resource: Resource<VmbusDeviceHandleKind>,
238) -> anyhow::Result<SpawnedUnit<ChannelUnit<dyn VmbusDevice>>> {
239    let channel = resolver
240        .resolve(resource, ResolveVmbusDeviceHandleParams { driver_source })
241        .await?;
242    let offer = channel.0.offer();
243    let name = format!("{}:{}", offer.interface_name, offer.instance_id);
244    let handle =
245        offer_generic_channel(&driver_source.simple(), vmbus.control.as_ref(), channel.0).await?;
246    let unit = state_units
247        .add(name)
248        .depends_on(vmbus.unit.handle())
249        .spawn(driver_source.simple(), |recv| {
250            run_async_unit(ChannelUnit(handle), recv)
251        })?;
252    Ok(unit)
253}