1use crate::bus::ChannelRequest;
7use crate::bus::ChannelServerRequest;
8use crate::bus::ModifyRequest;
9use crate::bus::OfferInput;
10use crate::bus::OfferParams;
11use crate::bus::OfferResources;
12use crate::bus::OpenRequest;
13use crate::bus::ParentBus;
14use crate::gpadl::GpadlMap;
15use crate::gpadl::GpadlMapView;
16use anyhow::Context;
17use async_trait::async_trait;
18use futures::StreamExt;
19use futures::future::join_all;
20use futures::stream::SelectAll;
21use futures::stream::select;
22use inspect::Inspect;
23use inspect::InspectMut;
24use mesh::RecvError;
25use mesh::rpc::FailableRpc;
26use mesh::rpc::Rpc;
27use mesh::rpc::RpcSend;
28use pal_async::task::Spawn;
29use pal_async::task::Task;
30use pal_event::Event;
31use std::any::Any;
32use std::collections::BTreeSet;
33use std::marker::PhantomData;
34use std::pin::pin;
35use std::sync::Arc;
36use thiserror::Error;
37use tracing::instrument;
38use vmbus_core::TaggedStream;
39use vmbus_core::protocol::GpadlId;
40use vmbus_ring::gparange::MultiPagedRangeBuf;
41use vmcore::notify::Notify;
42use vmcore::save_restore::RestoreError;
43use vmcore::save_restore::SaveError;
44use vmcore::save_restore::SavedStateBlob;
45use vmcore::slim_event::SlimEvent;
46
47pub type ChannelOpenError = anyhow::Error;
49
50#[async_trait]
52pub trait VmbusDevice: Send + Any + InspectMut {
53 fn offer(&self) -> OfferParams;
55
56 fn max_subchannels(&self) -> u16 {
58 0
59 }
60
61 fn install(&mut self, resources: DeviceResources);
63
64 async fn open(
66 &mut self,
67 channel_idx: u16,
68 open_request: &OpenRequest,
69 ) -> Result<(), ChannelOpenError>;
70
71 async fn close(&mut self, channel_idx: u16);
73
74 async fn retarget_vp(&mut self, channel_idx: u16, target_vp: u32);
76
77 fn start(&mut self);
79
80 async fn stop(&mut self);
82
83 fn supports_save_restore(&mut self) -> Option<&mut dyn SaveRestoreVmbusDevice>;
88}
89
90#[async_trait]
92pub trait SaveRestoreVmbusDevice: VmbusDevice {
93 async fn save(&mut self) -> Result<SavedStateBlob, SaveError>;
95
96 async fn restore(
101 &mut self,
102 control: RestoreControl<'_>,
103 state: SavedStateBlob,
104 ) -> Result<(), RestoreError>;
105}
106
107#[derive(Debug, Default)]
109pub struct DeviceResources {
110 pub offer_resources: OfferResources,
112 pub gpadl_map: GpadlMapView,
114 pub channel_control: ChannelControl,
116 pub channels: Vec<ChannelResources>,
118}
119
120#[derive(Debug)]
122pub struct ChannelResources {
123 pub event: Notify,
125}
126
127#[derive(Debug, Default, Clone)]
129pub struct ChannelControl {
130 send: Option<mesh::Sender<u16>>,
131 max: u16,
132}
133
134#[derive(Debug, Error)]
136#[error("too many subchannels requested")]
137pub struct TooManySubchannels;
138
139impl ChannelControl {
140 pub fn enable_subchannels(&self, count: u16) -> Result<(), TooManySubchannels> {
147 if count > self.max {
148 return Err(TooManySubchannels);
149 }
150 if let Some(send) = &self.send {
151 send.send(count);
152 }
153 Ok(())
154 }
155
156 pub fn max_subchannels(&self) -> u16 {
158 self.max
159 }
160}
161
162#[must_use]
166#[derive(Inspect)]
167pub(crate) struct GenericChannelHandle {
168 #[inspect(flatten, send = "StateRequest::Inspect")]
169 state_req: mesh::Sender<StateRequest>,
170 #[inspect(skip)]
171 task: Task<Box<dyn VmbusDevice>>,
172}
173
174#[derive(Debug)]
175enum StateRequest {
176 Start,
178 Stop(Rpc<(), ()>),
180
181 Reset(Rpc<(), ()>),
185
186 Save(FailableRpc<(), Option<SavedStateBlob>>),
190
191 Restore(FailableRpc<SavedStateBlob, ()>),
195
196 Inspect(inspect::Deferred),
198}
199
200impl std::fmt::Debug for GenericChannelHandle {
201 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
202 f.pad("ChannelHandle")
203 }
204}
205
206impl GenericChannelHandle {
207 pub async fn revoke(self) -> Option<Box<dyn VmbusDevice>> {
209 drop(self.state_req);
210 Some(self.task.await)
211 }
212
213 pub fn start(&self) {
214 self.state_req.send(StateRequest::Start);
215 }
216
217 pub async fn stop(&self) {
218 self.state_req
219 .call(StateRequest::Stop, ())
220 .await
221 .expect("critical channel failure")
222 }
223
224 pub async fn reset(&self) {
225 self.state_req
226 .call(StateRequest::Reset, ())
227 .await
228 .expect("critical channel failure")
229 }
230
231 pub async fn save(&self) -> anyhow::Result<Option<SavedStateBlob>> {
232 self.state_req
233 .call(StateRequest::Save, ())
234 .await
235 .expect("critical channel failure")
236 .map_err(|err| err.into())
237 }
238
239 pub async fn restore(&self, buffer: SavedStateBlob) -> anyhow::Result<()> {
240 self.state_req
241 .call(StateRequest::Restore, buffer)
242 .await
243 .expect("critical channel failure")
244 .map_err(|err| err.into())
245 }
246}
247
248#[must_use]
252#[derive(Inspect)]
253#[inspect(transparent)]
254pub struct ChannelHandle<T: ?Sized>(GenericChannelHandle, PhantomData<fn() -> Box<T>>);
255
256impl<T: ?Sized> std::fmt::Debug for ChannelHandle<T> {
257 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
258 f.pad("ChannelHandle")
259 }
260}
261
262impl<T: 'static + VmbusDevice> ChannelHandle<T> {
263 pub async fn revoke(self) -> Option<T> {
265 let device = self.0.revoke().await? as Box<dyn Any>;
266 Some(
267 *device
268 .downcast()
269 .expect("type must match the one used to create it"),
270 )
271 }
272}
273
274impl ChannelHandle<dyn VmbusDevice> {
275 pub async fn revoke(self) -> Option<Box<dyn VmbusDevice>> {
277 self.0.revoke().await
278 }
279}
280
281impl<T: 'static + VmbusDevice + ?Sized> ChannelHandle<T> {
282 pub fn start(&self) {
284 self.0.start()
285 }
286
287 pub async fn stop(&self) {
289 self.0.stop().await
290 }
291
292 pub async fn reset(&self) {
294 self.0.reset().await
295 }
296
297 pub async fn save(&self) -> anyhow::Result<Option<SavedStateBlob>> {
299 self.0.save().await
300 }
301
302 pub async fn restore(&self, buffer: SavedStateBlob) -> anyhow::Result<()> {
304 self.0.restore(buffer).await
305 }
306}
307
308async fn offer_generic(
309 driver: &impl Spawn,
310 bus: &(impl ParentBus + ?Sized),
311 mut channel: Box<dyn VmbusDevice>,
312) -> anyhow::Result<GenericChannelHandle> {
313 let offer = channel.offer();
314 let max_subchannels = channel.max_subchannels();
315 let instance_id = offer.instance_id;
316 let (request_send, request_recv) = mesh::channel();
317 let (server_request_send, server_request_recv) = mesh::channel();
318 let (state_req_send, state_req_recv) = mesh::channel();
319
320 let use_event = bus.use_event();
321
322 let events: Vec<_> = (0..max_subchannels + 1)
323 .map(|_| {
324 if use_event {
325 Notify::from_event(Event::new())
326 } else {
327 Notify::from_slim_event(Arc::new(SlimEvent::new()))
328 }
329 })
330 .collect();
331
332 let request = OfferInput {
333 params: offer,
334 event: events[0].clone().interrupt(),
335 request_send,
336 server_request_recv,
337 };
338
339 let gpadl_map = GpadlMap::new();
340
341 let offer_result = bus.add_child(request).await?;
342
343 let resources = events
344 .iter()
345 .map(|event| ChannelResources {
346 event: event.clone(),
347 })
348 .collect();
349
350 let (subchannel_enable_send, subchannel_enable_recv) = mesh::channel();
351 channel.install(DeviceResources {
352 offer_resources: offer_result,
353 gpadl_map: gpadl_map.clone().view(),
354 channels: resources,
355 channel_control: ChannelControl {
356 send: Some(subchannel_enable_send),
357 max: max_subchannels,
358 },
359 });
360
361 let bus = bus.clone_bus();
362 let task = driver.spawn(format!("vmbus offer {}", instance_id), async move {
363 let device = Device::new(
364 request_recv,
365 server_request_send,
366 events,
367 gpadl_map,
368 subchannel_enable_recv,
369 );
370 device
371 .run_channel(bus.as_ref(), channel.as_mut(), state_req_recv)
372 .await;
373 channel
374 });
375
376 Ok(GenericChannelHandle {
377 state_req: state_req_send,
378 task,
379 })
380}
381
382pub struct RestoreControl<'a> {
385 device: &'a mut Device,
386 bus: &'a dyn ParentBus,
387 offer: OfferParams,
388}
389
390impl RestoreControl<'_> {
391 pub async fn restore(
400 &mut self,
401 states: &[bool],
402 ) -> Result<Vec<Option<OpenRequest>>, ChannelRestoreError> {
403 self.device.restore(self.bus, &self.offer, states).await
404 }
405}
406
407#[derive(Debug, Error)]
409pub enum ChannelRestoreError {
410 #[error("failed to enable subchannels")]
412 EnablingSubchannels(#[source] anyhow::Error),
413 #[error("failed to restore vmbus channel")]
415 RestoreError(#[source] anyhow::Error),
416 #[error("failed to restore gpadl")]
418 GpadlError(#[source] vmbus_ring::gparange::Error),
419}
420
421impl From<ChannelRestoreError> for RestoreError {
422 fn from(err: ChannelRestoreError) -> Self {
423 RestoreError::Other(err.into())
424 }
425}
426
427enum DeviceState {
428 Running,
429 Stopped(Vec<(usize, ChannelRequest)>),
433}
434
435struct Device {
436 state: DeviceState,
437 server_requests: Vec<mesh::Sender<ChannelServerRequest>>,
438 open: Vec<bool>,
439 subchannel_gpadls: Vec<BTreeSet<GpadlId>>,
440 requests: SelectAll<TaggedStream<usize, mesh::Receiver<ChannelRequest>>>,
441 events: Vec<Notify>,
442 gpadl_map: Arc<GpadlMap>,
443 subchannel_enable_recv: mesh::Receiver<u16>,
444}
445
446impl Device {
447 fn new(
448 request_recv: mesh::Receiver<ChannelRequest>,
449 server_request_send: mesh::Sender<ChannelServerRequest>,
450 events: Vec<Notify>,
451 gpadl_map: Arc<GpadlMap>,
452 subchannel_enable_recv: mesh::Receiver<u16>,
453 ) -> Self {
454 let open: Vec<bool> = vec![false];
455 let subchannel_gpadls: Vec<BTreeSet<GpadlId>> = vec![];
456 let mut requests: SelectAll<TaggedStream<usize, mesh::Receiver<ChannelRequest>>> =
457 SelectAll::new();
458 requests.push(TaggedStream::new(0, request_recv));
459 Self {
460 state: DeviceState::Running,
461 server_requests: vec![server_request_send],
462 open,
463 subchannel_gpadls,
464 requests,
465 events,
466 gpadl_map,
467 subchannel_enable_recv,
468 }
469 }
470
471 async fn run_channel(
473 mut self,
474 bus: &dyn ParentBus,
475 channel: &mut dyn VmbusDevice,
476 state_req_recv: mesh::Receiver<StateRequest>,
477 ) {
478 enum Event {
479 Request(usize, Option<ChannelRequest>),
480 EnableSubchannels(u16),
481 StateRequest(Result<StateRequest, RecvError>),
482 }
483
484 let mut state_req_recv = pin!(futures::stream::unfold(state_req_recv, async |mut recv| {
485 Some((recv.recv().await, recv))
486 }));
487
488 let map_request = |(idx, req)| Event::Request(idx, req);
489 loop {
490 let mut s = select(
491 (&mut self.requests).map(map_request),
492 select(
493 (&mut self.subchannel_enable_recv).map(Event::EnableSubchannels),
494 (&mut state_req_recv).map(Event::StateRequest),
495 ),
496 );
497 if let Some(event) = s.next().await {
498 match event {
499 Event::Request(idx, Some(request)) => {
500 self.handle_channel_request(idx, request, channel).await;
501 }
502 Event::Request(_idx, None) => continue,
503 Event::EnableSubchannels(count) => {
504 let offer = channel.offer();
505 let _ = self.enable_channels(bus, &offer, count as usize + 1).await;
506 }
507 Event::StateRequest(Ok(request)) => {
508 self.handle_state_request(request, channel, bus).await;
509 }
510 Event::StateRequest(Err(_)) => {
511 break;
513 }
514 }
515 }
516 }
517
518 let subchannel_senders = self.server_requests.split_off(1);
525 join_all(
526 subchannel_senders
527 .into_iter()
528 .map(|s| s.call(ChannelServerRequest::Revoke, ())),
529 )
530 .await;
531
532 drop(self.server_requests);
534 for recv in self.requests.iter_mut() {
542 if recv.value().is_some() {
543 while recv.next().await.is_some() {}
544 }
545 }
546
547 for subchannel_idx in (0..self.open.len()).rev() {
548 if self.open[subchannel_idx] {
549 channel.close(subchannel_idx as u16).await;
550 }
551 }
552 }
553
554 #[instrument(level = "debug", skip_all, fields(channel_idx, ?request))]
555 async fn handle_channel_request(
556 &mut self,
557 channel_idx: usize,
558 request: ChannelRequest,
559 channel: &mut dyn VmbusDevice,
560 ) {
561 if matches!(request, ChannelRequest::Open(_) | ChannelRequest::Modify(_)) {
567 if let DeviceState::Stopped(pending_messages) = &mut self.state {
568 pending_messages.push((channel_idx, request));
569 return;
570 }
571 }
572
573 match request {
574 ChannelRequest::Open(rpc) => {
575 rpc.handle(async |open_request| {
576 self.handle_open(channel, channel_idx, open_request).await
577 })
578 .await
579 }
580 ChannelRequest::Close(rpc) => {
581 rpc.handle(async |()| {
582 self.handle_close(channel_idx, channel).await;
583 })
584 .await
585 }
586 ChannelRequest::Gpadl(rpc) => rpc.handle_sync(|gpadl| {
587 self.handle_gpadl(gpadl.id, gpadl.count, gpadl.buf, channel_idx);
588 true
589 }),
590 ChannelRequest::TeardownGpadl(rpc) => {
591 self.handle_teardown_gpadl(rpc, channel_idx);
592 }
593 ChannelRequest::Modify(rpc) => {
594 rpc.handle(async |req| {
595 self.handle_modify(channel, channel_idx, req).await;
596 0
597 })
598 .await
599 }
600 }
601 }
602
603 async fn handle_open(
604 &mut self,
605 channel: &mut dyn VmbusDevice,
606 channel_idx: usize,
607 open_request: OpenRequest,
608 ) -> bool {
609 assert!(!self.open[channel_idx]);
610 let opened = channel
613 .open(channel_idx as u16, &open_request)
614 .await
615 .inspect_err(|error| {
616 tracelimit::error_ratelimited!(
617 error = error.as_ref() as &dyn std::error::Error,
618 "failed to open channel"
619 );
620 })
621 .is_ok();
622 self.open[channel_idx] = opened;
623 opened
624 }
625
626 async fn handle_close(&mut self, channel_idx: usize, channel: &mut dyn VmbusDevice) {
627 assert!(self.open[channel_idx]);
628 if channel_idx == 0 {
629 self.server_requests.truncate(1);
631 for recv in self.requests.iter_mut() {
632 if let Some(&idx) = recv.value() {
633 if idx > 0 {
634 while recv.next().await.is_some() {}
635 }
636 }
637 }
638 for subchannel_idx in 1..self.open.len() {
639 if self.open[subchannel_idx] {
640 channel.close(subchannel_idx as u16).await;
641 }
642 for &gpadl_id in &self.subchannel_gpadls[subchannel_idx - 1] {
643 self.gpadl_map.remove(gpadl_id, Box::new(|| ()));
644 }
645 }
646 self.open.truncate(1);
647 self.subchannel_gpadls.clear();
648 }
649 channel.close(channel_idx as u16).await;
650 self.open[channel_idx] = false;
651 if channel_idx == 0 {
652 while self.subchannel_enable_recv.try_recv().is_ok() {}
654 }
655 }
656
657 fn handle_gpadl(&mut self, id: GpadlId, count: u16, buf: Vec<u64>, channel_idx: usize) {
658 self.gpadl_map.add(
659 id,
660 MultiPagedRangeBuf::from_range_buffer(count.into(), buf).unwrap(),
661 );
662 if channel_idx > 0 {
663 self.subchannel_gpadls[channel_idx - 1].insert(id);
664 }
665 }
666
667 fn handle_teardown_gpadl(&mut self, rpc: Rpc<GpadlId, ()>, channel_idx: usize) {
668 let id = *rpc.input();
669 if let Some(f) = self.gpadl_map.remove(
670 id,
671 Box::new(move || {
672 rpc.complete(());
673 }),
674 ) {
675 f()
676 }
677 if channel_idx > 0 {
678 assert!(self.subchannel_gpadls[channel_idx - 1].remove(&id));
679 }
680 }
681
682 async fn handle_modify(
683 &mut self,
684 channel: &mut dyn VmbusDevice,
685 channel_idx: usize,
686 req: ModifyRequest,
687 ) {
688 match req {
689 ModifyRequest::TargetVp { target_vp } => {
690 channel.retarget_vp(channel_idx as u16, target_vp).await
691 }
692 }
693 }
694
695 #[instrument(level = "debug", skip_all, fields(?request))]
696 async fn handle_state_request(
697 &mut self,
698 request: StateRequest,
699 channel: &mut dyn VmbusDevice,
700 bus: &dyn ParentBus,
701 ) {
702 match request {
703 StateRequest::Start => {
704 channel.start();
705 if let DeviceState::Stopped(pending_messages) =
706 std::mem::replace(&mut self.state, DeviceState::Running)
707 {
708 for (channel_idx, request) in pending_messages.into_iter() {
709 self.handle_channel_request(channel_idx, request, channel)
710 .await;
711 }
712 }
713 }
714 StateRequest::Stop(rpc) => {
715 if matches!(self.state, DeviceState::Running) {
716 self.state = DeviceState::Stopped(Vec::new());
717 rpc.handle(async |()| {
718 channel.stop().await;
719 })
720 .await;
721 } else {
722 rpc.complete(());
723 }
724 }
725 StateRequest::Reset(rpc) => {
726 if let DeviceState::Stopped(pending_messages) = &mut self.state {
727 pending_messages.clear();
728 }
729 rpc.complete(());
730 }
731 StateRequest::Save(rpc) => {
732 rpc.handle_failable(async |()| {
733 if let Some(channel) = channel.supports_save_restore() {
734 channel.save().await.map(Some)
735 } else {
736 Ok(None)
737 }
738 })
739 .await;
740 }
741 StateRequest::Restore(rpc) => {
742 rpc.handle_failable(async |buffer| {
743 let channel = channel
744 .supports_save_restore()
745 .context("saved state not supported")?;
746 let control = RestoreControl {
747 device: &mut *self,
748 offer: channel.offer(),
749 bus,
750 };
751 channel
752 .restore(control, buffer)
753 .await
754 .map_err(anyhow::Error::from)?;
755 anyhow::Ok(())
756 })
757 .await;
758 }
759 StateRequest::Inspect(deferred) => {
760 deferred.inspect(&mut *channel);
761 }
762 }
763 }
764
765 async fn enable_channels(
766 &mut self,
767 bus: &dyn ParentBus,
768 offer: &OfferParams,
769 count: usize,
770 ) -> anyhow::Result<()> {
771 let mut r = Ok(());
773 for subchannel_idx in self.server_requests.len()..count {
774 let (request_send, request_recv) = mesh::channel();
775 let (server_request_send, server_request_recv) = mesh::channel();
776 let request = OfferInput {
777 params: OfferParams {
778 subchannel_index: subchannel_idx as u16,
779 ..offer.clone()
780 },
781 event: self.events[subchannel_idx].clone().interrupt(),
782 request_send,
783 server_request_recv,
784 };
785 match bus.add_child(request).await {
786 Ok(_) => {
787 self.requests
788 .push(TaggedStream::new(subchannel_idx, request_recv));
789 self.server_requests.push(server_request_send);
790 self.subchannel_gpadls.push(BTreeSet::new());
791 self.open.push(false);
792 }
793 Err(err) => {
794 tracing::error!(
795 error = err.as_ref() as &dyn std::error::Error,
796 "could not offer subchannel"
797 );
798 if r.is_ok() {
799 r = Err(err);
800 }
801 }
802 }
803 }
804 r
805 }
806
807 pub async fn restore(
808 &mut self,
809 bus: &dyn ParentBus,
810 offer: &OfferParams,
811 states: &[bool],
812 ) -> Result<Vec<Option<OpenRequest>>, ChannelRestoreError> {
813 self.enable_channels(bus, offer, states.len())
814 .await
815 .map_err(ChannelRestoreError::EnablingSubchannels)?;
816
817 let mut results = Vec::with_capacity(states.len());
818 for (channel_idx, open) in states.iter().copied().enumerate() {
819 let result = self.server_requests[channel_idx]
820 .call_failable(ChannelServerRequest::Restore, open)
821 .await
822 .map_err(|err| ChannelRestoreError::RestoreError(err.into()))?;
823
824 assert!(open == result.open_request.is_some());
825
826 for gpadl in result.gpadls {
827 let buf = match MultiPagedRangeBuf::from_range_buffer(
828 gpadl.request.count.into(),
829 gpadl.request.buf,
830 ) {
831 Ok(buf) => buf,
832 Err(err) => {
833 if gpadl.accepted {
834 return Err(ChannelRestoreError::GpadlError(err));
835 } else {
836 continue;
839 }
840 }
841 };
842 self.gpadl_map.add(gpadl.request.id, buf);
843 if channel_idx > 0 {
844 self.subchannel_gpadls[channel_idx - 1].insert(gpadl.request.id);
845 }
846 }
847
848 results.push(result.open_request);
849 }
850 self.open.copy_from_slice(states);
851 Ok(results)
852 }
853}
854
855pub async fn offer_channel<T: 'static + VmbusDevice>(
858 driver: &impl Spawn,
859 bus: &(impl ParentBus + ?Sized),
860 channel: T,
861) -> anyhow::Result<ChannelHandle<T>> {
862 let handle = offer_generic(driver, bus, Box::new(channel)).await?;
863 Ok(ChannelHandle(handle, PhantomData))
864}
865
866pub async fn offer_generic_channel(
868 driver: &impl Spawn,
869 bus: &(impl ParentBus + ?Sized),
870 channel: Box<dyn VmbusDevice>,
871) -> anyhow::Result<ChannelHandle<dyn VmbusDevice>> {
872 let handle = offer_generic(driver, bus, channel).await?;
873 Ok(ChannelHandle(handle, PhantomData))
874}