intermediate changes

ref:cde9a383711a11544ce7e107a78147fb96cc4029

7 files changed, 252 insertions, 0 deletions

diff --git a/library/python/filelock/__init__.py b/library/python/filelock/__init__.py
new file mode 100644
index 00000000000..f81ff67f37f
--- /dev/null
+++ b/library/python/filelock/__init__.py
@@ -0,0 +1,122 @@
+import errno
+import logging
+import os
+import sys
+
+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)
+
+
+class AbstractFileLock(object):
+
+    def __init__(self, path):
+        self.path = path
+
+    def acquire(self, blocking=True):
+        raise NotImplementedError
+
+    def release(self):
+        raise NotImplementedError
+
+    def __enter__(self):
+        self.acquire()
+        return self
+
+    def __exit__(self, type, value, traceback):
+        self.release()
+
+
+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)
+
+    def acquire(self, blocking=True):
+        import errno
+        try:
+            self._locker(self._lock, blocking)
+        except IOError as e:
+            if e.errno in (errno.EAGAIN, errno.EACCES) and not blocking:
+                return False
+            raise
+        return True
+
+    def release(self):
+        self._unlocker(self._lock)
+
+    def __del__(self):
+        if hasattr(self, "_lock"):
+            self._lock.close()
+
+
+class _WinFileLock(AbstractFileLock):
+    """
+    Based on LockFile / UnlockFile from win32 API
+    https://msdn.microsoft.com/en-us/library/windows/desktop/aa365202(v=vs.85).aspx
+    """
+
+    _LOCKED_BYTES_NUM = 1
+
+    def __init__(self, path):
+        super(_WinFileLock, self).__init__(path)
+        self._lock = None
+        try:
+            with file(path, 'w') as lock_file:
+                lock_file.write(" " * 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)
+            else:
+                return False
+
+    def release(self):
+        if self._lock:
+            library.python.windows.unlock_file(self._lock, 0, self._LOCKED_BYTES_NUM, raises=False)
+            self._lock.close()
+            self._lock = None
+
+
+class FileLock(AbstractFileLock):
+
+    def __init__(self, path):
+        super(FileLock, self).__init__(path)
+
+        if sys.platform.startswith('win'):
+            self._lock = _WinFileLock(path)
+        else:
+            self._lock = _NixFileLock(path)
+
+    def acquire(self, blocking=True):
+        logger.debug('Acquiring filelock (blocking=%s): %s', blocking, self.path)
+        return self._lock.acquire(blocking)
+
+    def release(self):
+        logger.debug('Ensuring filelock released: %s', self.path)
+        return self._lock.release()
diff --git a/library/python/filelock/ut/lib/test_filelock.py b/library/python/filelock/ut/lib/test_filelock.py
new file mode 100644
index 00000000000..1b11d891231
--- /dev/null
+++ b/library/python/filelock/ut/lib/test_filelock.py
@@ -0,0 +1,83 @@
+import os
+import time
+import logging
+import multiprocessing
+import tempfile
+import threading
+
+import library.python.filelock
+
+
+def _acquire_lock(lock_path, out_file_path):
+    with library.python.filelock.FileLock(lock_path):
+        with open(out_file_path, "a") as out:
+            out.write("{}:{}\n".format(os.getpid(), time.time()))
+        time.sleep(2)
+
+
+def test_filelock():
+    temp_dir = tempfile.mkdtemp()
+    lock_path = os.path.join(temp_dir, "file.lock")
+    out_file_path = os.path.join(temp_dir, "out.txt")
+
+    process_count = 5
+    processes = []
+    for i in range(process_count):
+        process = multiprocessing.Process(target=_acquire_lock, args=(lock_path, out_file_path))
+        process.start()
+        processes.append(process)
+
+    for process in processes:
+        process.join()
+
+    pids = []
+    times = []
+    with open(out_file_path) as out:
+        content = out.read()
+        logging.info("Times:\n%s", content)
+        for line in content.strip().split("\n"):
+            pid, time_val = line.split(":")
+            pids.append(pid)
+            times.append(float(time_val))
+
+    assert len(set(pids)) == process_count
+    time1 = times.pop()
+    while times:
+        time2 = times.pop()
+        assert int(time1) - int(time2) >= 2
+        time1 = time2
+
+
+def test_filelock_init_acquired():
+    temp_dir = tempfile.mkdtemp()
+    lock_path = os.path.join(temp_dir, "file.lock")
+
+    with library.python.filelock.FileLock(lock_path):
+        sublock = library.python.filelock.FileLock(lock_path)
+        del sublock
+
+
+def test_concurrent_lock():
+    filename = 'con.lock'
+
+    def lock():
+        l = library.python.filelock.FileLock(filename)
+        time.sleep(1)
+        l.acquire()
+        l.release()
+        try:
+            os.unlink(filename)
+        except OSError:
+            pass
+
+    threads = []
+    for i in range(100):
+        t = threading.Thread(target=lock)
+        t.daemon = True
+        threads.append(t)
+
+    for t in threads:
+        t.start()
+
+    for t in threads:
+        t.join()
diff --git a/library/python/filelock/ut/lib/ya.make b/library/python/filelock/ut/lib/ya.make
new file mode 100644
index 00000000000..f3f9da5a67e
--- /dev/null
+++ b/library/python/filelock/ut/lib/ya.make
@@ -0,0 +1,11 @@
+OWNER(g:yatool)
+
+PY23_LIBRARY()
+
+TEST_SRCS(test_filelock.py)
+
+PEERDIR(
+    library/python/filelock
+)
+
+END()
diff --git a/library/python/filelock/ut/py2/ya.make b/library/python/filelock/ut/py2/ya.make
new file mode 100644
index 00000000000..30b54e0232a
--- /dev/null
+++ b/library/python/filelock/ut/py2/ya.make
@@ -0,0 +1,9 @@
+OWNER(g:yatool)
+
+PY2TEST()
+
+PEERDIR(
+    library/python/filelock/ut/lib
+)
+
+END()
diff --git a/library/python/filelock/ut/py3/ya.make b/library/python/filelock/ut/py3/ya.make
new file mode 100644
index 00000000000..05a856a19a2
--- /dev/null
+++ b/library/python/filelock/ut/py3/ya.make
@@ -0,0 +1,9 @@
+OWNER(g:yatool)
+
+PY3TEST()
+
+PEERDIR(
+    library/python/filelock/ut/lib
+)
+
+END()
diff --git a/library/python/filelock/ut/ya.make b/library/python/filelock/ut/ya.make
new file mode 100644
index 00000000000..9280ea415e1
--- /dev/null
+++ b/library/python/filelock/ut/ya.make
@@ -0,0 +1,7 @@
+OWNER(g:yatool)
+
+RECURSE(
+    lib
+    py2
+    py3
+)
diff --git a/library/python/filelock/ya.make b/library/python/filelock/ya.make
new file mode 100644
index 00000000000..958cc1866f6
--- /dev/null
+++ b/library/python/filelock/ya.make
@@ -0,0 +1,11 @@
+OWNER(g:yatool)
+
+PY23_LIBRARY()
+
+PY_SRCS(__init__.py)
+
+PEERDIR(
+    library/python/windows
+)
+
+END()