Re: [PATCH 3/8] bpf: Add bpf_cookie support to fprobe

[Jiri Olsa <olsajiri@gmail.com>]

On Tue, Feb 08, 2022 at 03:35:24PM -0800, Andrii Nakryiko wrote:
> On Tue, Feb 8, 2022 at 1:07 AM Jiri Olsa <jolsa@redhat.com> wrote:
> >
> > On Mon, Feb 07, 2022 at 10:59:21AM -0800, Andrii Nakryiko wrote:
> > > On Wed, Feb 2, 2022 at 5:54 AM Jiri Olsa <jolsa@redhat.com> wrote:
> > > >
> > > > Adding support to call bpf_get_attach_cookie helper from
> > > > kprobe program attached by fprobe link.
> > > >
> > > > The bpf_cookie is provided by array of u64 values, where
> > > > each value is paired with provided function address with
> > > > the same array index.
> > > >
> > > > Suggested-by: Andrii Nakryiko <andrii@kernel.org>
> > > > Signed-off-by: Jiri Olsa <jolsa@kernel.org>
> > > > ---
> > > >  include/linux/bpf.h            |  2 +
> > > >  include/uapi/linux/bpf.h       |  1 +
> > > >  kernel/bpf/syscall.c           | 83 +++++++++++++++++++++++++++++++++-
> > > >  kernel/trace/bpf_trace.c       | 16 ++++++-
> > > >  tools/include/uapi/linux/bpf.h |  1 +
> > > >  5 files changed, 100 insertions(+), 3 deletions(-)
> > > >
> > > > diff --git a/include/linux/bpf.h b/include/linux/bpf.h
> > > > index 6eb0b180d33b..7b65f05c0487 100644
> > > > --- a/include/linux/bpf.h
> > > > +++ b/include/linux/bpf.h
> > > > @@ -1301,6 +1301,8 @@ static inline void bpf_reset_run_ctx(struct bpf_run_ctx *old_ctx)
> > > >  #endif
> > > >  }
> > > >
> > > > +u64 bpf_fprobe_cookie(struct bpf_run_ctx *ctx, u64 ip);
> > > > +
> > > >  /* BPF program asks to bypass CAP_NET_BIND_SERVICE in bind. */
> > > >  #define BPF_RET_BIND_NO_CAP_NET_BIND_SERVICE                   (1 << 0)
> > > >  /* BPF program asks to set CN on the packet. */
> > > > diff --git a/include/uapi/linux/bpf.h b/include/uapi/linux/bpf.h
> > > > index c0912f0a3dfe..0dc6aa4f9683 100644
> > > > --- a/include/uapi/linux/bpf.h
> > > > +++ b/include/uapi/linux/bpf.h
> > > > @@ -1484,6 +1484,7 @@ union bpf_attr {
> > > >                                 __aligned_u64   addrs;
> > > >                                 __u32           cnt;
> > > >                                 __u32           flags;
> > > > +                               __aligned_u64   bpf_cookies;
> > >
> > > maybe put it right after addrs, they are closely related and cnt
> > > describes all of syms/addrs/cookies.
> >
> > ok
> >
> > >
> > > >                         } fprobe;
> > > >                 };
> > > >         } link_create;
> > > > diff --git a/kernel/bpf/syscall.c b/kernel/bpf/syscall.c
> > > > index 0cfbb112c8e1..6c5e74bc43b6 100644
> > > > --- a/kernel/bpf/syscall.c
> > > > +++ b/kernel/bpf/syscall.c
> > > > @@ -33,6 +33,8 @@
> > > >  #include <linux/rcupdate_trace.h>
> > > >  #include <linux/memcontrol.h>
> > > >  #include <linux/fprobe.h>
> > > > +#include <linux/bsearch.h>
> > > > +#include <linux/sort.h>
> > > >
> > > >  #define IS_FD_ARRAY(map) ((map)->map_type == BPF_MAP_TYPE_PERF_EVENT_ARRAY || \
> > > >                           (map)->map_type == BPF_MAP_TYPE_CGROUP_ARRAY || \
> > > > @@ -3025,10 +3027,18 @@ static int bpf_perf_link_attach(const union bpf_attr *attr, struct bpf_prog *pro
> > > >
> > > >  #ifdef CONFIG_FPROBE
> > > >
> > > > +struct bpf_fprobe_cookie {
> > > > +       unsigned long addr;
> > > > +       u64 bpf_cookie;
> > > > +};
> > > > +
> > > >  struct bpf_fprobe_link {
> > > >         struct bpf_link link;
> > > >         struct fprobe fp;
> > > >         unsigned long *addrs;
> > > > +       struct bpf_run_ctx run_ctx;
> > > > +       struct bpf_fprobe_cookie *bpf_cookies;
> > >
> > > you already have all the addrs above, why keeping a second copy of
> > > each addrs in bpf_fprobe_cookie. Let's have two arrays: addrs
> > > (unsigned long) and cookies (u64) and make sure that they are sorted
> > > together. Then lookup addrs, calculate index, use that index to fetch
> > > cookie.
> > >
> > > Seems like sort_r() provides exactly the interface you'd need to do
> > > this very easily. Having addrs separate from cookies also a bit
> > > advantageous in terms of TLB misses (if you need any more persuasion
> > > ;)
> >
> > no persuation needed, I actually tried that but it turned out sort_r
> > is not ready yet ;-)
> >
> > because you can't pass priv pointer to the swap callback, so we can't
> > swap the other array.. I did a change to allow that, but it's not trivial
> > and will need some bigger testing/review because the original sort
> > calls sort_r, and of course there are many 'sort' users ;-)
> 
> Big sigh... :( Did you do something similar to _CMP_WRAPPER? You don't
> need to change the interface of sort(), so it shouldn't require
> extensive code refactoring. You'll just need to adjust priv to be not
> just cmp_func, but cmp_func + swap_fun (need a small struct on the
> stack in sort, probably). Or you did something else?

I ended up with change below

jirka


---
 include/linux/sort.h  |  2 +-
 include/linux/types.h |  1 +
 lib/sort.c            | 44 +++++++++++++++++++++++++++++++++----------
 3 files changed, 36 insertions(+), 11 deletions(-)

diff --git a/include/linux/sort.h b/include/linux/sort.h
index b5898725fe9d..e163287ac6c1 100644
--- a/include/linux/sort.h
+++ b/include/linux/sort.h
@@ -6,7 +6,7 @@
 
 void sort_r(void *base, size_t num, size_t size,
 	    cmp_r_func_t cmp_func,
-	    swap_func_t swap_func,
+	    swap_r_func_t swap_func,
 	    const void *priv);
 
 void sort(void *base, size_t num, size_t size,
diff --git a/include/linux/types.h b/include/linux/types.h
index ac825ad90e44..ea8cf60a8a79 100644
--- a/include/linux/types.h
+++ b/include/linux/types.h
@@ -226,6 +226,7 @@ struct callback_head {
 typedef void (*rcu_callback_t)(struct rcu_head *head);
 typedef void (*call_rcu_func_t)(struct rcu_head *head, rcu_callback_t func);
 
+typedef void (*swap_r_func_t)(void *a, void *b, int size, const void *priv);
 typedef void (*swap_func_t)(void *a, void *b, int size);
 
 typedef int (*cmp_r_func_t)(const void *a, const void *b, const void *priv);
diff --git a/lib/sort.c b/lib/sort.c
index aa18153864d2..f65078608c16 100644
--- a/lib/sort.c
+++ b/lib/sort.c
@@ -122,16 +122,29 @@ static void swap_bytes(void *a, void *b, size_t n)
  * a pointer, but small integers make for the smallest compare
  * instructions.
  */
-#define SWAP_WORDS_64 (swap_func_t)0
-#define SWAP_WORDS_32 (swap_func_t)1
-#define SWAP_BYTES    (swap_func_t)2
+#define SWAP_WORDS_64 (swap_r_func_t)0
+#define SWAP_WORDS_32 (swap_r_func_t)1
+#define SWAP_BYTES    (swap_r_func_t)2
+#define SWAP_WRAPPER  (swap_r_func_t)3
+
+struct wrapper {
+	cmp_func_t cmp;
+	swap_func_t swap;
+};
 
 /*
  * The function pointer is last to make tail calls most efficient if the
  * compiler decides not to inline this function.
  */
-static void do_swap(void *a, void *b, size_t size, swap_func_t swap_func)
+static void do_swap(void *a, void *b, size_t size, swap_r_func_t swap_func, const void *priv)
 {
+	const struct wrapper *w = priv;
+
+	if (swap_func == SWAP_WRAPPER) {
+		w->swap(a, b, (int)size);
+		return;
+	}
+
 	if (swap_func == SWAP_WORDS_64)
 		swap_words_64(a, b, size);
 	else if (swap_func == SWAP_WORDS_32)
@@ -139,15 +152,17 @@ static void do_swap(void *a, void *b, size_t size, swap_func_t swap_func)
 	else if (swap_func == SWAP_BYTES)
 		swap_bytes(a, b, size);
 	else
-		swap_func(a, b, (int)size);
+		swap_func(a, b, (int)size, priv);
 }
 
 #define _CMP_WRAPPER ((cmp_r_func_t)0L)
 
 static int do_cmp(const void *a, const void *b, cmp_r_func_t cmp, const void *priv)
 {
+	const struct wrapper *w = priv;
+
 	if (cmp == _CMP_WRAPPER)
-		return ((cmp_func_t)(priv))(a, b);
+		return w->cmp(a, b);
 	return cmp(a, b, priv);
 }
 
@@ -198,16 +213,20 @@ static size_t parent(size_t i, unsigned int lsbit, size_t size)
  */
 void sort_r(void *base, size_t num, size_t size,
 	    cmp_r_func_t cmp_func,
-	    swap_func_t swap_func,
+	    swap_r_func_t swap_func,
 	    const void *priv)
 {
 	/* pre-scale counters for performance */
 	size_t n = num * size, a = (num/2) * size;
 	const unsigned int lsbit = size & -size;  /* Used to find parent */
+	const struct wrapper *w = priv;
 
 	if (!a)		/* num < 2 || size == 0 */
 		return;
 
+	if (swap_func == SWAP_WRAPPER && !w->swap)
+		swap_func = NULL;
+
 	if (!swap_func) {
 		if (is_aligned(base, size, 8))
 			swap_func = SWAP_WORDS_64;
@@ -230,7 +249,7 @@ void sort_r(void *base, size_t num, size_t size,
 		if (a)			/* Building heap: sift down --a */
 			a -= size;
 		else if (n -= size)	/* Sorting: Extract root to --n */
-			do_swap(base, base + n, size, swap_func);
+			do_swap(base, base + n, size, swap_func, priv);
 		else			/* Sort complete */
 			break;
 
@@ -257,7 +276,7 @@ void sort_r(void *base, size_t num, size_t size,
 		c = b;			/* Where "a" belongs */
 		while (b != a) {	/* Shift it into place */
 			b = parent(b, lsbit, size);
-			do_swap(base + b, base + c, size, swap_func);
+			do_swap(base + b, base + c, size, swap_func, priv);
 		}
 	}
 }
@@ -267,6 +286,11 @@ void sort(void *base, size_t num, size_t size,
 	  cmp_func_t cmp_func,
 	  swap_func_t swap_func)
 {
-	return sort_r(base, num, size, _CMP_WRAPPER, swap_func, cmp_func);
+	struct wrapper w = {
+		.cmp  = cmp_func,
+		.swap = swap_func,
+	};
+
+	return sort_r(base, num, size, _CMP_WRAPPER, SWAP_WRAPPER, &w);
 }
 EXPORT_SYMBOL(sort);
-- 
2.34.1