[PATCH 00/13] kprobes: Reduce registration and optimization lookup costs

[PATCH 00/13] kprobes: Reduce registration and optimization lookup costs

[Jiri Olsa]

hi,
this series reduces repeated lookup work when registering and preparing
large numbers of kprobes, particularly on x86 with optimized probes.

The kernel already provides interfaces for registering large numbers of
kprobes. Faster registration will become more important for our planned
inline function tracing support, which will use kprobes.

The changes address two kinds of lookup:

  - add indexes sorted by code address for core static-call and jump-label
    reservation checks. The original tables remain sorted by key for
    updates. Build the indexes during late init and retain linear scans
    if allocation fails. Module reservation checks remain linear.

  - add an address-ordered rbtree of registered kprobes and iterators for
    ranges and monotonically increasing addresses. Convert instruction
    recovery and optimization checks to reuse these iterators instead
    of repeatedly looking up individual addresses.

The series also determines x86 jump-label instruction sizes directly
from their opcode or NOP encoding, avoiding instruction decoding.

The final samples patch is included only to demonstrate the testing and
is not intended for merging. It adds 50k functions to the kernel, along
with sample module code that registers kprobes on them, simulating kprobes
on inline kernel functions.

There's overall speedup about 10x for registering the kprobe, please check
the perf stats below. Plus there's one more low hanging fruit leftover I
plan to send later for bug entry lookup that will add bit more speed up.

The test below measures perf stat for the probe_many_example.ko module to
attach kprobes to 20,000 functions.

  Performance counter stats for:
    'insmod ./samples/kprobes/kprobe_many_example.ko nr_functions=20000':

  Before:
    50,908,260,992      cycles:k                                                              
    53,886,723,615      instructions:k                                                        

      10.072721766 seconds time elapsed

  After:
     5,122,978,035      cycles:k                                                              
     5,620,945,319      instructions:k                                                        

       1.027289983 seconds time elapsed

thanks,
jirka


Assisted-by: LLM
---
Jiri Olsa (13):
      static_call: Index code text reservation lookups
      jump_label: Index text reservation lookups
      jump_label/x86: Avoid decoding instructions for jump entry size
      kprobes: Maintain an rbtree of registered probes
      kprobes: Add an iterator for address ranges
      kprobes: Add an address iterator for probe lookup
      kprobes: Use kprobe iterator in get_optimized_kprobe
      x86/kprobes: Use kprobe iterator in __recover_optprobed_insn
      x86/kprobes: Use kprobe iterator in arch_check_optimized_kprobe
      x86/kprobes: Use kprobe address iterator in can_optimize
      x86/kprobes: Use kprobe address iterator in can_probe
      x86/kprobes: Use kprobe address iterator in in copy_optimized_instructions
      samples/kprobes: Add batch probe example

 arch/x86/kernel/jump_label.c          |  14 +++++++----
 arch/x86/kernel/kprobes/common.h      |   7 +++---
 arch/x86/kernel/kprobes/core.c        |  26 +++++++++++--------
 arch/x86/kernel/kprobes/opt.c         |  32 ++++++++++++++---------
 include/linux/kprobes.h               |  25 ++++++++++++++++++
 kernel/Makefile                       |   3 +++
 kernel/jump_label.c                   | 115 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++-
 kernel/kprobes.c                      | 134 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++------
 kernel/kprobes_many.c                 |  81 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 kernel/static_call_inline.c           | 114 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++--
 samples/Kconfig                       |   8 ++++++
 samples/kprobes/Makefile              |   1 +
 samples/kprobes/kprobe_many_example.c |  80 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 13 files changed, 598 insertions(+), 42 deletions(-)
 create mode 100644 kernel/kprobes_many.c
 create mode 100644 samples/kprobes/kprobe_many_example.c

[PATCH 01/13] static_call: Index code text reservation lookups

[Jiri Olsa]

Text reservation checks scan every static call site. Build an address-sorted
index for core text sites and use binary search to find the starting point
near the requested range.

Module indexing is intentionally unsupported. Module reservation checks
continue to use a linear scan, avoiding per-module index allocation and
teardown complexity.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 kernel/static_call_inline.c | 114 +++++++++++++++++++++++++++++++++++-
 1 file changed, 112 insertions(+), 2 deletions(-)

diff --git a/kernel/static_call_inline.c b/kernel/static_call_inline.c
index 2b6a0d99cdbe..884dfc897a1f 100644
--- a/kernel/static_call_inline.c
+++ b/kernel/static_call_inline.c
@@ -15,6 +15,18 @@ extern struct static_call_site __start_static_call_sites[],
 extern struct static_call_tramp_key __start_static_call_tramp_key[],
 				    __stop_static_call_tramp_key[];
 
+struct static_call_site_index {
+	struct static_call_site *base;
+	unsigned int nr;
+	u32 entries[];
+};
+
+/*
+ * Index static call sites by code address for text reservation lookups.
+ * The site table itself is sorted by static-call-key address.
+ */
+static struct static_call_site_index *static_call_site_index;
+
 int static_call_initialized;
 
 /*
@@ -41,7 +53,7 @@ static void static_call_unlock(void)
 	mutex_unlock(&static_call_mutex);
 }
 
-static inline void *static_call_addr(struct static_call_site *site)
+static inline void *static_call_addr(const struct static_call_site *site)
 {
 	return (void *)((long)site->addr + (long)&site->addr);
 }
@@ -301,6 +313,79 @@ static int addr_conflict(struct static_call_site *site, void *start, void *end)
 	return 0;
 }
 
+static int static_call_site_index_cmp(const void *a, const void *b, const void *priv)
+{
+	const struct static_call_site *base = priv;
+	unsigned int ia = *(u32 *)a;
+	unsigned int ib = *(u32 *)b;
+	unsigned long ca = (unsigned long)static_call_addr(base + ia);
+	unsigned long cb = (unsigned long)static_call_addr(base + ib);
+
+	return cmp_int(ca, cb);
+}
+
+static struct static_call_site_index *
+static_call_site_index_alloc(struct static_call_site *start,
+			     struct static_call_site *stop)
+{
+	struct static_call_site_index *index;
+	unsigned int nr = stop - start;
+	unsigned int i;
+
+	if (!nr)
+		return NULL;
+
+	index = kvmalloc_flex(*index, entries, nr, GFP_KERNEL);
+	if (!index)
+		return NULL;
+
+	index->nr = nr;
+	index->base = start;
+	for (i = 0; i < nr; i++)
+		index->entries[i] = i;
+
+	sort_r(index->entries, nr, sizeof(index->entries[0]),
+	       static_call_site_index_cmp, NULL, start);
+
+	return index;
+}
+
+static int
+static_call_site_index_text_reserved(struct static_call_site_index *index,
+				     void *start, void *end, bool init)
+{
+	unsigned int low = 0, high = index->nr, i;
+
+	/* Find the first site after @start. */
+	while (low < high) {
+		unsigned int mid = low + (high - low) / 2;
+		void *addr = static_call_addr(index->base + index->entries[mid]);
+
+		if (addr <= start)
+			low = mid + 1;
+		else
+			high = mid;
+	}
+
+	/* The preceding site may overlap @start. */
+	if (low)
+		low--;
+
+	for (i = low; i < index->nr; i++) {
+		struct static_call_site *site = index->base + index->entries[i];
+		void *addr = static_call_addr(site);
+
+		/* Later sites cannot overlap once this address is past @end. */
+		if (addr > end)
+			break;
+		if ((init || !static_call_is_init(site)) &&
+		    addr_conflict(site, start, end))
+			return 1;
+	}
+
+	return 0;
+}
+
 static int __static_call_text_reserved(struct static_call_site *iter_start,
 				       struct static_call_site *iter_stop,
 				       void *start, void *end, bool init)
