1 files changed, 4 insertions, 4 deletions

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index bdeb82ca6c20..4c31a21a55be 100644
--- 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);
         }
 }