[PATCH v4] slab: clear OBJEXTS_ALLOC_FAIL when freeing a slab

[Hao Ge <hao.ge@linux.dev>]

From: Hao Ge <gehao@kylinos.cn>

If obj_exts allocation failed, slab->obj_exts is set to OBJEXTS_ALLOC_FAIL,
But we did not clear it when freeing the slab. Since OBJEXTS_ALLOC_FAIL and
MEMCG_DATA_OBJEXTS currently share the same bit position, during the
release of the associated folio, a VM_BUG_ON_FOLIO() check in
folio_memcg_kmem() is triggered because it was mistakenly assumed that
a valid folio->memcg_data was not cleared before freeing the folio.

When freeing a slab, we clear slab->obj_exts if the obj_ext array has been
successfully allocated. So let's clear OBJEXTS_ALLOC_FAIL when freeing
a slab if the obj_ext array allocated fail to allow them to be returned
to the buddy system more smoothly.

Fixes: 7612833192d5 ("slab: Reuse first bit for OBJEXTS_ALLOC_FAIL")
Suggested-by: Harry Yoo <harry.yoo@oracle.com>
Signed-off-by: Hao Ge <gehao@kylinos.cn>
Reviewed-by: Suren Baghdasaryan <surenb@google.com>
Acked-by: Shakeel Butt <shakeel.butt@linux.dev>
---
v4: Based on the discussion between Vlastimil and Harry,
    modify the solution to clear OBJEXTS_ALLOC_FAIL when freeing a slab.
    This does seem more reasonable. Thank you both.
---
 mm/slab.h | 26 ++++++++++++++++++++++++++
 mm/slub.c |  6 ++++++
 2 files changed, 32 insertions(+)

diff --git a/mm/slab.h b/mm/slab.h
index 078daecc7cf5..52424d6871bd 100644
--- a/mm/slab.h
+++ b/mm/slab.h
@@ -547,6 +547,28 @@ static inline struct slabobj_ext *slab_obj_exts(struct slab *slab)
 	return (struct slabobj_ext *)(obj_exts & ~OBJEXTS_FLAGS_MASK);
 }
 
+/*
+ * objexts_clear_alloc_fail -  Clear the OBJEXTS_ALLOC_FAIL for
+ * the slab object extension vector associated with a slab.
+ * @slab: a pointer to the slab struct
+ */
+static inline void objexts_clear_alloc_fail(struct slab *slab)
+{
+	unsigned long obj_exts = READ_ONCE(slab->obj_exts);
+
+#ifdef CONFIG_MEMCG
+	/*
+	 * obj_exts should be either NULL, a valid pointer with
+	 * MEMCG_DATA_OBJEXTS bit set or be equal to OBJEXTS_ALLOC_FAIL.
+	 */
+	VM_BUG_ON_PAGE(obj_exts && !(obj_exts & MEMCG_DATA_OBJEXTS) &&
+		       obj_exts != OBJEXTS_ALLOC_FAIL, slab_page(slab));
+	VM_BUG_ON_PAGE(obj_exts & MEMCG_DATA_KMEM, slab_page(slab));
+#endif
+
+	obj_exts &= ~OBJEXTS_ALLOC_FAIL;
+	WRITE_ONCE(slab->obj_exts, obj_exts);
+}
 int alloc_slab_obj_exts(struct slab *slab, struct kmem_cache *s,
                         gfp_t gfp, bool new_slab);
 
@@ -557,6 +579,10 @@ static inline struct slabobj_ext *slab_obj_exts(struct slab *slab)
 	return NULL;
 }
 
+static inline void objexts_clear_alloc_fail(struct slab *slab)
+{
+}
+
 #endif /* CONFIG_SLAB_OBJ_EXT */
 
 static inline enum node_stat_item cache_vmstat_idx(struct kmem_cache *s)
diff --git a/mm/slub.c b/mm/slub.c
index b1f15598fbfd..80166a4a62f9 100644
--- a/mm/slub.c
+++ b/mm/slub.c
@@ -2169,6 +2169,12 @@ static inline void free_slab_obj_exts(struct slab *slab)
 {
 	struct slabobj_ext *obj_exts;
 
+	/*
+	 * If obj_exts allocation failed, slab->obj_exts is set to OBJEXTS_ALLOC_FAIL,
+	 * Therefore, we should clear the OBJEXTS_ALLOC_FAIL flag first when freeing a slab.
+	 */
+	objexts_clear_alloc_fail(slab);
+
 	obj_exts = slab_obj_exts(slab);
 	if (!obj_exts)
 		return;
-- 
2.25.1