[PATCH 2/3] sched/core: Introduce nr_running percpu for each cookie

[Cruz Zhao]

Introduce a percpu count to struct sched_core_cookie, which indicates how
many tasks with this cookie in the runqueue of this cpu.

Signed-off-by: Cruz Zhao <CruzZhao@xxxxxxxxxxxxxxxxx>
---
 kernel/sched/core.c       |  7 +++++++
 kernel/sched/core_sched.c | 16 ++++++++--------
 kernel/sched/sched.h      |  9 +++++++++
 3 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 263d764..9f71042 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -235,21 +235,28 @@ static inline int rb_sched_core_cmp(const void *key, const struct rb_node *node)
 
 void sched_core_enqueue(struct rq *rq, struct task_struct *p)
 {
+	struct sched_core_cookie *ck = (struct sched_core_cookie *)p->core_cookie;
+
 	rq->core->core_task_seq++;
 
 	if (!p->core_cookie)
 		return;
 
 	rb_add(&p->core_node, &rq->core_tree, rb_sched_core_less);
+
+	*per_cpu_ptr(ck->nr_running, rq->cpu) += 1;
 }
 
 void sched_core_dequeue(struct rq *rq, struct task_struct *p, int flags)
 {
+	struct sched_core_cookie *ck = (struct sched_core_cookie *)p->core_cookie;
+
 	rq->core->core_task_seq++;
 
 	if (sched_core_enqueued(p)) {
 		rb_erase(&p->core_node, &rq->core_tree);
 		RB_CLEAR_NODE(&p->core_node);
+		*per_cpu_ptr(ck->nr_running, rq->cpu) -= 1;
 	}
 
 	/*
diff --git a/kernel/sched/core_sched.c b/kernel/sched/core_sched.c
index ba2466c..65ab9fcb 100644
--- a/kernel/sched/core_sched.c
+++ b/kernel/sched/core_sched.c
@@ -1,20 +1,19 @@
 // SPDX-License-Identifier: GPL-2.0-only
 
-/*
- * A simple wrapper around refcount. An allocated sched_core_cookie's
- * address is used to compute the cookie of the task.
- */
-struct sched_core_cookie {
-	refcount_t refcnt;
-};
-
 static unsigned long sched_core_alloc_cookie(void)
 {
 	struct sched_core_cookie *ck = kmalloc(sizeof(*ck), GFP_KERNEL);
+	int cpu;
+
 	if (!ck)
 		return 0;
 
 	refcount_set(&ck->refcnt, 1);
+
+	ck->nr_running = alloc_percpu(unsigned int);
+	for_each_possible_cpu(cpu)
+		*per_cpu_ptr(ck->nr_running, cpu) = 0;
+
 	sched_core_get();
 
 	return (unsigned long)ck;
@@ -25,6 +24,7 @@ static void sched_core_put_cookie(unsigned long cookie)
 	struct sched_core_cookie *ptr = (void *)cookie;
 
 	if (ptr && refcount_dec_and_test(&ptr->refcnt)) {
+		free_percpu(ptr->nr_running);
 		kfree(ptr);
 		sched_core_put();
 	}
diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index 5b14b6b..d852c67 100644
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -1186,6 +1186,15 @@ static inline raw_spinlock_t *__rq_lockp(struct rq *rq)
 bool cfs_prio_less(struct task_struct *a, struct task_struct *b, bool fi);
 
 /*
+ * A simple wrapper around refcount. An allocated sched_core_cookie's
+ * address is used to compute the cookie of the task.
+ */
+struct sched_core_cookie {
+	refcount_t refcnt;
+	unsigned int __percpu *nr_running;
+};
+
+/*
  * Helpers to check if the CPU's core cookie matches with the task's cookie
  * when core scheduling is enabled.
  * A special case is that the task's cookie always matches with CPU's core
-- 
1.8.3.1