[PATCH FIX v0] mm: memcg/slab: Uncharge during kmem_cache_free_bulk()

[Bharata B Rao]

Object cgroup charging is done for all the objects during
allocation, but during freeing, uncharging ends up happening
for only one object in the case of bulk allocation/freeing.

Fix this by having a separate call to uncharge all the
objects from kmem_cache_free_bulk() and by modifying
memcg_slab_free_hook() to take care of bulk uncharging.

Signed-off-by: Bharata B Rao <bhar...@linux.ibm.com>
---
 mm/slab.c |  2 +-
 mm/slab.h | 42 +++++++++++++++++++++++++++---------------
 mm/slub.c |  3 ++-
 3 files changed, 30 insertions(+), 17 deletions(-)

diff --git a/mm/slab.c b/mm/slab.c
index f658e86ec8cee..5c70600d8b1cc 100644
--- a/mm/slab.c
+++ b/mm/slab.c
@@ -3440,7 +3440,7 @@ void ___cache_free(struct kmem_cache *cachep, void *objp,
                memset(objp, 0, cachep->object_size);
        kmemleak_free_recursive(objp, cachep->flags);
        objp = cache_free_debugcheck(cachep, objp, caller);
-       memcg_slab_free_hook(cachep, virt_to_head_page(objp), objp);
+       memcg_slab_free_hook(cachep, &objp, 1);
 
        /*
         * Skip calling cache_free_alien() when the platform is not numa.
diff --git a/mm/slab.h b/mm/slab.h
index 6cc323f1313af..6dd4b702888a7 100644
--- a/mm/slab.h
+++ b/mm/slab.h
@@ -345,30 +345,42 @@ static inline void memcg_slab_post_alloc_hook(struct 
kmem_cache *s,
        obj_cgroup_put(objcg);
 }
 
-static inline void memcg_slab_free_hook(struct kmem_cache *s, struct page 
*page,
-                                       void *p)
+static inline void memcg_slab_free_hook(struct kmem_cache *s_orig,
+                                       void **p, int objects)
 {
+       struct kmem_cache *s;
        struct obj_cgroup *objcg;
+       struct page *page;
        unsigned int off;
+       int i;
 
        if (!memcg_kmem_enabled())
                return;
 
-       if (!page_has_obj_cgroups(page))
-               return;
+       for (i = 0; i < objects; i++) {
+               if (unlikely(!p[i]))
+                       continue;
 
-       off = obj_to_index(s, page, p);
-       objcg = page_obj_cgroups(page)[off];
-       page_obj_cgroups(page)[off] = NULL;
+               page = virt_to_head_page(p[i]);
+               if (!page_has_obj_cgroups(page))
+                       continue;
 
-       if (!objcg)
-               return;
+               if (!s_orig)
+                       s = page->slab_cache;
+               else
+                       s = s_orig;
 
-       obj_cgroup_uncharge(objcg, obj_full_size(s));
-       mod_objcg_state(objcg, page_pgdat(page), cache_vmstat_idx(s),
-                       -obj_full_size(s));
+               off = obj_to_index(s, page, p[i]);
+               objcg = page_obj_cgroups(page)[off];
+               if (!objcg)
+                       continue;
 
-       obj_cgroup_put(objcg);
+               page_obj_cgroups(page)[off] = NULL;
+               obj_cgroup_uncharge(objcg, obj_full_size(s));
+               mod_objcg_state(objcg, page_pgdat(page), cache_vmstat_idx(s),
+                               -obj_full_size(s));
+               obj_cgroup_put(objcg);
+       }
 }
 
 #else /* CONFIG_MEMCG_KMEM */
@@ -406,8 +418,8 @@ static inline void memcg_slab_post_alloc_hook(struct 
kmem_cache *s,
 {
 }
 
-static inline void memcg_slab_free_hook(struct kmem_cache *s, struct page 
*page,
-                                       void *p)
+static inline void memcg_slab_free_hook(struct kmem_cache *s,
+                                       void **p, int objects)
 {
 }
 #endif /* CONFIG_MEMCG_KMEM */
diff --git a/mm/slub.c b/mm/slub.c
index 6d3574013b2f8..0cbe67f13946e 100644
--- a/mm/slub.c
+++ b/mm/slub.c
@@ -3091,7 +3091,7 @@ static __always_inline void do_slab_free(struct 
kmem_cache *s,
        struct kmem_cache_cpu *c;
        unsigned long tid;
 
-       memcg_slab_free_hook(s, page, head);
+       memcg_slab_free_hook(s, &head, 1);
 redo:
        /*
         * Determine the currently cpus per cpu slab.
@@ -3253,6 +3253,7 @@ void kmem_cache_free_bulk(struct kmem_cache *s, size_t 
size, void **p)
        if (WARN_ON(!size))
                return;
 
+       memcg_slab_free_hook(s, p, size);
        do {
                struct detached_freelist df;
 
-- 
2.26.2