[PATCH slab/for-next v4 6/8] mm/slab: introduce struct kvfree_rcu_head for kvfree_rcu batching

["Harry Yoo (Oracle)" <harry@kernel.org>]

rcu_head is overkill for kvfree_rcu() because the callback
function is always either kfree(), vfree(), or free_large_kmalloc(),
and thus there is no need for a function pointer.

kvfree_rcu batching reuses the field to store the start address
of an object, however, this is not strictly needed because we can
calculate the start address in the slowpath. For the purpose of
kvfree_rcu batching, it is sufficient to implement a linked list using
a single pointer.

Introduce a new struct called kvfree_rcu_head (the name was suggested
by Vlastimil Babka), which is similar to rcu_head but is only a single
pointer to build a linked list, without a function pointer, when
CONFIG_KVFREE_RCU_BATCHED=y.

When kvfree_rcu is not batched, kvfree_rcu_head is the same size
as rcu_head. Note that shrinking struct kvfree_rcu_head on
CONFIG_KVFREE_RCU_BATCHED=n kernels would inevitably require additional
complexity and also some sort of batching (which defeats the purpose of
the config option) because it cannot fall back to call_rcu().

For now there are no user-visible changes to the API. k[v]free_rcu()
simply casts rcu_head to kvfree_rcu_head. While this does not affect
the API, it allows kfree_rcu_nolock() to reuse kvfree_rcu batching
as a fallback when trylock or sheaf allocation fails.

Stop storing the object pointer in rcu_head.func and instead calculate
the object's start address in kvfree_rcu_list(). Factor out the existing
logic to calculate the start address from kvfree_rcu_cb() to
kvmalloc_obj_start_addr().

Signed-off-by: Harry Yoo (Oracle) <harry@kernel.org>
---
 include/linux/rcupdate.h   | 10 ++++++----
 include/linux/types.h      |  8 ++++++++
 include/trace/events/rcu.h |  2 +-
 mm/slab.h                  | 28 ++++++++++++++++++++++++++++
 mm/slab_common.c           | 17 ++++++++---------
 mm/slub.c                  | 40 +++++++++-------------------------------
 6 files changed, 60 insertions(+), 45 deletions(-)

diff --git a/include/linux/rcupdate.h b/include/linux/rcupdate.h
index 5e95acc33989..ef5bb6981133 100644
--- a/include/linux/rcupdate.h
+++ b/include/linux/rcupdate.h
@@ -1098,19 +1098,21 @@ static inline void rcu_read_unlock_migrate(void)
 /*
  * In mm/slab_common.c, no suitable header to include here.
  */
-void kvfree_call_rcu(struct rcu_head *head, void *ptr);
+void kvfree_call_rcu(struct kvfree_rcu_head *head, void *ptr);
 
 /*
  * The BUILD_BUG_ON() makes sure the rcu_head offset can be handled. See the
  * comment of kfree_rcu() for details.
  */
-#define kvfree_rcu_arg_2(ptr, rhf)					\
+#define kvfree_rcu_arg_2(ptr, kvrhf)					\
 do {									\
 	typeof (ptr) ___p = (ptr);					\
+	struct kvfree_rcu_head *___head;				\
 									\
 	if (___p) {							\
-		BUILD_BUG_ON(offsetof(typeof(*(ptr)), rhf) >= 4096);	\
-		kvfree_call_rcu(&((___p)->rhf), (void *) (___p));	\
+		BUILD_BUG_ON(offsetof(typeof(*(ptr)), kvrhf) >= 4096);	\
+		___head = (struct kvfree_rcu_head *) &(___p)->kvrhf;	\
+		kvfree_call_rcu(___head, (void *) (___p));		\
 	}								\
 } while (0)
 
diff --git a/include/linux/types.h b/include/linux/types.h
index 93166b0b0617..8b0976669d66 100644
--- a/include/linux/types.h
+++ b/include/linux/types.h
@@ -255,6 +255,14 @@ struct callback_head {
 } __attribute__((aligned(sizeof(void *))));
 #define rcu_head callback_head
 
