[cli,lock] Allow wait_for_lock

dvc/cli/command.py

@@ -18,7 +18,10 @@ def __init__(self, args: Any):
 
         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
 
@@ -114,10 +115,20 @@ def _do_lock(self):
             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
+        retries = None if self._wait else default_retries
+        delay = DEFAULT_TIMEOUT / default_retries
+        attempts = 0
+
+        while retries is None or attempts < retries:
+            try:
+                self._do_lock()
+                return
+            except LockError:
+                attempts += 1
+                time.sleep(delay)
+        raise LockError(FAILED_TO_LOCK_MESSAGE)
 
     def unlock(self):
         if self._lock_failed:
@@ -150,7 +161,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 +184,16 @@ def __init__(self, lockfile, tmp_dir=None, **kwargs):
         self._separator = flufl.lock.SEP
         self._set_claimfile()
         self._owned = True
+        self._wait = wait
         self._retry_errnos = []
 
     def lock(self, timeout: Optional[Union[timedelta, int]] = None):
         try:
-            super().lock(timeout or timedelta(seconds=DEFAULT_TIMEOUT))
+            if timeout is None:
+                # Default timeout if not explicitly passed
+                timeout = timedelta(seconds=DEFAULT_TIMEOUT)
+
+            super().lock(timeout)
         except flufl.lock.TimeOutError:
             raise LockError(FAILED_TO_LOCK_MESSAGE)  # noqa: B904
 
@@ -191,6 +207,6 @@ def _set_claimfile(self):
             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 @@ def __init__(  # noqa: PLR0915, PLR0913
         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 @@ def __init__(  # noqa: PLR0915, PLR0913
         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 @@ def __init__(  # noqa: PLR0915, PLR0913
                     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_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,38 @@ 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(lockfile_path, hold_seconds):
+    lock = make_lock(lockfile_path)
+    with lock:
+        time.sleep(hold_seconds)
+
+
+def try_lock_with_wait(lockfile_path, wait, result_queue):
+    try:
+        lock = make_lock(lockfile_path, wait=wait)
+        with lock:
+            result_queue.put("acquired")
+    except LockError as e:
+        result_queue.put(f"error: {e}")
+
+
+def test_lock_waits_when_requested(tmp_path):
+    lockfile = tmp_path / "lock"
+
+    # Process 1 holds the lock for 2 seconds
+    p1 = multiprocessing.Process(target=hold_lock, args=(lockfile, 2))
+    p1.start()
+    time.sleep(0.2)  # let p1 acquire lock
+
+    # Process 2 will wait for the lock (should succeed)
+    q = multiprocessing.Queue()
+    p2 = multiprocessing.Process(target=try_lock_with_wait, args=(lockfile, True, q))
+    p2.start()
+    p2.join(timeout=4)
+
+    result = q.get(timeout=1)
+    assert result == "acquired"
+
+    p1.join()