From 08d435e6ddc8cc565636c3a46e1ed698b7da684d Mon Sep 17 00:00:00 2001
From: Ronen Friedman <rfriedma@redhat.com>
Date: Sat, 16 Sep 2023 11:53:35 -0500
Subject: [PATCH] osd/scrub: introduce OsdService::get_locked_pg()

which returns an RAII-wrapper around a locked PG.

Signed-off-by: Ronen Friedman <rfriedma@redhat.com>
---
 src/osd/OSD.cc                     | 11 +++++++++++
 src/osd/OSD.h                      |  6 ++++++
 src/osd/PG.cc                      |  8 ++++++++
 src/osd/PG.h                       | 18 ++++++++++++++++++
 src/osd/scrubber/osd_scrub_sched.h | 11 +++++++++--
 src/test/osd/test_scrub_sched.cc   |  6 ++++++
 6 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/src/osd/OSD.cc b/src/osd/OSD.cc
index 255976aba0ec1..4b7ecfc00cbe0 100644
--- a/src/osd/OSD.cc
+++ b/src/osd/OSD.cc
@@ -7595,6 +7595,17 @@ void OSD::handle_fast_scrub(MOSDScrub2 *m)
   m->put();
 }
 
+std::optional<PGLockWrapper> OSDService::get_locked_pg(spg_t pgid)
+{
+  auto pg = osd->lookup_lock_pg(pgid);
+  if (pg) {
+    return PGLockWrapper{std::move(pg)};
+  } else {
+    return std::nullopt;
+  }
+}
+
+
 MPGStats* OSD::collect_pg_stats()
 {
   dout(15) << __func__ << dendl;
diff --git a/src/osd/OSD.h b/src/osd/OSD.h
index 338b390dde184..efde588d939eb 100644
--- a/src/osd/OSD.h
+++ b/src/osd/OSD.h
@@ -260,6 +260,12 @@ public:
     spg_t pgid,
     bool allow_requested_repair_only) final;
 
+  /**
+   * locks the named PG, returning an RAII wrapper that unlocks upon
+   * destruction.
+   * returns nullopt if failing to lock.
+   */
+  std::optional<PGLockWrapper> get_locked_pg(spg_t pgid) final;
 
  private:
   // -- agent shared state --
diff --git a/src/osd/PG.cc b/src/osd/PG.cc
index 79758e22c3af0..25a8ac7197845 100644
--- a/src/osd/PG.cc
+++ b/src/osd/PG.cc
@@ -2836,3 +2836,11 @@ void PG::with_heartbeat_peers(std::function<void(int)>&& f)
 uint64_t PG::get_min_alloc_size() const {
   return osd->store->get_min_alloc_size();
 }
+
+PGLockWrapper::~PGLockWrapper()
+{
+  if (m_pg) {
+    // otherwise - we were 'moved from'
+    m_pg->unlock();
+  }
+}
diff --git a/src/osd/PG.h b/src/osd/PG.h
index 3cb22b73e534d..56ac9a65bf1a0 100644
--- a/src/osd/PG.h
+++ b/src/osd/PG.h
@@ -1450,4 +1450,22 @@ public:
  }
 };
 
+/**
+ * Initialized with a locked PG. That PG is unlocked in the
+ * destructor.
+ * Used by OsdScrub when initiating a scrub.
+ */
+class PGLockWrapper {
+ public:
+  explicit PGLockWrapper(PGRef locked_pg) : m_pg{locked_pg} {}
+  PGRef pg() { return m_pg; }
+  ~PGLockWrapper();
+  PGLockWrapper(PGLockWrapper&& rhs) : m_pg(std::move(rhs.m_pg)) {
+    rhs.m_pg = nullptr;
+  }
+  PGLockWrapper(const PGLockWrapper& rhs) = delete;
+ private:
+  PGRef m_pg;
+};
+
 #endif
diff --git a/src/osd/scrubber/osd_scrub_sched.h b/src/osd/scrubber/osd_scrub_sched.h
index 94cca0f38471b..2ec36bbc43031 100644
--- a/src/osd/scrubber/osd_scrub_sched.h
+++ b/src/osd/scrubber/osd_scrub_sched.h
@@ -109,8 +109,8 @@ ScrubQueue interfaces (main functions):
 
 #include <optional>
 #include "utime.h"
-#include "scrub_job.h"
-class PG;
+#include "osd/scrubber/scrub_job.h"
+#include "osd/PG.h"
 
 namespace Scrub {
 
@@ -132,6 +132,13 @@ class ScrubSchedListener {
  public:
   virtual int get_nodeid() const = 0;  // returns the OSD number ('whoami')
 
+  /**
+   * locks the named PG, returning an RAII wrapper that unlocks upon
+   * destruction.
+   * returns nullopt if failing to lock.
+   */
+  virtual std::optional<PGLockWrapper> get_locked_pg(spg_t pgid) = 0;
+
   /**
    * A callback used by the ScrubQueue object to initiate a scrub on a specific
    * PG.
diff --git a/src/test/osd/test_scrub_sched.cc b/src/test/osd/test_scrub_sched.cc
index a3f1a348584aa..e1847619db725 100644
--- a/src/test/osd/test_scrub_sched.cc
+++ b/src/test/osd/test_scrub_sched.cc
@@ -114,6 +114,12 @@ class FakeOsd : public Scrub::ScrubSchedListener {
     m_next_response[pgid] = result;
   }
 
+  std::optional<PGLockWrapper> get_locked_pg(spg_t pgid)
+  {
+    std::ignore = pgid;
+    return std::nullopt;
+  }
+
  private:
   int m_osd_num;
   std::map<spg_t, schedule_result_t> m_next_response;
-- 
2.39.5