Re: [PATCH mm v3] memcg: enable memory accounting in __alloc_pages_bulk

Enable memory accounting for bulk page allocator.

Fixes: 387ba26fb1cb ("mm/page_alloc: add a bulk page allocator")
Cc: <redacted>
Signed-off-by: Vasily Averin <redacted>
---
v3: added comments,
    removed call of post charge hook for nr_pages = 0
v2: modified according to Shakeel Butt's remarks
---
 include/linux/memcontrol.h | 11 +++++++++
 mm/memcontrol.c            | 48 +++++++++++++++++++++++++++++++++++++-
 mm/page_alloc.c            | 21 ++++++++++++++++-
 3 files changed, 78 insertions(+), 2 deletions(-)

diff --git a/include/linux/memcontrol.h b/include/linux/memcontrol.h
index 3096c9a0ee01..990acd70c846 100644
--- a/include/linux/memcontrol.h
+++ b/include/linux/memcontrol.h

@@ -810,6 +810,12 @@ static inline void obj_cgroup_put(struct obj_cgroup *objcg)
        percpu_ref_put(&objcg->refcnt);
 }

+static inline void obj_cgroup_put_many(struct obj_cgroup *objcg,
+                                      unsigned long nr)
+{
+       percpu_ref_put_many(&objcg->refcnt, nr);
+}
+
 static inline void mem_cgroup_put(struct mem_cgroup *memcg)
 {
        if (memcg)

@@ -1746,4 +1752,9 @@ static inline struct mem_cgroup *mem_cgroup_from_obj(void *p)

 #endif /* CONFIG_MEMCG_KMEM */

+bool memcg_bulk_pre_charge_hook(struct obj_cgroup **objcgp, gfp_t gfp,
+                               unsigned int nr_pages);
+void memcg_bulk_charge_hook(struct obj_cgroup *objcgp, struct page *page);
+void memcg_bulk_post_charge_hook(struct obj_cgroup *objcg,
+                                unsigned int nr_pages);
 #endif /* _LINUX_MEMCONTROL_H */

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 87e41c3cac10..16fe3384c12c 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c

@@ -3239,7 +3239,53 @@ void obj_cgroup_uncharge(struct obj_cgroup *objcg, size_t size)
        refill_obj_stock(objcg, size, true);
 }

-#endif /* CONFIG_MEMCG_KMEM */
+bool memcg_bulk_pre_charge_hook(struct obj_cgroup **objcgp, gfp_t gfp,
+                               unsigned int nr_pages)
+{
+       struct obj_cgroup *objcg = NULL;

+
+       if (!memcg_kmem_enabled() || !(gfp & __GFP_ACCOUNT))
+               return true;
+
+       objcg = get_obj_cgroup_from_current();
+
+       if (objcg && obj_cgroup_charge_pages(objcg, gfp, nr_pages)) {
+               obj_cgroup_put(objcg);
+               return false;
+       }

+       obj_cgroup_get_many(objcg, nr_pages - 1);
+       *objcgp = objcg;
+       return true;
+}
+
+void memcg_bulk_charge_hook(struct obj_cgroup *objcg, struct page *page)
+{
+       page->memcg_data = (unsigned long)objcg | MEMCG_DATA_KMEM;
+}
+
+void memcg_bulk_post_charge_hook(struct obj_cgroup *objcg,
+                                unsigned int nr_pages)
+{
+       obj_cgroup_uncharge_pages(objcg, nr_pages);
+       obj_cgroup_put_many(objcg, nr_pages);
+}

+#else /* !CONFIG_MEMCG_KMEM */
+bool memcg_bulk_pre_charge_hook(struct obj_cgroup **objcgp, gfp_t gfp,
+                               unsigned int nr_pages)
+{
+       return true;
+}
+
+void memcg_bulk_charge_hook(struct obj_cgroup *objcgp, struct page *page)
+{
+}
+
+void memcg_bulk_post_charge_hook(struct obj_cgroup *objcg,
+                                unsigned int nr_pages)
+{
+}
+#endif
+