From f9de0c40305823d043928964d00986a09b4e0cca 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 | 52 ++++++++++++++++++++-------
 ucollection/unified_collection_data.h |  7 ++--
 2 files changed, 44 insertions(+), 15 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,24 +162,17 @@ static long ioctrl_collect_process_count(void __user *argp)
 	return 0;
 }
 
-static long ioctrl_collect_thread_count(void __user *argp)
+static long read_thread_count_locked(struct ucollection_process_thread_count *kcount,
+	struct ucollection_process_thread_count __user *count)
 {
-	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;
-	}
-	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);
+	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);
+		pr_info("pid=%d is task NULL or not alive", kcount->pid);
 		rcu_read_unlock();
 		return -EINVAL;
 	}
-	uint32_t thread_count = 0;
+	unsigned int thread_count = 0;
 	struct task_struct *t = task;
 	do {
 		thread_count++;
@@ -189,6 +182,36 @@ static long ioctrl_collect_thread_count(void __user *argp)
 	return 0;
 }
 
+static long ioctrl_collect_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;
+	}
+	memset(&kcount, 0, sizeof(struct ucollection_process_thread_count));
+	(void)copy_from_user(&kcount, count, sizeof(struct ucollection_process_thread_count));
+	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;
+	}
+	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,
 	struct ucollection_thread_cpu_entry __user *entry)
 {
@@ -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)
-- 
Gitee