@@ -483,7 +568,15 @@ static inline int __static_call_mod_text_reserved(void *start, void *end)
 int static_call_text_reserved(void *start, void *end)
 {
 	bool init = system_state < SYSTEM_RUNNING;
-	int ret = __static_call_text_reserved(__start_static_call_sites,
+	struct static_call_site_index *index;
+	int ret;
+
+	/* Pairs with the release in static_call_site_index_build(). */
+	index = smp_load_acquire(&static_call_site_index);
+	if (index)
+		ret = static_call_site_index_text_reserved(index, start, end, init);
+	else
+		ret = __static_call_text_reserved(__start_static_call_sites,
 			__stop_static_call_sites, start, end, init);
 
 	if (ret)
@@ -522,6 +615,23 @@ int __init static_call_init(void)
 }
 early_initcall(static_call_init);
 
+static int __init static_call_site_index_build(void)
+{
+	struct static_call_site_index *index;
+
+	/*
+	 * Build the index in an initcall so normal memory allocation is
+	 * available, static_call_init() can run before that.
+	 */
+	index = static_call_site_index_alloc(__start_static_call_sites,
+					     __stop_static_call_sites);
+
+	/* Publish the initialized index to static_call_text_reserved(). */
+	smp_store_release(&static_call_site_index, index);
+	return 0;
+}
+late_initcall(static_call_site_index_build);
+
 #ifdef CONFIG_STATIC_CALL_SELFTEST
 
 static int func_a(int x)
-- 
2.43.0

[PATCH 02/13] jump_label: Index text reservation lookups

[Jiri Olsa]

Text reservation checks scan every jump entry. Build an index sorted by
code address for core entries and use binary search to find the starting
point near the requested range. Keep the original table sorted by
static-key address for updates.

Module reservation checks continue to use a linear scan, avoiding
per-module index allocation and teardown complexity.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 kernel/jump_label.c | 115 +++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 114 insertions(+), 1 deletion(-)

diff --git a/kernel/jump_label.c b/kernel/jump_label.c
index ab6b774bcfd7..64467b971efe 100644
--- a/kernel/jump_label.c
+++ b/kernel/jump_label.c
@@ -8,6 +8,7 @@
  */
 #include <linux/memory.h>
 #include <linux/uaccess.h>
+#include <linux/cache.h>
 #include <linux/module.h>
 #include <linux/list.h>
 #include <linux/slab.h>
@@ -22,6 +23,18 @@
 /* mutex to protect coming/going of the jump_label table */
 static DEFINE_MUTEX(jump_label_mutex);
 
+struct jump_entry_index {
+	struct jump_entry *base;
+	unsigned int nr;
+	u32 entries[];
+};
+
+/*
+ * Index jump entries by code address for text reservation lookups. The jump
+ * table itself is sorted by static-key address, then by code address per key.
+ */
+static struct jump_entry_index *jump_entry_index;
+
 void jump_label_lock(void)
 {
 	mutex_lock(&jump_label_mutex);
@@ -381,6 +394,81 @@ static int addr_conflict(struct jump_entry *entry, void *start, void *end)
 	return 0;
 }
 
+static int jump_entry_index_cmp(const void *a, const void *b, const void *priv)
+{
+	const struct jump_entry *base = priv;
+	unsigned int ia = *(u32 *)a;
+	unsigned int ib = *(u32 *)b;
+	unsigned long ca = jump_entry_code(base + ia);
+	unsigned long cb = jump_entry_code(base + ib);
+
+	if (ca < cb)
+		return -1;
+	if (ca > cb)
+		return 1;
+	return 0;
+}
+
+static struct jump_entry_index *
+jump_entry_index_alloc(struct jump_entry *start, struct jump_entry *stop)
+{
+	unsigned int i, nr = stop - start;
+	struct jump_entry_index *index;
+
+	if (!nr)
+		return NULL;
+
+	index = kvmalloc_flex(*index, entries, nr, GFP_KERNEL);
+	if (!index)
+		return NULL;
+
+	index->nr = nr;
+	index->base = start;
+	for (i = 0; i < nr; i++)
+		index->entries[i] = i;
+
+	sort_r(index->entries, nr, sizeof(index->entries[0]),
+	       jump_entry_index_cmp, NULL, start);
+	return index;
+}
+
+static int jump_entry_index_text_reserved(struct jump_entry_index *index,
+					  void *start, void *end, bool init)
+{
+	unsigned long start_addr = (unsigned long)start;
+	unsigned long end_addr = (unsigned long)end;
+	unsigned int low = 0, high = index->nr, i;
+
+	/* Find the first entry after @start. */
+	while (low < high) {
+		unsigned int mid = low + (high - low) / 2;
+		unsigned long addr = jump_entry_code(index->base + index->entries[mid]);
+
+		if (addr <= start_addr)
+			low = mid + 1;
+		else
+			high = mid;
+	}
+
+	/* The preceding entry may start before @start and overlap it. */
+	if (low)
+		low--;
+
+	for (i = low; i < index->nr; i++) {
+		struct jump_entry *entry = index->base + index->entries[i];
+		unsigned long addr = jump_entry_code(entry);
+
+		/* Later entries cannot overlap once this address is past @end. */
+		if (addr > end_addr)
+			break;
+		if ((init || !jump_entry_is_init(entry)) &&
+		    addr_conflict(entry, start, end))
+			return 1;
+	}
+
+	return 0;
+}
+
 static int __jump_label_text_reserved(struct jump_entry *iter_start,
 		struct jump_entry *iter_stop, void *start, void *end, bool init)
 {
@@ -869,8 +957,16 @@ early_initcall(jump_label_init_module);
  */
 int jump_label_text_reserved(void *start, void *end)
 {
+	struct jump_entry_index *index;
 	bool init = system_state < SYSTEM_RUNNING;
-	int ret = __jump_label_text_reserved(__start___jump_table,
+	int ret;
+
+	/* Pairs with the release in jump_entry_index_build(). */
+	index = smp_load_acquire(&jump_entry_index);
+	if (index)
+		ret = jump_entry_index_text_reserved(index, start, end, init);
+	else
+		ret = __jump_label_text_reserved(__start___jump_table,
 			__stop___jump_table, start, end, init);
 
 	if (ret)
@@ -909,6 +1005,23 @@ static void jump_label_update(struct static_key *key)
 		__jump_label_update(key, entry, stop, init);
 }
 
+static int __init jump_entry_index_build(void)
+{
+	struct jump_entry_index *index;
+
+	/*
+	 * Build the index in an initcall so normal memory allocation is
+	 * available; jump_label_init() runs before that.
+	 */
+	index = jump_entry_index_alloc(__start___jump_table, __stop___jump_table);
+
+	/* Publish the initialized index to jump_label_text_reserved(). */
+	smp_store_release(&jump_entry_index, index);
+
+	return 0;
+}
+late_initcall(jump_entry_index_build);
+
 #ifdef CONFIG_STATIC_KEYS_SELFTEST
 static DEFINE_STATIC_KEY_TRUE(sk_true);
 static DEFINE_STATIC_KEY_FALSE(sk_false);
-- 
2.43.0

[PATCH 03/13] jump_label/x86: Avoid decoding instructions for jump entry size

[Jiri Olsa]

Determine the jump entry size from its opcode or NOP encoding instead of
decoding the instruction. This reduces kprobe registration time when
checking jump label entries.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 arch/x86/kernel/jump_label.c | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/arch/x86/kernel/jump_label.c b/arch/x86/kernel/jump_label.c
index a7949a54a0ff..848a2aee05bd 100644
--- a/arch/x86/kernel/jump_label.c
+++ b/arch/x86/kernel/jump_label.c
@@ -15,16 +15,20 @@
 #include <asm/kprobes.h>
 #include <asm/alternative.h>
 #include <asm/text-patching.h>
-#include <asm/insn.h>
 
 int arch_jump_entry_size(struct jump_entry *entry)
 {
-	struct insn insn = {};
+	const u8 *code = (void *)jump_entry_code(entry);
 
-	insn_decode_kernel(&insn, (void *)jump_entry_code(entry));
-	BUG_ON(insn.length != 2 && insn.length != 5);
+	if (code[0] == JMP8_INSN_OPCODE ||
+	    !memcmp(code, x86_nops[JMP8_INSN_SIZE], JMP8_INSN_SIZE))
+		return JMP8_INSN_SIZE;
 
-	return insn.length;
+	if (code[0] == JMP32_INSN_OPCODE ||
+	    !memcmp(code, x86_nops[JMP32_INSN_SIZE], JMP32_INSN_SIZE))
+		return JMP32_INSN_SIZE;
+
+	BUG();
 }
 
 struct jump_label_patch {
-- 
2.43.0

[PATCH 04/13] kprobes: Maintain an rbtree of registered probes

[Jiri Olsa]

Add an address-ordered rbtree for registered probes. Keep it synchronized
with the hash table when probes are added, removed, or replaced by
aggregator probes.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 include/linux/kprobes.h |  2 ++
 kernel/kprobes.c        | 49 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git a/include/linux/kprobes.h b/include/linux/kprobes.h
index e6de7ae55bda..6f16cc01a16c 100644
--- a/include/linux/kprobes.h
+++ b/include/linux/kprobes.h
@@ -28,6 +28,7 @@
 #include <linux/ftrace.h>
 #include <linux/objpool.h>
 #include <linux/rethook.h>
+#include <linux/rbtree.h>
 #include <asm/kprobes.h>
 
 #ifdef CONFIG_KPROBES
@@ -58,6 +59,7 @@ typedef int (*kretprobe_handler_t) (struct kretprobe_instance *,
 
 struct kprobe {
 	struct hlist_node hlist;
+	struct rb_node rb_node;
 
 	/* list of kprobes for multi-handler support */
 	struct list_head list;
diff --git a/kernel/kprobes.c b/kernel/kprobes.c
index 2441b9e39343..6be1430a4ebe 100644
--- a/kernel/kprobes.c
+++ b/kernel/kprobes.c
@@ -64,6 +64,7 @@ static int kprobes_initialized;
  * - RCU hlist traversal under disabling preempt (breakpoint handlers)
  */
 static struct hlist_head kprobe_table[KPROBE_TABLE_SIZE];
+static struct rb_root kprobe_tree = RB_ROOT;
 
 /* NOTE: change this value only with 'kprobe_mutex' held */
 static bool kprobes_all_disarmed;
@@ -395,6 +396,45 @@ struct kprobe *get_kprobe(void *addr)
 }
 NOKPROBE_SYMBOL(get_kprobe);
 
+static void kprobe_tree_insert(struct kprobe *p)
+{
+	struct rb_node **link = &kprobe_tree.rb_node, *parent = NULL;
+	unsigned long addr = (unsigned long)p->addr;
+
+	lockdep_assert_held(&kprobe_mutex);
+	while (*link) {
+		struct kprobe *kp = rb_entry(*link, struct kprobe, rb_node);
+		unsigned long kp_addr = (unsigned long)kp->addr;
+
+		parent = *link;
+		if (addr < kp_addr)
+			link = &(*link)->rb_left;
+		else if (addr > kp_addr)
+			link = &(*link)->rb_right;
+		else {
+			WARN_ON_ONCE(1);
+			return;
+		}
+	}
+
+	rb_link_node(&p->rb_node, parent, link);
+	rb_insert_color(&p->rb_node, &kprobe_tree);
+}
+
+static void kprobe_tree_erase(struct kprobe *p)
+{
+	lockdep_assert_held(&kprobe_mutex);
+	rb_erase(&p->rb_node, &kprobe_tree);
+	RB_CLEAR_NODE(&p->rb_node);
+}
+
+static void kprobe_tree_replace(struct kprobe *old, struct kprobe *new)
+{
+	lockdep_assert_held(&kprobe_mutex);
+	rb_replace_node(&old->rb_node, &new->rb_node, &kprobe_tree);
+	RB_CLEAR_NODE(&old->rb_node);
+}
+
 static int aggr_pre_handler(struct kprobe *p, struct pt_regs *regs);
 
 /* Return true if 'p' is an aggregator */
@@ -643,6 +683,7 @@ static void do_unoptimize_kprobes(struct list_head *unopt_list,
 			 * (reclaiming is done by do_free_cleaned_kprobes().)
 			 */
 			hlist_del_init_rcu(&op->kp.hlist);
+			kprobe_tree_erase(&op->kp);
 		} else {
 			/*
 			 * Keep on cooling_list until the quiescence period
@@ -722,6 +763,7 @@ static void optprobe_finalize_generation(struct optprobe_generation *gen)
 			 * it will be reclaimed after the next quiescence period.
 			 */
 			hlist_del_init_rcu(&op->kp.hlist);
+			kprobe_tree_erase(&op->kp);
 			list_move(&op->list, &optprobe_gens[optprobe_cur_gen].freeing_list);
 			kick_kprobe_optimizer();
 		} else {
@@ -1006,6 +1048,7 @@ static void unoptimize_kprobe(struct kprobe *p, bool force)
 					if (!kprobe_gone(&op->kp))
 						arch_disarm_kprobe(&op->kp);
 					hlist_del_init_rcu(&op->kp.hlist);
+					kprobe_tree_erase(&op->kp);
 					list_move(&op->list, &cur->freeing_list);
 				} else {
 					list_move(&op->list, &cur->cooling_list);
@@ -1021,6 +1064,7 @@ static void unoptimize_kprobe(struct kprobe *p, bool force)
 				if (!kprobe_gone(&op->kp))
 					arch_disarm_kprobe(&op->kp);
 				hlist_del_init_rcu(&op->kp.hlist);
+				kprobe_tree_erase(&op->kp);
 				list_move(&op->list, &cur->freeing_list);
 				kick_kprobe_optimizer();
 			}
@@ -1076,6 +1120,7 @@ static void kill_optimized_kprobe(struct kprobe *p)
 		if (kprobe_unused(p)) {
 			if (optprobe_queued_unopt(op) || optprobe_queued_cooling(op)) {
 				hlist_del_init_rcu(&op->kp.hlist);
+				kprobe_tree_erase(&op->kp);
 				list_move(&op->list,
 					  &optprobe_gens[optprobe_cur_gen].freeing_list);
 				kick_kprobe_optimizer();
@@ -1599,6 +1644,7 @@ static void init_aggr_kprobe(struct kprobe *ap, struct kprobe *p)
 
 	list_add_rcu(&p->list, &ap->list);
 	hlist_replace_rcu(&p->hlist, &ap->hlist);
+	kprobe_tree_replace(p, ap);
 }
 
 /*
@@ -1940,11 +1986,13 @@ static int __register_kprobe(struct kprobe *p)
 	INIT_HLIST_NODE(&p->hlist);
 	hlist_add_head_rcu(&p->hlist,
 		       &kprobe_table[hash_ptr(p->addr, KPROBE_HASH_BITS)]);
+	kprobe_tree_insert(p);
 
 	if (!kprobes_all_disarmed && !kprobe_disabled(p)) {
 		ret = arm_kprobe(p);
 		if (ret) {
 			hlist_del_rcu(&p->hlist);
+			kprobe_tree_erase(p);
 			synchronize_rcu();
 		}
 	}
@@ -2075,6 +2123,7 @@ static int __unregister_kprobe_top(struct kprobe *p)
 		 * unoptimizing.
 		 */
 		hlist_del_rcu(&ap->hlist);
+		kprobe_tree_erase(ap);
 		return 0;
 	}
 
-- 
2.43.0

[PATCH 05/13] kprobes: Add an iterator for address ranges

[Jiri Olsa]

Add an iterator that walks registered probes in address order within
a specified range. Use the kprobe rbtree to find the first probe at
or after the start address, then advance through the tree.

Require kprobe_mutex to be held while iterating and provide a
for_each_kprobe() helper.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 include/linux/kprobes.h | 13 +++++++++++++
 kernel/kprobes.c        | 42 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/include/linux/kprobes.h b/include/linux/kprobes.h
index 6f16cc01a16c..bdba893cc0b2 100644
--- a/include/linux/kprobes.h
+++ b/include/linux/kprobes.h
@@ -95,6 +95,11 @@ struct kprobe {
 	u32 flags;
 };
 
+struct kprobe_iter {
+	struct kprobe *next;
+	unsigned long end;
+};
+
 /* Kprobe status flags */
 #define KPROBE_FLAG_GONE	1 /* breakpoint has already gone */
 #define KPROBE_FLAG_DISABLED	2 /* probe is temporarily disabled */
@@ -386,6 +391,14 @@ static inline void kprobe_ftrace_kill(void) {}
 /* Get the kprobe at this addr (if any) - called with preemption disabled */
 struct kprobe *get_kprobe(void *addr);
 
+struct kprobe *kprobe_iter_start(struct kprobe_iter *iter,
+				 unsigned long start, unsigned long end);
+struct kprobe *kprobe_iter_next(struct kprobe_iter *iter);
+
+#define for_each_kprobe(iter, p, start, end)			\
+	for ((p) = kprobe_iter_start(&(iter), (start), (end));	\
+	     (p); (p) = kprobe_iter_next(&(iter)))
+
 /* kprobe_running() will just return the current_kprobe on this CPU */
 static inline struct kprobe *kprobe_running(void)
 {
diff --git a/kernel/kprobes.c b/kernel/kprobes.c
index 6be1430a4ebe..bbd51e0fab7e 100644
--- a/kernel/kprobes.c
+++ b/kernel/kprobes.c
@@ -435,6 +435,48 @@ static void kprobe_tree_replace(struct kprobe *old, struct kprobe *new)
 	RB_CLEAR_NODE(&old->rb_node);
 }
 
+struct kprobe *kprobe_iter_start(struct kprobe_iter *iter, unsigned long start,
+				 unsigned long end)
+{
+	struct rb_node *node = kprobe_tree.rb_node;
+	struct kprobe *kp, *next = NULL;
+
+	lockdep_assert_held(&kprobe_mutex);
+	iter->next = NULL;
+	iter->end = end;
+	if (start > end)
+		return NULL;
+
+	while (node) {
+		kp = rb_entry(node, struct kprobe, rb_node);
+		if ((unsigned long)kp->addr >= start) {
+			next = kp;
+			node = node->rb_left;
+		} else {
+			node = node->rb_right;
+		}
+	}
+
+	iter->next = next;
+	return kprobe_iter_next(iter);
+}
+
+struct kprobe *kprobe_iter_next(struct kprobe_iter *iter)
+{
+	struct kprobe *kp = iter->next;
+	struct rb_node *node;
+
+	lockdep_assert_held(&kprobe_mutex);
+	if (!kp || (unsigned long)kp->addr > iter->end) {
+		iter->next = NULL;
+		return NULL;
+	}
+
+	node = rb_next(&kp->rb_node);
+	iter->next = node ? rb_entry(node, struct kprobe, rb_node) : NULL;
+	return kp;
+}
+
 static int aggr_pre_handler(struct kprobe *p, struct pt_regs *regs);
 
 /* Return true if 'p' is an aggregator */
-- 
2.43.0

[PATCH 06/13] kprobes: Add an address iterator for probe lookup

[Jiri Olsa]

Add kprobe_addr_iter_start() and kprobe_addr_iter_next() to look up
probes at nondecreasing addresses while retaining the next candidate
in the range iterator.

Provide for_each_addr_kprobe() for callers that advance the address
in the loop body. Return NULL for addresses without a registered probe
and require kprobe_mutex to remain held during iteration.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 include/linux/kprobes.h | 10 ++++++++++
 kernel/kprobes.c        | 25 +++++++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/include/linux/kprobes.h b/include/linux/kprobes.h
index bdba893cc0b2..30d0de29ceb5 100644
--- a/include/linux/kprobes.h
+++ b/include/linux/kprobes.h
@@ -394,11 +394,21 @@ struct kprobe *get_kprobe(void *addr);
 struct kprobe *kprobe_iter_start(struct kprobe_iter *iter,
 				 unsigned long start, unsigned long end);
 struct kprobe *kprobe_iter_next(struct kprobe_iter *iter);
+struct kprobe *kprobe_addr_iter_start(struct kprobe_iter *iter,
+				      unsigned long start, unsigned long end);
+struct kprobe *kprobe_addr_iter_next(struct kprobe_iter *iter,
+				     unsigned long addr);
 
 #define for_each_kprobe(iter, p, start, end)			\
 	for ((p) = kprobe_iter_start(&(iter), (start), (end));	\
 	     (p); (p) = kprobe_iter_next(&(iter)))
 
+/* The caller holds kprobe_mutex and initializes and advances addr. */
+#define for_each_addr_kprobe(iter, kp, addr, end)			\
+	for ((kp) = kprobe_addr_iter_start(&(iter), (addr), (end));	\
+	     (addr) <= (end);						\
+	     (kp) = kprobe_addr_iter_next(&(iter), (addr)))
+
 /* kprobe_running() will just return the current_kprobe on this CPU */
 static inline struct kprobe *kprobe_running(void)
 {
diff --git a/kernel/kprobes.c b/kernel/kprobes.c
index bbd51e0fab7e..6655c24bb66f 100644
--- a/kernel/kprobes.c
+++ b/kernel/kprobes.c
@@ -477,6 +477,31 @@ struct kprobe *kprobe_iter_next(struct kprobe_iter *iter)
 	return kp;
 }
 
+/* Address iteration retains the candidate in iter->next until addr passes it. */
+struct kprobe *kprobe_addr_iter_start(struct kprobe_iter *iter,
+				      unsigned long start, unsigned long end)
+{
+	iter->next = kprobe_iter_start(iter, start, end);
+	return kprobe_addr_iter_next(iter, start);
+}
+
+/* Addresses must be nondecreasing; return NULL at addresses without a probe. */
+struct kprobe *kprobe_addr_iter_next(struct kprobe_iter *iter,
+				     unsigned long addr)
+{
+	struct kprobe *kp = iter->next;
+
+	lockdep_assert_held(&kprobe_mutex);
+	while (kp && (unsigned long)kp->addr < addr) {
+		kprobe_iter_next(iter);
+		kp = iter->next;
+	}
+	if (!kp || (unsigned long)kp->addr > iter->end)
+		return NULL;
+
+	return (unsigned long)kp->addr == addr ? kp : NULL;
+}
+
 static int aggr_pre_handler(struct kprobe *p, struct pt_regs *regs);
 
 /* Return true if 'p' is an aggregator */
-- 
2.43.0

[PATCH 07/13] kprobes: Use kprobe iterator in get_optimized_kprobe

[Jiri Olsa]

Use the kprobe tree iterator to find optimized probes whose instruction
ranges may contain the given address. This replaces repeated hash lookups
for each preceding instruction address.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 kernel/kprobes.c | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/kernel/kprobes.c b/kernel/kprobes.c
index 6655c24bb66f..371ecc7d5b54 100644
--- a/kernel/kprobes.c
+++ b/kernel/kprobes.c
@@ -591,23 +591,25 @@ bool kprobe_disarmed(struct kprobe *p)
  */
 static struct kprobe *get_optimized_kprobe(kprobe_opcode_t *addr)
 {
-	int i;
-	struct kprobe *p;
+	unsigned long offset = MAX_OPTIMIZED_LENGTH / sizeof(kprobe_opcode_t) - 1;
+	unsigned long addr_val = (unsigned long)addr;
+	struct kprobe *p, *candidate = NULL;
 	struct optimized_kprobe *op;
+	struct kprobe_iter iter;
 
-	/* Don't check i == 0, since that is a breakpoint case. */
-	for (i = 1; i < MAX_OPTIMIZED_LENGTH / sizeof(kprobe_opcode_t); i++) {
-		p = get_kprobe(addr - i);
+	/* Don't check addr itself, since that is a breakpoint case. */
+	for_each_kprobe(iter, p, addr_val - offset, addr_val - 1) {
 		/* A disabled probe can have prepared, but inactive, optinsns. */
-		if (!p || !kprobe_optready(p) || kprobe_disarmed(p))
+		if (!kprobe_optready(p) || kprobe_disarmed(p))
 			continue;
 
+		/* The iterator walks low to high; keep the nearest matching probe. */
 		op = container_of(p, struct optimized_kprobe, kp);
 		if (arch_within_optimized_kprobe(op, addr))
-			return p;
+			candidate = p;
 	}
 
-	return NULL;
+	return candidate;
 }
 
 #define OPTPROBE_GEN_MAX 2
-- 
2.43.0

[PATCH 08/13] x86/kprobes: Use kprobe iterator in __recover_optprobed_insn

[Jiri Olsa]

Use the kprobe range iterator to find an optimized probe covering the
address being recovered. This replaces repeated hash lookups across the
JMP32 instruction range.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 arch/x86/kernel/kprobes/opt.c | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/arch/x86/kernel/kprobes/opt.c b/arch/x86/kernel/kprobes/opt.c
index 3f8fea52619f..9da41424be77 100644
--- a/arch/x86/kernel/kprobes/opt.c
+++ b/arch/x86/kernel/kprobes/opt.c
@@ -37,12 +37,11 @@
 unsigned long __recover_optprobed_insn(kprobe_opcode_t *buf, unsigned long addr)
 {
 	struct optimized_kprobe *op;
+	struct kprobe_iter iter;
 	struct kprobe *kp;
 	long offs;
-	int i;
 
-	for (i = 0; i < JMP32_INSN_SIZE; i++) {
-		kp = get_kprobe((void *)addr - i);
+	for_each_kprobe(iter, kp, addr - (JMP32_INSN_SIZE - 1), addr) {
 		/* This function only handles jump-optimized kprobe */
 		if (kp && kprobe_optimized(kp)) {
 			op = container_of(kp, struct optimized_kprobe, kp);
-- 
2.43.0

[PATCH 09/13] x86/kprobes: Use kprobe iterator in arch_check_optimized_kprobe

[Jiri Olsa]

arch_check_optimized_kprobe() checks each byte of the optimized range
with a separate get_kprobe() lookup. Iterate over registered probes in
the range instead, retaining the check for probes that are still armed.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 arch/x86/kernel/kprobes/opt.c | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/arch/x86/kernel/kprobes/opt.c b/arch/x86/kernel/kprobes/opt.c
index 9da41424be77..cdf69c939631 100644
--- a/arch/x86/kernel/kprobes/opt.c
+++ b/arch/x86/kernel/kprobes/opt.c
@@ -334,12 +334,13 @@ static int can_optimize(unsigned long paddr)
 /* Check optimized_kprobe can actually be optimized. */
 int arch_check_optimized_kprobe(struct optimized_kprobe *op)
 {
-	int i;
+	unsigned long start = (unsigned long)op->kp.addr;
+	unsigned long end = start + op->optinsn.size - 1;
+	struct kprobe_iter iter;
 	struct kprobe *p;
 
-	for (i = 1; i < op->optinsn.size; i++) {
-		p = get_kprobe(op->kp.addr + i);
-		if (p && !kprobe_disarmed(p))
+	for_each_kprobe(iter, p, start + 1, end) {
+		if (!kprobe_disarmed(p))
 			return -EEXIST;
 	}
 
-- 
2.43.0

[PATCH 10/13] x86/kprobes: Use kprobe address iterator in can_optimize

[Jiri Olsa]

can_optimize() already iterates over kprobes in the function. Pass the
iterator's kprobe to recover_probed_instruction() so it doesn't need to
look up the kprobe again for each address. Keep the lookup for callers
that pass NULL.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 arch/x86/kernel/kprobes/common.h |  4 ++--
 arch/x86/kernel/kprobes/core.c   | 16 ++++++++--------
 arch/x86/kernel/kprobes/opt.c    |  9 ++++++---
 3 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/arch/x86/kernel/kprobes/common.h b/arch/x86/kernel/kprobes/common.h
index e772276f5aa9..062746e20c35 100644
--- a/arch/x86/kernel/kprobes/common.h
+++ b/arch/x86/kernel/kprobes/common.h
@@ -79,9 +79,9 @@
 
 /* Ensure if the instruction can be boostable */
 extern bool can_boost(struct insn *insn, void *orig_addr);
-/* Recover instruction if given address is probed */
+/* Recover instruction; @kp is NULL if the caller needs a kprobe lookup. */
 extern unsigned long recover_probed_instruction(kprobe_opcode_t *buf,
-					 unsigned long addr);
+					 unsigned long addr, struct kprobe **kp);
 /*
  * Copy an instruction and adjust the displacement if the instruction
  * uses the %rip-relative addressing mode.
diff --git a/arch/x86/kernel/kprobes/core.c b/arch/x86/kernel/kprobes/core.c
index 133ff20caccd..6ad8b89a9ba1 100644
--- a/arch/x86/kernel/kprobes/core.c
+++ b/arch/x86/kernel/kprobes/core.c
@@ -201,12 +201,11 @@ bool can_boost(struct insn *insn, void *addr)
 }
 
 static unsigned long
-__recover_probed_insn(kprobe_opcode_t *buf, unsigned long addr)
+__recover_probed_insn(kprobe_opcode_t *buf, unsigned long addr, struct kprobe **kpp)
 {
-	struct kprobe *kp;
+	struct kprobe *kp = kpp ? *kpp : get_kprobe((void *)addr);
 	bool faddr;
 
-	kp = get_kprobe((void *)addr);
 	faddr = ftrace_location(addr) == addr;
 	/*
 	 * Use the current code if it is not modified by Kprobe
@@ -251,7 +250,8 @@ __recover_probed_insn(kprobe_opcode_t *buf, unsigned long addr)
  * for preventing to release referencing kprobes.
  * Returns zero if the instruction can not get recovered (or access failed).
  */
-unsigned long recover_probed_instruction(kprobe_opcode_t *buf, unsigned long addr)
+unsigned long recover_probed_instruction(kprobe_opcode_t *buf, unsigned long addr,
+					 struct kprobe **kp)
 {
 	unsigned long __addr;
 
@@ -259,7 +259,7 @@ unsigned long recover_probed_instruction(kprobe_opcode_t *buf, unsigned long add
 	if (__addr != addr)
 		return __addr;
 
-	return __recover_probed_insn(buf, addr);
+	return __recover_probed_insn(buf, addr, kp);
 }
 
 /* Check if insn is INT or UD */
@@ -304,7 +304,7 @@ static bool can_probe(unsigned long paddr)
 		 * relative-jump. Since the relative-jump itself is
 		 * normally used, we just go through if there is no kprobe.
 		 */
-		__addr = recover_probed_instruction(buf, addr);
+		__addr = recover_probed_instruction(buf, addr, NULL);
 		if (!__addr)
 			return false;
 
@@ -327,7 +327,7 @@ static bool can_probe(unsigned long paddr)
 	if (addr != paddr)
 		return false;
 
-	__addr = recover_probed_instruction(buf, addr);
+	__addr = recover_probed_instruction(buf, addr, NULL);
 	if (!__addr)
 		return false;
 
@@ -394,7 +394,7 @@ kprobe_opcode_t *arch_adjust_kprobe_addr(unsigned long addr, unsigned long offse
 int __copy_instruction(u8 *dest, u8 *src, u8 *real, struct insn *insn)
 {
 	kprobe_opcode_t buf[MAX_INSN_SIZE];
-	unsigned long recovered_insn = recover_probed_instruction(buf, (unsigned long)src);
+	unsigned long recovered_insn = recover_probed_instruction(buf, (unsigned long)src, NULL);
 	int ret;
 
 	if (!recovered_insn || !insn)
diff --git a/arch/x86/kernel/kprobes/opt.c b/arch/x86/kernel/kprobes/opt.c
index cdf69c939631..9f0af24155ba 100644
--- a/arch/x86/kernel/kprobes/opt.c
+++ b/arch/x86/kernel/kprobes/opt.c
@@ -257,7 +257,9 @@ static int insn_jump_into_range(struct insn *insn, unsigned long start, int len)
 /* Decode whole function to ensure any instructions don't jump into target */
 static int can_optimize(unsigned long paddr)
 {
-	unsigned long addr, size = 0, offset = 0;
+	unsigned long addr, end, size = 0, offset = 0;
+	struct kprobe_iter iter;
+	struct kprobe *kp;
 	struct insn insn;
 	kprobe_opcode_t buf[MAX_INSN_SIZE];
 
@@ -279,7 +281,8 @@ static int can_optimize(unsigned long paddr)
 
 	/* Decode instructions */
 	addr = paddr - offset;
-	while (addr < paddr - offset + size) { /* Decode until function end */
+	end = paddr - offset + size - 1;
+	for_each_addr_kprobe(iter, kp, addr, end) { /* Decode until function end */
 		unsigned long recovered_insn;
 		int ret;
 
@@ -289,7 +292,7 @@ static int can_optimize(unsigned long paddr)
 			 * we can't optimize kprobe in this function.
 			 */
 			return 0;
-		recovered_insn = recover_probed_instruction(buf, addr);
+		recovered_insn = recover_probed_instruction(buf, addr, &kp);
 		if (!recovered_insn)
 			return 0;
 
-- 
2.43.0

[PATCH 11/13] x86/kprobes: Use kprobe address iterator in can_probe

[Jiri Olsa]

Walk the instructions before the target with for_each_addr_kprobe() and
pass each matching kprobe to recover_probed_instruction(). Reuse the
iterator’s kprobe for the target instruction as well.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 arch/x86/kernel/kprobes/core.c | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/arch/x86/kernel/kprobes/core.c b/arch/x86/kernel/kprobes/core.c
index 6ad8b89a9ba1..4d63826f6c04 100644
--- a/arch/x86/kernel/kprobes/core.c
+++ b/arch/x86/kernel/kprobes/core.c
@@ -287,6 +287,8 @@ static inline bool is_exception_insn(struct insn *insn)
 static bool can_probe(unsigned long paddr)
 {
 	unsigned long addr, __addr, offset = 0;
+	struct kprobe_iter iter;
+	struct kprobe *kp;
 	struct insn insn;
 	kprobe_opcode_t buf[MAX_INSN_SIZE];
 
@@ -295,7 +297,9 @@ static bool can_probe(unsigned long paddr)
 
 	/* Decode instructions */
 	addr = paddr - offset;
-	while (addr < paddr) {
+	for_each_addr_kprobe(iter, kp, addr, paddr) {
+		if (addr == paddr)
+			break;
 		/*
 		 * Check if the instruction has been modified by another
 		 * kprobe, in which case we replace the breakpoint by the
@@ -304,7 +308,7 @@ static bool can_probe(unsigned long paddr)
 		 * relative-jump. Since the relative-jump itself is
 		 * normally used, we just go through if there is no kprobe.
 		 */
-		__addr = recover_probed_instruction(buf, addr, NULL);
+		__addr = recover_probed_instruction(buf, addr, &kp);
 		if (!__addr)
 			return false;
 
@@ -327,7 +331,7 @@ static bool can_probe(unsigned long paddr)
 	if (addr != paddr)
 		return false;
 
-	__addr = recover_probed_instruction(buf, addr, NULL);
+	__addr = recover_probed_instruction(buf, addr, &kp);
 	if (!__addr)
 		return false;
 
-- 
2.43.0

[PATCH 12/13] x86/kprobes: Use kprobe address iterator in in copy_optimized_instructions

[Jiri Olsa]

copy_optimized_instructions() copies whole instructions until it covers
the five-byte jump. Pass the iterator's kprobe to instruction recovery
at each instruction start, avoiding a separate lookup. Other
__copy_instruction() callers retain the lookup behavior by passing NULL.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 arch/x86/kernel/kprobes/common.h | 3 ++-
 arch/x86/kernel/kprobes/core.c   | 6 +++---
 arch/x86/kernel/kprobes/opt.c    | 9 +++++++--
 3 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/arch/x86/kernel/kprobes/common.h b/arch/x86/kernel/kprobes/common.h
index 062746e20c35..c04d01fe075a 100644
--- a/arch/x86/kernel/kprobes/common.h
+++ b/arch/x86/kernel/kprobes/common.h
@@ -86,7 +86,8 @@ extern unsigned long recover_probed_instruction(kprobe_opcode_t *buf,
  * Copy an instruction and adjust the displacement if the instruction
  * uses the %rip-relative addressing mode.
  */
-extern int __copy_instruction(u8 *dest, u8 *src, u8 *real, struct insn *insn);
+extern int __copy_instruction(u8 *dest, u8 *src, u8 *real, struct insn *insn,
+			      struct kprobe **kp);
 
 /* Generate a relative-jump/call instruction */
 extern void synthesize_reljump(void *dest, void *from, void *to);
diff --git a/arch/x86/kernel/kprobes/core.c b/arch/x86/kernel/kprobes/core.c
index 4d63826f6c04..ae6611fdb19a 100644
--- a/arch/x86/kernel/kprobes/core.c
+++ b/arch/x86/kernel/kprobes/core.c
@@ -395,10 +395,10 @@ kprobe_opcode_t *arch_adjust_kprobe_addr(unsigned long addr, unsigned long offse
  * instruction, displacement must be adjust by @real, not @dest.
  * This returns the length of copied instruction, or 0 if it has an error.
  */
-int __copy_instruction(u8 *dest, u8 *src, u8 *real, struct insn *insn)
+int __copy_instruction(u8 *dest, u8 *src, u8 *real, struct insn *insn, struct kprobe **kp)
 {
 	kprobe_opcode_t buf[MAX_INSN_SIZE];
-	unsigned long recovered_insn = recover_probed_instruction(buf, (unsigned long)src, NULL);
+	unsigned long recovered_insn = recover_probed_instruction(buf, (unsigned long)src, kp);
 	int ret;
 
 	if (!recovered_insn || !insn)
@@ -735,7 +735,7 @@ static int arch_copy_kprobe(struct kprobe *p)
 	int ret, len;
 
 	/* Copy an instruction with recovering if other optprobe modifies it.*/
-	len = __copy_instruction(buf, p->addr, p->ainsn.insn, &insn);
+	len = __copy_instruction(buf, p->addr, p->ainsn.insn, &insn, NULL);
 	if (!len)
 		return -EINVAL;
 
diff --git a/arch/x86/kernel/kprobes/opt.c b/arch/x86/kernel/kprobes/opt.c
index 9f0af24155ba..058b32564ea0 100644
--- a/arch/x86/kernel/kprobes/opt.c
+++ b/arch/x86/kernel/kprobes/opt.c
@@ -201,14 +201,19 @@ NOKPROBE_SYMBOL(optimized_callback);
 
 static int copy_optimized_instructions(u8 *dest, u8 *src, u8 *real)
 {
+	struct kprobe_iter iter;
+	unsigned long addr;
+	struct kprobe *kp;
 	struct insn insn;
 	int len = 0, ret;
 
-	while (len < JMP32_INSN_SIZE) {
-		ret = __copy_instruction(dest + len, src + len, real + len, &insn);
+	addr = (unsigned long)src;
+	for_each_addr_kprobe(iter, kp, addr, (unsigned long)src + JMP32_INSN_SIZE - 1) {
+		ret = __copy_instruction(dest + len, src + len, real + len, &insn, &kp);
 		if (!ret || !can_boost(&insn, src + len))
 			return -EINVAL;
 		len += ret;
+		addr += ret;
 	}
 	/* Check whether the address range is reserved */
 	if (ftrace_text_reserved(src, src + len - 1) ||
-- 
2.43.0

[PATCH 13/13] samples/kprobes: Add batch probe example

[Jiri Olsa]

Add a sample module that registers kprobes on a batch of
generated functions and calls each function to trigger
the probes.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 kernel/Makefile                       |  3 +
 kernel/kprobes_many.c                 | 81 ++++++++++++++++++++++++
 samples/Kconfig                       |  8 +++
 samples/kprobes/Makefile              |  1 +
 samples/kprobes/kprobe_many_example.c | 91 +++++++++++++++++++++++++++
 5 files changed, 184 insertions(+)
 create mode 100644 kernel/kprobes_many.c
 create mode 100644 samples/kprobes/kprobe_many_example.c

diff --git a/kernel/Makefile b/kernel/Makefile
index 4ff92e24680c..f307b29d4fef 100644
--- a/kernel/Makefile
+++ b/kernel/Makefile
@@ -99,6 +99,9 @@ obj-$(CONFIG_AUDITSYSCALL) += auditsc.o audit_watch.o audit_fsnotify.o audit_tre
 obj-$(CONFIG_GCOV_KERNEL) += gcov/
 obj-$(CONFIG_KCOV) += kcov.o
 obj-$(CONFIG_KPROBES) += kprobes.o
+ifneq ($(CONFIG_SAMPLE_KPROBES_MANY),)
+obj-y += kprobes_many.o
+endif
 obj-$(CONFIG_FAIL_FUNCTION) += fail_function.o
 obj-$(CONFIG_KGDB) += debug/
 obj-$(CONFIG_DETECT_HUNG_TASK) += hung_task.o
diff --git a/kernel/kprobes_many.c b/kernel/kprobes_many.c
new file mode 100644
index 000000000000..1c31a63309bf
--- /dev/null
+++ b/kernel/kprobes_many.c
@@ -0,0 +1,81 @@
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+/*
+ * Generated functions for the batch kprobe sample.
+ */
+
+#include <linux/annotate.h>
+#include <linux/export.h>
+#include <linux/linkage.h>
+#include <linux/stringify.h>
+#include <asm/asm.h>
+#include <asm/ibt.h>
+
+#define KPROBE_MANY_MAX_FUNCS	50000
+#define KPROBE_MANY_FUNC_SIZE	24
+#define KPROBE_MANY_PROBE_OFF	8
+
+#if defined(CONFIG_MITIGATION_RETHUNK)
+#define KPROBE_MANY_RET		ASM_RET
+#define KPROBE_MANY_RET_SIZE	5
+#elif defined(CONFIG_MITIGATION_SLS)
+#define KPROBE_MANY_RET		ASM_RET
+#define KPROBE_MANY_RET_SIZE	2
+#else
+#define KPROBE_MANY_RET		ASM_RET
+#define KPROBE_MANY_RET_SIZE	1
+#endif
+
+#define KPROBE_MANY_TAIL_NOPS \
+	(KPROBE_MANY_FUNC_SIZE - KPROBE_MANY_PROBE_OFF - 7 - \
+	 KPROBE_MANY_RET_SIZE)
+
+extern char kprobes_many_func_0[];
+EXPORT_SYMBOL_GPL(kprobes_many_func_0);
+
+/*
+ * Keep these functions at a fixed size so the sample module can place
+ * probes at the same offset in every function.
+ *
+ * For example, KPROBE_MANY_FUNC kprobes_many_func_0 emits:
+ *
+ * kprobes_many_func_0:		ENDBR
+ * 				NOPS to offset 8
+ * 1:				probe site (offset 8), annotated as NOENDBR
+ * 				7 NOPS for the kprobe optimizer
+ * 				NOPS to fill the function
+ * 				RET
+ *
+ * The NOPS and return sequence keep each function 24 bytes long, including
+ * the ENDBR instruction when enabled.
+ */
+asm (
+	".pushsection .text, \"ax\", @progbits\n"
+	".macro KPROBE_MANY_FUNC name\n"
+	".type \\name, @function\n"
+	"\\name:\n"
+	ASM_ENDBR
+	IBT_NOSEAL("\\name")
+	".rept " __stringify(KPROBE_MANY_PROBE_OFF - ENDBR_INSN_SIZE) "\n"
+	"	nop\n"
+	".endr\n"
+	"1:\n"
+	ASM_ANNOTATE_LABEL(1b, ANNOTYPE_NOENDBR) "\n"
+	".rept 7\n"
+	"	nop\n"
+	".endr\n"
+	".rept " __stringify(KPROBE_MANY_TAIL_NOPS) "\n"
+	"	nop\n"
+	".endr\n"
+	KPROBE_MANY_RET
+	".size \\name, .-\\name\n"
+	".endm\n"
+	".macro KPROBE_MANY_FUNC_LOCAL\n"
+	"KPROBE_MANY_FUNC kprobes_many_func_\\@\n"
+	".endm\n"
+	".globl kprobes_many_func_0\n"
+	"KPROBE_MANY_FUNC kprobes_many_func_0\n"
+	".rept " __stringify(KPROBE_MANY_MAX_FUNCS - 1) "\n"
+	"	KPROBE_MANY_FUNC_LOCAL\n"
+	".endr\n"
+	".popsection\n");
diff --git a/samples/Kconfig b/samples/Kconfig
index 00eed83c07df..92a41c5fdbfd 100644
--- a/samples/Kconfig
+++ b/samples/Kconfig
@@ -80,6 +80,14 @@ config SAMPLE_KRETPROBES
 	default m
 	depends on SAMPLE_KPROBES && KRETPROBES
 
+config SAMPLE_KPROBES_MANY
+	tristate "Build batch kprobe example -- x86 loadable module only"
+	depends on SAMPLE_KPROBES && X86 && m
+	help
+	  This builds a sample module that registers kprobes in the middle of
+	  50000 generated x86 kernel functions using the batch kprobe
+	  interface.
+
 config SAMPLE_HW_BREAKPOINT
 	tristate "Build kernel hardware breakpoint examples -- loadable module only"
 	depends on HAVE_HW_BREAKPOINT && m
diff --git a/samples/kprobes/Makefile b/samples/kprobes/Makefile
index e774592718d6..2fd2150e523c 100644
--- a/samples/kprobes/Makefile
+++ b/samples/kprobes/Makefile
@@ -4,3 +4,4 @@
 
 obj-$(CONFIG_SAMPLE_KPROBES) += kprobe_example.o
 obj-$(CONFIG_SAMPLE_KRETPROBES) += kretprobe_example.o
+obj-$(CONFIG_SAMPLE_KPROBES_MANY) += kprobe_many_example.o
diff --git a/samples/kprobes/kprobe_many_example.c b/samples/kprobes/kprobe_many_example.c
new file mode 100644
index 000000000000..5a58908c840c
--- /dev/null
+++ b/samples/kprobes/kprobe_many_example.c
@@ -0,0 +1,91 @@
+// SPDX-License-Identifier: GPL-2.0-only
+/*
+ * Sample module that registers a batch of kprobes in generated kernel
+ * functions.
+ */
+
+#include <linux/kernel.h>
+#include <linux/kprobes.h>
+#include <linux/module.h>
+
+#define KPROBE_MANY_MAX_FUNCS 50000
+#define KPROBE_MANY_FUNC_SIZE 24
+#define KPROBE_MANY_PROBE_OFF 8
+
+extern char kprobes_many_func_0[];
+
+static unsigned int nr_functions = 20000;
+module_param(nr_functions, uint, 0444);
+MODULE_PARM_DESC(nr_functions, "number of generated functions to probe (1-50000)");
+
+static struct kprobe probes[KPROBE_MANY_MAX_FUNCS];
+static struct kprobe *probe_ptrs[KPROBE_MANY_MAX_FUNCS];
+static unsigned int kprobe_many_pre_handler_count;
+
+static int __kprobes kprobe_many_pre_handler(struct kprobe *p, struct pt_regs *regs)
+{
+	kprobe_many_pre_handler_count++;
+	return 0;
+}
+
+/* The generated assembly functions have no CFI type metadata. */
+static noinline __nocfi void
+kprobe_many_call_function(char *functions, unsigned int index)
+{
+	void (*function)(void);
+
+	function = (void (*)(void))(functions + index * KPROBE_MANY_FUNC_SIZE);
+	function();
+}
+
+static int __init kprobe_many_init(void)
+{
+	char *functions;
+	unsigned int i;
+	unsigned int hit_count;
+	int ret;
+
+	if (!nr_functions || nr_functions > KPROBE_MANY_MAX_FUNCS)
+		return -EINVAL;
+
+	functions = kprobes_many_func_0;
+
+	for (i = 0; i < nr_functions; i++) {
+		probes[i].addr = (void *)(functions + i * KPROBE_MANY_FUNC_SIZE +
+					  KPROBE_MANY_PROBE_OFF);
+		probes[i].pre_handler = kprobe_many_pre_handler;
+		probe_ptrs[i] = &probes[i];
+	}
+
+	ret = register_kprobes(probe_ptrs, nr_functions);
+	if (ret) {
+		pr_err("register_kprobes failed, returned %d\n", ret);
+		return ret;
+	}
+
+	pr_info("registered %u probes\n", nr_functions);
+	for (i = 0; i < nr_functions; i++)
+		kprobe_many_call_function(functions, i);
+
+	hit_count = kprobe_many_pre_handler_count;
+	if (hit_count != nr_functions) {
+		pr_err("pre-handler triggered %u times, expected %u\n",
+		       hit_count, nr_functions);
+		unregister_kprobes(probe_ptrs, nr_functions);
+		return -EINVAL;
+	}
+
+	return 0;
+}
+
+static void __exit kprobe_many_exit(void)
+{
+	unregister_kprobes(probe_ptrs, nr_functions);
+	pr_info("unregistered %u probes\n", nr_functions);
+}
+
+module_init(kprobe_many_init);
+module_exit(kprobe_many_exit);
+
+MODULE_DESCRIPTION("sample module registering kprobes in generated functions");
+MODULE_LICENSE("GPL");
-- 
2.43.0