From 2260d31be951024d30ddfa75225ea25142a60ae2 Mon Sep 17 00:00:00 2001
From: secwall <secwall@yandex-team.ru>
Date: Wed, 11 Mar 2026 15:10:41 +0100
Subject: [PATCH] Make lock cache expirable

---
 internal/dcs/config.go |  2 ++
 internal/dcs/zk.go     | 12 +++++++-----
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/internal/dcs/config.go b/internal/dcs/config.go
index f51baee4..de64d9db 100644
--- a/internal/dcs/config.go
+++ b/internal/dcs/config.go
@@ -11,6 +11,7 @@ import (
 type ZookeeperConfig struct {
 	Hostname              string                   `config:"hostname" yaml:"hostname"`
 	SessionTimeout        time.Duration            `config:"session_timeout" yaml:"session_timeout"`
+	LockHeldTTL           time.Duration            `config:"lock_held_ttl" yaml:"lock_held_ttl"`
 	Namespace             string                   `config:"namespace,required" yaml:"namespace"`
 	Hosts                 []string                 `config:"hosts,required" yaml:"hosts"`
 	BackoffInterval       time.Duration            `config:"backoff_interval" yaml:"backoff_interval"`
@@ -57,6 +58,7 @@ func DefaultZookeeperConfig() (ZookeeperConfig, error) {
 	config := ZookeeperConfig{
 		Hostname:              hostname,
 		SessionTimeout:        2 * time.Second,
+		LockHeldTTL:           30 * time.Second,
 		BackoffInterval:       backoff.DefaultInitialInterval,
 		BackoffRandFactor:     backoff.DefaultRandomizationFactor,
 		BackoffMultiplier:     backoff.DefaultMultiplier,
diff --git a/internal/dcs/zk.go b/internal/dcs/zk.go
index f1a96b97..1286f7dd 100644
--- a/internal/dcs/zk.go
+++ b/internal/dcs/zk.go
@@ -340,9 +340,11 @@ func (z *zkDCS) retryDelete(path string, version int32) (err error) {
 
 func (z *zkDCS) AcquireLock(path string) bool {
 	fullPath := z.buildFullPath(path)
-	_, hasLock := z.lockHeld.Load(fullPath)
-	if hasLock {
-		return true
+	if cached, hasLock := z.lockHeld.Load(fullPath); hasLock {
+		if time.Since(cached.(time.Time)) < z.config.LockHeldTTL {
+			return true
+		}
+		z.lockHeld.Delete(fullPath)
 	}
 	self := z.getSelfLockOwner()
 	data, _, err := z.retryGet(fullPath)
@@ -362,7 +364,7 @@ func (z *zkDCS) AcquireLock(path string) bool {
 			}
 			return false
 		}
-		z.lockHeld.Store(fullPath, struct{}{})
+		z.lockHeld.Store(fullPath, time.Now())
 		return true
 	}
 	owner := LockOwner{}
@@ -371,7 +373,7 @@ func (z *zkDCS) AcquireLock(path string) bool {
 		return false
 	}
 	if owner == self {
-		z.lockHeld.Store(fullPath, struct{}{})
+		z.lockHeld.Store(fullPath, time.Now())
 		return true
 	}
 	return false