Merge pull request #2584 from efiop/755

dvc: release locks when running a command

dvc/lock.py

@@ -7,10 +7,9 @@ import time
 from datetime import timedelta
 
 import zc.lockfile
-from funcy.py3 import lkeep
 
 from dvc.exceptions import DvcException
-from dvc.utils import makedirs, format_link
+from dvc.utils import format_link
 from dvc.utils.compat import is_py3, is_py2
 from dvc.progress import Tqdm
 
@@ -93,8 +92,6 @@ if is_py3:
             import socket
 
             self._tmp_dir = tmp_dir
-            if self._tmp_dir is not None:
-                makedirs(self._tmp_dir, exist_ok=True)
 
             # NOTE: this is basically Lock.__init__ copy-paste, except that
             # instead of using `socket.getfqdn()` we use `socket.gethostname()`
@@ -120,10 +117,6 @@ if is_py3:
         def lockfile(self):
             return self._lockfile
 
-        @property
-        def files(self):
-            return lkeep([self._lockfile, self._tmp_dir])
-
         def lock(self):
             try:
                 super(Lock, self).lock(timedelta(seconds=DEFAULT_TIMEOUT))

dvc/path_info.py

@@ -23,7 +23,16 @@ if is_py2:
     fs_encoding = "utf-8"
 
 
-class PathInfo(pathlib.PurePath):
+class _BasePath(object):
+    def overlaps(self, other):
+        if isinstance(other, basestring):
+            other = self.__class__(other)
+        elif self.__class__ != other.__class__:
+            return False
+        return self == other or self.isin(other) or other.isin(self)
+
+
+class PathInfo(pathlib.PurePath, _BasePath):
     # Use __slots__ in PathInfo objects following PurePath implementation.
     # This makes objects smaller and speeds up attribute access.
     # We don't add any fields so it's empty.
@@ -138,7 +147,7 @@ class _URLPathParents(object):
         return "<{}.parents>".format(self.src)
 
 
-class URLInfo(object):
+class URLInfo(_BasePath):
     DEFAULT_PORTS = {"http": 80, "https": 443, "ssh": 22, "hdfs": 0}
 
     def __init__(self, url):

dvc/repo/__init__.py

@@ -76,6 +76,7 @@ class Repo(object):
         from dvc.repo.metrics import Metrics
         from dvc.scm.tree import WorkingTree
         from dvc.repo.tag import Tag
+        from dvc.utils import makedirs
 
         root_dir = self.find_root(root_dir)
 
@@ -88,6 +89,9 @@ class Repo(object):
 
         self.tree = WorkingTree(self.root_dir)
 
+        self.tmp_dir = os.path.join(self.dvc_dir, "tmp")
+        makedirs(self.tmp_dir, exist_ok=True)
+
         hardlink_lock = self.config.config["core"].get("hardlink_lock", False)
         self.lock = make_lock(
             os.path.join(self.dvc_dir, "lock"),
@@ -95,6 +99,7 @@ class Repo(object):
             hardlink_lock=hardlink_lock,
             friendly=True,
         )
+
         # NOTE: storing state and link_state in the repository itself to avoid
         # any possible state corruption in 'shared cache dir' scenario.
         self.state = State(self, self.config.config)
@@ -165,9 +170,8 @@ class Repo(object):
 
         flist = (
             [self.config.config_local_file, updater.updater_file]
+            + [self.lock.lockfile, updater.lock.lockfile, self.tmp_dir]
             + self.state.files
-            + self.lock.files
-            + updater.lock.files
         )
 
         if path_isin(self.cache.local.cache_dir, self.root_dir):
@@ -352,11 +356,7 @@ class Repo(object):
                     continue
 
                 for out in outs:
-                    if (
-                        out == dep.path_info
-                        or dep.path_info.isin(out)
-                        or out.isin(dep.path_info)
-                    ):
+                    if out.overlaps(dep.path_info):
                         dep_stage = outs[out].stage
                         dep_node = relpath(dep_stage.path, self.root_dir)
                         G.add_node(dep_node, stage=dep_stage)

