From be760356b866d3b5a4d426cd951b4b74c87616c6 Mon Sep 17 00:00:00 2001
From: Jason Dillaman <dillaman@redhat.com>
Date: Thu, 30 Apr 2020 21:13:08 -0400
Subject: [PATCH] librbd: support a finish handler for image IO dispatch layers

This will be faster than allocating per-IO Context callbacks for
layers that need to track the in-flight state of particular IOs.

Signed-off-by: Jason Dillaman <dillaman@redhat.com>
---
 src/librbd/io/ImageDispatchSpec.cc            |  2 ++
 src/librbd/io/ImageDispatcher.cc              | 32 +++++++++++++++++++
 src/librbd/io/ImageDispatcher.h               |  3 ++
 src/librbd/io/ImageDispatcherInterface.h      |  5 +++
 src/test/librbd/mock/io/MockImageDispatcher.h |  1 +
 5 files changed, 43 insertions(+)

diff --git a/src/librbd/io/ImageDispatchSpec.cc b/src/librbd/io/ImageDispatchSpec.cc
index 285d462737d4f..4d2f6497b195e 100644
--- a/src/librbd/io/ImageDispatchSpec.cc
+++ b/src/librbd/io/ImageDispatchSpec.cc
@@ -150,11 +150,13 @@ void ImageDispatchSpec<I>::send() {
 
 template <typename I>
 void ImageDispatchSpec<I>::finish(int r) {
+  image_dispatcher->finish(r, dispatch_layer, tid);
   delete this;
 }
 
 template <typename I>
 void ImageDispatchSpec<I>::fail(int r) {
+  dispatch_result = DISPATCH_RESULT_COMPLETE;
   aio_comp->fail(r);
 }
 
diff --git a/src/librbd/io/ImageDispatcher.cc b/src/librbd/io/ImageDispatcher.cc
index 19fdd295ac79d..5ea46822bc48e 100644
--- a/src/librbd/io/ImageDispatcher.cc
+++ b/src/librbd/io/ImageDispatcher.cc
@@ -105,6 +105,38 @@ ImageDispatcher<I>::ImageDispatcher(I* image_ctx)
   this->register_dispatch(image_dispatch);
 }
 
+template <typename I>
+void ImageDispatcher<I>::finish(int r, ImageDispatchLayer image_dispatch_layer,
+                                uint64_t tid) {
+  auto cct = this->m_image_ctx->cct;
+  ldout(cct, 20) << "tid=" << tid << dendl;
+
+  // loop in reverse order from last invoked dispatch layer calling its
+  // handle_finished method
+  while (image_dispatch_layer != IMAGE_DISPATCH_LAYER_NONE) {
+    std::shared_lock locker{this->m_lock};
+    auto it = this->m_dispatches.find(image_dispatch_layer);
+    image_dispatch_layer = static_cast<ImageDispatchLayer>(
+      image_dispatch_layer - 1);
+
+    if (it == this->m_dispatches.end()) {
+      continue;
+    }
+
+    // track callback while lock is dropped so that the layer cannot be shutdown
+    auto& dispatch_meta = it->second;
+    auto dispatch = dispatch_meta.dispatch;
+    auto async_op_tracker = dispatch_meta.async_op_tracker;
+    async_op_tracker->start_op();
+    locker.unlock();
+
+    dispatch->handle_finished(r, tid);
+
+    // safe since dispatch_meta cannot be deleted until ops finished
+    async_op_tracker->finish_op();
+  }
+}
+
 template <typename I>
 bool ImageDispatcher<I>::send_dispatch(
     ImageDispatchInterface* image_dispatch,
diff --git a/src/librbd/io/ImageDispatcher.h b/src/librbd/io/ImageDispatcher.h
index 1516a76daed62..1cda3bdf5aafe 100644
--- a/src/librbd/io/ImageDispatcher.h
+++ b/src/librbd/io/ImageDispatcher.h
@@ -26,6 +26,9 @@ class ImageDispatcher : public Dispatcher<ImageCtxT, ImageDispatcherInterface> {
 public:
   ImageDispatcher(ImageCtxT* image_ctx);
 
+  void finish(int r, ImageDispatchLayer image_dispatch_layer,
+              uint64_t tid) override;
+
 protected:
   bool send_dispatch(
     ImageDispatchInterface* object_dispatch,
diff --git a/src/librbd/io/ImageDispatcherInterface.h b/src/librbd/io/ImageDispatcherInterface.h
index 7628def5ddffd..5ba8ee47d51e4 100644
--- a/src/librbd/io/ImageDispatcherInterface.h
+++ b/src/librbd/io/ImageDispatcherInterface.h
@@ -7,6 +7,7 @@
 #include "include/int_types.h"
 #include "librbd/io/DispatcherInterface.h"
 #include "librbd/io/ImageDispatchInterface.h"
+#include "librbd/io/Types.h"
 
 struct Context;
 
@@ -16,6 +17,10 @@ namespace io {
 struct ImageDispatcherInterface
   : public DispatcherInterface<ImageDispatchInterface> {
 public:
+
+  virtual void finish(int r, ImageDispatchLayer image_dispatch_layer,
+                      uint64_t tid) = 0;
+
 };
 
 } // namespace io
diff --git a/src/test/librbd/mock/io/MockImageDispatcher.h b/src/test/librbd/mock/io/MockImageDispatcher.h
index 5eefde7156a6a..90b2f072885c2 100644
--- a/src/test/librbd/mock/io/MockImageDispatcher.h
+++ b/src/test/librbd/mock/io/MockImageDispatcher.h
@@ -25,6 +25,7 @@ public:
   MOCK_METHOD2(shut_down_dispatch, void(ImageDispatchLayer, Context*));
 
   MOCK_METHOD1(send, void(ImageDispatchSpec<>*));
+  MOCK_METHOD3(finish, void(int r, ImageDispatchLayer, uint64_t));
 };
 
 } // namespace io
-- 
2.39.5