percpu-refcount: implement percpu_ref_tryget()

Implement percpu_ref_tryget() which fails if the refcnt already
reached zero.  Note that this is different from the recently renamed
percpu_ref_tryget_live() which fails if the refcnt has been killed and
is draining the remaining references.  percpu_ref_tryget() succeeds on
a killed refcnt as long as its current refcnt is above zero.

Signed-off-by: Tejun Heo <tj@kernel.org>
Acked-by: Kent Overstreet <kmo@daterainc.com>

diff --git a/include/linux/percpu-refcount.h b/include/linux/percpu-refcount.h
index e22d15597cc37deafc079a66f07283a7112bf81b..dba35c411e8c06ab9b3cd5732d3258e4339b55c0 100644 (file)
--- a/include/linux/percpu-refcount.h
+++ b/include/linux/percpu-refcount.h
@@ -117,6 +117,36 @@ static inline void percpu_ref_get(struct percpu_ref *ref)
        rcu_read_unlock_sched();
 }
 
+/**
+ * percpu_ref_tryget - try to increment a percpu refcount
+ * @ref: percpu_ref to try-get
+ *
+ * Increment a percpu refcount unless its count already reached zero.
+ * Returns %true on success; %false on failure.
+ *
+ * The caller is responsible for ensuring that @ref stays accessible.
+ */
+static inline bool percpu_ref_tryget(struct percpu_ref *ref)
+{
+       unsigned __percpu *pcpu_count;
+       int ret = false;
+
+       rcu_read_lock_sched();
+
+       pcpu_count = ACCESS_ONCE(ref->pcpu_count);
+
+       if (likely(REF_STATUS(pcpu_count) == PCPU_REF_PTR)) {
+               __this_cpu_inc(*pcpu_count);
+               ret = true;
+       } else {
+               ret = atomic_inc_not_zero(&ref->count);
+       }
+
+       rcu_read_unlock_sched();
+
+       return ret;
+}
+
 /**
  * percpu_ref_tryget_live - try to increment a live percpu refcount
  * @ref: percpu_ref to try-get
@@ -128,6 +158,8 @@ static inline void percpu_ref_get(struct percpu_ref *ref)
  * will fail.  For such guarantee, percpu_ref_kill_and_confirm() should be
  * used.  After the confirm_kill callback is invoked, it's guaranteed that
  * no new reference will be given out by percpu_ref_tryget().
+ *
+ * The caller is responsible for ensuring that @ref stays accessible.
  */
 static inline bool percpu_ref_tryget_live(struct percpu_ref *ref)
 {