commit 3b7ed473b3e4c1449811d26f0458c5cebe3733a1
Author: Guido Trotter <ultrotter@google.com>
Date:   Tue Mar 4 13:18:22 2008 +0000

    LockSet: make acquire() able to get the whole set
    
    This new functionality makes it possible to acquire a whole set, by passing
    "None" to the acquire() function as the list of elements. This will avoid new
    additions to the set, and then acquire all the current elements. The list of
    all elements acquired will be returned at the end.
    
    Deletions can still happen during the acquire process and we'll deal with it by
    just skipping the deleted elements: it's effectively as if they were deleted
    before we called the function. After we've finished though we hold all the
    elements, so no more deletes can be performed before we release them.
    
    Any call to release() will then first of all release the "set-level" lock if
    we're holding it, and then all or some of the locks we have.
    
    Some new tests checks that this feature works as intended.
    
    Reviewed-by: imsnah

diff --git a/lib/locking.py b/lib/locking.py
index 7c886a5..b251f30 100644
--- a/lib/locking.py
+++ b/lib/locking.py
@@ -372,11 +372,25 @@ class LockSet:
     # Check we don't already own locks at this level
     assert not self._is_owned(), "Cannot acquire locks in the same set twice"
 
+    if names is None:
+      # If no names are given acquire the whole set by not letting new names
+      # being added before we release, and getting the current list of names.
+      # Some of them may then be deleted later, but we'll cope with this.
+      #
+      # We'd like to acquire this lock in a shared way, as it's nice if
+      # everybody else can use the instances at the same time. If are acquiring
+      # them exclusively though they won't be able to do this anyway, though,
+      # so we'll get the list lock exclusively as well in order to be able to
+      # do add() on the set while owning it.
+      self.__lock.acquire(shared=shared)
+
     try:
       # Support passing in a single resource to acquire rather than many
       if isinstance(names, basestring):
         names = [names]
       else:
+        if names is None:
+          names = self.__names()
         names.sort()
 
       acquire_list = []
@@ -388,7 +402,12 @@ class LockSet:
           lock = self.__lockdict[lname] # raises KeyError if the lock is not there
           acquire_list.append((lname, lock))
         except (KeyError):
-          raise errors.LockError('non-existing lock in set (%s)' % lname)
+          if self.__lock._is_owned():
+            # We are acquiring all the set, it doesn't matter if this particular
+            # element is not there anymore.
+            continue
+          else:
+            raise errors.LockError('non-existing lock in set (%s)' % lname)
 
       # This will hold the locknames we effectively acquired.
       acquired = set()
@@ -411,13 +430,21 @@ class LockSet:
             raise
 
         except (errors.LockError):
-          name_fail = lname
-          for lname in self._list_owned():
-            self.__lockdict[lname].release()
-            self._del_owned(lname)
-          raise errors.LockError('non-existing lock in set (%s)' % name_fail)
+          if self.__lock._is_owned():
+            # We are acquiring all the set, it doesn't matter if this particular
+            # element is not there anymore.
+            continue
+          else:
+            name_fail = lname
+            for lname in self._list_owned():
+              self.__lockdict[lname].release()
+              self._del_owned(lname)
+            raise errors.LockError('non-existing lock in set (%s)' % name_fail)
 
     except:
+      # If something went wrong and we had the set-lock let's release it...
+      if self.__lock._is_owned():
+        self.__lock.release()
       raise
 
     return acquired
@@ -448,6 +475,11 @@ class LockSet:
                "release() on unheld resources %s" %
                names.difference(self._list_owned()))
 
+    # First of all let's release the "all elements" lock, if set.
+    # After this 'add' can work again
+    if self.__lock._is_owned():
+      self.__lock.release()
+
     for lockname in names:
       # If we are sure the lock doesn't leave __lockdict without being
       # exclusively held we can do this...
