Intermediate changes

commit_hash:afe2bc6dcfc79e6d34dabab9a0c92d3fa18bc87d

4 files changed, 151 insertions, 25 deletions

diff --git a/library/python/filelock/__init__.py b/library/python/filelock/__init__.py
index fae89713d3..545844b598 100644
--- a/library/python/filelock/__init__.py
+++ b/library/python/filelock/__init__.py
@@ -1,19 +1,17 @@
+import collections
 import errno
 import logging
 import os
+import struct
 import sys
+import time
 
 import library.python.windows
 
 logger = logging.getLogger(__name__)
 
-
-def set_close_on_exec(stream):
-    if library.python.windows.on_win():
-        library.python.windows.set_handle_information(stream, inherit=False)
-    else:
-        import fcntl
-        fcntl.fcntl(stream, fcntl.F_SETFD, fcntl.FD_CLOEXEC)
+# python2 compat
+os_O_CLOEXEC = getattr(os, 'O_CLOEXEC', 1 << 19)
 
 
 class AbstractFileLock(object):
@@ -40,13 +38,15 @@ class _NixFileLock(AbstractFileLock):
     def __init__(self, path):
         super(_NixFileLock, self).__init__(path)
         from fcntl import flock, LOCK_EX, LOCK_UN, LOCK_NB
+
         self._locker = lambda lock, blocking: flock(lock, LOCK_EX if blocking else LOCK_EX | LOCK_NB)
         self._unlocker = lambda lock: flock(lock, LOCK_UN)
-        self._lock = open(self.path, 'a')
-        set_close_on_exec(self._lock)
+        # nonbuffered random access rw mode
+        self._lock = os.fdopen(os.open(self.path, os.O_RDWR | os.O_CREAT | os_O_CLOEXEC), 'r+b', 0)
 
     def acquire(self, blocking=True):
         import errno
+
         try:
             self._locker(self._lock, blocking)
         except IOError as e:
@@ -73,34 +73,32 @@ class _WinFileLock(AbstractFileLock):
 
     def __init__(self, path):
         super(_WinFileLock, self).__init__(path)
-        self._lock = None
+        # nonbuffered random access rw mode
+        self._lock = os.fdopen(os.open(self.path, os.O_RDWR | os.O_CREAT | os.O_BINARY | os.O_NOINHERIT), 'r+b', 0)
         try:
-            with open(path, 'w') as lock_file:
-                lock_file.write(" " * self._LOCKED_BYTES_NUM)
+            self._lock.write(b' ' * self._LOCKED_BYTES_NUM)
         except IOError as e:
             if e.errno != errno.EACCES or not os.path.isfile(path):
                 raise
 
     def acquire(self, blocking=True):
-        self._lock = open(self.path)
-        set_close_on_exec(self._lock)
-
-        import time
         locked = False
         while not locked:
             locked = library.python.windows.lock_file(self._lock, 0, self._LOCKED_BYTES_NUM, raises=False)
             if locked:
                 return True
             if blocking:
-                time.sleep(.5)
+                time.sleep(0.5)
             else:
                 return False
 
     def release(self):
         if self._lock:
             library.python.windows.unlock_file(self._lock, 0, self._LOCKED_BYTES_NUM, raises=False)
+
+    def __del__(self):
+        if getattr(self, '_lock', False):
             self._lock.close()
-            self._lock = None
 
 
 class FileLock(AbstractFileLock):
@@ -114,9 +112,64 @@ class FileLock(AbstractFileLock):
             self._lock = _NixFileLock(path)
 
     def acquire(self, blocking=True):
-        logger.debug('Acquiring filelock (blocking=%s): %s', blocking, self.path)
+        logger.debug('Acquiring %s (blocking=%s): %s', type(self).__name__, blocking, self.path)
         return self._lock.acquire(blocking)
 
     def release(self):
-        logger.debug('Ensuring filelock released: %s', self.path)
+        logger.debug('Ensuring %s released: %s', type(self).__name__, self.path)
         return self._lock.release()