dvc/rwlock.py

@@ -0,0 +1,191 @@
+from __future__ import unicode_literals
+
+import os
+import json
+
+from collections import defaultdict
+from contextlib import contextmanager
+
+from voluptuous import Schema, Invalid
+from funcy.py3 import lmap, lfilter, lmapcat
+
+from .exceptions import DvcException
+from .lock import LockError
+from .utils.compat import (
+    convert_to_unicode,
+    FileNotFoundError,
+    JSONDecodeError,
+    str,
+)
+from .utils.fs import relpath
+
+INFO_SCHEMA = {"pid": int, "cmd": str}
+
+SCHEMA = Schema({"write": {str: INFO_SCHEMA}, "read": {str: [INFO_SCHEMA]}})
+
+
+class RWLockFileCorruptedError(DvcException):
+    def __init__(self, path, cause):
+        super(RWLockFileCorruptedError, self).__init__(
+            "Unable to read RWLock-file '{}'. JSON structure is "
+            "corrupted".format(relpath(path)),
+            cause=cause,
+        )
+
+
+class RWLockFileFormatError(DvcException):
+    def __init__(self, path, cause):
+        super(RWLockFileFormatError, self).__init__(
+            "RWLock-file '{}' format error.".format(relpath(path)), cause=cause
+        )
+
+
+@contextmanager
+def _edit_rwlock(lock_dir):
+    path = os.path.join(lock_dir, "rwlock")
+    try:
+        with open(path, "r") as fobj:
+            lock = json.load(fobj)
+        lock = SCHEMA(convert_to_unicode(lock))
+    except FileNotFoundError:
+        lock = SCHEMA({})
+    except JSONDecodeError as exc:
+        raise RWLockFileCorruptedError(path, cause=exc)
+    except Invalid as exc:
+        raise RWLockFileFormatError(path, cause=exc)
+    lock = defaultdict(dict, lock)
+    lock["read"] = defaultdict(list, lock["read"])
+    lock["write"] = defaultdict(dict, lock["write"])
+    yield lock
+    with open(path, "w+") as fobj:
+        json.dump(lock, fobj)
+
+
+def _infos_to_str(infos):
+    return "\n".join(
+        "  (PID {}): {}".format(info["pid"], info["cmd"]) for info in infos
+    )
+
+
+def _check_no_writers(lock, info, path_infos):
+    for path_info in path_infos:
+        blocking_urls = lfilter(path_info.overlaps, lock["write"])
+        if not blocking_urls:
+            continue
+
+        writers = lmap(lock["write"].get, blocking_urls)
+        writers = lfilter(lambda i: info != i, writers)
+        if not writers:
+            continue
+
+        raise LockError(
+            "'{}' is busy, it is being written to by:\n{}".format(
+                str(path_info), _infos_to_str(writers)
+            )
+        )
+
+
+def _check_no_readers(lock, info, path_infos):
+    for path_info in path_infos:
+        blocking_urls = lfilter(path_info.overlaps, lock["read"])
+        if not blocking_urls:
+            continue
+
+        readers = lmapcat(lock["read"].get, blocking_urls)
+        readers = lfilter(lambda i: info != i, readers)
+        if not readers:
+            continue
+
+        raise LockError(
+            "'{}' is busy, it is being read by:\n{}".format(
+                str(path_info), _infos_to_str(readers)
+            )
+        )
+
+
+def _acquire_read(lock, info, path_infos):
+    changes = []
+
+    for path_info in path_infos:
+        url = path_info.url
+
+        readers = lock["read"][url]
+        if info in readers:
+            continue
+
+        changes.append(url)
+        readers.append(info)
+
+    return changes
+
+
+def _acquire_write(lock, info, path_infos):
+    changes = []
+
+    for path_info in path_infos:
+        url = path_info.url
+
+        if lock["write"][url] == info:
+            continue
+
+        changes.append(url)
+        lock["write"][url] = info
+
+    return changes
+
+
+def _release_write(lock, info, changes):
+    for url in changes:
+        assert "write" in lock
+        assert url in lock["write"]
+        assert lock["write"][url] == info
+        del lock["write"][url]
+        if not lock["write"]:
+            del lock["write"]
+
+
+def _release_read(lock, info, changes):
+    for url in changes:
+        assert "read" in lock
+        assert url in lock["read"]
+        assert info in lock["read"][url]
+        lock["read"][url].remove(info)
+        if not lock["read"][url]:
+            del lock["read"][url]
+        if not lock["read"]:
+            del lock["read"]
+
+
+@contextmanager
+def rwlock(tmp_dir, cmd, read, write):
+    """Create non-thread-safe RWLock for PathInfos.
+
+    Args:
+        tmp_dir (str): existing directory where to create the rwlock file.
+        cmd (str): command that will be working on these PathInfos.
+        read ([PathInfo]): PathInfos that are going to be read.
+        write ([PathInfo]): PathInfos that are going to be written.
+
+    Raises:
+        LockError: raised if PathInfo we want to read is being written to by
+            another command or if PathInfo we want to write is being written
+            to or read from by another command.
+        RWLockFileCorruptedError: raised if rwlock file is not a valid JSON.
+        RWLockFileFormatError: raised if rwlock file is a valid JSON, but
+            has internal format that doesn't pass our schema validation.
+    """
+    info = {"pid": os.getpid(), "cmd": cmd}
+
+    with _edit_rwlock(tmp_dir) as lock:
+        _check_no_writers(lock, info, read + write)
+        _check_no_readers(lock, info, write)
+
+        rchanges = _acquire_read(lock, info, read)
+        wchanges = _acquire_write(lock, info, write)
+
+    try:
+        yield
+    finally:
+        with _edit_rwlock(tmp_dir) as lock:
+            _release_write(lock, info, wchanges)
+            _release_read(lock, info, rchanges)

