feat core: add prolong() to dist lock strategy for cheaper TTL refresh

commit_hash:177bb59df1c01236aa4a5228ed331745f2a4b8e0

Author: oigolovanova

core/include/userver/dist_lock/dist_lock_strategy.hpp

@@ -34,6 +34,24 @@ class DistLockStrategyBase {
     /// cleanup, Release won't be invoked.
     virtual void Acquire(std::chrono::milliseconds lock_ttl, const std::string& locker_id) = 0;
 
+    /// Prolongs (refreshes the TTL of) a lock already held by @a locker_id.
+    ///
+    /// Called instead of Acquire() on every refresh while the lock is held, so
+    /// backends may implement a cheaper query than the full Acquire().
+    ///
+    /// @param lock_ttl The new duration for which the lock must be held.
+    /// @param locker_id Globally unique ID of the locking entity, must be the
+    /// same as in Acquire().
+    /// @throws LockIsAcquiredByAnotherHostException when the lock is no longer
+    /// held by @a locker_id (ownership was lost)
+    /// @throws anything else when the prolongation fails, strategy is responsible
+    /// for cleanup, Release won't be invoked.
+    /// @note The default implementation simply calls Acquire(), preserving the
+    /// legacy behaviour for strategies that do not override it.
+    virtual void Prolong(std::chrono::milliseconds lock_ttl, const std::string& locker_id) {
+        Acquire(lock_ttl, locker_id);
+    }
+
     /// Releases the lock.
     ///
     /// @param locker_id Globally unique ID of the locking entity, must be the

core/src/dist_lock/dist_lock_test.cpp

@@ -48,6 +48,11 @@ class MockDistLockStrategy final : public dist_lock::DistLockStrategyBase {
         *locked_by = locker_id;
     }
 
+    void Prolong(std::chrono::milliseconds lock_ttl, const std::string& locker_id) override {
+        Acquire(lock_ttl, locker_id);
+        prolongs_++;
+    }
+
     void Release(const std::string& locker_id) override {
         auto locked_by = locked_by_var_.Lock();
         if (*locked_by == locker_id) {
@@ -69,15 +74,51 @@ class MockDistLockStrategy final : public dist_lock::DistLockStrategyBase {
 
     size_t GetAttemptsCount() const { return attempts_; }
 
+    size_t GetProlongsCount() const { return prolongs_; }
+
 private:
     concurrent::Variable<std::string> locked_by_var_;
     std::atomic<bool> allowed_{false};
     std::atomic<size_t> attempts_{0};
+    std::atomic<size_t> prolongs_{0};
 };
 /// [Sample distlock strategy]
 
 auto MakeMockStrategy() { return std::make_shared<MockDistLockStrategy>(); }
 
+// A strategy that intentionally does NOT override Prolong(), so refreshes fall
+// back to the default DistLockStrategyBase::Prolong() (which calls Acquire()).
+class DefaultProlongStrategy final : public dist_lock::DistLockStrategyBase {
+public:
+    ~DefaultProlongStrategy() override {
+        auto locked_by = locked_by_var_.Lock();
+        EXPECT_TRUE(locked_by->empty());
+    }
+
+    void Acquire(std::chrono::milliseconds, const std::string& locker_id) override {
+        UASSERT(!locker_id.empty());
+        acquires_++;
+        auto locked_by = locked_by_var_.Lock();
+        if (!locked_by->empty() && *locked_by != locker_id) {
+            throw dist_lock::LockIsAcquiredByAnotherHostException();
+        }
+        *locked_by = locker_id;
+    }
+
+    void Release(const std::string& locker_id) override {
+        auto locked_by = locked_by_var_.Lock();
+        if (*locked_by == locker_id) {
+            locked_by->clear();
+        }
+    }
+
+    size_t GetAcquiresCount() const { return acquires_; }
+
+private:
+    concurrent::Variable<std::string> locked_by_var_;
+    std::atomic<size_t> acquires_{0};
+};
+
 class DistLockWorkload {
 public:
     explicit DistLockWorkload(bool abort_on_cancel = false)
@@ -470,4 +511,40 @@ UTEST_MT(LockedTask, SingleAttemptExample, 3) {
     /// [Sample distributed locked task SingleAttempt]
 }
 
+UTEST_MT(LockedWorker, ProlongUsedForRefresh, 3) {
+    auto strategy = MakeMockStrategy();
+    DistLockWorkload work;
+    dist_lock::DistLockedWorker locked_worker(kWorkerName, [&] { work.Work(); }, strategy, MakeSettings());
+
+    locked_worker.Start();
+    strategy->Allow(true);
+    EXPECT_TRUE(work.WaitForLocked(true, utest::kMaxTestWaitTime));
+
+    for (int i = 0; i < 100 && strategy->GetProlongsCount() < 2; ++i) {
+        engine::InterruptibleSleepFor(kAttemptInterval);
+    }
+
+    EXPECT_LE(2U, strategy->GetProlongsCount());
+    EXPECT_LT(strategy->GetProlongsCount(), strategy->GetAttemptsCount());
+
+    locked_worker.Stop();
+}
+
+UTEST_MT(LockedWorker, DefaultProlongRoutesThroughAcquire, 3) {
+    auto strategy = std::make_shared<DefaultProlongStrategy>();
+    DistLockWorkload work;
+    dist_lock::DistLockedWorker locked_worker(kWorkerName, [&] { work.Work(); }, strategy, MakeSettings());
+
+    locked_worker.Start();
+    EXPECT_TRUE(work.WaitForLocked(true, utest::kMaxTestWaitTime));
+
+    const auto acquires_after_take = strategy->GetAcquiresCount();
+    for (int i = 0; i < 100 && strategy->GetAcquiresCount() <= acquires_after_take; ++i) {
+        engine::InterruptibleSleepFor(kAttemptInterval);
+    }
+    EXPECT_LT(acquires_after_take, strategy->GetAcquiresCount());
+
+    locked_worker.Stop();
+}
+
 USERVER_NAMESPACE_END

core/src/dist_lock/impl/locker.cpp

@@ -111,7 +111,11 @@ void Locker::Run(LockerMode mode, dist_lock::DistLockWaitingMode waiting_mode, t
 
         try {
             if (settings.is_enabled || is_locked_) {
-                strategy_->Acquire(settings.lock_ttl, Id());
+                if (is_locked_) {
+                    strategy_->Prolong(settings.lock_ttl, Id());
+                } else {
+                    strategy_->Acquire(settings.lock_ttl, Id());
+                }
                 stats_.lock_successes++;
                 if (!ExchangeLockState(true, attempt_start)) {
                     LOG_DEBUG() << "Starting watchdog task";