sbitmap: add helpers for add/del wait queue handling

After commit 5d2ee7122c73, users of sbitmap that need wait queue
handling must use the provided helpers. But we only added
prepare_to_wait()/finish_wait() style helpers, add the equivalent
add_wait_queue/list_del wrappers as we..

This is needed to ensure kyber plays by the sbitmap waitqueue
rules.

Tested-by: Ming Lei <ming.lei@redhat.com>
Reviewed-by: Omar Sandoval <osandov@fb.com>
Signed-off-by: Jens Axboe <axboe@kernel.dk>

2 files changed, 40 insertions, 6 deletions

diff --git a/include/linux/sbitmap.h b/include/linux/sbitmap.h
index 03f50fcedc79..14d558146aea 100644
--- a/include/linux/sbitmap.h
+++ b/include/linux/sbitmap.h
@@ -560,13 +560,13 @@ void sbitmap_queue_wake_up(struct sbitmap_queue *sbq);
 void sbitmap_queue_show(struct sbitmap_queue *sbq, struct seq_file *m);
 
 struct sbq_wait {
-        int accounted;
+        struct sbitmap_queue *sbq;      /* if set, sbq_wait is accounted */
         struct wait_queue_entry wait;
 };
 
 #define DEFINE_SBQ_WAIT(name)                                                   \
         struct sbq_wait name = {                                                \
-                .accounted = 0,                                                 \
+                .sbq = NULL,                                                    \
                 .wait = {                                                       \
                         .private        = current,                              \
                         .func           = autoremove_wake_function,             \
@@ -588,4 +588,16 @@ void sbitmap_prepare_to_wait(struct sbitmap_queue *sbq,
 void sbitmap_finish_wait(struct sbitmap_queue *sbq, struct sbq_wait_state *ws,
                                 struct sbq_wait *sbq_wait);
 
+/*
+ * Wrapper around add_wait_queue(), which maintains some extra internal state
+ */
+void sbitmap_add_wait_queue(struct sbitmap_queue *sbq,
+                            struct sbq_wait_state *ws,
+                            struct sbq_wait *sbq_wait);
+
+/*
+ * Must be paired with sbitmap_add_wait_queue()
+ */
+void sbitmap_del_wait_queue(struct sbq_wait *sbq_wait);
+
 #endif /* __LINUX_SCALE_BITMAP_H */
diff --git a/lib/sbitmap.c b/lib/sbitmap.c
index 5b3e56d68dab..65c2d06250a6 100644
--- a/lib/sbitmap.c
+++ b/lib/sbitmap.c
@@ -671,13 +671,35 @@ void sbitmap_queue_show(struct sbitmap_queue *sbq, struct seq_file *m)
 }
 EXPORT_SYMBOL_GPL(sbitmap_queue_show);
 
+void sbitmap_add_wait_queue(struct sbitmap_queue *sbq,
+                            struct sbq_wait_state *ws,
+                            struct sbq_wait *sbq_wait)
+{
+        if (!sbq_wait->sbq) {
+                sbq_wait->sbq = sbq;
+                atomic_inc(&sbq->ws_active);
+        }
+        add_wait_queue(&ws->wait, &sbq_wait->wait);
+}
+EXPORT_SYMBOL_GPL(sbitmap_add_wait_queue);
+
+void sbitmap_del_wait_queue(struct sbq_wait *sbq_wait)
+{
+        list_del_init(&sbq_wait->wait.entry);
+        if (sbq_wait->sbq) {
+                atomic_dec(&sbq_wait->sbq->ws_active);
+                sbq_wait->sbq = NULL;
+        }
+}
+EXPORT_SYMBOL_GPL(sbitmap_del_wait_queue);
+
 void sbitmap_prepare_to_wait(struct sbitmap_queue *sbq,
                              struct sbq_wait_state *ws,
                              struct sbq_wait *sbq_wait, int state)
 {
-        if (!sbq_wait->accounted) {
+        if (!sbq_wait->sbq) {
                 atomic_inc(&sbq->ws_active);
-                sbq_wait->accounted = 1;
+                sbq_wait->sbq = sbq;
         }
         prepare_to_wait_exclusive(&ws->wait, &sbq_wait->wait, state);
 }
@@ -687,9 +709,9 @@ void sbitmap_finish_wait(struct sbitmap_queue *sbq, struct sbq_wait_state *ws,
                          struct sbq_wait *sbq_wait)
 {
         finish_wait(&ws->wait, &sbq_wait->wait);
-        if (sbq_wait->accounted) {
+        if (sbq_wait->sbq) {
                 atomic_dec(&sbq->ws_active);
-                sbq_wait->accounted = 0;
+                sbq_wait->sbq = NULL;
         }
 }
 EXPORT_SYMBOL_GPL(sbitmap_finish_wait);