[PATCH] sched/cputime: Account cgroup fields to the scheduling context

[Hui Su]

Proxy execution separates the scheduling context from the execution
context. Scheduler runtime accounting charges cgroup time to the donor,
while tick and vtime accounting use the execution task when updating
cgroup fields.

When the donor and execution task belong to different cgroups, this can
charge cpu.stat usage to the donor cgroup while charging cpu.stat user
and system fields to the execution task's cgroup.

Before this fix, a reproducer with donor and execution tasks in separate
cgroups showed the donor cgroup gaining usage time while the execution
cgroup gained system time.

Account cgroup fields to the scheduling context as well. Separate
per-CPU cpustat accounting from cgroup field accounting so the former
continues to follow the execution context while the latter follows the
scheduling context.

Keep a per-rq cputime owner so delayed virtual cputime accounting is
charged to the scheduling context that owned the interval.

A wakeup can reset rq->donor while holding an rq lock. If the reset is
remote or occurs outside task context, defer the cputime boundary until
the target CPU next enters the scheduler. This avoids accessing remote
or interrupt-context vtime and per-CPU cpustat state. Hold a reference
to the old cputime owner while the deferred boundary is pending.

Keep CPU cpustat classification based on the execution task and cgroup
field classification based on the accounting owner, including nice_usec.
Generic vtime residuals are flushed at proxy boundaries so they cannot
span two scheduling contexts.

Leave per-task, thread-group, and force-idle accounting unchanged.

Tested: proxy execution cgroup reproducer with
CONFIG_SCHED_PROXY_EXEC=y and CONFIG_VIRT_CPU_ACCOUNTING_GEN=y.
Tested: CONFIG_SCHED_PROXY_EXEC=y with CONFIG_VIRT_CPU_ACCOUNTING=n.
Built with CONFIG_SCHED_PROXY_EXEC=n and CONFIG_VIRT_CPU_ACCOUNTING_GEN=y.

Fixes: aa4f74dfd42b ("sched: Fix runtime accounting w/ split exec & sched contexts")
Signed-off-by: Hui Su <sh_def@xxxxxxx>
---
 kernel/sched/core.c    |  17 ++++
 kernel/sched/cputime.c | 189 +++++++++++++++++++++++++++++++++++++----
 kernel/sched/sched.h   |   8 ++
 3 files changed, 197 insertions(+), 17 deletions(-)

diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index f78275192036..43c09110cc98 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -3759,6 +3759,11 @@ static inline void proxy_reset_donor(struct rq *rq)
 {
 	WARN_ON_ONCE(rq->donor == rq->curr);
 
+	/* Establish the cputime boundary locally from task context. */
+	if (rq == this_rq() && in_task())
+		sched_cputime_donor_changed(rq, rq->curr);
+	else
+		sched_cputime_donor_defer(rq);
 	put_prev_set_next_task(rq, rq->donor, rq->curr);
 	rq_set_donor(rq, rq->curr);
 	zap_balance_callbacks(rq);
@@ -5363,6 +5368,10 @@ static struct rq *finish_task_switch(struct task_struct *prev)
 	 */
 	prev_state = READ_ONCE(prev->__state);
 	vtime_task_switch(prev);
+#ifdef CONFIG_SCHED_PROXY_EXEC
+	if (sched_proxy_exec())
+		sched_cputime_switch(rq, prev);
+#endif
 	perf_event_task_sched_in(prev, current);
 	finish_task(prev);
 	tick_nohz_task_switch();
@@ -7227,6 +7236,11 @@ static void __sched notrace __schedule(int sched_mode)
 keep_resched:
 	rq->last_seen_need_resched_ns = 0;
 
+#ifdef CONFIG_SCHED_PROXY_EXEC
+	if (prev == next && sched_proxy_exec())
+		sched_cputime_donor_changed(rq, rq->donor);
+#endif
+
 	is_switch = prev != next;
 	if (likely(is_switch)) {
 		rq->nr_switches++;
@@ -8326,6 +8340,9 @@ void __init init_idle(struct task_struct *idle, int cpu)
 
 	rq->idle = idle;
 	rq_set_donor(rq, idle);
+#ifdef CONFIG_SCHED_PROXY_EXEC
+	rq->cputime_donor = idle;
+#endif
 	rcu_assign_pointer(rq->curr, idle);
 	idle->on_rq = TASK_ON_RQ_QUEUED;
 	idle->on_cpu = 1;
diff --git a/kernel/sched/cputime.c b/kernel/sched/cputime.c
index 06bddaa738e5..5cda78496c23 100644
--- a/kernel/sched/cputime.c
+++ b/kernel/sched/cputime.c
@@ -103,17 +103,29 @@ static u64 irqtime_tick_accounted(u64 dummy)
 
 #endif /* !CONFIG_IRQ_TIME_ACCOUNTING */
 
-static inline void task_group_account_field(struct task_struct *p, int index,
-					    u64 tmp)
+static inline struct task_struct *cgroup_account_task(struct task_struct *p)
+{
+#ifdef CONFIG_SCHED_PROXY_EXEC
+	struct task_struct *donor;
+
+	/* Cgroup fields follow the scheduling context, not the execution task. */
+	if (sched_proxy_exec()) {
+		donor = this_rq()->cputime_donor;
+		if (donor)
+			return donor;
+	}
+#endif
+	return p;
+}
+
+static inline void account_cpustat_field(int index, u64 tmp)
 {
-	/*
-	 * Since all updates are sure to touch the root cgroup, we
-	 * get ourselves ahead and touch it first. If the root cgroup
-	 * is the only cgroup, then nothing else should be necessary.
-	 *
-	 */
 	__this_cpu_add(kernel_cpustat.cpustat[index], tmp);
+}
 
+static inline void account_cgroup_field(struct task_struct *p, int index,
+					u64 tmp)
+{
 	cgroup_account_cputime_field(p, index, tmp);
 }
 
@@ -124,16 +136,23 @@ static inline void task_group_account_field(struct task_struct *p, int index,
  */
 void account_user_time(struct task_struct *p, u64 cputime)
 {
-	int index;
+	struct task_struct *cgroup_task = cgroup_account_task(p);
+	int index, cgroup_index;
 
 	/* Add user time to process. */
 	p->utime += cputime;
 	account_group_user_time(p, cputime);
 
 	index = (task_nice(p) > 0) ? CPUTIME_NICE : CPUTIME_USER;
+	if (cgroup_task == p)
+		cgroup_index = index;
+	else
+		cgroup_index = (task_nice(cgroup_task) > 0) ?
+			CPUTIME_NICE : CPUTIME_USER;
 
-	/* Add user time to cpustat. */
-	task_group_account_field(p, index, cputime);
+	/* Add user time to cpustat and the cgroup. */
+	account_cpustat_field(index, cputime);
+	account_cgroup_field(cgroup_task, cgroup_index, cputime);
 
 	/* Account for user time used */
 	acct_account_cputime(p);
@@ -146,7 +165,9 @@ void account_user_time(struct task_struct *p, u64 cputime)
  */
 void account_guest_time(struct task_struct *p, u64 cputime)
 {
+	struct task_struct *cgroup_task = cgroup_account_task(p);
 	u64 *cpustat = kcpustat_this_cpu->cpustat;
+	int index, cgroup_index;
 
 	/* Add guest time to process. */
 	p->utime += cputime;
@@ -154,11 +175,19 @@ void account_guest_time(struct task_struct *p, u64 cputime)
 	p->gtime += cputime;
 
 	/* Add guest time to cpustat. */
-	if (task_nice(p) > 0) {
-		task_group_account_field(p, CPUTIME_NICE, cputime);
+	index = (task_nice(p) > 0) ? CPUTIME_NICE : CPUTIME_USER;
+	if (cgroup_task == p)
+		cgroup_index = index;
+	else
+		cgroup_index = (task_nice(cgroup_task) > 0) ?
+			CPUTIME_NICE : CPUTIME_USER;
+	if (index == CPUTIME_NICE) {
+		account_cpustat_field(index, cputime);
+		account_cgroup_field(cgroup_task, cgroup_index, cputime);
 		cpustat[CPUTIME_GUEST_NICE] += cputime;
 	} else {
-		task_group_account_field(p, CPUTIME_USER, cputime);
+		account_cpustat_field(index, cputime);
+		account_cgroup_field(cgroup_task, cgroup_index, cputime);
 		cpustat[CPUTIME_GUEST] += cputime;
 	}
 }
@@ -172,12 +201,15 @@ void account_guest_time(struct task_struct *p, u64 cputime)
 void account_system_index_time(struct task_struct *p,
 			       u64 cputime, enum cpu_usage_stat index)
 {
+	struct task_struct *cgroup_task = cgroup_account_task(p);
+
 	/* Add system time to process. */
 	p->stime += cputime;
 	account_group_system_time(p, cputime);
 
-	/* Add system time to cpustat. */
-	task_group_account_field(p, index, cputime);
+	/* Add system time to cpustat and the cgroup. */
+	account_cpustat_field(index, cputime);
+	account_cgroup_field(cgroup_task, index, cputime);
 
 	/* Account for system time used */
 	acct_account_cputime(p);
@@ -245,7 +277,8 @@ void __account_forceidle_time(struct task_struct *p, u64 delta)
 {
 	__schedstat_add(p->stats.core_forceidle_sum, delta);
 
-	task_group_account_field(p, CPUTIME_FORCEIDLE, delta);
+	account_cpustat_field(CPUTIME_FORCEIDLE, delta);
+	account_cgroup_field(p, CPUTIME_FORCEIDLE, delta);
 }
 #endif /* CONFIG_SCHED_CORE */
 
@@ -988,6 +1021,30 @@ void vtime_task_switch_generic(struct task_struct *prev)
 	write_seqcount_end(&vtime->seqcount);
 }
 
+#if defined(CONFIG_SCHED_PROXY_EXEC)
+static void vtime_flush_pending(struct task_struct *tsk)
+{
+	struct vtime *vtime = &tsk->vtime;
+
+	write_seqcount_begin(&vtime->seqcount);
+	if (vtime->utime) {
+		account_user_time(tsk, vtime->utime);
+		vtime->utime = 0;
+	}
+
+	if (vtime->gtime) {
+		account_guest_time(tsk, vtime->gtime);
+		vtime->gtime = 0;
+	}
+
+	if (vtime->stime) {
+		account_system_time(tsk, irq_count(), vtime->stime);
+		vtime->stime = 0;
+	}
+	write_seqcount_end(&vtime->seqcount);
+}
+#endif /* CONFIG_SCHED_PROXY_EXEC */
+
 void vtime_init_idle(struct task_struct *t, int cpu)
 {
 	struct vtime *vtime = &t->vtime;
@@ -1300,3 +1357,101 @@ void kcpustat_cpu_fetch(struct kernel_cpustat *dst, int cpu)
 EXPORT_SYMBOL_GPL(kcpustat_cpu_fetch);
 
 #endif /* CONFIG_VIRT_CPU_ACCOUNTING_GEN */
+
+#ifdef CONFIG_SCHED_PROXY_EXEC
+static void __sched_cputime_release_donor(struct rq *rq)
+{
+	struct task_struct *hold = rq->cputime_donor_hold;
+
+	rq->cputime_donor_hold = NULL;
+	if (hold)
+		put_task_struct(hold);
+}
+
+static void sched_cputime_release_donor(struct rq *rq)
+{
+	lockdep_assert_rq_held(rq);
+	__sched_cputime_release_donor(rq);
+}
+
+static void sched_cputime_release_donor_switch(struct rq *rq)
+{
+	/*
+	 * The rq lock is carried across the context switch physically,
+	 * but its lockdep state has been released.
+	 */
+	lockdep_assert_irqs_disabled();
+	__sched_cputime_release_donor(rq);
+}
+
+void sched_cputime_donor_defer(struct rq *rq)
+{
+	lockdep_assert_rq_held(rq);
+
+	if (rq->cputime_donor_hold)
+		WARN_ON_ONCE(rq->cputime_donor_hold != rq->cputime_donor);
+	if (rq->cputime_donor_hold)
+		return;
+
+	get_task_struct(rq->cputime_donor);
+	rq->cputime_donor_hold = rq->cputime_donor;
+}
+
+void sched_cputime_donor_changed(struct rq *rq,
+				 struct task_struct *new_donor)
+{
+	lockdep_assert_rq_held(rq);
+
+	if (rq->cputime_donor == new_donor) {
+		sched_cputime_release_donor(rq);
+		return;
+	}
+
+	/* Commit the current interval before changing its cgroup owner. */
+	if (!sched_proxy_exec())
+		goto update_owner;
+
+	if (is_idle_task(rq->curr))
+		goto update_owner;
+
+#ifdef CONFIG_VIRT_CPU_ACCOUNTING_GEN
+	if (vtime_accounting_enabled_this_cpu()) {
+		vtime_account_kernel(rq->curr);
+		vtime_flush_pending(rq->curr);
+	}
+#elif defined(CONFIG_VIRT_CPU_ACCOUNTING_NATIVE)
+	vtime_account_kernel(rq->curr);
+	vtime_flush(rq->curr);
+#endif
+
+update_owner:
+	rq->cputime_donor = new_donor;
+	sched_cputime_release_donor(rq);
+}
+
+void sched_cputime_switch(struct rq *rq, struct task_struct *prev)
+{
+	/* The rq lock is carried across the context switch physically. */
+	lockdep_assert_irqs_disabled();
+
+	/* rq->donor already names the next scheduling context here. */
+	if (!sched_proxy_exec())
+		goto update_owner;
+
+#ifdef CONFIG_VIRT_CPU_ACCOUNTING_GEN
+	if (vtime_accounting_enabled_this_cpu()) {
+		if (rq->cputime_donor != prev)
+			vtime_flush_pending(prev);
+
+		if (rq->donor != current) {
+			rq->cputime_donor = current;
+			vtime_flush_pending(current);
+		}
+	}
+#endif
+
+update_owner:
+	rq->cputime_donor = rq->donor;
+	sched_cputime_release_donor_switch(rq);
+}
+#endif /* CONFIG_SCHED_PROXY_EXEC */
diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index e656c7059bf8..bb268e8160c7 100644
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -1175,6 +1175,10 @@ struct rq {
 #ifdef CONFIG_SCHED_PROXY_EXEC
 	struct task_struct __rcu	*donor;  /* Scheduling context */
 	struct task_struct __rcu	*curr;   /* Execution context */
+	/* Cgroup cputime owner; lags donor while the previous task is flushed. */
+	struct task_struct	*cputime_donor;
+	/* Reference held while a remote donor change awaits local accounting. */
+	struct task_struct	*cputime_donor_hold;
 #else
 	union {
 		struct task_struct __rcu *donor; /* Scheduler context */
@@ -1466,6 +1470,10 @@ static inline bool available_idle_cpu(int cpu)
 }
 
 #ifdef CONFIG_SCHED_PROXY_EXEC
+void sched_cputime_donor_changed(struct rq *rq, struct task_struct *new_donor);
+void sched_cputime_donor_defer(struct rq *rq);
+void sched_cputime_switch(struct rq *rq, struct task_struct *prev);
+
 static inline void rq_set_donor(struct rq *rq, struct task_struct *t)
 {
 	rcu_assign_pointer(rq->donor, t);
-- 
2.54.0