[PATCH] percpu-ref: correctly get percpu pointer

[Shaohua Li <shli@kernel.org>]

I saw randam system hang testing virtio with blk-mq enabled and cpu hotplug
runing in the background. It turns out __ref_is_percpu() doesn't always return
correct percpu pointer. percpu_ref_put() calls __ref_is_percpu(), which checks
__PERCPU_REF_ATOMIC. After this check, the __PERCPU_REF_ATOMIC or
__PERCPU_REF_DEAD might be set, so we must exclude the two bits from the percpu
pointer. Fortunately we can still use percpu data for percpu_ref_put() even
this happens, because the final transistion from percpu to atomic occurs at rcu
context while __ref_is_percpu() is always called with rcu read lock protected.

CC: Jens Axboe <axboe@fb.com>
CC: Tejun Heo <tj@kernel.org>
CC: Kent Overstreet <kmo@daterainc.com>
Signed-off-by: Shaohua Li <shli@fb.com>
---
 include/linux/percpu-refcount.h | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/include/linux/percpu-refcount.h b/include/linux/percpu-refcount.h
index d5c89e0..6beee08 100644
--- a/include/linux/percpu-refcount.h
+++ b/include/linux/percpu-refcount.h
@@ -136,7 +136,14 @@ static inline bool __ref_is_percpu(struct percpu_ref *ref,
 	if (unlikely(percpu_ptr & __PERCPU_REF_ATOMIC))
 		return false;
 
-	*percpu_countp = (unsigned long __percpu *)percpu_ptr;
+	/*
+	 * At this point ATOMIC or DEAD might be set when percpu_ref_kill() is
+	 * running. It's still safe to use percpu here, because the final
+	 * transition from percpu to atomic occurs at rcu context while this
+	 * routine is protected with rcu read lock.
+	 */
+	*percpu_countp = (unsigned long __percpu *)(percpu_ptr &
+		~__PERCPU_REF_ATOMIC_DEAD);
 	return true;
 }
 
-- 
1.8.3.2