From 84dec4729a4168f131f4c758b6dfd916c0c0dfe4 Mon Sep 17 00:00:00 2001
From: Jason Dillaman <dillaman@redhat.com>
Date: Tue, 13 Sep 2016 16:38:51 -0400
Subject: [PATCH] librbd: helper class for quiescing in-flight async ops

Signed-off-by: Jason Dillaman <dillaman@redhat.com>
---
 src/librbd/Utils.h | 36 ++++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/src/librbd/Utils.h b/src/librbd/Utils.h
index 6bd4320073b9f..b098881d6f5a9 100644
--- a/src/librbd/Utils.h
+++ b/src/librbd/Utils.h
@@ -161,6 +161,42 @@ inline ImageCtx *get_image_ctx(ImageCtx *image_ctx) {
   return image_ctx;
 }
 
+/// helper for tracking in-flight async ops when coordinating
+/// a shut down of the invoking class instance
+class AsyncOpTracker {
+public:
+  AsyncOpTracker() : m_refs(0) {
+  }
+
+  void start_op() {
+    m_refs.inc();
+  }
+
+  void finish_op() {
+    if (m_refs.dec() == 0 && m_on_finish != nullptr) {
+      Context *on_finish = nullptr;
+      std::swap(on_finish, m_on_finish);
+      on_finish->complete(0);
+    }
+  }
+
+  template <typename I>
+  void wait(I &image_ctx, Context *on_finish) {
+    assert(m_on_finish == nullptr);
+
+    on_finish = create_async_context_callback(image_ctx, on_finish);
+    if (m_refs.read() == 0) {
+      on_finish->complete(0);
+      return;
+    }
+    m_on_finish = on_finish;
+  }
+
+private:
+  atomic_t m_refs;
+  Context *m_on_finish = nullptr;
+};
+
 } // namespace util
 } // namespace librbd