static_branch_enable() does not work from a __init function?

[Dexuan Cui <decui@microsoft.com>]

Hi,
The below init_module() prints "foo: false". This is strange since
static_branch_enable() is called before the static_branch_unlikely().
This strange behavior happens to v5.10 and an old v5.4 kernel.

If I remove the "__init" marker from the init_module() function, then
I get the expected output of "foo: true"! I guess here I'm missing
something with Static Keys?

#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/jump_label.h>

static DEFINE_STATIC_KEY_FALSE(enable_foo);

int __init init_module(void)
{
        static_branch_enable(&enable_foo);

        if (static_branch_unlikely(&enable_foo))
                printk("foo: true\n");
        else
                printk("foo: false\n");

        return 0;
}

void cleanup_module(void)
{
        static_branch_disable(&enable_foo);
}

MODULE_LICENSE("GPL");


PS, I originally found: in arch/x86/kvm/vmx/vmx.c: vmx_init(), it looks
like the line "static_branch_enable(&enable_evmcs);" does not take effect
in a v5.4-based kernel, but does take effect in the v5.10 kernel in the
same x86-64 virtual machine on Hyper-V, so I made the above test module
to test static_branch_enable(), and found that static_branch_enable() in
the test module does not work with both v5.10 and my v5.4 kernel, if the
__init marker is used.

Thanks,
-- Dexuan