[cli,lock] Allow wait_for_lock (#10784)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Saugat Pachhai (सौगात) <[email protected]>

dvc/cli/command.py

@@ -18,7 +18,10 @@ class CmdBase(ABC):
 
         os.chdir(args.cd)
 
-        self.repo: Repo = Repo(uninitialized=self.UNINITIALIZED)
+        self.repo: Repo = Repo(
+            uninitialized=self.UNINITIALIZED,
+            _wait_for_lock=args.wait_for_lock,
+        )
         self.config: Config = self.repo.config
         self.args = args
 

dvc/cli/parser.py

@@ -200,6 +200,14 @@ def get_main_parser():
         type=str,
     )
 
+    parser.add_argument(
+        "--wait-for-lock",
+        action="store_true",
+        default=False,
+        help="Wait for the lock if it is already held by another process, instead of"
+        " failing immediately.",
+    )
+
     # Sub commands
     subparsers = parser.add_subparsers(
         title="Available Commands",

dvc/lock.py

@@ -2,13 +2,13 @@
 
 import hashlib
 import os
+import time
 from abc import ABC, abstractmethod
 from datetime import timedelta
 from typing import Optional, Union
 
 import flufl.lock
 import zc.lockfile
-from funcy import retry
 
 from dvc.exceptions import DvcException
 from dvc.progress import Tqdm
@@ -88,9 +88,10 @@ class Lock(LockBase):
     Uses zc.lockfile as backend.
     """
 
-    def __init__(self, lockfile, friendly=False, **kwargs):
+    def __init__(self, lockfile, friendly=False, wait=False, **kwargs):
         super().__init__(lockfile)
         self._friendly = friendly
+        self._wait = wait
         self._lock = None
         self._lock_failed = False
 
@@ -101,23 +102,38 @@ class Lock(LockBase):
     def _do_lock(self):
         try:
             self._lock_failed = False
-            with Tqdm(
-                bar_format="{desc}",
-                disable=not self._friendly,
-                desc="If DVC froze, see `hardlink_lock` in {}".format(
-                    format_link("https://man.dvc.org/config#core")
-                ),
-            ):
-                self._lock = zc.lockfile.LockFile(self._lockfile)
+            self._lock = zc.lockfile.LockFile(self._lockfile)
         except zc.lockfile.LockError:
             self._lock_failed = True
             raise LockError(FAILED_TO_LOCK_MESSAGE)  # noqa: B904
 
     def lock(self):
-        retries = 6
-        delay = DEFAULT_TIMEOUT / retries
-        lock_retry = retry(retries, LockError, timeout=delay)(self._do_lock)
-        lock_retry()
+        """Acquire the lock, either waiting forever, or after default_retries."""
+        default_retries = 6
+        delay = DEFAULT_TIMEOUT / default_retries
+        attempts = 0
+
+        max_retries = float("inf") if self._wait else default_retries
+
+        with Tqdm(
+            bar_format="{desc}",
+            disable=not self._friendly,
+            desc="Waiting to acquire lock. "
+            "If DVC froze, see `hardlink_lock` in {}.".format(
+                format_link("https://man.dvc.org/config#core")
+            ),
+        ) as pbar:
+            while True:
+                try:
+                    self._do_lock()
+                    return
+                except LockError:
+                    attempts += 1
+                    if attempts > max_retries:
+                        raise
+                    time.sleep(delay)
+                finally:
+                    pbar.update()
 
     def unlock(self):
         if self._lock_failed:
@@ -150,7 +166,7 @@ class HardlinkLock(flufl.lock.Lock, LockBase):
         tmp_dir (str): a directory to store claim files.
     """
 
-    def __init__(self, lockfile, tmp_dir=None, **kwargs):
+    def __init__(self, lockfile, tmp_dir=None, wait=False, **kwargs):
         import socket
 
         self._tmp_dir = tmp_dir
@@ -173,11 +189,21 @@ class HardlinkLock(flufl.lock.Lock, LockBase):
         self._separator = flufl.lock.SEP
         self._set_claimfile()
         self._owned = True
+        self._wait = wait
         self._retry_errnos = []
+        self._friendly = kwargs.get("friendly", False)
 
     def lock(self, timeout: Optional[Union[timedelta, int]] = None):
         try:
-            super().lock(timeout or timedelta(seconds=DEFAULT_TIMEOUT))
+            if not self._wait:
+                timeout = timeout or timedelta(seconds=DEFAULT_TIMEOUT)
+
+            with Tqdm(
+                bar_format="{desc}",
+                disable=not (self._wait and self._friendly),
+                desc="Waiting to acquire lock",
+            ):
+                super().lock(timeout)
         except flufl.lock.TimeOutError:
             raise LockError(FAILED_TO_LOCK_MESSAGE)  # noqa: B904
 
@@ -191,6 +217,6 @@ class HardlinkLock(flufl.lock.Lock, LockBase):
             self._claimfile = os.path.join(self._tmp_dir, filename + ".lock")
 
 
-def make_lock(lockfile, tmp_dir=None, friendly=False, hardlink_lock=False):
+def make_lock(lockfile, tmp_dir=None, friendly=False, hardlink_lock=False, wait=False):
     cls = HardlinkLock if hardlink_lock else Lock
-    return cls(lockfile, tmp_dir=tmp_dir, friendly=friendly)
+    return cls(lockfile, tmp_dir=tmp_dir, friendly=friendly, wait=wait)

dvc/repo/__init__.py

@@ -144,6 +144,7 @@ class Repo:
         scm: Optional[Union["Git", "NoSCM"]] = None,
         remote: Optional[str] = None,
         remote_config: Optional["DictStrAny"] = None,
+        _wait_for_lock: bool = False,
     ):
         from dvc.cachemgr import CacheManager
         from dvc.data_cloud import DataCloud
