From 72ecf3c8e3b45058ce7e4d3e6bfbdfa3894e2860 Mon Sep 17 00:00:00 2001
From: I Hsin Cheng <richard120310@gmail.com>
Date: Tue, 19 Nov 2024 21:36:12 +0800
Subject: [PATCH] scx_rusty: Temporary fix of duplicate active tptr

Under severe load unbalance scenario such as mixtures of CPU-insensive
workload and I/O-intensive worload, same tptr may be written into the
same dom_active_tptrs's array.

It will lead to load balancer's failure because when the tptr task
contains large enough load, it tends be to selected so warnings about
same tptr being set in "lb_data" will continue to pop up.

Use a workaround for now , which is to keep a HashSet in userspace
recording the current active tptr under a domain, and do not generate
the same task repeatedly.

Signed-off-by: I Hsin Cheng <richard120310@gmail.com>
---
 scheds/rust/scx_rusty/src/load_balance.rs | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/scheds/rust/scx_rusty/src/load_balance.rs b/scheds/rust/scx_rusty/src/load_balance.rs
index c8d5b26..8ed1bd5 100644
--- a/scheds/rust/scx_rusty/src/load_balance.rs
+++ b/scheds/rust/scx_rusty/src/load_balance.rs
@@ -133,6 +133,7 @@
 use core::cmp::Ordering;
 use std::cell::Cell;
 use std::collections::BTreeMap;
+use std::collections::HashSet;
 use std::collections::VecDeque;
 use std::fmt;
 use std::sync::Arc;
@@ -343,6 +344,7 @@ struct Domain {
     queried_tasks: bool,
     load: LoadEntity,
     tasks: SortedVec<TaskInfo>,
+    active_tptr_set: HashSet<u64>,
 }
 
 impl Domain {
@@ -362,6 +364,7 @@ impl Domain {
                 load_avg,
             ),
             tasks: SortedVec::new(),
+            active_tptr_set: HashSet::new(),
         }
     }
 
@@ -680,6 +683,10 @@ impl<'a, 'b> LoadBalancer<'a, 'b> {
             let tptr = active_tptrs.tptrs[(idx % MAX_TPTRS) as usize];
             let key = unsafe { std::mem::transmute::<u64, [u8; 8]>(tptr) };
 
+            if dom.active_tptr_set.contains(&tptr) {
+                continue;
+            }
+
             if let Some(task_data_elem) = task_data.lookup(&key, libbpf_rs::MapFlags::ANY)? {
                 let task_ctx =
                     unsafe { &*(task_data_elem.as_slice().as_ptr() as *const bpf_intf::task_ctx) };
@@ -705,6 +712,7 @@ impl<'a, 'b> LoadBalancer<'a, 'b> {
                 };
                 load *= weight;
 
+                dom.active_tptr_set.insert(tptr);
                 dom.tasks.insert(TaskInfo {
                     tptr,
                     load: OrderedFloat(load),