1"""holds locking functionality that works across processes"""
2from __future__ import absolute_import, unicode_literals
3
4import logging
5import os
6from abc import ABCMeta, abstractmethod
7from contextlib import contextmanager
8from threading import Lock, RLock
9
10from filelock import FileLock, Timeout
11from six import add_metaclass
12
13from virtualenv.util.path import Path
14
15
16class _CountedFileLock(FileLock):
17    def __init__(self, lock_file):
18        parent = os.path.dirname(lock_file)
19        if not os.path.isdir(parent):
20            try:
21                os.makedirs(parent)
22            except OSError:
23                pass
24        super(_CountedFileLock, self).__init__(lock_file)
25        self.count = 0
26        self.thread_safe = RLock()
27
28    def acquire(self, timeout=None, poll_intervall=0.05):
29        with self.thread_safe:
30            if self.count == 0:
31                super(_CountedFileLock, self).acquire(timeout=timeout, poll_intervall=poll_intervall)
32            self.count += 1
33
34    def release(self, force=False):
35        with self.thread_safe:
36            if self.count == 1:
37                super(_CountedFileLock, self).release(force=force)
38            self.count = max(self.count - 1, 0)
39
40
41_lock_store = {}
42_store_lock = Lock()
43
44
45@add_metaclass(ABCMeta)
46class PathLockBase(object):
47    def __init__(self, folder):
48        path = Path(folder)
49        self.path = path.resolve() if path.exists() else path
50
51    def __repr__(self):
52        return "{}({})".format(self.__class__.__name__, self.path)
53
54    def __div__(self, other):
55        return type(self)(self.path / other)
56
57    def __truediv__(self, other):
58        return self.__div__(other)
59
60    @abstractmethod
61    def __enter__(self):
62        raise NotImplementedError
63
64    @abstractmethod
65    def __exit__(self, exc_type, exc_val, exc_tb):
66        raise NotImplementedError
67
68    @abstractmethod
69    @contextmanager
70    def lock_for_key(self, name, no_block=False):
71        raise NotImplementedError
72
73    @abstractmethod
74    @contextmanager
75    def non_reentrant_lock_for_key(name):
76        raise NotImplementedError
77
78
79class ReentrantFileLock(PathLockBase):
80    def __init__(self, folder):
81        super(ReentrantFileLock, self).__init__(folder)
82        self._lock = None
83
84    def _create_lock(self, name=""):
85        lock_file = str(self.path / "{}.lock".format(name))
86        with _store_lock:
87            if lock_file not in _lock_store:
88                _lock_store[lock_file] = _CountedFileLock(lock_file)
89            return _lock_store[lock_file]
90
91    @staticmethod
92    def _del_lock(lock):
93        with _store_lock:
94            if lock is not None:
95                with lock.thread_safe:
96                    if lock.count == 0:
97                        _lock_store.pop(lock.lock_file, None)
98
99    def __del__(self):
100        self._del_lock(self._lock)
101
102    def __enter__(self):
103        self._lock = self._create_lock()
104        self._lock_file(self._lock)
105
106    def __exit__(self, exc_type, exc_val, exc_tb):
107        self._release(self._lock)
108
109    def _lock_file(self, lock, no_block=False):
110        # multiple processes might be trying to get a first lock... so we cannot check if this directory exist without
111        # a lock, but that lock might then become expensive, and it's not clear where that lock should live.
112        # Instead here we just ignore if we fail to create the directory.
113        try:
114            os.makedirs(str(self.path))
115        except OSError:
116            pass
117        try:
118            lock.acquire(0.0001)
119        except Timeout:
120            if no_block:
121                raise
122            logging.debug("lock file %s present, will block until released", lock.lock_file)
123            lock.release()  # release the acquire try from above
124            lock.acquire()
125
126    @staticmethod
127    def _release(lock):
128        lock.release()
129
130    @contextmanager
131    def lock_for_key(self, name, no_block=False):
132        lock = self._create_lock(name)
133        try:
134            try:
135                self._lock_file(lock, no_block)
136                yield
137            finally:
138                self._release(lock)
139        finally:
140            self._del_lock(lock)
141
142    @contextmanager
143    def non_reentrant_lock_for_key(self, name):
144        with _CountedFileLock(str(self.path / "{}.lock".format(name))):
145            yield
146
147
148class NoOpFileLock(PathLockBase):
149    def __enter__(self):
150        raise NotImplementedError
151
152    def __exit__(self, exc_type, exc_val, exc_tb):
153        raise NotImplementedError
154
155    @contextmanager
156    def lock_for_key(self, name, no_block=False):
157        yield
158
159    @contextmanager
160    def non_reentrant_lock_for_key(self, name):
161        yield
162
163
164__all__ = (
165    "NoOpFileLock",
166    "ReentrantFileLock",
167    "Timeout",
168)
169