diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index 39024b7677ea754cbb3a45fdc3d0b5c02a56c939..d076c1644ea3b14a36d5e888a25370a8929fd4b1 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -1172,8 +1172,10 @@ _out:							\
 		_array = rcu_dereference(array);	\
 		_item = &_array->items[0];		\
 		while ((_prog = READ_ONCE(_item->prog))) {		\
-			bpf_cgroup_storage_set(_item->cgroup_storage);	\
+			if (unlikely(bpf_cgroup_storage_set(_item->cgroup_storage)))	\
+				break;			\
 			ret = func(_prog, ctx);		\
+			bpf_cgroup_storage_unset();	\
 			_ret &= (ret & 1);		\
 			_cn |= (ret & 2);		\
 			_item++;			\