1#![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
37pub struct VmbusServerHandle {
41 unit: SpawnedUnit<VmbusServerUnit>,
42 control: Arc<VmbusServerControl>,
43}
44
45impl VmbusServerHandle {
46 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 pub fn control(&self) -> &Arc<VmbusServerControl> {
61 &self.control
62 }
63
64 pub fn unit_handle(&self) -> &UnitHandle {
66 self.unit.handle()
67 }
68
69 pub async fn remove(self) -> VmbusServer {
71 self.unit.remove().await.0
72 }
73}
74
75#[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#[must_use]
108#[derive(Debug, Inspect)]
109#[inspect(transparent)]
110pub struct ChannelUnit<T: ?Sized>(ChannelHandle<T>);
111
112pub 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 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#[must_use]
162#[derive(Debug)]
163pub struct SimpleChannelUnit<T: SimpleVmbusDevice>(SimpleDeviceHandle<T>);
164
165pub 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 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
231pub 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}