[PATCH v3] jump_label: reduce the size of struct static_key

[PATCH v3] jump_label: reduce the size of struct static_key

[Jason Baron]

The static_key->next field goes mostly unused. The field is used for
associating module uses with a static key. Most uses of struct static_key
define a static key in the core kernel and make use of it entirely within
the core kernel, or define the static key in a module and make use of it
only from within that module. In fact, of the ~3,000 static keys defined,
I found only about 5 or so that did not fit this pattern.

Thus, we can remove the static_key->next field entirely and overload
the static_key->entries field. That is, when all the static_key uses
are contained within the same module, static_key->entries continues
to point to those uses. However, if the static_key uses are not contained
within the module where the static_key is defined, then we allocate a
struct static_key_mod, store a pointer to the uses within that
struct static_key_mod, and have the static key point at the static_key_mod.
This does incur some extra memory usage when a static_key is used in a
module that does not define it, but since there are only a handful of such
cases there is a net savings.

In order to identify if the static_key->entries pointer contains a
struct static_key_mod or a struct jump_entry pointer, bit 1 of
static_key->entries is set to 1 if it points to a struct static_key_mod and
is 0 if it points to a struct jump_entry. We were already using bit 0 in a
similar way to store the initial value of the static_key. This does mean
that allocations of struct static_key_mod and that the struct jump_entry
tables need to be at least 4-byte aligned in memory. As far as I can tell
all arches meet this criteria.

For my .config, the patch increased the text by 267 bytes, but reduced
the data + bss size by 14,976, for a net savings of 14,709 bytes.

   text	   data	    bss	    dec	    hex	filename
8092427	5016512	 790528	13899467	 d416cb	vmlinux.pre
8092694	5001536	 790528	13884758	 d3dd56	vmlinux.post

Cc: Peter Zijlstra <peterz@infradead.org>
Cc: Steven Rostedt <rostedt@goodmis.org>
Cc: Ingo Molnar <mingo@kernel.org>
Cc: Joe Perches <joe@perches.com>
Signed-off-by: Jason Baron <jbaron@akamai.com>
---
Changed in v3:
-added static_key_set_[entries|mod]() to clean up type casting (Ingo Molnar)
Changed in v2:
-Replace static_key->entries with union (Steven Rostedt)
---
 Documentation/static-keys.txt |   4 +-
 include/linux/jump_label.h    |  23 +++++---
 kernel/jump_label.c           | 126 ++++++++++++++++++++++++++++++++++--------
 3 files changed, 120 insertions(+), 33 deletions(-)

diff --git a/Documentation/static-keys.txt b/Documentation/static-keys.txt
index ea8d7b4e53f0..32a25fad0c1b 100644
--- a/Documentation/static-keys.txt
+++ b/Documentation/static-keys.txt
@@ -155,7 +155,9 @@ or:
 
 There are a few functions and macros that architectures must implement in order
 to take advantage of this optimization. If there is no architecture support, we
-simply fall back to a traditional, load, test, and jump sequence.
+simply fall back to a traditional, load, test, and jump sequence. Also, the
+struct jump_entry table must be at least 4-byte aligned because the
+static_key->entry field makes use of the two least significant bits.
 
 * select HAVE_ARCH_JUMP_LABEL, see: arch/x86/Kconfig
 
diff --git a/include/linux/jump_label.h b/include/linux/jump_label.h
index a0547c571800..680c98b2f41c 100644
--- a/include/linux/jump_label.h
+++ b/include/linux/jump_label.h
@@ -89,11 +89,17 @@ extern bool static_key_initialized;
 
 struct static_key {
 	atomic_t enabled;
-/* Set lsb bit to 1 if branch is default true, 0 ot */
-	struct jump_entry *entries;
-#ifdef CONFIG_MODULES
-	struct static_key_mod *next;
-#endif
+/*
+ * bit 0 => 1 if key is initially true
+ *	    0 if initially false
+ * bit 1 => 1 if points to struct static_key_mod
+ *	    0 if points to struct jump_entry
+ */
+	union {
+		unsigned long type;
+		struct jump_entry *entries;
+		struct static_key_mod *next;
+	};
 };
 
 #else
@@ -118,9 +124,10 @@ struct module;
 
 #ifdef HAVE_JUMP_LABEL
 
