Skip to content

Commit

Permalink
scx_rusty: refactor mempolicy handling bpf code and load balancing
Browse files Browse the repository at this point in the history 
This change refactors some of the helper methods for getting the
preferred node for tasks using mempolicy. The load balancing logic in
try_find_move_task is updated to allow for a filter, which is used to
filter for tasks with a preferred mempolicy.

Signed-off-by: Daniel Hodges <[email protected]>
  • Loading branch information
@hodgesds
hodgesds committed Jul 15, 2024
1 parent 5c38f97 commit 7eca676
Show file tree
Hide file tree
Showing 2 changed files with 47 additions and 28 deletions.
23 changes: 12 additions & 11 deletions scheds/rust/scx_rusty/src/bpf/main.bpf.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1203,7 +1203,7 @@ static u64 node_dom_mask(u32 node_id)
u32 dom_id = 0;

bpf_for(dom_id, 0, nr_doms) {
if (!(dom_node_id(dom_id) == node_id))
if (dom_node_id(dom_id) != node_id)
continue;

mask |= 1LLU << dom_id;
Expand All @@ -1216,20 +1216,22 @@ static u64 node_dom_mask(u32 node_id)
* Sets the preferred domain mask according to the mempolicy. See man(2)
* set_mempolicy for more details on mempolicy.
*/
static int task_set_preferred_mempolicy_dom_mask(struct task_struct *p,
struct task_ctx *taskc)
static void task_set_preferred_mempolicy_dom_mask(struct task_struct *p,
struct task_ctx *taskc)
{
u32 node_id;
u32 val = 0;
nodemask_t *node_mask = &p->mempolicy->nodes;
void *mask;

taskc->preferred_dom_mask = 0;

if (!mempolicy_affinity || !bpf_core_field_exists(p->mempolicy) ||
!p->mempolicy || !taskc->cpumask)
return -1;
return;

if (!(p->mempolicy->mode & (MPOL_BIND|MPOL_PREFERRED|MPOL_PREFERRED_MANY)))
return -1;
return;

// MPOL_BIND and MPOL_PREFERRED_MANY use the home_node field on the
// mempolicy struct, so use that for now. In the future the memory
Expand All @@ -1238,12 +1240,12 @@ static int task_set_preferred_mempolicy_dom_mask(struct task_struct *p,
if ((int)p->mempolicy->home_node >= 0) {
taskc->preferred_dom_mask =
node_dom_mask((u32)p->mempolicy->home_node);
return 0;
return;
}

mask = BPF_CORE_READ(node_mask, bits);
if (bpf_core_read(&val, sizeof(val), mask))
return -1;
return;

bpf_for(node_id, 0, nr_nodes) {
if (!(val & 1 << node_id))
Expand All @@ -1252,7 +1254,7 @@ static int task_set_preferred_mempolicy_dom_mask(struct task_struct *p,
taskc->preferred_dom_mask |= node_dom_mask(node_id);
}

return 0;
return;
}

void BPF_STRUCT_OPS(rusty_dispatch, s32 cpu, struct task_struct *prev)
Expand Down Expand Up @@ -1511,10 +1513,9 @@ static u32 task_pick_domain(struct task_ctx *taskc, struct task_struct *p,
return NO_DOM_FOUND;

taskc->dom_mask = 0;
taskc->preferred_dom_mask = 0;

dom = pcpu_ctx[cpu].dom_rr_cur++;
has_preferred_dom = task_set_preferred_mempolicy_dom_mask(p, taskc);
task_set_preferred_mempolicy_dom_mask(p, taskc);
bpf_repeat(nr_doms) {
dom = (dom + 1) % nr_doms;

Expand All @@ -1527,7 +1528,7 @@ static u32 task_pick_domain(struct task_ctx *taskc, struct task_struct *p,
if (first_dom == NO_DOM_FOUND)
first_dom = dom;

if (has_preferred_dom < 0)
if (taskc->preferred_dom_mask == 0)
continue;

if (((1LLU << dom) & taskc->preferred_dom_mask))
Expand Down
52 changes: 35 additions & 17 deletions scheds/rust/scx_rusty/src/load_balance.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -755,18 +755,8 @@ impl<'a, 'b> LoadBalancer<'a, 'b> {
skip_kworkers: bool,
) -> Option<&'d TaskInfo>
where
I: IntoIterator<Item = &'d TaskInfo> + Clone,
I: IntoIterator<Item = &'d TaskInfo>,
{
// First try to find a task in the preferred domain mask.
if let Some(task) = tasks_by_load.clone().into_iter().take_while(
|task| {
!task.migrated.get()
|| (task.preferred_dom_mask & (1 << pull_dom) > 0)
}).next()
{
return Some(task);
}

match tasks_by_load
.into_iter()
.skip_while(|task| {
Expand All @@ -788,6 +778,7 @@ impl<'a, 'b> LoadBalancer<'a, 'b> {
&mut self,
(push_dom, to_push): (&mut Domain, f64),
(pull_dom, to_pull): (&mut Domain, f64),
task_filter: impl Fn(&TaskInfo, u32) -> bool,
to_xfer: f64,
) -> Result<Option<f64>> {
let to_pull = to_pull.abs();
Expand All @@ -803,22 +794,29 @@ impl<'a, 'b> LoadBalancer<'a, 'b> {
// counterpart while scanning right and picking the better of the
// two.
let tasks = std::mem::take(&mut push_dom.tasks).into_vec();
let pull_dom_id = pull_dom.id.try_into().unwrap(),
let (task, new_imbal) = match (
Self::find_first_candidate(
tasks
.as_slice()
.iter()
.filter(|x| x.load <= OrderedFloat(to_xfer))
.filter(|x| {
x.load <= OrderedFloat(to_xfer)
&& task_filter(x, pull_dom_id)
})
.rev(),
pull_dom.id.try_into().unwrap(),
pull_dom_id,
self.skip_kworkers,
),
Self::find_first_candidate(
tasks
.as_slice()
.iter()
.filter(|x| x.load >= OrderedFloat(to_xfer)),
pull_dom.id.try_into().unwrap(),
.filter(|x| {
x.load >= OrderedFloat(to_xfer)
&& task_filter(x, pull_dom_id)
}),
pull_dom_id,
self.skip_kworkers,
),
) {
Expand Down Expand Up @@ -885,9 +883,19 @@ impl<'a, 'b> LoadBalancer<'a, 'b> {
pull_node.domains.insert(pull_dom);
break;
}
let transferred = self.try_find_move_task((&mut push_dom, push_imbal),
let mut transferred = self.try_find_move_task((&mut push_dom, push_imbal),
(&mut pull_dom, pull_imbal),
|task: &TaskInfo, pull_dom:u32| -> bool {
(task.preferred_dom_mask & (1 << pull_dom)) > 0
},
xfer)?;
if transferred.is_none() {
transferred = self.try_find_move_task((&mut push_dom, push_imbal),
(&mut pull_dom, pull_imbal),
|_task: &TaskInfo, _pull_dom:u32| -> bool {true},
xfer)?;
}

pullers.push(pull_dom);
if let Some(transferred) = transferred {
pushed = transferred;
Expand Down Expand Up @@ -1047,9 +1055,19 @@ impl<'a, 'b> LoadBalancer<'a, 'b> {
bail!("Node {} pull dom {} had imbal {}", node.id, pull_dom.id, pull_imbal);
}
let xfer = push_dom.xfer_between(&pull_dom);
let transferred = self.try_find_move_task((&mut push_dom, push_imbal),
let mut transferred = self.try_find_move_task((&mut push_dom, push_imbal),
(&mut pull_dom, pull_imbal),
|task: &TaskInfo, pull_dom:u32| -> bool {
(task.preferred_dom_mask & (1 << pull_dom)) > 0
},
xfer)?;
if transferred.is_none() {
transferred = self.try_find_move_task((&mut push_dom, push_imbal),
(&mut pull_dom, pull_imbal),
|_task: &TaskInfo, _pull_dom:u32| -> bool {true},
xfer)?;
}

if let Some(transferred) = transferred {
if transferred <= 0.0f64 {
bail!("Expected nonzero load transfer")
Expand Down

0 comments on commit 7eca676

Please sign in to comment.