1 files changed, 44 insertions, 2 deletions

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 9fbcff71245e..3de3901ae0a7 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -379,7 +379,48 @@ enum mem_type {
 
 static void mem_cgroup_get(struct mem_cgroup *memcg);
 static void mem_cgroup_put(struct mem_cgroup *memcg);
-static struct mem_cgroup *parent_mem_cgroup(struct mem_cgroup *memcg);
+
+/* Writing them here to avoid exposing memcg's inner layout */
+#ifdef CONFIG_CGROUP_MEM_RES_CTLR_KMEM
+#ifdef CONFIG_INET
+#include <net/sock.h>
+
+static bool mem_cgroup_is_root(struct mem_cgroup *memcg);
+void sock_update_memcg(struct sock *sk)
+{
+        /* A socket spends its whole life in the same cgroup */
+        if (sk->sk_cgrp) {
+                WARN_ON(1);
+                return;
+        }
+        if (static_branch(&memcg_socket_limit_enabled)) {
+                struct mem_cgroup *memcg;
+
+                BUG_ON(!sk->sk_prot->proto_cgroup);
+
+                rcu_read_lock();
+                memcg = mem_cgroup_from_task(current);
+                if (!mem_cgroup_is_root(memcg)) {
+                        mem_cgroup_get(memcg);
+                        sk->sk_cgrp = sk->sk_prot->proto_cgroup(memcg);
+                }
+                rcu_read_unlock();
+        }
+}
+EXPORT_SYMBOL(sock_update_memcg);
+
+void sock_release_memcg(struct sock *sk)
+{
+        if (static_branch(&memcg_socket_limit_enabled) && sk->sk_cgrp) {
+                struct mem_cgroup *memcg;
+                WARN_ON(!sk->sk_cgrp->memcg);
+                memcg = sk->sk_cgrp->memcg;
+                mem_cgroup_put(memcg);
+        }
+}
+#endif /* CONFIG_INET */
+#endif /* CONFIG_CGROUP_MEM_RES_CTLR_KMEM */
+
 static void drain_all_stock_async(struct mem_cgroup *memcg);
 
 static struct mem_cgroup_per_zone *
@@ -4932,12 +4973,13 @@ static void mem_cgroup_put(struct mem_cgroup *memcg)
 /*
  * Returns the parent mem_cgroup in memcgroup hierarchy with hierarchy enabled.
  */
-static struct mem_cgroup *parent_mem_cgroup(struct mem_cgroup *memcg)
+struct mem_cgroup *parent_mem_cgroup(struct mem_cgroup *memcg)
 {
         if (!memcg->res.parent)
                 return NULL;
         return mem_cgroup_from_res_counter(memcg->res.parent, res);
 }
+EXPORT_SYMBOL(parent_mem_cgroup);
 
 #ifdef CONFIG_CGROUP_MEM_RES_CTLR_SWAP
 static void __init enable_swap_cgroup(void)