+
+
+_LockInfo = collections.namedtuple('LockInfo', ['pid', 'time'])
+
+
+class _PidLockMixin(object):
+    _LockedBytes = 0
+    _InfoFormat = 'QQ'
+    _InfoFmtSize = struct.calcsize(_InfoFormat)
+
+    def _register_lock(self):
+        self._lock.seek(self._LockedBytes, os.SEEK_SET)
+        self._lock.write(struct.pack(self._InfoFormat, os.getpid(), int(time.time())))
+
+    @property
+    def info(self):
+        self._lock.seek(self._LockedBytes, os.SEEK_SET)
+        try:
+            data = struct.unpack(self._InfoFormat, self._lock.read(self._InfoFmtSize))
+        except struct.error:
+            data = 0, 0
+        return _LockInfo(*data)
+
+
+class _NixPidFileLock(_NixFileLock, _PidLockMixin):
+    def acquire(self, blocking=True):
+        if super(_NixPidFileLock, self).acquire(blocking):
+            self._register_lock()
+            return True
+        return False
+
+
+class _WinPidFileLock(_WinFileLock, _PidLockMixin):
+    _LockedBytes = _WinFileLock._LOCKED_BYTES_NUM
+
+    def acquire(self, blocking=True):
+        if super(_WinPidFileLock, self).acquire(blocking):
+            self._register_lock()
+            return True
+        return False
+
+
+class PidFileLock(FileLock):
+
+    def __init__(self, path):
+        AbstractFileLock.__init__(self, path)
+
+        if sys.platform.startswith('win'):
+            self._lock = _WinPidFileLock(path)
+        else:
+            self._lock = _NixPidFileLock(path)
+
+    @property
+    def info(self):
+        return self._lock.info
diff --git a/library/python/filelock/ut/test_filelock.py b/library/python/filelock/ut/test_filelock.py
index 60a3722390..018acf290b 100644
--- a/library/python/filelock/ut/test_filelock.py
+++ b/library/python/filelock/ut/test_filelock.py
@@ -58,12 +58,12 @@ def test_filelock_init_acquired():
 
 
 def test_concurrent_lock():
-    filename = 'con.lock'
+    filename = 'thread.lock'
 
     def lock():
         lock = library.python.filelock.FileLock(filename)
         time.sleep(1)
-        lock.acquire()
+        assert lock.acquire()
         lock.release()
         try:
             os.unlink(filename)
@@ -73,7 +73,6 @@ def test_concurrent_lock():
     threads = []
     for i in range(100):
         t = threading.Thread(target=lock)
-        t.daemon = True
         threads.append(t)
 
     for t in threads:
@@ -81,3 +80,59 @@ def test_concurrent_lock():
 
     for t in threads:
         t.join()
+
+
+def test_pidfilelock():
+    lock_file = 'pidfile.lock'
+    # there should be no info
+    lock = library.python.filelock.PidFileLock(lock_file)
+    assert lock.info.pid == 0
+    assert lock.info.time == 0
+
+    with library.python.filelock.PidFileLock(lock_file) as lock:
+        assert lock.info.pid == os.getpid()
+        assert lock.info.time <= time.time()
+        assert lock.info.time > time.time() - 2
+
+        newlock = library.python.filelock.PidFileLock(lock_file)
+        # info shouldn't require locking
+        assert newlock.info.pid == os.getpid()
+        assert not newlock.acquire(blocking=False)
+
+    newlock = library.python.filelock.PidFileLock(lock_file)
+    # info is still accessible
+    assert newlock.info.pid == os.getpid()
+    t = newlock.info.time
+    # info is updated
+    time.sleep(1)
+    with newlock as lock:
+        assert lock.info.time > t
+
+
+def _try_acquire_pidlock(lock_file, out_file, lock_pid=None):
+    lock = library.python.filelock.PidFileLock(lock_file)
+    with open(out_file, "w") as afile:
+        afile.write("1" if lock.acquire(blocking=False) else "0")
+
+    if lock_pid is not None:
+        assert lock.info.pid == lock_pid
+
+
+def test_pidfilelock_multiprocessing():
+    lock_file = 'mp_pidfile.lock'
+    out_file = lock_file + ".out"
+
+    # subprocess can aquire lock
+    proc = multiprocessing.Process(target=_try_acquire_pidlock, args=(lock_file, out_file))
+    proc.start()
+    proc.join()
+    with open(out_file) as afile:
+        assert "1" == afile.read()
+
+    # subprocess can't aquire lock
+    with library.python.filelock.PidFileLock(lock_file) as lock:
+        proc = multiprocessing.Process(target=_try_acquire_pidlock, args=(lock_file, out_file, lock.info.pid))
+        proc.start()
+        proc.join()
+        with open(out_file) as afile:
+            assert "0" == afile.read()
diff --git a/library/python/filelock/ut/ya.make b/library/python/filelock/ut/ya.make
index a62699640f..60108f73c6 100644
--- a/library/python/filelock/ut/ya.make
+++ b/library/python/filelock/ut/ya.make
@@ -1,9 +1,25 @@
 PY23_TEST()
 
-TEST_SRCS(test_filelock.py)
+TEST_SRCS(
+    test_filelock.py
+)
 
 PEERDIR(
     library/python/filelock
 )
 
+IF (OS_DARWIN)
+    SIZE(LARGE)
+    TAG(
+        ya:fat
+        ya:exotic_platform
+    )
+ELSEIF (OS_WINDOWS)
+    SIZE(LARGE)
+    TAG(
+        ya:fat
+        sb:ssd&~MULTISLOT&WINDOWS
+    )
+ENDIF()
+
 END()
diff --git a/library/python/filelock/ya.make b/library/python/filelock/ya.make
index 55c6c23ef3..408961d20f 100644
--- a/library/python/filelock/ya.make
+++ b/library/python/filelock/ya.make
@@ -1,6 +1,8 @@
 PY23_LIBRARY()
 
-PY_SRCS(__init__.py)
+PY_SRCS(
+    __init__.py
+)
 
 PEERDIR(
     library/python/windows