dvc/stage.py

@@ -7,7 +7,10 @@ import re
 import signal
 import subprocess
 import threading
+
+from functools import wraps
 from itertools import chain
+from funcy import decorator
 
 from voluptuous import Any, Schema, MultipleInvalid
 
@@ -129,6 +132,52 @@ class MissingDataSource(DvcException):
         super(MissingDataSource, self).__init__(msg)
 
 
+@decorator
+def rwlocked(call, read=None, write=None):
+    import sys
+    from dvc.rwlock import rwlock
+    from dvc.dependency.repo import DependencyREPO
+
+    if read is None:
+        read = []
+
+    if write is None:
+        write = []
+
+    stage = call._args[0]
+
+    def _chain(names):
+        return [
+            item.path_info
+            for attr in names
+            for item in getattr(stage, attr)
+            # There is no need to lock DependencyREPO deps, as there is no
+            # corresponding OutputREPO, so we can't even write it.
+            if not isinstance(item, DependencyREPO)
+        ]
+
+    cmd = " ".join(sys.argv)
+
+    with rwlock(stage.repo.tmp_dir, cmd, _chain(read), _chain(write)):
+        return call()
+
+
+def unlocked_repo(f):
+    @wraps(f)
+    def wrapper(stage, *args, **kwargs):
+        stage.repo.state.dump()
+        stage.repo.lock.unlock()
+        stage.repo._reset()
+        try:
+            ret = f(stage, *args, **kwargs)
+        finally:
+            stage.repo.lock.lock()
+            stage.repo.state.load()
+        return ret
+
+    return wrapper
+
+
 class Stage(object):
     STAGE_FILE = "Dvcfile"
     STAGE_FILE_SUFFIX = ".dvc"
@@ -286,6 +335,7 @@ class Stage(object):
             return True
         return False
 
