Add methods to lock/unlock shared directories

2 files changed, 43 insertions, 0 deletions

diff --git a/obnamlib/repo.py b/obnamlib/repo.py
index be9e081c..fe9e345b 100644
--- a/obnamlib/repo.py
+++ b/obnamlib/repo.py
@@ -145,6 +145,7 @@ class Repository(object):
         self.clientlist = obnamlib.ClientList(self.fs, node_size, 
                                               upload_queue_size, 
                                               lru_size, self)
+        self.got_shared_lock = False
         self.got_client_lock = False
         self.current_client = None
         self.current_client_id = None
@@ -214,6 +215,11 @@ class Repository(object):
         if not self.got_root_lock:
             raise LockFail('have not got lock on root node')
 
+    def require_shared_lock(self):
+        '''Ensure we have the lock on the shared B-trees except clientlist.'''
+        if not self.got_shared_lock:
+            raise LockFail('have not got lock on shared B-trees')
+
     def require_client_lock(self):
         '''Ensure we have the lock on the currently open client.'''
         if not self.got_client_lock:
@@ -337,6 +343,30 @@ class Repository(object):
         if client_name not in self.list_clients():
             raise obnamlib.Error('client %s does not exist' % client_name)
         self.removed_clients.append(client_name)
+
+    @property
+    def shared_dirs(self):
+        return [self.chunklist.dirname, self.chunksums.dirname]
+        
+    def lock_shared(self):
+        '''Lock a client for exclusive write access.
+        
+        Raise obnamlib.LockFail if locking fails. Lock will be released
+        by commit_client() or unlock_client().
+
+        '''
+
+        tracing.trace('locking shared')
+        self.check_format_version()
+        self.lockmgr.lock(self.shared_dirs)
+        self.got_shared_lock = True
+
+    def unlock_shared(self):
+        '''Unlock currently locked shared B-trees.'''
+        tracing.trace('unlocking shared')
+        self.require_shared_lock()
+        self.lockmgr.unlock(self.shared_dirs)
+        self.got_shared_lock = False
         
     def lock_client(self, client_name):
         '''Lock a client for exclusive write access.
diff --git a/obnamlib/repo_tests.py b/obnamlib/repo_tests.py
index ef28ab51..1da63fd0 100644
--- a/obnamlib/repo_tests.py
+++ b/obnamlib/repo_tests.py
@@ -127,6 +127,19 @@ class RepositoryRootNodeTests(unittest.TestCase):
         self.repo._write_format_version(0)
         self.assertRaises(obnamlib.BadFormat, self.repo.list_clients)
 
+    def test_locks_shared(self):
+        self.repo.lock_shared()
+        self.assertTrue(self.repo.got_shared_lock)
+
+    def test_unlocks_shared(self):
+        self.repo.lock_shared()
+        self.repo.unlock_shared()
+        self.assertFalse(self.repo.got_shared_lock)
+
+    def test_unlock_shared_when_locked_by_other_fails(self):
+        self.other.lock_shared()
+        self.assertRaises(obnamlib.LockFail, self.repo.unlock_shared)
+
     def test_lock_client_fails_if_format_is_incompatible(self):
         self.repo._write_format_version(0)
         self.assertRaises(obnamlib.BadFormat, self.repo.lock_client, 'foo')