lunarkernel_sched_extention/lse_task_struct_ext.c at dev · LunarKernel-Dev/lunarkernel_sched_extention · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
// SPDX-License-Identifier: GPL-2.0-only
/*
 * Copyright (C) 2025, LunarKernel Project. All rights reserved.
 *
 * This code is part of LunarKernel.
 * task_struct extension.
 *
 * File name: lse_task_struct_ext.c
 * Author: Cloud_Yun <1770669041@qq.com>
 * Version: v251003_Dev
 * Date: 2025/10/3 Friday
 */

#include <linux/smp.h>
#include <linux/rwsem.h>
#include <trace/hooks/sched.h>
#include <../kernel/sched/sched.h>

#include "lse_main.h"

static struct kmem_cache *lse_task_struct_cachep;

static void init_lse_task_struct(struct lse_task_struct *lts, struct task_struct *tsk)
{
    memset(lts, 0, sizeof(struct lse_task_struct));
    lts->task = tsk;
}

static void alloc_lse_task_struct(void *unused, struct task_struct *tsk,
                                     struct task_struct *orig)
{
    struct lse_task_struct *lts;

    if (!tsk)
        return;

    if (smp_load_acquire(&tsk->android_vendor_data1[LTS_IDX]) != 0)
        return;

    lts = kmem_cache_alloc(lse_task_struct_cachep, GFP_ATOMIC);
    if (!lts)
        return;

    init_lse_task_struct(lts, tsk);

    smp_store_release(&tsk->android_vendor_data1[LTS_IDX], (u64)lts);
}

static void free_lse_task_struct(void *unused, struct task_struct *tsk)
{
    struct lse_task_struct *lts;

    if (!tsk)
        return;

    lts = (struct lse_task_struct *)smp_load_acquire(&tsk->android_vendor_data1[LTS_IDX]);
    if (!lts)
        return;

    smp_store_release(&tsk->android_vendor_data1[LTS_IDX], 0);

    kmem_cache_free(lse_task_struct_cachep, lts);
}

static void alloc_lts_mem_for_all_threads(void)
{
    struct task_struct *p, *g;
    u32 iter_cpu;

    read_lock(&tasklist_lock);
    for_each_process_thread(g, p) {
        struct lse_task_struct *lts = NULL;

        lts = (struct lse_task_struct *)smp_load_acquire(&p->android_vendor_data1[LTS_IDX]);

        if (!lts) {
            lts = kmem_cache_alloc(lse_task_struct_cachep, GFP_ATOMIC);

            if (lts) {
                init_lse_task_struct(lts, p);
                smp_store_release(&p->android_vendor_data1[LTS_IDX], (u64)lts);
            }
        }
    }

    for_each_possible_cpu(iter_cpu) {
        struct lse_task_struct *lts = NULL;

        p = cpu_rq(iter_cpu)->idle;
        lts = (struct lse_task_struct *)smp_load_acquire(&p->android_vendor_data1[LTS_IDX]);

        if (!lts) {
            lts = kmem_cache_alloc(lse_task_struct_cachep, GFP_ATOMIC);

            if (lts) {
                init_lse_task_struct(lts, p);
                smp_store_release(&p->android_vendor_data1[LTS_IDX], (u64)lts);
            }
        }
    }
    read_unlock(&tasklist_lock);
}

int lse_task_struct_ext_init(void)
{
    lse_task_struct_cachep = kmem_cache_create("lse_task_struct",
            sizeof(struct lse_task_struct), 0,
            SLAB_PANIC | SLAB_ACCOUNT, NULL);

    if (!lse_task_struct_cachep)
        return -ENOMEM;

    alloc_lts_mem_for_all_threads();

    register_trace_android_vh_dup_task_struct(alloc_lse_task_struct, NULL);
    register_trace_android_vh_free_task(free_lse_task_struct, NULL);

    return 0;
}