@@ -167,6 +168,7 @@ class Repo:
         self._remote = remote
         self._remote_config = remote_config
         self._data_index = None
+        self._wait_for_lock = _wait_for_lock
 
         if rev and not fs:
             self._scm = scm = SCM(root_dir or os.curdir)
@@ -204,6 +206,7 @@ class Repo:
                     tmp_dir=self.tmp_dir,
                     hardlink_lock=self.config["core"].get("hardlink_lock", False),
                     friendly=True,
+                    wait=self._wait_for_lock,
                 )
                 os.makedirs(self.site_cache_dir, exist_ok=True)
                 if not fs and (

tests/func/test_cli.py

@@ -124,6 +124,7 @@ def test_find_root(dvc):
     class A:
         quiet = False
         verbose = True
+        wait_for_lock = False
         cd = os.path.pardir
 
     args = A()
@@ -139,6 +140,7 @@ def test_cd(dvc):
     class A:
         quiet = False
         verbose = True
+        wait_for_lock = False
         cd = os.path.pardir
 
     parent_dir = os.path.realpath(os.path.pardir)

tests/func/test_lock.py

@@ -1,8 +1,11 @@
+import multiprocessing
+import time
+
 import pytest
 
 from dvc.cli import main
 from dvc.exceptions import DvcException
-from dvc.lock import Lock, LockError
+from dvc.lock import Lock, LockError, make_lock
 
 
 def test_with(tmp_dir, dvc, mocker):
@@ -52,3 +55,63 @@ def test_cli(tmp_dir, dvc, mocker, caplog):
     with Lock(tmp_dir / dvc.tmp_dir / "lock"):
         assert main(["add", "foo"]) == 1
     assert expected_error_msg in caplog.text
+
+
+def hold_lock_until_signaled(lockfile_path, result_queue, release_signal):
+    lock = make_lock(lockfile_path)
+    with lock:
+        result_queue.put("p1_acquired")
+        release_signal.wait()
+    result_queue.put("p1_released")
+
+
+def try_lock_with_wait(lockfile_path, wait, result_queue):
+    result_queue.put("p2_starting")
+    try:
+        lock = make_lock(lockfile_path, wait=wait)
+        with lock:
+            result_queue.put("p2_acquired")
+    except LockError as e:
+        result_queue.put(f"error: {e}")
+    else:
+        result_queue.put("p2_released")
+
+
+def test_lock_waits_when_requested(request, tmp_path):
+    lockfile = tmp_path / "lock"
+
+    q: multiprocessing.Queue[str] = multiprocessing.Queue()
+    release_signal = multiprocessing.Event()
+    # Process 1 holds the lock until signaled to release it
+    p1 = multiprocessing.Process(
+        target=hold_lock_until_signaled, args=(lockfile, q, release_signal)
+    )
+    p2 = multiprocessing.Process(target=try_lock_with_wait, args=(lockfile, True, q))
+
+    p1.start()
+    request.addfinalizer(p1.kill)
+
+    assert q.get(timeout=1) == "p1_acquired"
+
+    # Process 2 will wait for the lock (should succeed)
+    p2.start()
+    request.addfinalizer(p2.kill)
+
+    assert q.get(timeout=1) == "p2_starting"
+    assert q.empty()
+
+    # sleep to ensure Process 2 is waiting for the lock
+    time.sleep(1)
+    release_signal.set()  # release the lock
+
+    p1.join(timeout=4)
+
+    events = [q.get(timeout=2), q.get(timeout=2), q.get(timeout=2)]
+    # we still can't be sure of the order of events.
+    assert "p1_released" in events
+    assert "p2_acquired" in events
+    assert "p2_released" in events
+
+    p2.join(timeout=1)
+
+    assert q.empty()