prevent multiple instance problem

stockroom/main.py

@@ -79,16 +79,19 @@ def get_hangar_checkout(self, write: bool = False) -> Any:
         return self._repo.hangar_repository.checkout(write=write)
 
     @contextmanager
-    def optimize(self):
+    def optimize(self, write=False):
         """
         This context manager, on `enter`, asks the :class:`StockRepository` object to
         open the global checkout. Global checkout is being stored as property of the
         repository singleton. Hence all the downstream tasks will get this opened
         checkout until it is closed. This global checkout will be closed on the `exit` of
         this context manager
         """
+        if self._repo.is_optimized:
+            raise RuntimeError("Attempt to open one optimized checkout while another is "
+                               "in action in the same process")
         try:
-            self._repo.open_global_checkout()
+            self._repo.open_global_checkout(write)
             yield None
         finally:
             self._repo.close_global_checkout()
@@ -100,7 +103,7 @@ def commit(self, message: str) -> str:
         close after the commit. Which means, no other write operations should be running
         while stock commit is in progress
         """
-        with self._repo.checkout(write=True) as co:
+        with self._repo.write_checkout() as co:
             digest = co.commit(message)
         set_current_head(self._repo.stockroot, digest)
         return digest

stockroom/repository.py

@@ -37,60 +37,76 @@ def __init__(self, root):
         self._hangar_repo = Repository(root)
         self._optimized_Rcheckout = None
         self._optimized_Wcheckout = None
-        self._has_optimized = False
+        self._has_optimized = {'R': False, 'W': False}
 
     @property
     def hangar_repository(self):
         return self._hangar_repo
 
-    def open_global_checkout(self):
+    @property
+    def is_optimized(self):
+        return any(self._has_optimized.values())
+
+    def open_global_checkout(self, write):
         head_commit = get_current_head(self._root)
         self._optimized_Rcheckout = self._hangar_repo.checkout(commit=head_commit)
-        self._optimized_Wcheckout = self._hangar_repo.checkout(write=True)
-        self._optimized_Wcheckout.__enter__()
         self._optimized_Rcheckout.__enter__()
-        self._has_optimized = True
+        self._has_optimized['R'] = True
+        if write:
+            self._optimized_Wcheckout = self._hangar_repo.checkout(write=True)
+            self._optimized_Wcheckout.__enter__()
+            self._has_optimized['W'] = True
 
     def close_global_checkout(self):
-        self._has_optimized = False
-        self._optimized_Wcheckout.__exit__()
+        self._has_optimized['R'] = False
         self._optimized_Rcheckout.__exit__()
-        self._optimized_Wcheckout.close()
         self._optimized_Rcheckout.close()
-        self._optimized_Wcheckout = None
         self._optimized_Rcheckout = None
+        if self._has_optimized['W']:
+            self._has_optimized['W'] = False
+            self._optimized_Wcheckout.__exit__()
+            self._optimized_Wcheckout.close()
+            self._optimized_Wcheckout = None
 
     @contextmanager
-    def checkout(self, write=False):
+    def read_checkout(self):
         """
-        An api similar to hangar checkout but creates the checkout object using the
-        commit hash from stock file instead of user supplying one. This enables users
-        to rely on git checkout for hangar checkout as well. This checkout is being
+        An api similar to hangar checkout in read mode but creates the checkout object
+        using the commit hash from stock file instead of user supplying one. This enables
+        users to rely on git checkout for hangar checkout as well. This checkout is being
         designed as a context manager that makes sure the checkout is closed. On entry
-        and exit, the CM checks the existence of a global checkout. On entry, ff global
+        and exit, the CM checks the existence of a global checkout. On entry, if global
         checkout exists, it returns the that instead of creating a new checkout. On exit,
         it doesn't close in case of global checkout instead it lets the CM do the closure
         """
-        if write:
-            if self._has_optimized:
-                co = self._optimized_Wcheckout
-            else:
-                if self._hangar_repo.writer_lock_held:
-                    raise PermissionError("Another write operation is in progress. "
-                                          "Could not acquire the lock")
-                co = self._hangar_repo.checkout(write=True)
+        if self._has_optimized['R']:
+            co = self._optimized_Rcheckout
         else:
-            if self._has_optimized:
-                co = self._optimized_Rcheckout
-            else:
-                head_commit = get_current_head(self._root)
-                co = self._hangar_repo.checkout(commit=head_commit)
+            head_commit = get_current_head(self._root)
+            co = self._hangar_repo.checkout(commit=head_commit)
         try:
             yield co
         finally:
-            if not self._has_optimized:
+            if not self._has_optimized['R']:
                 co.close()
-
+
+    @contextmanager
+    def write_checkout(self):
+        """
+        An API similar to hangar checkout in write mode but does the closure of checkout
+        on the exit of CM. It also monitors the existence of global checkout and open
+        or close a local checkout if the global checkout doesn't exist
+        """
+        if self._has_optimized['W']:
+            co = self._optimized_Wcheckout
+        else:
+            co = self._hangar_repo.checkout(write=True)
+        try:
+            yield co
+        finally:
+            if not self._has_optimized['W']:
+                co.close()
+
     @property
     def stockroot(self) -> Path:
         """
@@ -117,7 +133,7 @@ def init_repo(name=None, email=None, overwrite=False):
             raise ValueError("Both ``name`` and ``email`` cannot be None")
         commit_hash = ''
         repo.init(user_name=name, user_email=email, remove_old=overwrite)