+    @rwlocked(read=["deps", "outs"])
     def changed(self):
         # Short-circuit order: stage md5 is fast, deps are expected to change
         ret = (
@@ -299,6 +349,7 @@ class Stage(object):
 
         return ret
 
+    @rwlocked(write=["outs"])
     def remove_outs(self, ignore_remove=False, force=False):
         """Used mainly for `dvc remove --outs` and :func:`Stage.reproduce`."""
         for out in self.outs:
@@ -316,6 +367,7 @@ class Stage(object):
         for out in self.outs:
             out.unprotect()
 
+    @rwlocked(write=["outs"])
     def remove(self, force=False, remove_outs=True):
         if remove_outs:
             self.remove_outs(ignore_remove=True, force=force)
@@ -323,6 +375,7 @@ class Stage(object):
             self.unprotect_outs()
         os.unlink(self.path)
 
+    @rwlocked(read=["deps"], write=["outs"])
     def reproduce(self, interactive=False, **kwargs):
 
         if not kwargs.get("force", False) and not self.changed():
@@ -755,6 +808,7 @@ class Stage(object):
                 )
             self.save()
 
+    @rwlocked(write=["outs"])
     def commit(self):
         for out in self.outs:
             out.commit()
@@ -801,6 +855,7 @@ class Stage(object):
             if occurrence > 1:
                 raise ArgumentDuplicationError(str(path))
 
+    @unlocked_repo
     def _run(self):
         self._check_missing_deps()
 
@@ -850,6 +905,7 @@ class Stage(object):
         if (p is None) or (p.returncode != 0):
             raise StageCmdFailedError(self)
 
+    @rwlocked(read=["deps"], write=["outs"])
     def run(self, dry=False, no_commit=False, force=False):
         if (self.cmd or self.is_import) and not self.locked and not dry:
             self.remove_outs(ignore_remove=False, force=False)
@@ -874,7 +930,6 @@ class Stage(object):
                     self.outs[0].checkout()
                 else:
                     self.deps[0].download(self.outs[0])
-
         elif self.is_data_source:
             msg = "Verifying data sources in '{}'".format(self.relpath)
             logger.info(msg)
@@ -904,6 +959,7 @@ class Stage(object):
         if paths:
             raise MissingDataSource(paths)
 
+    @rwlocked(write=["outs"])
     def checkout(self, force=False, progress_callback=None):
         failed_checkouts = []
         for out in self.outs:
@@ -923,6 +979,7 @@ class Stage(object):
 
         return ret
 
+    @rwlocked(read=["deps", "outs"])
     def status(self):
         ret = []
 

dvc/utils/compat.py

@@ -123,6 +123,7 @@ if is_py2:
     makedirs = _makedirs
     range = xrange  # noqa: F821
     FileNotFoundError = IOError
+    JSONDecodeError = ValueError
 
     import StringIO
     import io
@@ -164,6 +165,7 @@ elif is_py3:
     import configparser as ConfigParser  # noqa: F401
     from collections.abc import Mapping  # noqa: F401
     from contextlib import ExitStack  # noqa: F401
+    from json.decoder import JSONDecodeError  # noqa: F401
 
     builtin_str = str  # noqa: F821
     str = str  # noqa: F821

tests/func/test_run.py

@@ -893,7 +893,7 @@ class TestShouldNotCheckoutUponCorruptedLocalHardlinkCache(TestDvc):
         )
         patch_run = mock.patch.object(stage, "_run", wraps=stage._run)
 
-        with self.dvc.state:
+        with self.dvc.lock, self.dvc.state:
             with patch_checkout as mock_checkout:
                 with patch_run as mock_run:
                     stage.run()

tests/unit/test_rwlock.py

