underhill_core/
threadpool_vm_task_backend.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

use inspect::Inspect;
use pal_async::driver::Driver;
use underhill_threadpool::AffinitizedThreadpool;
use underhill_threadpool::RetargetableDriver;
use vmcore::vm_task::BuildVmTaskDriver;
use vmcore::vm_task::TargetedDriver;

/// A VM task driver backend backed by the threadpool.
#[derive(Debug)]
pub struct ThreadpoolBackend(AffinitizedThreadpool);

impl ThreadpoolBackend {
    pub fn new(tp: AffinitizedThreadpool) -> Self {
        Self(tp)
    }
}

impl BuildVmTaskDriver for ThreadpoolBackend {
    type Driver = ThreadpoolDriver;

    fn build(&self, name: String, target_vp: Option<u32>, run_on_target: bool) -> Self::Driver {
        let _ = name;
        ThreadpoolDriver {
            spawn_target: Target::new(&self.0, if run_on_target { target_vp } else { None }),
            io_target: Target::new(&self.0, target_vp),
        }
    }
}

#[derive(Debug, Inspect)]
pub struct ThreadpoolDriver {
    spawn_target: Target,
    io_target: Target,
}

#[derive(Debug)]
enum Target {
    Untargeted(AffinitizedThreadpool),
    Targeted(RetargetableDriver),
}

impl Inspect for Target {
    fn inspect(&self, req: inspect::Request<'_>) {
        let v = match self {
            Target::Untargeted(_) => "any".into(),
            Target::Targeted(driver) => driver.current_target_cpu().into(),
        };
        req.value(v)
    }
}

impl Target {
    fn new(tp: &AffinitizedThreadpool, vp: Option<u32>) -> Self {
        vp.map_or_else(
            || Target::Untargeted(tp.clone()),
            |vp| Target::Targeted(RetargetableDriver::new(tp.clone(), vp)),
        )
    }
}

impl TargetedDriver for ThreadpoolDriver {
    fn spawner(&self) -> &dyn pal_async::task::Spawn {
        match &self.spawn_target {
            Target::Untargeted(x) => x,
            Target::Targeted(x) => x,
        }
    }

    fn driver(&self) -> &dyn Driver {
        match &self.io_target {
            // TODO: for consistency with existing behavior, initiate IO to VP0.
            // Consider using the current initiator instead.
            Target::Untargeted(x) => x.driver(0),
            Target::Targeted(x) => x,
        }
    }

    fn retarget_vp(&self, target_vp: u32) {
        match &self.spawn_target {
            Target::Untargeted(_) => {}
            Target::Targeted(x) => x.retarget(target_vp),
        }
        match &self.io_target {
            Target::Untargeted(_) => {}
            Target::Targeted(x) => x.retarget(target_vp),
        }
    }

    fn is_target_vp_ready(&self) -> bool {
        match &self.io_target {
            Target::Untargeted(_) => true,
            Target::Targeted(x) => x.current_driver().is_affinity_set(),
        }
    }

    async fn wait_target_vp_ready(&self) {
        match &self.io_target {
            Target::Untargeted(_) => {}
            Target::Targeted(x) => x.current_driver().wait_for_affinity().await,
        }
    }
}