-#define JUMP_TYPE_FALSE	0UL
-#define JUMP_TYPE_TRUE	1UL
-#define JUMP_TYPE_MASK	1UL
+#define JUMP_TYPE_FALSE		0UL
+#define JUMP_TYPE_TRUE		1UL
+#define JUMP_TYPE_LINKED	2UL
+#define JUMP_TYPE_MASK		3UL
 
 static __always_inline bool static_key_false(struct static_key *key)
 {
diff --git a/kernel/jump_label.c b/kernel/jump_label.c
index a9b8cf500591..2eb9e80dc691 100644
--- a/kernel/jump_label.c
+++ b/kernel/jump_label.c
@@ -236,12 +236,27 @@ void __weak __init_or_module arch_jump_label_transform_static(struct jump_entry
 
 static inline struct jump_entry *static_key_entries(struct static_key *key)
 {
-	return (struct jump_entry *)((unsigned long)key->entries & ~JUMP_TYPE_MASK);
+	return (struct jump_entry *)(key->type & ~JUMP_TYPE_MASK);
 }
 
 static inline bool static_key_type(struct static_key *key)
 {
-	return (unsigned long)key->entries & JUMP_TYPE_MASK;
+	return key->type & JUMP_TYPE_TRUE;
+}
+
+static inline bool static_key_linked(struct static_key *key)
+{
+	return key->type & JUMP_TYPE_LINKED;
+}
+
+static inline void static_key_clear_linked(struct static_key *key)
+{
+	key->type &= ~JUMP_TYPE_LINKED;
+}
+
+static inline void static_key_set_linked(struct static_key *key)
+{
+	key->type |= JUMP_TYPE_LINKED;
 }
 
 static inline struct static_key *jump_entry_key(struct jump_entry *entry)
@@ -254,6 +269,25 @@ static bool jump_entry_branch(struct jump_entry *entry)
 	return (unsigned long)entry->key & 1UL;
 }
 
+/***
+ * A 'struct static_key' uses a union such that it either points directly
+ * to a table of 'struct jump_entry' or to a linked list of modules, which in
+ * turn point to 'struct jump_entry' tables. Since most static_key uses occur
+ * within the module in which they are defined, this saves space. Since the
+ * two lower bits of the pointer are used to keep track of which pointer type
+ * is in use and to store the initial branch direction, we need a special
+ * access function which preserves these bits.
+ */
+static void static_key_set_entries(struct static_key *key,
+				   struct jump_entry *entries)
+{
+	unsigned long type;
+
+	type = key->type & JUMP_TYPE_MASK;
+	key->entries = entries;
+	key->type |= type;
+}
+
 static enum jump_label_type jump_label_type(struct jump_entry *entry)
 {
 	struct static_key *key = jump_entry_key(entry);
@@ -313,13 +347,7 @@ void __init jump_label_init(void)
 			continue;
 
 		key = iterk;
-		/*
-		 * Set key->entries to iter, but preserve JUMP_LABEL_TRUE_BRANCH.
-		 */
-		*((unsigned long *)&key->entries) += (unsigned long)iter;
-#ifdef CONFIG_MODULES
-		key->next = NULL;
-#endif
+		static_key_set_entries(key, iter);
 	}
 	static_key_initialized = true;
 	jump_label_unlock();
@@ -343,6 +371,22 @@ struct static_key_mod {
 	struct module *mod;
 };
 
+static inline struct static_key_mod *static_key_mod(struct static_key *key)
+{
+	return (struct static_key_mod *)(key->type & ~JUMP_TYPE_MASK);
+}
+
+/* See comments above static_key_set_entries() */
+static void static_key_set_mod(struct static_key *key,
+			       struct static_key_mod *mod)
+{
+	unsigned long type;
+
+	type = static_key_type(key);
+	key->next = mod;
+	key->type |= type;
+}
+
 static int __jump_label_mod_text_reserved(void *start, void *end)
 {
 	struct module *mod;
@@ -363,13 +407,19 @@ static int __jump_label_mod_text_reserved(void *start, void *end)
 
 static void __jump_label_mod_update(struct static_key *key)
 {
-	struct static_key_mod *mod;
+	struct static_key_mod *mod = static_key_mod(key);
 
-	for (mod = key->next; mod; mod = mod->next) {
+	while (mod) {
+		struct jump_entry *stop;
 		struct module *m = mod->mod;
 
-		__jump_label_update(key, mod->entries,
-				    m->jump_entries + m->num_jump_entries);
+		if (!m)
+			stop = __stop___jump_table;
+		else
+			stop = m->jump_entries + m->num_jump_entries;
+		if (mod->entries)
+			__jump_label_update(key, mod->entries, stop);
+		mod = mod->next;
 	}
 }
 
@@ -404,7 +454,7 @@ static int jump_label_add_module(struct module *mod)
 	struct jump_entry *iter_stop = iter_start + mod->num_jump_entries;
 	struct jump_entry *iter;
 	struct static_key *key = NULL;
-	struct static_key_mod *jlm;
+	struct static_key_mod *jlm, *jlm2;
 
 	/* if the module doesn't have jump label entries, just return */
 	if (iter_start == iter_stop)
@@ -421,20 +471,32 @@ static int jump_label_add_module(struct module *mod)
 
 		key = iterk;
 		if (within_module(iter->key, mod)) {
-			/*
-			 * Set key->entries to iter, but preserve JUMP_LABEL_TRUE_BRANCH.
-			 */
-			*((unsigned long *)&key->entries) += (unsigned long)iter;
-			key->next = NULL;
+			static_key_set_entries(key, iter);
 			continue;
 		}
 		jlm = kzalloc(sizeof(struct static_key_mod), GFP_KERNEL);
 		if (!jlm)
 			return -ENOMEM;
+		if (!static_key_linked(key)) {
+			jlm2 = kzalloc(sizeof(struct static_key_mod),
+				       GFP_KERNEL);
+			if (!jlm2) {
+				kfree(jlm);
+				return -ENOMEM;
+			}
+			preempt_disable();
+			jlm2->mod = __module_address((unsigned long)key);
+			preempt_enable();
+			jlm2->entries = static_key_entries(key);
+			jlm2->next = NULL;
+			static_key_set_mod(key, jlm2);
+			static_key_set_linked(key);
+		}
 		jlm->mod = mod;
 		jlm->entries = iter;
-		jlm->next = key->next;
-		key->next = jlm;
+		jlm->next = static_key_mod(key);
+		static_key_set_mod(key, jlm);
+		static_key_set_linked(key);
 
 		/* Only update if we've changed from our initial state */
 		if (jump_label_type(iter) != jump_label_init_type(iter))
@@ -462,15 +524,28 @@ static void jump_label_del_module(struct module *mod)
 			continue;
 
 		prev = &key->next;
-		jlm = key->next;
+		jlm = static_key_mod(key);
 
 		while (jlm && jlm->mod != mod) {
 			prev = &jlm->next;
 			jlm = jlm->next;
 		}
 
-		if (jlm) {
+		if (!jlm)
+			continue;
+
+		if (prev == &key->next)
+			static_key_set_mod(key, jlm->next);
+		else
 			*prev = jlm->next;
+
+		kfree(jlm);
+
+		jlm = static_key_mod(key);
+		/* if only one etry is left, fold it back into the static_key */
+		if (jlm->next == NULL) {
+			static_key_set_entries(key, jlm->entries);
+			static_key_clear_linked(key);
 			kfree(jlm);
 		}
 	}
@@ -565,7 +640,10 @@ static void jump_label_update(struct static_key *key)
 #ifdef CONFIG_MODULES
 	struct module *mod;
 
-	__jump_label_mod_update(key);
+	if (static_key_linked(key)) {
+		__jump_label_mod_update(key);
+		return;
+	}
 
 	preempt_disable();
 	mod = __module_address((unsigned long)key);
-- 
2.6.1

Re: [PATCH v3] jump_label: reduce the size of struct static_key

[Steven Rostedt]

On Tue, 31 Jan 2017 17:00:43 -0500
Jason Baron <jbaron@akamai.com> wrote:

> ---
>  Documentation/static-keys.txt |   4 +-
>  include/linux/jump_label.h    |  23 +++++---
>  kernel/jump_label.c           | 126 ++++++++++++++++++++++++++++++++++--------
>  3 files changed, 120 insertions(+), 33 deletions(-)
> 
> diff --git a/Documentation/static-keys.txt b/Documentation/static-keys.txt
> index ea8d7b4e53f0..32a25fad0c1b 100644
> --- a/Documentation/static-keys.txt
> +++ b/Documentation/static-keys.txt
> @@ -155,7 +155,9 @@ or:
>  
>  There are a few functions and macros that architectures must implement in order
>  to take advantage of this optimization. If there is no architecture support, we
> -simply fall back to a traditional, load, test, and jump sequence.
> +simply fall back to a traditional, load, test, and jump sequence. Also, the
> +struct jump_entry table must be at least 4-byte aligned because the
> +static_key->entry field makes use of the two least significant bits.
>  
>  * select HAVE_ARCH_JUMP_LABEL, see: arch/x86/Kconfig
>  
> diff --git a/include/linux/jump_label.h b/include/linux/jump_label.h
> index a0547c571800..680c98b2f41c 100644
> --- a/include/linux/jump_label.h
> +++ b/include/linux/jump_label.h
> @@ -89,11 +89,17 @@ extern bool static_key_initialized;
>  
>  struct static_key {
>  	atomic_t enabled;
> -/* Set lsb bit to 1 if branch is default true, 0 ot */
> -	struct jump_entry *entries;
> -#ifdef CONFIG_MODULES
> -	struct static_key_mod *next;
> -#endif
> +/*
> + * bit 0 => 1 if key is initially true
> + *	    0 if initially false
> + * bit 1 => 1 if points to struct static_key_mod
> + *	    0 if points to struct jump_entry
> + */
> +	union {
> +		unsigned long type;
> +		struct jump_entry *entries;
> +		struct static_key_mod *next;
> +	};
>  };
>  
>  #else
> @@ -118,9 +124,10 @@ struct module;
>  
>  #ifdef HAVE_JUMP_LABEL
>  
> -#define JUMP_TYPE_FALSE	0UL
> -#define JUMP_TYPE_TRUE	1UL
> -#define JUMP_TYPE_MASK	1UL
> +#define JUMP_TYPE_FALSE		0UL
> +#define JUMP_TYPE_TRUE		1UL
> +#define JUMP_TYPE_LINKED	2UL
> +#define JUMP_TYPE_MASK		3UL
>  
>  static __always_inline bool static_key_false(struct static_key *key)
>  {
> diff --git a/kernel/jump_label.c b/kernel/jump_label.c
> index a9b8cf500591..2eb9e80dc691 100644
> --- a/kernel/jump_label.c
> +++ b/kernel/jump_label.c
> @@ -236,12 +236,27 @@ void __weak __init_or_module arch_jump_label_transform_static(struct jump_entry
>  
>  static inline struct jump_entry *static_key_entries(struct static_key *key)
>  {
> -	return (struct jump_entry *)((unsigned long)key->entries & ~JUMP_TYPE_MASK);
> +	return (struct jump_entry *)(key->type & ~JUMP_TYPE_MASK);
>  }
>  
>  static inline bool static_key_type(struct static_key *key)
>  {
> -	return (unsigned long)key->entries & JUMP_TYPE_MASK;
> +	return key->type & JUMP_TYPE_TRUE;
> +}
> +
> +static inline bool static_key_linked(struct static_key *key)
> +{
> +	return key->type & JUMP_TYPE_LINKED;
> +}
> +
> +static inline void static_key_clear_linked(struct static_key *key)
> +{
> +	key->type &= ~JUMP_TYPE_LINKED;
> +}
> +
> +static inline void static_key_set_linked(struct static_key *key)
> +{
> +	key->type |= JUMP_TYPE_LINKED;
>  }
>  
>  static inline struct static_key *jump_entry_key(struct jump_entry *entry)
> @@ -254,6 +269,25 @@ static bool jump_entry_branch(struct jump_entry *entry)
>  	return (unsigned long)entry->key & 1UL;
>  }
>  
> +/***
> + * A 'struct static_key' uses a union such that it either points directly
> + * to a table of 'struct jump_entry' or to a linked list of modules, which in
> + * turn point to 'struct jump_entry' tables. Since most static_key uses occur

I read that three times thinking that the grammar was incorrect, but I
finally figured out that it's not. But the issue is that darn comma
between "modules" and "which". With the comma, it causes one to think
the subject of the sentence is "a linked list of modules" (singular)
where we would use "in turn points to". But if you take out that comma,
it becomes obvious that the "point" is for the "modules" (plural), and
then it makes sense again.

"or a linked list of modules which in turn point to"

See, it reads better that way. Please nuke that comma.

> + * within the module in which they are defined, this saves space. Since the

Also, please nuke the "Since most static_key uses occur..." sentence.
It just reads funny, and I don't think it adds much new info. It's
pretty obvious that we do things like this to save space.

> + * two lower bits of the pointer are used to keep track of which pointer type
> + * is in use and to store the initial branch direction, we need a special
> + * access function which preserves these bits.
> + */

Also, make the "Since the two lower" a new paragraph. It's a separate
topic. And also nuke the "Since", it's bad form ;-)

"The two lower bits of the pointer are used to keep track of which
pointer type is in use and to store the initial branch direction, we
use a special access function which preserves these bits."


> +static void static_key_set_entries(struct static_key *key,
> +				   struct jump_entry *entries)
> +{
> +	unsigned long type;
> +
> +	type = key->type & JUMP_TYPE_MASK;
> +	key->entries = entries;
> +	key->type |= type;
> +}
> +
>  static enum jump_label_type jump_label_type(struct jump_entry *entry)
>  {
>  	struct static_key *key = jump_entry_key(entry);
> @@ -313,13 +347,7 @@ void __init jump_label_init(void)
>  			continue;
>  
>  		key = iterk;
> -		/*
> -		 * Set key->entries to iter, but preserve JUMP_LABEL_TRUE_BRANCH.
> -		 */
> -		*((unsigned long *)&key->entries) += (unsigned long)iter;
> -#ifdef CONFIG_MODULES
> -		key->next = NULL;
> -#endif
> +		static_key_set_entries(key, iter);

The new code does make sense, but I'm trying to see how the old code
works. How was the "+=" an assignment? I'm assuming that as this is
init code, that the key->entries was only assigned as 0 or 1 (depending
on if it was true or false). Looking around, this does seem to be the
case. But as a reviewer, I still like to have that warm cozy feeling of
confirmation ;-)


>  	}
>  	static_key_initialized = true;
>  	jump_label_unlock();
> @@ -343,6 +371,22 @@ struct static_key_mod {
>  	struct module *mod;
>  };
>  
> +static inline struct static_key_mod *static_key_mod(struct static_key *key)
> +{
> +	return (struct static_key_mod *)(key->type & ~JUMP_TYPE_MASK);
> +}
> +
> +/* See comments above static_key_set_entries() */
> +static void static_key_set_mod(struct static_key *key,
> +			       struct static_key_mod *mod)
> +{
> +	unsigned long type;

Even though you have the "see comments", it may still be helpful to
have something simple here. Like:

	/*
	 * key->type and key->next are the same via union.
	 * This sets key->next and preserves the type bits.
	 */

> +
> +	type = static_key_type(key);
> +	key->next = mod;
> +	key->type |= type;
> +}
> +
>  static int __jump_label_mod_text_reserved(void *start, void *end)
>  {
>  	struct module *mod;
> @@ -363,13 +407,19 @@ static int __jump_label_mod_text_reserved(void *start, void *end)
>  
>  static void __jump_label_mod_update(struct static_key *key)
>  {
> -	struct static_key_mod *mod;
> +	struct static_key_mod *mod = static_key_mod(key);
>  
> -	for (mod = key->next; mod; mod = mod->next) {
> +	while (mod) {

Maybe keep the for loop?

	for (mod = static_key_mod(key); mod; mod->next) {


It's still nicely compact.

Also, with a for loop, this can be changed to:

	if (!mod->entries)
		continue;

	if (!m)
		....

Speaking of which ...

> +		struct jump_entry *stop;
>  		struct module *m = mod->mod;
>  
> -		__jump_label_update(key, mod->entries,
> -				    m->jump_entries + m->num_jump_entries);
> +		if (!m)
> +			stop = __stop___jump_table;
> +		else
> +			stop = m->jump_entries + m->num_jump_entries;

How did this work before, if mod->entries was NULL? Or does this have
to do with this being now called by the module itself if linked. But
I'm not sure how that works either.

> +		if (mod->entries)
> +			__jump_label_update(key, mod->entries, stop);
> +		mod = mod->next;
>  	}
>  }
>  
> @@ -404,7 +454,7 @@ static int jump_label_add_module(struct module *mod)
>  	struct jump_entry *iter_stop = iter_start + mod->num_jump_entries;
>  	struct jump_entry *iter;
>  	struct static_key *key = NULL;
> -	struct static_key_mod *jlm;
> +	struct static_key_mod *jlm, *jlm2;
>  
>  	/* if the module doesn't have jump label entries, just return */
>  	if (iter_start == iter_stop)
> @@ -421,20 +471,32 @@ static int jump_label_add_module(struct module *mod)
>  
>  		key = iterk;
>  		if (within_module(iter->key, mod)) {
> -			/*
> -			 * Set key->entries to iter, but preserve JUMP_LABEL_TRUE_BRANCH.
> -			 */
> -			*((unsigned long *)&key->entries) += (unsigned long)iter;
> -			key->next = NULL;
> +			static_key_set_entries(key, iter);
>  			continue;
>  		}
>  		jlm = kzalloc(sizeof(struct static_key_mod), GFP_KERNEL);
>  		if (!jlm)
>  			return -ENOMEM;
> +		if (!static_key_linked(key)) {
> +			jlm2 = kzalloc(sizeof(struct static_key_mod),
> +				       GFP_KERNEL);
> +			if (!jlm2) {
> +				kfree(jlm);
> +				return -ENOMEM;
> +			}
> +			preempt_disable();
> +			jlm2->mod = __module_address((unsigned long)key);
> +			preempt_enable();
> +			jlm2->entries = static_key_entries(key);
> +			jlm2->next = NULL;
> +			static_key_set_mod(key, jlm2);
> +			static_key_set_linked(key);
> +		}
>  		jlm->mod = mod;
>  		jlm->entries = iter;
> -		jlm->next = key->next;
> -		key->next = jlm;

I wonder if a comment would be useful here:

		/*
		 * jlm->next = key->next;
		 * key->next = jlm;
		 */

Just because the encapsulated functions don't show this basic exchange
well. It tripped me up for a moment.

> +		jlm->next = static_key_mod(key);
> +		static_key_set_mod(key, jlm);
> +		static_key_set_linked(key);
>  
>  		/* Only update if we've changed from our initial state */
>  		if (jump_label_type(iter) != jump_label_init_type(iter))
> @@ -462,15 +524,28 @@ static void jump_label_del_module(struct module *mod)
>  			continue;
>  
>  		prev = &key->next;
> -		jlm = key->next;
> +		jlm = static_key_mod(key);
>  
>  		while (jlm && jlm->mod != mod) {
>  			prev = &jlm->next;
>  			jlm = jlm->next;
>  		}
>  

BTW, can !jlm ever happen? I guess it can if it failed to allocate :-/
Maybe a comment about that? And maybe even a WARN()?

-- Steve

> -		if (jlm) {
> +		if (!jlm)
> +			continue;
> +
> +		if (prev == &key->next)
> +			static_key_set_mod(key, jlm->next);
> +		else
>  			*prev = jlm->next;
> +
> +		kfree(jlm);
> +
> +		jlm = static_key_mod(key);
> +		/* if only one etry is left, fold it back into the static_key */
> +		if (jlm->next == NULL) {
> +			static_key_set_entries(key, jlm->entries);
> +			static_key_clear_linked(key);
>  			kfree(jlm);
>  		}
>  	}
> @@ -565,7 +640,10 @@ static void jump_label_update(struct static_key *key)
>  #ifdef CONFIG_MODULES
>  	struct module *mod;
>  
> -	__jump_label_mod_update(key);
> +	if (static_key_linked(key)) {
> +		__jump_label_mod_update(key);
> +		return;
> +	}
>  
>  	preempt_disable();
>  	mod = __module_address((unsigned long)key);

Re: [PATCH v3] jump_label: reduce the size of struct static_key

[Jason Baron]

Hi Steve,

On 01/31/2017 07:32 PM, Steven Rostedt wrote:
> On Tue, 31 Jan 2017 17:00:43 -0500
> Jason Baron <jbaron@akamai.com> wrote:
> 
>> ---
>>  Documentation/static-keys.txt |   4 +-
>>  include/linux/jump_label.h    |  23 +++++---
>>  kernel/jump_label.c           | 126 ++++++++++++++++++++++++++++++++++--------
>>  3 files changed, 120 insertions(+), 33 deletions(-)
>>
>> diff --git a/Documentation/static-keys.txt b/Documentation/static-keys.txt
>> index ea8d7b4e53f0..32a25fad0c1b 100644
>> --- a/Documentation/static-keys.txt
>> +++ b/Documentation/static-keys.txt
>> @@ -155,7 +155,9 @@ or:
>>  
>>  There are a few functions and macros that architectures must implement in order
>>  to take advantage of this optimization. If there is no architecture support, we
>> -simply fall back to a traditional, load, test, and jump sequence.
>> +simply fall back to a traditional, load, test, and jump sequence. Also, the
>> +struct jump_entry table must be at least 4-byte aligned because the
>> +static_key->entry field makes use of the two least significant bits.
>>  
>>  * select HAVE_ARCH_JUMP_LABEL, see: arch/x86/Kconfig
>>  
>> diff --git a/include/linux/jump_label.h b/include/linux/jump_label.h
>> index a0547c571800..680c98b2f41c 100644
>> --- a/include/linux/jump_label.h
>> +++ b/include/linux/jump_label.h
>> @@ -89,11 +89,17 @@ extern bool static_key_initialized;
>>  
>>  struct static_key {
>>  	atomic_t enabled;
>> -/* Set lsb bit to 1 if branch is default true, 0 ot */
>> -	struct jump_entry *entries;
>> -#ifdef CONFIG_MODULES
>> -	struct static_key_mod *next;
>> -#endif
>> +/*
>> + * bit 0 => 1 if key is initially true
>> + *	    0 if initially false
>> + * bit 1 => 1 if points to struct static_key_mod
>> + *	    0 if points to struct jump_entry
>> + */
>> +	union {
>> +		unsigned long type;
>> +		struct jump_entry *entries;
>> +		struct static_key_mod *next;
>> +	};
>>  };
>>  
>>  #else
>> @@ -118,9 +124,10 @@ struct module;
>>  
>>  #ifdef HAVE_JUMP_LABEL
>>  
>> -#define JUMP_TYPE_FALSE	0UL
>> -#define JUMP_TYPE_TRUE	1UL
>> -#define JUMP_TYPE_MASK	1UL
>> +#define JUMP_TYPE_FALSE		0UL
>> +#define JUMP_TYPE_TRUE		1UL
>> +#define JUMP_TYPE_LINKED	2UL
>> +#define JUMP_TYPE_MASK		3UL
>>  
>>  static __always_inline bool static_key_false(struct static_key *key)
>>  {
>> diff --git a/kernel/jump_label.c b/kernel/jump_label.c
>> index a9b8cf500591..2eb9e80dc691 100644
>> --- a/kernel/jump_label.c
>> +++ b/kernel/jump_label.c
>> @@ -236,12 +236,27 @@ void __weak __init_or_module arch_jump_label_transform_static(struct jump_entry
>>  
>>  static inline struct jump_entry *static_key_entries(struct static_key *key)
>>  {
>> -	return (struct jump_entry *)((unsigned long)key->entries & ~JUMP_TYPE_MASK);
>> +	return (struct jump_entry *)(key->type & ~JUMP_TYPE_MASK);
>>  }
>>  
>>  static inline bool static_key_type(struct static_key *key)
>>  {
>> -	return (unsigned long)key->entries & JUMP_TYPE_MASK;
>> +	return key->type & JUMP_TYPE_TRUE;
>> +}
>> +
>> +static inline bool static_key_linked(struct static_key *key)
>> +{
>> +	return key->type & JUMP_TYPE_LINKED;
>> +}
>> +
>> +static inline void static_key_clear_linked(struct static_key *key)
>> +{
>> +	key->type &= ~JUMP_TYPE_LINKED;
>> +}
>> +
>> +static inline void static_key_set_linked(struct static_key *key)
>> +{
>> +	key->type |= JUMP_TYPE_LINKED;
>>  }
>>  
>>  static inline struct static_key *jump_entry_key(struct jump_entry *entry)
>> @@ -254,6 +269,25 @@ static bool jump_entry_branch(struct jump_entry *entry)
>>  	return (unsigned long)entry->key & 1UL;
>>  }
>>  
>> +/***
>> + * A 'struct static_key' uses a union such that it either points directly
>> + * to a table of 'struct jump_entry' or to a linked list of modules, which in
>> + * turn point to 'struct jump_entry' tables. Since most static_key uses occur
> 
> I read that three times thinking that the grammar was incorrect, but I
> finally figured out that it's not. But the issue is that darn comma
> between "modules" and "which". With the comma, it causes one to think
> the subject of the sentence is "a linked list of modules" (singular)
> where we would use "in turn points to". But if you take out that comma,
> it becomes obvious that the "point" is for the "modules" (plural), and
> then it makes sense again.
> 
> "or a linked list of modules which in turn point to"
> 
> See, it reads better that way. Please nuke that comma.
> 
>> + * within the module in which they are defined, this saves space. Since the
> 
> Also, please nuke the "Since most static_key uses occur..." sentence.
> It just reads funny, and I don't think it adds much new info. It's
> pretty obvious that we do things like this to save space.
> 
>> + * two lower bits of the pointer are used to keep track of which pointer type
>> + * is in use and to store the initial branch direction, we need a special
>> + * access function which preserves these bits.
>> + */
> 
> Also, make the "Since the two lower" a new paragraph. It's a separate
> topic. And also nuke the "Since", it's bad form ;-)
> 
> "The two lower bits of the pointer are used to keep track of which
> pointer type is in use and to store the initial branch direction, we
> use a special access function which preserves these bits."
> 
>

ok


>> +static void static_key_set_entries(struct static_key *key,
>> +				   struct jump_entry *entries)
>> +{
>> +	unsigned long type;
>> +
>> +	type = key->type & JUMP_TYPE_MASK;
>> +	key->entries = entries;
>> +	key->type |= type;
>> +}
>> +
>>  static enum jump_label_type jump_label_type(struct jump_entry *entry)
>>  {
>>  	struct static_key *key = jump_entry_key(entry);
>> @@ -313,13 +347,7 @@ void __init jump_label_init(void)
>>  			continue;
>>  
>>  		key = iterk;
>> -		/*
>> -		 * Set key->entries to iter, but preserve JUMP_LABEL_TRUE_BRANCH.
>> -		 */
>> -		*((unsigned long *)&key->entries) += (unsigned long)iter;
>> -#ifdef CONFIG_MODULES
>> -		key->next = NULL;
>> -#endif
>> +		static_key_set_entries(key, iter);
> 
> The new code does make sense, but I'm trying to see how the old code
> works. How was the "+=" an assignment? I'm assuming that as this is
> init code, that the key->entries was only assigned as 0 or 1 (depending
> on if it was true or false). Looking around, this does seem to be the
> case. But as a reviewer, I still like to have that warm cozy feeling of
> confirmation ;-)
> 

Yes, it's either 0 or 1 initially depending on the initial value (true
or false, the 'linked' bit can not be set at this point). So, I could
have left it, but since we have this fancy new static_key_set_entries(),
I thought it would look better :)

> 
>>  	}
>>  	static_key_initialized = true;
>>  	jump_label_unlock();
>> @@ -343,6 +371,22 @@ struct static_key_mod {
>>  	struct module *mod;
>>  };
>>  
>> +static inline struct static_key_mod *static_key_mod(struct static_key *key)
>> +{
>> +	return (struct static_key_mod *)(key->type & ~JUMP_TYPE_MASK);
>> +}
>> +
>> +/* See comments above static_key_set_entries() */
>> +static void static_key_set_mod(struct static_key *key,
>> +			       struct static_key_mod *mod)
>> +{
>> +	unsigned long type;
> 
> Even though you have the "see comments", it may still be helpful to
> have something simple here. Like:
> 
> 	/*
> 	 * key->type and key->next are the same via union.
> 	 * This sets key->next and preserves the type bits.
> 	 */
> 

ok

>> +
>> +	type = static_key_type(key);
>> +	key->next = mod;
>> +	key->type |= type;
>> +}
>> +
>>  static int __jump_label_mod_text_reserved(void *start, void *end)
>>  {
>>  	struct module *mod;
>> @@ -363,13 +407,19 @@ static int __jump_label_mod_text_reserved(void *start, void *end)
>>  
>>  static void __jump_label_mod_update(struct static_key *key)
>>  {
>> -	struct static_key_mod *mod;
>> +	struct static_key_mod *mod = static_key_mod(key);
>>  
>> -	for (mod = key->next; mod; mod = mod->next) {
>> +	while (mod) {
> 
> Maybe keep the for loop?
> 
> 	for (mod = static_key_mod(key); mod; mod->next) {
> 
> 
> It's still nicely compact.
> 
> Also, with a for loop, this can be changed to:
> 
> 	if (!mod->entries)
> 		continue;
> 
> 	if (!m)
> 		....
> 

ok, let's go back to the 'for' loop.

> Speaking of which ...
> 
>> +		struct jump_entry *stop;
>>  		struct module *m = mod->mod;
>>  
>> -		__jump_label_update(key, mod->entries,
>> -				    m->jump_entries + m->num_jump_entries);
>> +		if (!m)
>> +			stop = __stop___jump_table;
>> +		else
>> +			stop = m->jump_entries + m->num_jump_entries;
> 
> How did this work before, if mod->entries was NULL? Or does this have
> to do with this being now called by the module itself if linked. But
> I'm not sure how that works either.
> 

So previously, mod->entries could not have been NULL. When a static_key
is defined in a module or the core kernel but has no users in the module
in which it is defined then static_key->entries can be NULL. However,
mod->entries could not be NULL, because we only allocated mod if there
were uses.

With this patch though in the case where static_key->entries is NULL,
and a new module is added that references the static_key, we now
allocate a new 'mod', to point the entries that may be associated with
the initial static key. In the case where it is NULL, we really didn't
need to allocate this new module (and in fact in an earlier version I
didn't). I ended up doing it this way because it made 'add_module' and
'del_module' code slightly simpler and because it would be a little
strange for this case to exist (I'm not sure it does), but it is possible.

Note that 'jump_label_update()' already does:

        /* if there are no users, entry can be NULL */
        if (entry)
                __jump_label_update(key, entry, stop);

Thus, prior to this patch, we had this case for the static_key itself,
and with the introduction of this patch, we now also have it for the 'mod'.


>> +		if (mod->entries)
>> +			__jump_label_update(key, mod->entries, stop);
>> +		mod = mod->next;
>>  	}
>>  }
>>  
>> @@ -404,7 +454,7 @@ static int jump_label_add_module(struct module *mod)
>>  	struct jump_entry *iter_stop = iter_start + mod->num_jump_entries;
>>  	struct jump_entry *iter;
>>  	struct static_key *key = NULL;
>> -	struct static_key_mod *jlm;
>> +	struct static_key_mod *jlm, *jlm2;
>>  
>>  	/* if the module doesn't have jump label entries, just return */
>>  	if (iter_start == iter_stop)
>> @@ -421,20 +471,32 @@ static int jump_label_add_module(struct module *mod)
>>  
>>  		key = iterk;
>>  		if (within_module(iter->key, mod)) {
>> -			/*
>> -			 * Set key->entries to iter, but preserve JUMP_LABEL_TRUE_BRANCH.
>> -			 */
>> -			*((unsigned long *)&key->entries) += (unsigned long)iter;
>> -			key->next = NULL;
>> +			static_key_set_entries(key, iter);
>>  			continue;
>>  		}
>>  		jlm = kzalloc(sizeof(struct static_key_mod), GFP_KERNEL);
>>  		if (!jlm)
>>  			return -ENOMEM;
>> +		if (!static_key_linked(key)) {
>> +			jlm2 = kzalloc(sizeof(struct static_key_mod),
>> +				       GFP_KERNEL);
>> +			if (!jlm2) {
>> +				kfree(jlm);
>> +				return -ENOMEM;
>> +			}
>> +			preempt_disable();
>> +			jlm2->mod = __module_address((unsigned long)key);
>> +			preempt_enable();
>> +			jlm2->entries = static_key_entries(key);
>> +			jlm2->next = NULL;
>> +			static_key_set_mod(key, jlm2);
>> +			static_key_set_linked(key);
>> +		}
>>  		jlm->mod = mod;
>>  		jlm->entries = iter;
>> -		jlm->next = key->next;
>> -		key->next = jlm;
> 
> I wonder if a comment would be useful here:
> 
> 		/*
> 		 * jlm->next = key->next;
> 		 * key->next = jlm;
> 		 */
> 
> Just because the encapsulated functions don't show this basic exchange
> well. It tripped me up for a moment.
> 

I feel like this is more an artifact of the transformation here, but I
can add a comment here, sure.

>> +		jlm->next = static_key_mod(key);
>> +		static_key_set_mod(key, jlm);
>> +		static_key_set_linked(key);
>>  
>>  		/* Only update if we've changed from our initial state */
>>  		if (jump_label_type(iter) != jump_label_init_type(iter))
>> @@ -462,15 +524,28 @@ static void jump_label_del_module(struct module *mod)
>>  			continue;
>>  
>>  		prev = &key->next;
>> -		jlm = key->next;
>> +		jlm = static_key_mod(key);
>>  
>>  		while (jlm && jlm->mod != mod) {
>>  			prev = &jlm->next;
>>  			jlm = jlm->next;
>>  		}
>>  
> 
> BTW, can !jlm ever happen? I guess it can if it failed to allocate :-/
> Maybe a comment about that? And maybe even a WARN()?
> 

Yes, but I think it makes more sense to add that in the allocation path,
so that if a branch is not being switched we have some hint.

Actually, I think the better thing is that
blocking_notifier_call_chain() in kernel/module.c should check its
return value and refuse to load the module. It looks like
prepare_coming_module() already deals with returning an error, but
the return of blocking_notifier_call_chain() is not part of that error
path currently. Although this could be a can of worms since other
users on this chain may not want this...

Thanks,

-Jason


> -- Steve
> 
>> -		if (jlm) {
>> +		if (!jlm)
>> +			continue;
>> +
>> +		if (prev == &key->next)
>> +			static_key_set_mod(key, jlm->next);
>> +		else
>>  			*prev = jlm->next;
>> +
>> +		kfree(jlm);
>> +
>> +		jlm = static_key_mod(key);
>> +		/* if only one etry is left, fold it back into the static_key */
>> +		if (jlm->next == NULL) {
>> +			static_key_set_entries(key, jlm->entries);
>> +			static_key_clear_linked(key);
>>  			kfree(jlm);
>>  		}
>>  	}
>> @@ -565,7 +640,10 @@ static void jump_label_update(struct static_key *key)
>>  #ifdef CONFIG_MODULES
>>  	struct module *mod;
>>  
>> -	__jump_label_mod_update(key);
>> +	if (static_key_linked(key)) {
>> +		__jump_label_mod_update(key);
>> +		return;
>> +	}
>>  
>>  	preempt_disable();
>>  	mod = __module_address((unsigned long)key);
> 

Re: [PATCH v3] jump_label: reduce the size of struct static_key

[Steven Rostedt]

On Tue, 31 Jan 2017 17:00:43 -0500
Jason Baron <jbaron@akamai.com> wrote:

> --- a/kernel/jump_label.c
> +++ b/kernel/jump_label.c
> @@ -236,12 +236,27 @@ void __weak __init_or_module arch_jump_label_transform_static(struct jump_entry
>  
>  static inline struct jump_entry *static_key_entries(struct static_key *key)
>  {
> -	return (struct jump_entry *)((unsigned long)key->entries & ~JUMP_TYPE_MASK);

I wonder if we should add:

	WARN_ON_ONCE(key->type & JUMP_TYPE_LINKED);

> +	return (struct jump_entry *)(key->type & ~JUMP_TYPE_MASK);
>  }
>  

[...]

>  
> +static inline struct static_key_mod *static_key_mod(struct static_key *key)
> +{

	WARN_ON_ONCE(!(key->type & JUMP_TYPE_LINKED));

> +	return (struct static_key_mod *)(key->type & ~JUMP_TYPE_MASK);
> +}
> +

-- Steve

[PATCH v4] jump_label: reduce the size of struct static_key

[Jason Baron]

The static_key->next field goes mostly unused. The field is used for
associating module uses with a static key. Most uses of struct static_key
define a static key in the core kernel and make use of it entirely within
the core kernel, or define the static key in a module and make use of it
only from within that module. In fact, of the ~3,000 static keys defined,
I found only about 5 or so that did not fit this pattern.

Thus, we can remove the static_key->next field entirely and overload
the static_key->entries field. That is, when all the static_key uses
are contained within the same module, static_key->entries continues
to point to those uses. However, if the static_key uses are not contained
within the module where the static_key is defined, then we allocate a
struct static_key_mod, store a pointer to the uses within that
struct static_key_mod, and have the static key point at the static_key_mod.
This does incur some extra memory usage when a static_key is used in a
module that does not define it, but since there are only a handful of such
cases there is a net savings.

In order to identify if the static_key->entries pointer contains a
struct static_key_mod or a struct jump_entry pointer, bit 1 of
static_key->entries is set to 1 if it points to a struct static_key_mod and
is 0 if it points to a struct jump_entry. We were already using bit 0 in a
similar way to store the initial value of the static_key. This does mean
that allocations of struct static_key_mod and that the struct jump_entry
tables need to be at least 4-byte aligned in memory. As far as I can tell
all arches meet this criteria.

For my .config, the patch increased the text by 778 bytes, but reduced
the data + bss size by 14912, for a net savings of 14,134 bytes.

   text	   data	    bss	    dec	    hex	filename
8092427	5016512	 790528	13899467	 d416cb	vmlinux.pre
8093205	5001600	 790528	13885333	 d3df95	vmlinux.post

Cc: Peter Zijlstra <peterz@infradead.org>
Cc: Steven Rostedt <rostedt@goodmis.org>
Cc: Ingo Molnar <mingo@kernel.org>
Cc: Joe Perches <joe@perches.com>
Signed-off-by: Jason Baron <jbaron@akamai.com>
---
Changed in v4:
-added WARN*() invariants
-fixed up comments
-fixed bug in static_key_set_mod()
Changed in v3:
-added static_key_set_[entries|mod]() to clean up type casting (Ingo Molnar)
Changed in v2:
-Replace static_key->entries with union (Steven Rostedt)
---
 Documentation/static-keys.txt |   4 +-
 include/linux/jump_label.h    |  23 ++++---
 kernel/jump_label.c           | 153 +++++++++++++++++++++++++++++++++++-------
 3 files changed, 145 insertions(+), 35 deletions(-)

diff --git a/Documentation/static-keys.txt b/Documentation/static-keys.txt
index ea8d7b4e53f0..32a25fad0c1b 100644
--- a/Documentation/static-keys.txt
+++ b/Documentation/static-keys.txt
@@ -155,7 +155,9 @@ or:
 
 There are a few functions and macros that architectures must implement in order
 to take advantage of this optimization. If there is no architecture support, we
-simply fall back to a traditional, load, test, and jump sequence.
+simply fall back to a traditional, load, test, and jump sequence. Also, the
+struct jump_entry table must be at least 4-byte aligned because the
+static_key->entry field makes use of the two least significant bits.
 
 * select HAVE_ARCH_JUMP_LABEL, see: arch/x86/Kconfig
 
diff --git a/include/linux/jump_label.h b/include/linux/jump_label.h
index a0547c571800..680c98b2f41c 100644
--- a/include/linux/jump_label.h
+++ b/include/linux/jump_label.h
@@ -89,11 +89,17 @@ extern bool static_key_initialized;
 
 struct static_key {
 	atomic_t enabled;
-/* Set lsb bit to 1 if branch is default true, 0 ot */
-	struct jump_entry *entries;
-#ifdef CONFIG_MODULES
-	struct static_key_mod *next;
-#endif
+/*
+ * bit 0 => 1 if key is initially true
+ *	    0 if initially false
+ * bit 1 => 1 if points to struct static_key_mod
+ *	    0 if points to struct jump_entry
+ */
+	union {
+		unsigned long type;
+		struct jump_entry *entries;
+		struct static_key_mod *next;
+	};
 };
 
 #else
@@ -118,9 +124,10 @@ struct module;
 
 #ifdef HAVE_JUMP_LABEL
 
-#define JUMP_TYPE_FALSE	0UL
-#define JUMP_TYPE_TRUE	1UL
-#define JUMP_TYPE_MASK	1UL
+#define JUMP_TYPE_FALSE		0UL
+#define JUMP_TYPE_TRUE		1UL
+#define JUMP_TYPE_LINKED	2UL
+#define JUMP_TYPE_MASK		3UL
 
 static __always_inline bool static_key_false(struct static_key *key)
 {
diff --git a/kernel/jump_label.c b/kernel/jump_label.c
index a9b8cf500591..6c9cb208ac48 100644
--- a/kernel/jump_label.c
+++ b/kernel/jump_label.c
@@ -236,12 +236,28 @@ void __weak __init_or_module arch_jump_label_transform_static(struct jump_entry
 
 static inline struct jump_entry *static_key_entries(struct static_key *key)
 {
-	return (struct jump_entry *)((unsigned long)key->entries & ~JUMP_TYPE_MASK);
+	WARN_ON_ONCE(key->type & JUMP_TYPE_LINKED);
+	return (struct jump_entry *)(key->type & ~JUMP_TYPE_MASK);
 }
 
 static inline bool static_key_type(struct static_key *key)
 {
-	return (unsigned long)key->entries & JUMP_TYPE_MASK;
+	return key->type & JUMP_TYPE_TRUE;
+}
+
+static inline bool static_key_linked(struct static_key *key)
+{
+	return key->type & JUMP_TYPE_LINKED;
+}
+
+static inline void static_key_clear_linked(struct static_key *key)
+{
+	key->type &= ~JUMP_TYPE_LINKED;
+}
+
+static inline void static_key_set_linked(struct static_key *key)
+{
+	key->type |= JUMP_TYPE_LINKED;
 }
 
 static inline struct static_key *jump_entry_key(struct jump_entry *entry)
@@ -254,6 +270,26 @@ static bool jump_entry_branch(struct jump_entry *entry)
 	return (unsigned long)entry->key & 1UL;
 }
 
+/***
+ * A 'struct static_key' uses a union such that it either points directly
+ * to a table of 'struct jump_entry' or to a linked list of modules which in
+ * turn point to 'struct jump_entry' tables.
+ *
+ * The two lower bits of the pointer are used to keep track of which pointer
+ * type is in use and to store the initial branch direction, we use an access
+ * function which preserves these bits.
+ */
+static void static_key_set_entries(struct static_key *key,
+				   struct jump_entry *entries)
+{
+	unsigned long type;
+
+	WARN_ON_ONCE((unsigned long)entries & JUMP_TYPE_MASK);
+	type = key->type & JUMP_TYPE_MASK;
+	key->entries = entries;
+	key->type |= type;
+}
+
 static enum jump_label_type jump_label_type(struct jump_entry *entry)
 {
 	struct static_key *key = jump_entry_key(entry);
@@ -313,13 +349,7 @@ void __init jump_label_init(void)
 			continue;
 
 		key = iterk;
-		/*
-		 * Set key->entries to iter, but preserve JUMP_LABEL_TRUE_BRANCH.
-		 */
-		*((unsigned long *)&key->entries) += (unsigned long)iter;
-#ifdef CONFIG_MODULES
-		key->next = NULL;
-#endif
+		static_key_set_entries(key, iter);
 	}
 	static_key_initialized = true;
 	jump_label_unlock();
@@ -343,6 +373,29 @@ struct static_key_mod {
 	struct module *mod;
 };
 
+static inline struct static_key_mod *static_key_mod(struct static_key *key)
+{
+	WARN_ON_ONCE(!(key->type & JUMP_TYPE_LINKED));
+	return (struct static_key_mod *)(key->type & ~JUMP_TYPE_MASK);
+}
+
+/***
+ * key->type and key->next are the same via union.
+ * This sets key->next and preserves the type bits.
+ *
+ * See additional comments above static_key_set_entries().
+ */
+static void static_key_set_mod(struct static_key *key,
+			       struct static_key_mod *mod)
+{
+	unsigned long type;
+
+	WARN_ON_ONCE((unsigned long)mod & JUMP_TYPE_MASK);
+	type = key->type & JUMP_TYPE_MASK;
+	key->next = mod;
+	key->type |= type;
+}
+
 static int __jump_label_mod_text_reserved(void *start, void *end)
 {
 	struct module *mod;
@@ -365,11 +418,23 @@ static void __jump_label_mod_update(struct static_key *key)
 {
 	struct static_key_mod *mod;
 
-	for (mod = key->next; mod; mod = mod->next) {
-		struct module *m = mod->mod;
+	for (mod = static_key_mod(key); mod; mod = mod->next) {
+		struct jump_entry *stop;
+		struct module *m;
+
+		/*
+		 * NULL if the static_key is defined in a module
+		 * that does not use it
+		 */
+		if (!mod->entries)
+			continue;
 
-		__jump_label_update(key, mod->entries,
-				    m->jump_entries + m->num_jump_entries);
+		m = mod->mod;
+		if (!m)
+			stop = __stop___jump_table;
+		else
+			stop = m->jump_entries + m->num_jump_entries;
+		__jump_label_update(key, mod->entries, stop);
 	}
 }
 
@@ -404,7 +469,7 @@ static int jump_label_add_module(struct module *mod)
 	struct jump_entry *iter_stop = iter_start + mod->num_jump_entries;
 	struct jump_entry *iter;
 	struct static_key *key = NULL;
-	struct static_key_mod *jlm;
+	struct static_key_mod *jlm, *jlm2;
 
 	/* if the module doesn't have jump label entries, just return */
 	if (iter_start == iter_stop)
@@ -421,20 +486,32 @@ static int jump_label_add_module(struct module *mod)
 
 		key = iterk;
 		if (within_module(iter->key, mod)) {
-			/*
-			 * Set key->entries to iter, but preserve JUMP_LABEL_TRUE_BRANCH.
-			 */
-			*((unsigned long *)&key->entries) += (unsigned long)iter;
-			key->next = NULL;
+			static_key_set_entries(key, iter);
 			continue;
 		}
 		jlm = kzalloc(sizeof(struct static_key_mod), GFP_KERNEL);
 		if (!jlm)
 			return -ENOMEM;
+		if (!static_key_linked(key)) {
+			jlm2 = kzalloc(sizeof(struct static_key_mod),
+				       GFP_KERNEL);
+			if (!jlm2) {
+				kfree(jlm);
+				return -ENOMEM;
+			}
+			preempt_disable();
+			jlm2->mod = __module_address((unsigned long)key);
+			preempt_enable();
+			jlm2->entries = static_key_entries(key);
+			jlm2->next = NULL;
+			static_key_set_mod(key, jlm2);
+			static_key_set_linked(key);
+		}
 		jlm->mod = mod;
 		jlm->entries = iter;
-		jlm->next = key->next;
-		key->next = jlm;
+		jlm->next = static_key_mod(key);
+		static_key_set_mod(key, jlm);
+		static_key_set_linked(key);
 
 		/* Only update if we've changed from our initial state */
 		if (jump_label_type(iter) != jump_label_init_type(iter))
@@ -461,16 +538,34 @@ static void jump_label_del_module(struct module *mod)
 		if (within_module(iter->key, mod))
 			continue;
 
+		/* No memory during module load */
+		if (WARN_ON(!static_key_linked(key)))
+			continue;
+
 		prev = &key->next;
-		jlm = key->next;
+		jlm = static_key_mod(key);
 
 		while (jlm && jlm->mod != mod) {
 			prev = &jlm->next;
 			jlm = jlm->next;
 		}
 
-		if (jlm) {
+		/* No memory during module load */
+		if (WARN_ON(!jlm))
+			continue;
+
+		if (prev == &key->next)
+			static_key_set_mod(key, jlm->next);
+		else
 			*prev = jlm->next;
+
+		kfree(jlm);
+
+		jlm = static_key_mod(key);
+		/* if only one etry is left, fold it back into the static_key */
+		if (jlm->next == NULL) {
+			static_key_set_entries(key, jlm->entries);
+			static_key_clear_linked(key);
 			kfree(jlm);
 		}
 	}
@@ -499,8 +594,10 @@ jump_label_module_notify(struct notifier_block *self, unsigned long val,
 	case MODULE_STATE_COMING:
 		jump_label_lock();
 		ret = jump_label_add_module(mod);
-		if (ret)
+		if (ret) {
+			WARN(1, "Failed to allocatote memory: jump_label may not work properly.\n");
 			jump_label_del_module(mod);
+		}
 		jump_label_unlock();
 		break;
 	case MODULE_STATE_GOING:
@@ -561,11 +658,14 @@ 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;
-	struct jump_entry *entry = static_key_entries(key);
+	struct jump_entry *entry;
 #ifdef CONFIG_MODULES
 	struct module *mod;
 
-	__jump_label_mod_update(key);
+	if (static_key_linked(key)) {
+		__jump_label_mod_update(key);
+		return;
+	}
 
 	preempt_disable();
 	mod = __module_address((unsigned long)key);
@@ -573,6 +673,7 @@ static void jump_label_update(struct static_key *key)
 		stop = mod->jump_entries + mod->num_jump_entries;
 	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);
-- 
2.6.1