Search code examples
clinux-kernelprocess-migration

How to migrate an arbitrary process in a kernel module?


I'm writing a kernel module that iterates over all processes (using the for_each_process macro), now, I would like to migrate ones that meet certain criteria to a new NUMA node or processor. I found a few functions defined as extern in the sched.h:

extern void sched_setnuma(struct task_struct *p, int node);
extern int migrate_task_to(struct task_struct *p, int cpu);
extern int migrate_swap(struct task_struct *, struct task_struct *);

But can't use them cause the kernel is already compiled.

Is there any way to use these functions or similar ones to achieve a task migration?

EDIT

#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/init.h>
#include <linux/sched/signal.h>
#include <linux/sched.h>
#include <linux/kprobes.h>
#include <linux/kallsyms.h>


MODULE_LICENSE("GPL");
MODULE_AUTHOR("Mohammad Siavashi");
MODULE_DESCRIPTION("This is a test module");

int (*migrate_task_to_func)(struct task_struct *p, int cpu);

void start(void)
{
    struct task_struct *t;
    for_each_process(t)
    {
        if (strcmp(t->comm, "stream_task") == 0)
        {
            migrate_task_to_func(t, 2);
        }
    }
}

static int __init hello_init(void)
{
    migrate_task_to_func = (void *)kallsyms_lookup_name("migrate_task_to");
    start();
    return 0;
}

static void __exit hello_cleanup(void)
{
    printk(KERN_INFO "Cleaning up module.\n");
}

module_init(hello_init);
module_exit(hello_cleanup);


Solution

  • You can use kallsyms_lookup_name to find it out and then use it again. If kallsyms_lookup_name is not exported (kernel>=5.7), you can use kprobe to find it first

    e.g:

    #include <linux/kernel.h>
    #include <linux/module.h>
    #include <linux/kprobes.h>
    
    void (*sched_setnuma_func)(struct task_struct *p, int node);
    int (*migrate_task_to_func)(struct task_struct *p, int cpu);
    int (*migrate_swap_func)(struct task_struct *, struct task_struct *);
    
    static int __init test_init(void)
    {
        unsigned long (*kallsyms_lookup_name_func)(const char *name);
        struct kprobe kp = { .symbol_name    = "kallsyms_lookup_name", };
    
        register_kprobe(&kp);
        kallsyms_lookup_name_func = (void *)kp.addr;
        unregister_kprobe(&kp);
    
        if (kallsyms_lookup_name_func) {
            sched_setnuma_func = (void *)kallsyms_lookup_name_func("sched_setnuma");
            migrate_task_to_func = (void *)kallsyms_lookup_name_func("migrate_task_to");
            migrate_swap_func = (void *)kallsyms_lookup_name_func("migrate_swap");
        }
        printk("sched_setnuma_func: %p\n", sched_setnuma_func);
        printk("migrate_task_to_func: %p\n", migrate_task_to_func);
        printk("migrate_swap_func: %p\n", migrate_swap_func);
        return 0;
    }
    
    static void __exit test_exit(void)
    {
    }
    
    module_init(test_init);
    module_exit(test_exit);
    MODULE_LICENSE("GPL");