@@ -463,13 +495,21 @@ class LockSet:
       shared: is the pre-acquisition shared?
 
     """
+
+    assert not self.__lock._is_owned(shared=1), (
+           "Cannot add new elements while sharing the set-lock")
+
     # Support passing in a single resource to add rather than many
     if isinstance(names, basestring):
       names = [names]
 
-    # Acquire the internal lock in an exclusive way, so there cannot be a
-    # conflicting add()
-    self.__lock.acquire()
+    # If we don't already own the set-level lock acquire it in an exclusive way
+    # we'll get it and note we need to release it later.
+    release_lock = False
+    if not self.__lock._is_owned():
+      release_lock = True
+      self.__lock.acquire()
+
     try:
       invalid_names = set(self.__names()).intersection(names)
       if invalid_names:
@@ -499,7 +539,9 @@ class LockSet:
         self.__lockdict[lockname] = lock
 
     finally:
-      self.__lock.release()
+      # Only release __lock if we were not holding it previously.
+      if release_lock:
+        self.__lock.release()
 
     return True
 
diff --git a/test/ganeti.locking_unittest.py b/test/ganeti.locking_unittest.py
index 74524ed..996c8c2 100755
--- a/test/ganeti.locking_unittest.py
+++ b/test/ganeti.locking_unittest.py
@@ -331,6 +331,19 @@ class TestLockSet(unittest.TestCase):
     # Cannot remove 'three' as we are sharing it
     self.assertRaises(AssertionError, self.ls.remove, 'three')
 
+  def testAcquireSetLock(self):
+    # acquire the set-lock exclusively
+    self.assertEquals(self.ls.acquire(None), set(['one', 'two', 'three']))
+    # I can still add/remove elements...
+    self.assertEquals(self.ls.remove(['two', 'three']), ['two', 'three'])
+    self.assert_(self.ls.add('six'))
+    self.ls.release()
+    # share the set-lock
+    self.assertEquals(self.ls.acquire(None, shared=1), set(['one', 'six']))
+    # adding new elements is not possible
+    self.assertRaises(AssertionError, self.ls.add, 'five')
+    self.ls.release()
+
   def _doLockSet(self, set, shared):
     try:
       self.ls.acquire(set, shared=shared)
@@ -339,6 +352,14 @@ class TestLockSet(unittest.TestCase):
     except errors.LockError:
       self.done.put('ERR')
 
+  def _doAddSet(self, set):
+    try:
+      self.ls.add(set, acquired=1)
+      self.done.put('DONE')
+      self.ls.release()
+    except errors.LockError:
+      self.done.put('ERR')
+
   def _doRemoveSet(self, set):
     self.done.put(self.ls.remove(set))
 
@@ -412,6 +433,42 @@ class TestLockSet(unittest.TestCase):
     self.assertEqual(self.done.get(True, 1), ['two'])
     self.ls.release()
 
+  def testConcurrentSharedSetLock(self):
+    # share the set-lock...
+    self.ls.acquire(None, shared=1)
+    # ...another thread can share it too
+    Thread(target=self._doLockSet, args=(None, 1)).start()
+    self.assertEqual(self.done.get(True, 1), 'DONE')
+    # ...or just share some elements
+    Thread(target=self._doLockSet, args=(['one', 'three'], 1)).start()
+    self.assertEqual(self.done.get(True, 1), 'DONE')
+    # ...but not add new ones or remove any
+    Thread(target=self._doAddSet, args=(['nine'])).start()
+    Thread(target=self._doRemoveSet, args=(['two'], )).start()
+    self.assertRaises(Queue.Empty, self.done.get, True, 0.2)
+    # this just releases the set-lock
+    self.ls.release([])
+    self.assertEqual(self.done.get(True, 1), 'DONE')
+    # release the lock on the actual elements so remove() can proceed too
+    self.ls.release()
+    self.assertEqual(self.done.get(True, 1), ['two'])
+
+  def testConcurrentExclusiveSetLock(self):
+    # acquire the set-lock...
+    self.ls.acquire(None, shared=0)
+    # ...no one can do anything else
+    Thread(target=self._doLockSet, args=(None, 1)).start()
+    Thread(target=self._doLockSet, args=(None, 0)).start()
+    Thread(target=self._doLockSet, args=(['three'], 0)).start()
+    Thread(target=self._doLockSet, args=(['two'], 1)).start()
+    Thread(target=self._doAddSet, args=(['nine'])).start()
+    self.ls.release()
+    self.assertEqual(self.done.get(True, 1), 'DONE')
+    self.assertEqual(self.done.get(True, 1), 'DONE')
+    self.assertEqual(self.done.get(True, 1), 'DONE')
+    self.assertEqual(self.done.get(True, 1), 'DONE')
+    self.assertEqual(self.done.get(True, 1), 'DONE')
+
 
 class TestGanetiLockManager(unittest.TestCase):
 
