We are able to match a cache allocation to a particular memcg.  If the
task doesn't change groups during the allocation itself - a rare event,
this will give us a good picture about who is the first group to touch a
cache page.

This patch uses the now available infrastructure by calling
memcg_kmem_get_cache() before all the cache allocations.

Signed-off-by: Glauber Costa <redacted>
CC: Christoph Lameter <redacted>
CC: Pekka Enberg <penberg-bbCR+/B0CizivPeTLB3BmA@public.gmane.org>
CC: Michal Hocko <redacted>
CC: Kamezawa Hiroyuki <redacted>
CC: Johannes Weiner <redacted>
CC: Suleiman Souhlal <redacted>
CC: Tejun Heo <redacted>
---
 include/linux/slub_def.h | 15 ++++++++++-----
 mm/memcontrol.c          |  3 +++
 mm/slab.c                |  6 +++++-
 mm/slub.c                |  5 +++--
 4 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/include/linux/slub_def.h b/include/linux/slub_def.h
index 961e72e..ed330df 100644
--- a/include/linux/slub_def.h
+++ b/include/linux/slub_def.h

@@ -13,6 +13,8 @@
 #include <linux/kobject.h>

 #include <linux/kmemleak.h>
+#include <linux/memcontrol.h>
+#include <linux/mm.h>

 enum stat_item {
        ALLOC_FASTPATH,         /* Allocation from cpu slab */

@@ -209,14 +211,14 @@ static __always_inline int kmalloc_index(size_t size)
  * This ought to end up with a global pointer to the right cache
  * in kmalloc_caches.
  */
-static __always_inline struct kmem_cache *kmalloc_slab(size_t size)
+static __always_inline struct kmem_cache *kmalloc_slab(gfp_t flags, size_t size)
 {
        int index = kmalloc_index(size);

        if (index == 0)
                return NULL;

-       return kmalloc_caches[index];
+       return memcg_kmem_get_cache(kmalloc_caches[index], flags);
 }