From 768b13b1f9b07586506d935b47bf73e7e8219e82 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9D=A8=E4=B8=96=E7=90=A6=5Fhw?= <670097973@qq.com>
Date: Fri, 10 May 2024 21:11:04 +0800
Subject: [PATCH] add app thread count ioctrol
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: 杨世琦_hw <670097973@qq.com>
---
 ucollection/ucollection_process_cpu.c | 54 ++++++++++++++++++++-------
 ucollection/unified_collection_data.h |  7 +++-
 2 files changed, 45 insertions(+), 16 deletions(-)

diff --git a/ucollection/ucollection_process_cpu.c b/ucollection/ucollection_process_cpu.c
index 9fa5d7f..db12fbe 100644
--- a/ucollection/ucollection_process_cpu.c
+++ b/ucollection/ucollection_process_cpu.c
@@ -162,6 +162,26 @@ static long ioctrl_collect_process_count(void __user *argp)
 	return 0;
 }
 
+static long read_thread_count_locked(struct ucollection_process_thread_count *kcount,
+	struct ucollection_process_thread_count __user *count)
+{
+	rcu_read_lock();
+	struct task_struct *task = get_alive_task_by_pid(kcount->pid);
+	if (task == NULL) {
+		pr_info("pid=%d is task NULL or not alive", kcount->pid);
+		rcu_read_unlock();
+		return -EINVAL;
+	}
+	unsigned int thread_count = 0;
+	struct task_struct *t = task;
+	do {
+		thread_count++;
+	} while_each_thread(task, t);
+	put_user(thread_count, &count->thread_count);
+	rcu_read_unlock();
+	return 0;
+}
+
 static long ioctrl_collect_thread_count(void __user *argp)
 {
 	struct ucollection_process_thread_count kcount;
@@ -172,21 +192,24 @@ static long ioctrl_collect_thread_count(void __user *argp)
 	}
 	memset(&kcount, 0, sizeof(struct ucollection_process_thread_count));
 	(void)copy_from_user(&kcount, count, sizeof(struct ucollection_process_thread_count));
-	rcu_read_lock();
-	struct task_struct *task = get_alive_task_by_pid(kcount.pid);
-	if (task == NULL) {
-		pr_info("pid=%d is task NULL or not alive", kcount.pid);
-		rcu_read_unlock();
+	return read_thread_count_locked(&kcount, count);
+}
+
+static long ioctrl_collect_app_thread_count(void __user *argp)
+{
+	struct ucollection_process_thread_count kcount;
+	struct ucollection_process_thread_count __user *count = argp;
+	if (count == NULL) {
+		pr_err("cpu entry is null");
 		return -EINVAL;
 	}
-	uint32_t thread_count = 0;
-	struct task_struct *t = task;
-	do {
-		thread_count++;
-	} while_each_thread(task, t);
-	put_user(thread_count, &count->thread_count);
-	rcu_read_unlock();
-	return 0;
+	memset(&kcount, 0, sizeof(struct ucollection_process_thread_count));
+	(void)copy_from_user(&kcount, count, sizeof(struct ucollection_process_thread_count));
+	if (current->tgid != kcount.pid) {
+		pr_err("pid=%d is not self current tgid:%d", kcount.pid, current->tgid);
+		return -EINVAL;
+	}
+	return read_thread_count_locked(&kcount, count);
 }
 
 static long read_thread_info_locked(struct ucollection_thread_cpu_entry *kentry,
@@ -199,7 +222,7 @@ static long read_thread_info_locked(struct ucollection_thread_cpu_entry *kentry,
 		rcu_read_unlock();
 		return -EINVAL;
 	}
-	uint32_t thread_count = 0;
+	unsigned int thread_count = 0;
 	struct task_struct *t = task;
 	do {
 		if (thread_count >= kentry->total_count) {
@@ -295,6 +318,9 @@ long unified_collection_collect_process_cpu(unsigned int cmd, void __user *argp)
 	case IOCTRL_COLLECT_THREAD_COUNT:
 		ret = ioctrl_collect_thread_count(argp);
 		break;
+	case IOCTRL_COLLECT_APP_THREAD_COUNT:
+		ret = ioctrl_collect_app_thread_count(argp);
+		break;
 	case IOCTRL_COLLECT_APP_THREAD:
 		ret = ioctrl_collect_app_thread_cpu(argp);
 		break;
diff --git a/ucollection/unified_collection_data.h b/ucollection/unified_collection_data.h
index 1913b30..67e270f 100644
--- a/ucollection/unified_collection_data.h
+++ b/ucollection/unified_collection_data.h
@@ -33,8 +33,8 @@ struct ucollection_process_cpu_entry {
 };
 
 struct ucollection_process_thread_count {
-	uint32_t pid;
-	uint32_t thread_count;
+	int pid;
+	unsigned int thread_count;
 };
 
 struct ucollection_thread_cpu_item {
@@ -66,6 +66,7 @@ enum collection_type {
 	COLLECT_THREAD_COUNT,
 	COLLECT_APP_THREAD,
 	COLLECT_THE_THREAD,
+	COLLECT_APP_THREAD_COUNT,
 };
 
 #define DMIPS_NUM 128
@@ -76,6 +77,8 @@ enum collection_type {
 	struct ucollection_process_cpu_entry)
 #define IOCTRL_COLLECT_THREAD_COUNT _IOR(IOCTRL_COLLECT_CPU_BASE, COLLECT_THREAD_COUNT, \
 	struct ucollection_process_thread_count)
+#define IOCTRL_COLLECT_APP_THREAD_COUNT _IOR(IOCTRL_COLLECT_CPU_BASE, COLLECT_APP_THREAD_COUNT, \
+	struct ucollection_process_thread_count)
 #define IOCTRL_COLLECT_APP_THREAD _IOR(IOCTRL_COLLECT_CPU_BASE, COLLECT_APP_THREAD, struct ucollection_thread_cpu_entry)
 #define IOCTRL_COLLECT_THE_THREAD _IOR(IOCTRL_COLLECT_CPU_BASE, COLLECT_THE_THREAD, struct ucollection_thread_cpu_entry)
 #define IOCTRL_COLLECT_PROC_COUNT _IOR(IOCTRL_COLLECT_CPU_BASE, COLLECT_PROC_COUNT, unsigned int)