From e870a13ea09728e1b271358534be403b2276fdef Mon Sep 17 00:00:00 2001
From: Ivo Dilov <iddilov@gmail.com>
Date: Wed, 4 Dec 2024 15:15:43 +0200
Subject: [PATCH] Allow setting the on_lock_lost after construction

This will be useful if we want to construct a ReliableStorageLockGuard
but decide later what to do in case a lock is lost.

Also changes ReliableStorageLockGuard to hold a copy of the lock instead
of just a reference. The lock is a cheap object to copy, so it makes
easier to avoid errors where lock is destructed before the guard for it.
---
 cpp/arcticdb/util/reliable_storage_lock-inl.hpp | 14 ++++++++++++--
 cpp/arcticdb/util/reliable_storage_lock.hpp     | 10 +++++++---
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/cpp/arcticdb/util/reliable_storage_lock-inl.hpp b/cpp/arcticdb/util/reliable_storage_lock-inl.hpp
index 3ae517051d..83e7bf5985 100644
--- a/cpp/arcticdb/util/reliable_storage_lock-inl.hpp
+++ b/cpp/arcticdb/util/reliable_storage_lock-inl.hpp
@@ -189,7 +189,7 @@ std::optional<AcquiredLockId> ReliableStorageLock<ClockType>::try_take_next_id(c
     return lock_id;
 }
 
-inline ReliableStorageLockGuard::ReliableStorageLockGuard(const ReliableStorageLock<> &lock, AcquiredLockId acquired_lock, folly::Func&& on_lost_lock) :
+inline ReliableStorageLockGuard::ReliableStorageLockGuard(const ReliableStorageLock<> &lock, AcquiredLockId acquired_lock, std::optional<folly::Func>&& on_lost_lock) :
         lock_(lock), acquired_lock_(std::nullopt), on_lost_lock_(std::move(on_lost_lock)) {
     acquired_lock_ = acquired_lock;
     // We heartbeat 5 times per lock timeout to extend the lock.
@@ -211,7 +211,9 @@ inline ReliableStorageLockGuard::ReliableStorageLockGuard(const ReliableStorageL
 inline void ReliableStorageLockGuard::cleanup_on_lost_lock() {
     // We do not use shutdown because we don't want to run it from within a FunctionScheduler thread to avoid a deadlock
     extend_lock_heartbeat_.cancelAllFunctions();
-    on_lost_lock_();
+    if (on_lost_lock_.has_value()){
+        on_lost_lock_.value()();
+    }
 }
 
 inline ReliableStorageLockGuard::~ReliableStorageLockGuard() {
@@ -221,6 +223,14 @@ inline ReliableStorageLockGuard::~ReliableStorageLockGuard() {
     }
 }
 
+inline void ReliableStorageLockGuard::set_on_lost_lock(folly::Func &&on_lost_lock) {
+    on_lost_lock_ = std::make_optional<folly::Func>(std::move(on_lost_lock));
+    if (!acquired_lock_.has_value()) {
+        // Lock was lost before we set on_lost_lock. Running callback immediately.
+        on_lost_lock_.value()();
+    }
+}
+
 inline void ReliableStorageLockManager::take_lock_guard(const ReliableStorageLock<> &lock) {
     auto acquired = lock.retry_until_take_lock();
     guard = std::make_shared<ReliableStorageLockGuard>(lock, acquired, [](){
diff --git a/cpp/arcticdb/util/reliable_storage_lock.hpp b/cpp/arcticdb/util/reliable_storage_lock.hpp
index 8008a97f79..4628a67b2b 100644
--- a/cpp/arcticdb/util/reliable_storage_lock.hpp
+++ b/cpp/arcticdb/util/reliable_storage_lock.hpp
@@ -33,6 +33,7 @@ template <class ClockType = util::SysClock>
 class ReliableStorageLock {
 public:
     ReliableStorageLock(const std::string& base_name, const std::shared_ptr<Store> store, timestamp timeout);
+    ReliableStorageLock(const ReliableStorageLock<ClockType>& other) = default;
 
     AcquiredLockId retry_until_take_lock() const;
     std::optional<AcquiredLockId> try_take_lock() const;
@@ -55,14 +56,17 @@ class ReliableStorageLock {
 // via the on_lock_lost.
 class ReliableStorageLockGuard {
 public:
-    ReliableStorageLockGuard(const ReliableStorageLock<> &lock, AcquiredLockId acquired_lock, folly::Func&& on_lost_lock);
+    ReliableStorageLockGuard(const ReliableStorageLock<>& lock, AcquiredLockId acquired_lock, std::optional<folly::Func>&& on_lost_lock);
 
     ~ReliableStorageLockGuard();
+
+    // Will immediately trigger [on_lost_lock] if lock is already lost.
+    void set_on_lost_lock(folly::Func&& on_lost_lock);
 private:
     void cleanup_on_lost_lock();
-    const ReliableStorageLock<> &lock_;
+    const ReliableStorageLock<> lock_;
     std::optional<AcquiredLockId> acquired_lock_;
-    folly::Func on_lost_lock_;
+    std::optional<folly::Func> on_lost_lock_;
     folly::FunctionScheduler extend_lock_heartbeat_;
 };