-    # TODO: It's better to have the `close_environment` as public attribute in hangar
+    # closing the environment for avoiding issues in windows
     repo._env._close_environments()
 
     stock_file = Path.cwd()/'head.stock'

stockroom/storages/data.py

@@ -28,9 +28,9 @@ def __init__(self, repo):
         self._repo = repo
 
     def __setitem__(self, key, value):
-        with self._repo.checkout(write=True) as co:
+        with self._repo.write_checkout() as co:
             co[key] = value
 
     def __getitem__(self, key):
-        with self._repo.checkout() as co:
+        with self._repo.read_checkout() as co:
             return co[key]

stockroom/storages/model.py

@@ -52,7 +52,7 @@ def __setitem__(self, name, weights):
         longest = max([len(x.reshape(-1)) for x in weights])
         dtypes = [w.dtype.name for w in weights]
 
-        with self._repo.checkout(write=True) as co:
+        with self._repo.write_checkout() as co:
             co.metadata[parser.model_metakey(name, 'library')] = library
             co.metadata[parser.model_metakey(name, 'libraryVersion')] = library_version
             co.metadata[parser.model_metakey(name, 'longest')] = str(longest)
@@ -84,7 +84,7 @@ def __setitem__(self, name, weights):
                     shape_aset[i] = np.array(()).astype(shape_typ)
 
     def __getitem__(self, name):
-        with self._repo.checkout() as co:
+        with self._repo.read_checkout() as co:
             try:
                 library = co.metadata[parser.model_metakey(name, 'library')]
             except KeyError:

stockroom/storages/tag.py

@@ -23,7 +23,7 @@ def __init__(self, repo):
         self._repo = repo
 
     def __setitem__(self, key, value):
-        with self._repo.checkout(write=True) as co:
+        with self._repo.write_checkout() as co:
             if isinstance(value, int):
                 value_type = 'int'
             elif isinstance(value, float):
@@ -36,7 +36,7 @@ def __setitem__(self, key, value):
             co.metadata[parser.tag_typekey(key)] = value_type
 
     def __getitem__(self, key):
-        with self._repo.checkout() as co:
+        with self._repo.read_checkout() as co:
             try:
                 value = co.metadata[parser.tagkey(key)]
                 value_type = co.metadata[parser.tag_typekey(key)]

tests/test_multiple_instances.py

@@ -0,0 +1,78 @@
+from stockroom import StockRoom
+from random import randint
+import numpy as np
+import pytest
+
+
+class TestSameProcess:
+
+    def test_opening_two_instances(self, repo_with_aset):
+        # todo: should we allow this?
+        stk1 = StockRoom()
+        stk2 = StockRoom()
+        arr = np.arange(20).reshape(4, 5)
+        oldarr = arr * randint(1, 100)
+        newarr = arr * randint(1, 100)
+
+        stk1.data['aset', 1] = oldarr
+        stk2.data['aset', 1] = newarr
+        stk1.commit('added data')
+
+        assert np.allclose(stk2.data['aset', 1], newarr)
+        assert not np.allclose(stk2.data['aset', 1], oldarr)
+
+    def test_operating_one_in_another_write_contextmanager(self, repo_with_aset):
+        stk1 = StockRoom()
+        stk2 = StockRoom()
+        arr = np.arange(20).reshape(4, 5)
+        oldarr = arr * randint(1, 100)
+
+        with stk1.optimize(write=True):
+            assert stk1._repo._optimized_Rcheckout is not None
+            assert stk1._repo._optimized_Wcheckout is not None
+            assert stk2._repo._optimized_Rcheckout is not None
+            assert stk2._repo._optimized_Wcheckout is not None
+            stk2.data['aset', 1] = oldarr
+            stk1.commit('adding data inside cm')
+            with pytest.raises(KeyError):
+                # TODO: document this scenario
+                data = stk1.data['aset', 1]
+
+            stk3 = StockRoom()
+            assert stk3._repo._optimized_Rcheckout is not None
+            assert stk3._repo._optimized_Wcheckout is not None
+
+        assert np.allclose(oldarr, stk1.data['aset', 1])
+
+    def test_opening_one_contextmanager_in_another(self, repo_with_aset):
+        stk1 = StockRoom()
+        stk2 = StockRoom()
+
+        with stk1.optimize(write=True):
+            with pytest.raises(RuntimeError):
+                with stk2.optimize():
+                    pass
+            assert stk2._repo._optimized_Rcheckout is not None
+            assert stk2._repo._optimized_Wcheckout is not None
+        assert stk1._repo._optimized_Rcheckout is None
+        assert stk1._repo._optimized_Wcheckout is None
+        assert stk2._repo._optimized_Rcheckout is None
+        assert stk2._repo._optimized_Wcheckout is None
+
+    def test_one_inside_another_read_contextmanager(self, repo_with_aset):
+        stk1 = StockRoom()
+        stk2 = StockRoom()
+        arr = np.arange(20).reshape(4, 5)
+
+        with stk1.optimize():
+            # non-optimized write inside read CM
+            assert stk2._repo._optimized_Wcheckout is None
+            stk2.data['aset', 1] = arr
+            stk2.commit('adding data')
+
+            with pytest.raises(RuntimeError):
+                with stk2.optimize(write=True):
+                    pass
+
+        with stk1.optimize():
+            assert np.allclose(stk2.data['aset', 1], arr)