+#ifdef CONFIG_KVFREE_RCU_BATCHED
+struct kvfree_rcu_head {
+	struct kvfree_rcu_head *next;
+};
+#else
+#define kvfree_rcu_head rcu_head
+#endif
+
 typedef void (*rcu_callback_t)(struct rcu_head *head);
 typedef void (*call_rcu_func_t)(struct rcu_head *head, rcu_callback_t func);
 
diff --git a/include/trace/events/rcu.h b/include/trace/events/rcu.h
index 5fbdabe3faea..a74027126a2f 100644
--- a/include/trace/events/rcu.h
+++ b/include/trace/events/rcu.h
@@ -625,7 +625,7 @@ TRACE_EVENT_RCU(rcu_invoke_callback,
  */
 TRACE_EVENT_RCU(rcu_invoke_kvfree_callback,
 
-	TP_PROTO(const char *rcuname, struct rcu_head *rhp, unsigned long offset),
+	TP_PROTO(const char *rcuname, struct kvfree_rcu_head *rhp, unsigned long offset),
 
 	TP_ARGS(rcuname, rhp, offset),
 
diff --git a/mm/slab.h b/mm/slab.h
index fdd293b3efa5..1efeaceb7db3 100644
--- a/mm/slab.h
+++ b/mm/slab.h
@@ -351,6 +351,33 @@ static inline int objs_per_slab(const struct kmem_cache *cache,
 	return slab->objects;
 }
 
+/* kvfree_rcu_head offset can be only less than page size */
+static inline void *kvmalloc_obj_start_addr(void *head)
+{
+	void *obj = head;
+
+	if (unlikely(is_vmalloc_addr(obj))) {
+		obj = (void *) PAGE_ALIGN_DOWN((unsigned long)obj);
+	} else {
+		struct page *page = virt_to_page(obj);
+		struct slab *slab = page_slab(page);
+
+		if (!slab) {
+			obj = (void *) PAGE_ALIGN_DOWN((unsigned long)obj);
+		} else if (is_kfence_address(obj)) {
+			obj = kfence_object_start(obj);
+		} else {
+			struct kmem_cache *s = slab->slab_cache;
+			unsigned int idx = __obj_to_index(s, slab_address(slab), obj);
+
+			obj = slab_address(slab) + s->size * idx;
+			obj = fixup_red_left(s, obj);
+		}
+	}
+
+	return obj;
+}
+
 /*
  * State of the slab allocator.
  *
@@ -761,6 +788,7 @@ void __check_heap_object(const void *ptr, unsigned long n,
 			 const struct slab *slab, bool to_user);
 
 void deferred_work_barrier(void);
+void defer_kfree_rcu(struct kvfree_rcu_head *head);
 
 static inline bool slub_debug_orig_size(struct kmem_cache *s)
 {
diff --git a/mm/slab_common.c b/mm/slab_common.c
index 8c631bf97cd5..c1c32909fa52 100644
--- a/mm/slab_common.c
+++ b/mm/slab_common.c
@@ -1346,7 +1346,7 @@ struct kvfree_rcu_bulk_data {
 
 struct kfree_rcu_cpu_work {
 	struct rcu_work rcu_work;
-	struct rcu_head *head_free;
+	struct kvfree_rcu_head *head_free;
 	struct rcu_gp_oldstate head_free_gp_snap;
 	struct list_head bulk_head_free[FREE_N_CHANNELS];
 	struct kfree_rcu_cpu *krcp;
@@ -1382,7 +1382,7 @@ struct kfree_rcu_cpu_work {
 struct kfree_rcu_cpu {
 	// Objects queued on a linked list
 	// through their rcu_head structures.
-	struct rcu_head *head;
+	struct kvfree_rcu_head *head;
 	unsigned long head_gp_snap;
 	atomic_t head_count;
 
@@ -1523,12 +1523,12 @@ kvfree_rcu_bulk(struct kfree_rcu_cpu *krcp,
 }
 
 static void
-kvfree_rcu_list(struct rcu_head *head)
+kvfree_rcu_list(struct kvfree_rcu_head *head)
 {
-	struct rcu_head *next;
+	struct kvfree_rcu_head *next;
 
 	for (; head; head = next) {
-		void *ptr = (void *) head->func;
+		void *ptr = kvmalloc_obj_start_addr(head);
 		unsigned long offset = (void *) head - ptr;
 
 		next = head->next;
@@ -1552,7 +1552,7 @@ static void kfree_rcu_work(struct work_struct *work)
 	unsigned long flags;
 	struct kvfree_rcu_bulk_data *bnode, *n;
 	struct list_head bulk_head[FREE_N_CHANNELS];
-	struct rcu_head *head;
+	struct kvfree_rcu_head *head;
 	struct kfree_rcu_cpu *krcp;
 	struct kfree_rcu_cpu_work *krwp;
 	struct rcu_gp_oldstate head_gp_snap;
@@ -1683,7 +1683,7 @@ kvfree_rcu_drain_ready(struct kfree_rcu_cpu *krcp)
 {
 	struct list_head bulk_ready[FREE_N_CHANNELS];
 	struct kvfree_rcu_bulk_data *bnode, *n;
-	struct rcu_head *head_ready = NULL;
+	struct kvfree_rcu_head *head_ready = NULL;
 	unsigned long flags;
 	int i;
 
@@ -1946,7 +1946,7 @@ void __init kfree_rcu_scheduler_running(void)
  * be free'd in workqueue context. This allows us to: batch requests together to
  * reduce the number of grace periods during heavy kfree_rcu()/kvfree_rcu() load.
  */
-void kvfree_call_rcu(struct rcu_head *head, void *ptr)
+void kvfree_call_rcu(struct kvfree_rcu_head *head, void *ptr)
 {
 	unsigned long flags;
 	struct kfree_rcu_cpu *krcp;
@@ -1984,7 +1984,6 @@ void kvfree_call_rcu(struct rcu_head *head, void *ptr)
 			// Inline if kvfree_rcu(one_arg) call.
 			goto unlock_return;
 
-		head->func = ptr;
 		head->next = krcp->head;
 		WRITE_ONCE(krcp->head, head);
 		atomic_inc(&krcp->head_count);
diff --git a/mm/slub.c b/mm/slub.c
index deac315d0f23..d0976396d703 100644
--- a/mm/slub.c
+++ b/mm/slub.c
@@ -6699,43 +6699,21 @@ static void free_large_kmalloc(struct page *page, void *object)
  */
 void kvfree_rcu_cb(struct rcu_head *head)
 {
-	void *obj = head;
-	struct page *page;
-	struct slab *slab;
-	struct kmem_cache *s;
-	void *slab_addr;
+	void *obj;
+
+	obj = kvmalloc_obj_start_addr(head);
 
 	if (is_vmalloc_addr(obj)) {
-		obj = (void *) PAGE_ALIGN_DOWN((unsigned long)obj);
 		vfree(obj);
-		return;
-	}
-
-	page = virt_to_page(obj);
-	slab = page_slab(page);
-	if (!slab) {
-		/*
-		 * rcu_head offset can be only less than page size so no need to
-		 * consider allocation order
-		 */
-		obj = (void *) PAGE_ALIGN_DOWN((unsigned long)obj);
-		free_large_kmalloc(page, obj);
-		return;
-	}
-
-	s = slab->slab_cache;
-	slab_addr = slab_address(slab);
-
-	if (is_kfence_address(obj)) {
-		obj = kfence_object_start(obj);
 	} else {
-		unsigned int idx = __obj_to_index(s, slab_addr, obj);
+		struct page *page = virt_to_page(obj);
+		struct slab *slab = page_slab(page);
 
-		obj = slab_addr + s->size * idx;
-		obj = fixup_red_left(s, obj);
+		if (slab)
+			slab_free(slab->slab_cache, slab, obj, _RET_IP_);
+		else
+			free_large_kmalloc(page, obj);
 	}
-
-	slab_free(s, slab, obj, _RET_IP_);
 }
 
 /**

-- 
2.53.0