memcg: use css_get() in sock_update_memcg()

Use css_get/css_put instead of mem_cgroup_get/put.

Note, if at the same time someone is moving @current to a different
cgroup and removing the old cgroup, css_tryget() may return false, and
sock->sk_cgrp won't be initialized, which is fine.

Signed-off-by: Li Zefan <[email protected]>
Acked-by: KAMEZAWA Hiroyuki <[email protected]>
Acked-by: Michal Hocko <[email protected]>
Cc: Hugh Dickins <[email protected]>
Cc: Tejun Heo <[email protected]>
Cc: Glauber Costa <[email protected]>
Cc: Johannes Weiner <[email protected]>
Signed-off-by: Andrew Morton <[email protected]>
Signed-off-by: Linus Torvalds <[email protected]>

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index bdeb82ca6c20fe0cfd0a725066f2812f1c3beb90..4c31a21a55be87a1cd7b056aea1bb8de2e97c7b6 100644 (file)
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -551,15 +551,15 @@ void sock_update_memcg(struct sock *sk)
                 */
                if (sk->sk_cgrp) {
                        BUG_ON(mem_cgroup_is_root(sk->sk_cgrp->memcg));
-                       mem_cgroup_get(sk->sk_cgrp->memcg);
+                       css_get(&sk->sk_cgrp->memcg->css);
                        return;
                }
 
                rcu_read_lock();
                memcg = mem_cgroup_from_task(current);
                cg_proto = sk->sk_prot->proto_cgroup(memcg);
-               if (!mem_cgroup_is_root(memcg) && memcg_proto_active(cg_proto)) {
-                       mem_cgroup_get(memcg);
+               if (!mem_cgroup_is_root(memcg) &&
+                   memcg_proto_active(cg_proto) && css_tryget(&memcg->css)) {
                        sk->sk_cgrp = cg_proto;
                }
                rcu_read_unlock();
@@ -573,7 +573,7 @@ void sock_release_memcg(struct sock *sk)
                struct mem_cgroup *memcg;
                WARN_ON(!sk->sk_cgrp->memcg);
                memcg = sk->sk_cgrp->memcg;
-               mem_cgroup_put(memcg);
+               css_put(&sk->sk_cgrp->memcg->css);
        }
 }