[PATCH] taskstats: Avoid counting an exiting thread twice in TGID queries

["Baek Kyu, Kim" <pocj8ur4in@gmail.com>]

fill_tgid_exit() adds a thread's contribution to signal->stats before
exit_notify() sets exit_state. A TGID query in that interval copies the
cached contribution, then scans the same thread because exit_state is
still clear. The thread's contribution is counted twice.

Mark a task when its contribution reaches the cache, under the same
siglock used by the TGID query. Exclude marked tasks from the live scan.
Leave a task eligible for the scan if no cache exists, and clear the
marker when a new task is copied during fork.

Fixes: ad4ecbcba728 ("[PATCH] delay accounting taskstats interface send tgid once")
Signed-off-by: Baek Kyu, Kim <pocj8ur4in@gmail.com>
---
In an arm64 QEMU guest, holding the interval after cache accounting but
before exit_state changed made a 100 ms worker appear twice in a TGID CPU
query. The baseline reported 208000 us for 8000 us of the querying thread
plus 100000 us of the worker; the change reported 112000 us for 12000 us
plus 100000 us. Natural race frequency and other architectures were not
measured; KASAN and KCSAN were not enabled.

 include/linux/sched.h | 5 +++++
 kernel/fork.c         | 3 +++
 kernel/taskstats.c    | 7 ++++---
 3 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/include/linux/sched.h b/include/linux/sched.h
index d35ae49a991f..359ee982323c 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -1457,6 +1457,11 @@ struct task_struct {
 	struct task_delay_info		*delays;
 #endif
 
+#ifdef CONFIG_TASKSTATS
+	/* Protected by sighand->siglock. */
+	bool				taskstats_tgid_accounted;
+#endif
+
 #ifdef CONFIG_FAULT_INJECTION
 	int				make_it_fail;
 	unsigned int			fail_nth;
diff --git a/kernel/fork.c b/kernel/fork.c
index 33760ece3bb3..76d9abcc5ca8 100644
--- a/kernel/fork.c
+++ b/kernel/fork.c
@@ -2202,6 +2202,9 @@ __latent_entropy struct task_struct *copy_process(
 		goto bad_fork_cleanup_count;
 
 	delayacct_tsk_init(p);	/* Must remain after dup_task_struct() */
+#ifdef CONFIG_TASKSTATS
+	p->taskstats_tgid_accounted = false;
+#endif
 	p->flags &= ~(PF_SUPERPRIV | PF_WQ_WORKER | PF_IDLE | PF_NO_SETAFFINITY);
 	p->flags |= PF_FORKNOEXEC;
 	INIT_LIST_HEAD(&p->children);
diff --git a/kernel/taskstats.c b/kernel/taskstats.c
index f31df72f0e9d..dd9054de761f 100644
--- a/kernel/taskstats.c
+++ b/kernel/taskstats.c
@@ -245,8 +245,8 @@ static int fill_stats_for_tgid(pid_t tgid, struct taskstats *stats)
 	u64 now_ns;
 
 	/*
-	 * Add additional stats from live tasks except zombie thread group
-	 * leaders who are already counted with the dead tasks
+	 * Add stats from tasks not yet included in the dead-task cache.
+	 * Accounting can finish before exit_state is set.
 	 */
 	rcu_read_lock();
 	first = find_task_by_vpid(tgid);
@@ -261,7 +261,7 @@ static int fill_stats_for_tgid(pid_t tgid, struct taskstats *stats)
 
 	now_ns = ktime_get_ns();
 	for_each_thread(first, tsk) {
-		if (tsk->exit_state)
+		if (tsk->exit_state || tsk->taskstats_tgid_accounted)
 			continue;
 
 		tgid_stats_add_task(stats, tsk, now_ns);
@@ -291,6 +291,7 @@ static void fill_tgid_exit(struct task_struct *tsk)
 
 	now_ns = ktime_get_ns();
 	tgid_stats_add_task(tsk->signal->stats, tsk, now_ns);
+	tsk->taskstats_tgid_accounted = true;
 ret:
 	spin_unlock_irqrestore(&tsk->sighand->siglock, flags);
 	return;
-- 
2.55.0