diff --git a/mm/backing-dev.c b/mm/backing-dev.c
index cd15bc27dc0620b71e07b30f3bb5212a278b7927..c260d1248e814bb9240e5f23b76927afcd99f44b 100644
--- a/mm/backing-dev.c
+++ b/mm/backing-dev.c
@@ -578,8 +578,8 @@ static struct cgroup_subsys_state *cgwbv1_get_blkcss(struct mem_cgroup *memcg)
 	struct cgroup_subsys_state *blkcg_css;
 
 	rcu_read_lock();
-	blkcg_css = memcg->wb_blk_css;
-	if (!css_tryget_online(blkcg_css)) {
+	blkcg_css = READ_ONCE(memcg->wb_blk_css);
+	if (!blkcg_css || !css_tryget_online(blkcg_css)) {
 		blkcg_css = blkcg_root_css;
 		css_get(blkcg_css);
 	}
@@ -1180,6 +1180,7 @@ static void wb_kill_memcg(struct cgroup_subsys_state *memcg_css)
 
 	list_del_init(&memcg->memcg_node);
 	css_put(memcg->wb_blk_css);
+	memcg->wb_blk_css = NULL;
 }
 
 static void wb_kill_blkcg(struct cgroup_subsys_state *blkcg_css)
@@ -1221,6 +1222,8 @@ void wb_attach_memcg_to_blkcg(struct cgroup_subsys_state *memcg_css,
 	if (!cgroup1_writeback)
 		return;
 
+	lockdep_assert_held(&cgroup_mutex);
+
 	css_get(blkcg_css);
 	memcg->wb_blk_css = blkcg_css;
 	if (pre_blkcss == NULL)
diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 70223fc704c0a6685dbe206bdbb6a0c777dea05a..d21fed66101348229ec15586279847370e625877 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -6029,6 +6029,7 @@ static ssize_t memcg_high_async_ratio_write(struct kernfs_open_file *of,
 
 static int wb_blkio_show(struct seq_file *m, void *v)
 {
+	struct kernfs_open_file *of = m->private;
 	char *path;
 	ino_t blkcg_id;
 	struct cgroup *blkcg_cgroup;
@@ -6042,22 +6043,16 @@ static int wb_blkio_show(struct seq_file *m, void *v)
 	if (!path)
 		return -ENOMEM;
 
-	rcu_read_lock();
+	if (!cgroup_kn_lock_live(of->kn, false))
+		return -ENODEV;
 	blkcg_css = memcg->wb_blk_css;
-	if (!css_tryget_online(blkcg_css)) {
-		kfree(path);
-		rcu_read_unlock();
-
-		return -EINVAL;
-	}
 	blkcg_cgroup = blkcg_css->cgroup;
 	blkcg_id = cgroup_ino(blkcg_cgroup);
 	cgroup_path(blkcg_cgroup, path, PATH_MAX);
+	cgroup_kn_unlock(of->kn);
 	seq_printf(m, "wb_blkio_path:%s\n", path);
 	seq_printf(m, "wb_blkio_ino:%lu\n", blkcg_id);
 	kfree(path);
-	css_put(blkcg_css);
-	rcu_read_unlock();
 
 	return 0;
 }
@@ -6080,24 +6075,23 @@ static ssize_t wb_blkio_write(struct kernfs_open_file *of, char *buf,
 	if (ret)
 		return ret;
 
-	rcu_read_lock();
+	if (!cgroup_kn_lock_live(of->kn, false))
+		return -ENODEV;
 	root = blkcg_root_css->cgroup->root;
 	blk_cgroup = __cgroup_get_from_id(root, cgrp_id);
 	if (IS_ERR(blk_cgroup)) {
-		rcu_read_unlock();
+		cgroup_kn_unlock(of->kn);
 		return -EINVAL;
 	}
 	blkcg_css = cgroup_tryget_css(blk_cgroup, &io_cgrp_subsys);
-	if (!blkcg_css) {
-		ret = -EINVAL;
+	if (!blkcg_css)
 		goto out_unlock;
-	}
 	wb_attach_memcg_to_blkcg(memcg_css, blkcg_css);
 	css_put(blkcg_css);
 
 out_unlock:
 	cgroup_put(blk_cgroup);
-	rcu_read_unlock();
+	cgroup_kn_unlock(of->kn);
 
 	return ret < 0 ? ret : nbytes;
 }