@@ -0,0 +1,98 @@
+from __future__ import unicode_literals
+
+import os
+
+import pytest
+
+from dvc.rwlock import (
+    rwlock,
+    _edit_rwlock,
+    RWLockFileFormatError,
+    RWLockFileCorruptedError,
+)
+from dvc.lock import LockError
+from dvc.utils.compat import fspath
+from dvc.path_info import PathInfo
+
+
+def test_rwlock(tmp_path):
+    path = fspath(tmp_path)
+    foo = PathInfo("foo")
+
+    with rwlock(path, "cmd1", [foo], []):
+        with pytest.raises(LockError):
+            with rwlock(path, "cmd2", [], [foo]):
+                pass
+
+    with rwlock(path, "cmd1", [], [foo]):
+        with pytest.raises(LockError):
+            with rwlock(path, "cmd2", [foo], []):
+                pass
+
+    with rwlock(path, "cmd1", [], [foo]):
+        with pytest.raises(LockError):
+            with rwlock(path, "cmd2", [], [foo]):
+                pass
+
+
+def test_rwlock_reentrant(tmp_path):
+    path = fspath(tmp_path)
+    foo = PathInfo("foo")
+
+    with rwlock(path, "cmd1", [], [foo]):
+        with rwlock(path, "cmd1", [], [foo]):
+            pass
+        with _edit_rwlock(path) as lock:
+            assert lock == {
+                "read": {},
+                "write": {"foo": {"cmd": "cmd1", "pid": os.getpid()}},
+            }
+
+    with rwlock(path, "cmd", [foo], []):
+        with rwlock(path, "cmd", [foo], []):
+            pass
+        with _edit_rwlock(path) as lock:
+            assert lock == {
+                "read": {"foo": [{"cmd": "cmd", "pid": os.getpid()}]},
+                "write": {},
+            }
+
+
+def test_rwlock_subdirs(tmp_path):
+    path = fspath(tmp_path)
+    foo = PathInfo("foo")
+    subfoo = PathInfo("foo/subfoo")
+
+    with rwlock(path, "cmd1", [foo], []):
+        with pytest.raises(LockError):
+            with rwlock(path, "cmd2", [], [subfoo]):
+                pass
+
+    with rwlock(path, "cmd1", [], [subfoo]):
+        with pytest.raises(LockError):
+            with rwlock(path, "cmd2", [foo], []):
+                pass
+
+    with rwlock(path, "cmd1", [], [subfoo]):
+        with pytest.raises(LockError):
+            with rwlock(path, "cmd2", [], [foo]):
+                pass
+
+    with rwlock(path, "cmd1", [subfoo], []):
+        with rwlock(path, "cmd2", [foo], []):
+            pass
+
+
+def test_broken_rwlock(tmp_path):
+    dir_path = fspath(tmp_path)
+    path = tmp_path / "rwlock"
+
+    path.write_text('{"broken": "format"}', encoding="utf-8")
+    with pytest.raises(RWLockFileFormatError):
+        with _edit_rwlock(dir_path):
+            pass
+
+    path.write_text("{broken json", encoding="utf-8")
+    with pytest.raises(RWLockFileCorruptedError):
+        with _edit_rwlock(dir_path):
+            pass

tests/unit/test_stage.py

@@ -93,16 +93,14 @@ def test_stage_update(mocker):
     not isinstance(threading.current_thread(), threading._MainThread),
     reason="Not running in the main thread.",
 )
-def test_stage_run_ignore_sigint(mocker):
-    stage = Stage(None, "path")
-
+def test_stage_run_ignore_sigint(dvc_repo, mocker):
     proc = mocker.Mock()
     communicate = mocker.Mock()
     proc.configure_mock(returncode=0, communicate=communicate)
     popen = mocker.patch.object(subprocess, "Popen", return_value=proc)
     signal_mock = mocker.patch("signal.signal")
 
-    stage._run()
+    dvc_repo.run(cmd="path")
 
     assert popen.called_once()
     assert communicate.called_once_with()
@@ -110,8 +108,8 @@ def test_stage_run_ignore_sigint(mocker):
     assert signal.getsignal(signal.SIGINT) == signal.default_int_handler
 
 
-def test_always_changed():
-    stage = Stage(None, "path", always_changed=True)
+def test_always_changed(dvc_repo):
+    stage = Stage(dvc_repo, "path", always_changed=True)
     stage.save()
     assert stage.changed()
     assert stage.status()["path"] == ["always changed"]