Re: static_branch_enable() does not work from a __init function?

[Peter Zijlstra <peterz@infradead.org>]

On Wed, Dec 16, 2020 at 10:26:49AM +0100, Peter Zijlstra wrote:
> On Wed, Dec 16, 2020 at 03:54:29AM +0000, Dexuan Cui wrote:
> > 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?
> 
> *groan*... I think this is because __init is ran with
> MODULE_STATE_COMING, we only switch to MODULE_STATE_LIVE later.
> 
> Let me see if there's a sane way to untangle that.
> 
> > #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.

So I think the reason your above module doesn't work, while the one in
vmx_init() does work (for 5.10) should be fixed by the completely
untested below.

I've no clue about 5.4 and no desire to investigate. That's what distro
people are for.

Can you verify?

---
diff --git a/kernel/jump_label.c b/kernel/jump_label.c
index 015ef903ce8c..c6a39d662935 100644
--- a/kernel/jump_label.c
+++ b/kernel/jump_label.c
@@ -793,6 +793,7 @@ int jump_label_text_reserved(void *start, void *end)
 static void jump_label_update(struct static_key *key)
 {
 	struct jump_entry *stop = __stop___jump_table;
+	bool init = system_state < SYSTEM_RUNNING;
 	struct jump_entry *entry;
 #ifdef CONFIG_MODULES
 	struct module *mod;
@@ -804,15 +805,16 @@ static void jump_label_update(struct static_key *key)
 
 	preempt_disable();
 	mod = __module_address((unsigned long)key);
-	if (mod)
+	if (mod) {
 		stop = mod->jump_entries + mod->num_jump_entries;
+		init = mod->state == MODULE_STATE_COMING;
+	}
 	preempt_enable();
 #endif
 	entry = static_key_entries(key);
 	/* if there are no users, entry can be NULL */
 	if (entry)
-		__jump_label_update(key, entry, stop,
-				    system_state < SYSTEM_RUNNING);
+		__jump_label_update(key, entry, stop, init);
 }
 
 #ifdef CONFIG_STATIC_KEYS_SELFTEST