From b4478d3498056633aa0280ba7afcf5b52e91b2ef Mon Sep 17 00:00:00 2001
From: Michael Hanselmann <hansmi@google.com>
Date: Fri, 26 Feb 2010 18:35:59 +0100
Subject: [PATCH] Support passing in file object in utils.FileLock

This way we can re-use file objects opened in other places. Also add more
unittests.

Signed-off-by: Michael Hanselmann <hansmi@google.com>
Reviewed-by: Guido Trotter <ultrotter@google.com>
Reviewed-by: Iustin Pop <iustin@google.com>
---
 lib/jstore.py                 |  2 +-
 lib/ssconf.py                 |  2 +-
 lib/utils.py                  | 24 +++++++++---
 test/ganeti.utils_unittest.py | 70 ++++++++++++++++++++++++++++++++---
 4 files changed, 86 insertions(+), 12 deletions(-)

diff --git a/lib/jstore.py b/lib/jstore.py
index 723cda17f..a4e51fc49 100644
--- a/lib/jstore.py
+++ b/lib/jstore.py
@@ -88,7 +88,7 @@ def InitAndVerifyQueue(must_lock):
         raise
 
   # Lock queue
-  queue_lock = utils.FileLock(constants.JOB_QUEUE_LOCK_FILE)
+  queue_lock = utils.FileLock.Open(constants.JOB_QUEUE_LOCK_FILE)
   try:
     # The queue needs to be locked in exclusive mode to write to the serial and
     # version files.
diff --git a/lib/ssconf.py b/lib/ssconf.py
index baa7f1906..07a264a06 100644
--- a/lib/ssconf.py
+++ b/lib/ssconf.py
@@ -307,7 +307,7 @@ class SimpleStore(object):
     @param values: Dictionary of (name, value)
 
     """
-    ssconf_lock = utils.FileLock(constants.SSCONF_LOCK_FILE)
+    ssconf_lock = utils.FileLock.Open(constants.SSCONF_LOCK_FILE)
 
     # Get lock while writing files
     ssconf_lock.Exclusive(blocking=True, timeout=SSCONF_LOCK_TIMEOUT)
diff --git a/lib/utils.py b/lib/utils.py
index bb8b98613..a6b677c0b 100644
--- a/lib/utils.py
+++ b/lib/utils.py
@@ -2446,17 +2446,31 @@ class FileLock(object):
   """Utility class for file locks.
 
   """
-  def __init__(self, filename):
+  def __init__(self, fd, filename):
     """Constructor for FileLock.
 
-    This will open the file denoted by the I{filename} argument.
-
+    @type fd: file
+    @param fd: File object
     @type filename: str
-    @param filename: path to the file to be locked
+    @param filename: Path of the file opened at I{fd}
 
     """
+    self.fd = fd
     self.filename = filename
-    self.fd = open(self.filename, "w")
+
+  @classmethod
+  def Open(cls, filename):
+    """Creates and opens a file to be used as a file-based lock.
+
+    @type filename: string
+    @param filename: path to the file to be locked
+
+    """
+    # Using "os.open" is necessary to allow both opening existing file
+    # read/write and creating if not existing. Vanilla "open" will truncate an
+    # existing file -or- allow creating if not existing.
+    return cls(os.fdopen(os.open(filename, os.O_RDWR | os.O_CREAT), "w+"),
+               filename)
 
   def __del__(self):
     self.Close()
diff --git a/test/ganeti.utils_unittest.py b/test/ganeti.utils_unittest.py
index 14954684a..f1c665940 100755
--- a/test/ganeti.utils_unittest.py
+++ b/test/ganeti.utils_unittest.py
@@ -992,13 +992,9 @@ class TestTailFile(testutils.GanetiTestCase):
       self.failUnlessEqual(TailFile(fname, lines=i), data[-i:])
 
 
-class TestFileLock(unittest.TestCase):
+class _BaseFileLockTest:
   """Test case for the FileLock class"""
 
-  def setUp(self):
-    self.tmpfile = tempfile.NamedTemporaryFile()
-    self.lock = utils.FileLock(self.tmpfile.name)
-
   def testSharedNonblocking(self):
     self.lock.Shared(blocking=False)
     self.lock.Close()
@@ -1035,6 +1031,45 @@ class TestFileLock(unittest.TestCase):
     self.lock.Unlock(blocking=False)
     self.lock.Close()
 
+  def testSimpleTimeout(self):
+    # These will succeed on the first attempt, hence a short timeout
+    self.lock.Shared(blocking=True, timeout=10.0)
+    self.lock.Exclusive(blocking=False, timeout=10.0)
+    self.lock.Unlock(blocking=True, timeout=10.0)
+    self.lock.Close()
+
+  @staticmethod
+  def _TryLockInner(filename, shared, blocking):
+    lock = utils.FileLock.Open(filename)
+
+    if shared:
+      fn = lock.Shared
+    else:
+      fn = lock.Exclusive
+
+    try:
+      # The timeout doesn't really matter as the parent process waits for us to
+      # finish anyway.
+      fn(blocking=blocking, timeout=0.01)
+    except errors.LockError, err:
+      return False
+
+    return True
+
+  def _TryLock(self, *args):
+    return utils.RunInSeparateProcess(self._TryLockInner, self.tmpfile.name,
+                                      *args)
+
+  def testTimeout(self):
+    for blocking in [True, False]:
+      self.lock.Exclusive(blocking=True)
+      self.failIf(self._TryLock(False, blocking))
+      self.failIf(self._TryLock(True, blocking))
+
+      self.lock.Shared(blocking=True)
+      self.assert_(self._TryLock(True, blocking))
+      self.failIf(self._TryLock(False, blocking))
+
   def testCloseShared(self):
     self.lock.Close()
     self.assertRaises(AssertionError, self.lock.Shared, blocking=False)
@@ -1048,6 +1083,31 @@ class TestFileLock(unittest.TestCase):
     self.assertRaises(AssertionError, self.lock.Unlock, blocking=False)
 
 
+class TestFileLockWithFilename(testutils.GanetiTestCase, _BaseFileLockTest):
+  TESTDATA = "Hello World\n" * 10
+
+  def setUp(self):
+    testutils.GanetiTestCase.setUp(self)
+
+    self.tmpfile = tempfile.NamedTemporaryFile()
+    utils.WriteFile(self.tmpfile.name, data=self.TESTDATA)
+    self.lock = utils.FileLock.Open(self.tmpfile.name)
+
+    # Ensure "Open" didn't truncate file
+    self.assertFileContent(self.tmpfile.name, self.TESTDATA)
+
+  def tearDown(self):
+    self.assertFileContent(self.tmpfile.name, self.TESTDATA)
+
+    testutils.GanetiTestCase.tearDown(self)
+
+
+class TestFileLockWithFileObject(unittest.TestCase, _BaseFileLockTest):
+  def setUp(self):
+    self.tmpfile = tempfile.NamedTemporaryFile()
+    self.lock = utils.FileLock(open(self.tmpfile.name, "w"), self.tmpfile.name)
+
+
 class TestTimeFunctions(unittest.TestCase):
   """Test case for time functions"""
 
-- 
GitLab