diff --git a/.pylintrc b/.pylintrc index d6c2690..6f61736 100644 --- a/.pylintrc +++ b/.pylintrc @@ -340,7 +340,7 @@ indent-string=' ' max-line-length=100 # Maximum number of lines in a module. -max-module-lines=3000 +max-module-lines=4000 # Allow the body of a class to be on the same line as the declaration if body # contains single statement. diff --git a/diskcache/core.py b/diskcache/core.py index f52a445..044ddc0 100644 --- a/diskcache/core.py +++ b/diskcache/core.py @@ -23,6 +23,7 @@ import time import uuid import warnings +import weakref import zlib @@ -795,6 +796,207 @@ def resolve_ignore(func, ignore): return frozenset(ignore) | positions +class _ThreadSentinel: # pylint: disable=too-few-public-methods + """Marker kept in a Cache's thread-local state next to its connection. + + Thread-local state is destroyed when its thread exits, so once a weak + reference to the sentinel is dead the connection's thread is gone and + the connection may be closed from another thread. + """ + + __slots__ = ('__weakref__',) + + +class _Connections: + """Registry of the SQLite connections one Cache opened in this process. + + Connections are opened per thread (see ``Cache._con``). The registry + lets ``Cache.close`` and garbage collection close the connections of + threads that have already exited, and lets a forked child process + forget the connections it inherited without closing them. + """ + + def __init__(self): + self._lock = threading.Lock() + self._items = {} + self.pid = os.getpid() + + def add(self, ident, sentinel, con): + """Register `con` for thread `ident`; return connections to close. + + The returned connections belong to threads that have exited, or are + an earlier connection of the same thread whose thread-local state + was lost (a foreign thread that got a new interpreter thread state). + """ + with self._lock: + stale = self._pop_stale() + previous = self._items.pop(ident, None) + self._items[ident] = (weakref.ref(sentinel), con) + if previous is not None and previous[1] is not con: + stale.append(previous[1]) + return stale + + def pop(self, ident): + """Remove the connection of thread `ident` and return it with the + connections of threads that have exited, all to be closed.""" + with self._lock: + stale = self._pop_stale() + item = self._items.pop(ident, None) + if item is not None: + stale.append(item[1]) + return stale + + def _pop_stale(self): + # The caller holds the lock. + stale = [] + for ident, (ref, con) in list(self._items.items()): + if ref() is None: + del self._items[ident] + stale.append(con) + return stale + + def close_all(self): + """Close every registered connection. + + Runs when the Cache is garbage collected. In a child process the + connections were inherited and are discarded instead. + """ + if os.getpid() != self.pid: + self.discard() + return + with self._lock: + cons = [con for _, con in self._items.values()] + self._items.clear() + for con in cons: + con.close() + + def discard(self): + """Forget the connections after a fork without closing them. + + SQLite forbids a child process from using, or even closing, a + connection the parent opened: the child inherits the parent's lock + state, and closing may roll back or unlock underneath the parent + and corrupt the database. The connections stay referenced for the + life of the child so garbage collection never closes them either; + their file descriptors leak on purpose, as SQLite recommends. + """ + _INHERITED_CONNECTIONS.extend(con for _, con in self._items.values()) + self._items = {} + self._lock = threading.Lock() # The parent's lock may be held. + self.pid = os.getpid() + + +class _ForkGate: + """Delay ``os.fork`` while another thread is inside a transaction. + + A child process inherits the parent's SQLite state as of the fork. A + transaction in flight in another thread leaves the child's copy of that + state holding a write lock that nothing in the child ever releases, so + the child's own writes to that database would fail with ``Timeout``. + The gate holds new transactions back and makes fork wait, up to + `timeout` seconds, until the transactions in flight in other threads + finish; the held-back transactions proceed once the fork is done. The + forking thread's own transaction cannot be waited for and is not + counted; forking inside a transaction is unsupported. + """ + + timeout = 60 + + def __init__(self): + self._cond = threading.Condition(threading.Lock()) + self._active = 0 + self._forking = False + self._local = threading.local() + + def __enter__(self): + depth = getattr(self._local, 'depth', 0) + if not depth: + with self._cond: + while self._forking: + self._cond.wait() + self._active += 1 + self._local.depth = depth + 1 + + def __exit__(self, *exception): + depth = getattr(self._local, 'depth', 0) - 1 + if depth < 0: + return # A child process reset the gate during a transaction. + self._local.depth = depth + if not depth: + with self._cond: + self._active -= 1 + if not self._active: + self._cond.notify_all() + + def before_fork(self): + """Hold new transactions back, then wait for the active ones.""" + own = 1 if getattr(self._local, 'depth', 0) else 0 + deadline = time.monotonic() + self.timeout + with self._cond: + while self._forking: # Another thread is forking. + self._cond.wait() + self._forking = True + while self._active > own: + remaining = deadline - time.monotonic() + if remaining <= 0: + warnings.warn( + 'forking while a Cache transaction is in progress in' + ' another thread; the child process may be unable' + ' to write to that cache', + RuntimeWarning, + stacklevel=2, + ) + break + self._cond.wait(remaining) + + def after_fork_parent(self): + """Let transactions proceed again.""" + with self._cond: + self._forking = False + self._cond.notify_all() + + def after_fork_child(self): + """Start over: the child has one thread and no transactions.""" + self._cond = threading.Condition(threading.Lock()) + self._active = 0 + self._forking = False + self._local = threading.local() + + +_FORK_GATE = _ForkGate() + +# Every live Cache, so a child process can forget inherited connections. +_CACHES: 'weakref.WeakSet[Cache]' = weakref.WeakSet() + +# Connections inherited from a parent process, kept open on purpose. +_INHERITED_CONNECTIONS: list[sqlite3.Connection] = [] + + +def _before_fork(): + # SQLite advises closing every connection before fork. The gate waits + # for other threads' transactions; then each cache closes the forking + # thread's own connection and those of exited threads, which the + # parent reopens on next use. Connections held by other live threads + # cannot be closed here and are discarded in the child. + _FORK_GATE.before_fork() + for cache in list(_CACHES): + cache._close_before_fork() # pylint: disable=protected-access + + +def _after_fork_in_child(): + _FORK_GATE.after_fork_child() + for cache in list(_CACHES): + cache._discard_connections() # pylint: disable=protected-access + + +if hasattr(os, 'register_at_fork'): + os.register_at_fork( + before=_before_fork, + after_in_parent=_FORK_GATE.after_fork_parent, + after_in_child=_after_fork_in_child, + ) + + class Cache: """Disk and file backed cache.""" @@ -826,7 +1028,11 @@ def __init__(self, directory=None, timeout=60, disk=Disk, **settings): self._directory = directory self._timeout = 0 # Manually handle retries during initialization. self._local = threading.local() + self._cons = _Connections() self._txn_id = None + _CACHES.add(self) + finalizer = weakref.finalize(self, self._cons.close_all) + finalizer.atexit = False if not op.isdir(directory): try: @@ -962,42 +1168,71 @@ def disk(self): @property def _con(self): - # Check process ID to support process forking. If the process - # ID changes, close the connection and update the process ID. - - local_pid = getattr(self._local, 'pid', None) - pid = os.getpid() + # Check the process ID to support forking. If it changed, the + # connections were inherited from the parent process: forget them + # without closing them (see _Connections.discard) and open new ones. + # The at-fork handler normally does this before the cache is used. - if local_pid != pid: - self.close() - self._local.pid = pid + if self._cons.pid != os.getpid(): + self._discard_connections() con = getattr(self._local, 'con', None) if con is None: - con = self._local.con = sqlite3.connect( - op.join(self._directory, DBNAME), - timeout=self._timeout, - isolation_level=None, - ) + con = self._connect() - # Some SQLite pragmas work on a per-connection basis so - # query the Settings table and reset the pragmas. The - # Settings table may not exist so catch and ignore the - # OperationalError that may occur. + return con - try: - select = 'SELECT key, value FROM Settings' - settings = con.execute(select).fetchall() - except sqlite3.OperationalError: - pass - else: - for key, value in settings: - if key.startswith('sqlite_'): - self.reset(key, value, update=False) + def _connect(self): + # check_same_thread=False lets close() and the connection registry + # close a connection from a thread other than the one that opened + # it, once that thread has exited. Each connection is still used by + # one thread at a time, which is all SQLite's multi-thread mode asks. + + con = sqlite3.connect( + op.join(self._directory, DBNAME), + timeout=self._timeout, + isolation_level=None, + check_same_thread=False, + ) + sentinel = _ThreadSentinel() + stale = self._cons.add(threading.get_ident(), sentinel, con) + self._local.con = con + self._local.sentinel = sentinel + + for old in stale: + old.close() + + # Some SQLite pragmas work on a per-connection basis so + # query the Settings table and reset the pragmas. The + # Settings table may not exist so catch and ignore the + # OperationalError that may occur. + + try: + select = 'SELECT key, value FROM Settings' + settings = con.execute(select).fetchall() + except sqlite3.OperationalError: + pass + else: + for key, value in settings: + if key.startswith('sqlite_'): + self.reset(key, value, update=False) return con + def _close_before_fork(self): + # Called in the parent right before fork. Inside a transaction the + # connection cannot be closed; forking there is unsupported. + if self._txn_id != threading.get_ident(): + self.close() + + def _discard_connections(self): + # In a child process: forget the connections inherited from the + # parent without closing them, and start with fresh per-thread state. + self._cons.discard() + self._local = threading.local() + self._txn_id = None + @property def _sql(self): return self._con.execute @@ -1110,44 +1345,45 @@ def transact(self, retry=False): @cl.contextmanager def _transact(self, retry=False, filename=None): - sql = self._sql - filenames = [] - _disk_remove = self._disk.remove - tid = threading.get_ident() - txn_id = self._txn_id + with _FORK_GATE: + sql = self._sql + filenames = [] + _disk_remove = self._disk.remove + tid = threading.get_ident() + txn_id = self._txn_id - if tid == txn_id: - begin = False - else: - while True: - try: - sql('BEGIN IMMEDIATE') - begin = True - self._txn_id = tid - break - except sqlite3.OperationalError: - if retry: - continue - if filename is not None: - _disk_remove(filename) - raise Timeout from None + if tid == txn_id: + begin = False + else: + while True: + try: + sql('BEGIN IMMEDIATE') + begin = True + self._txn_id = tid + break + except sqlite3.OperationalError: + if retry: + continue + if filename is not None: + _disk_remove(filename) + raise Timeout from None - try: - yield sql, filenames.append - except BaseException: - if begin: - assert self._txn_id == tid - self._txn_id = None - sql('ROLLBACK') - raise - else: - if begin: - assert self._txn_id == tid - self._txn_id = None - sql('COMMIT') - for name in filenames: - if name is not None: - _disk_remove(name) + try: + yield sql, filenames.append + except BaseException: + if begin: + assert self._txn_id == tid + self._txn_id = None + sql('ROLLBACK') + raise + else: + if begin: + assert self._txn_id == tid + self._txn_id = None + sql('COMMIT') + for name in filenames: + if name is not None: + _disk_remove(name) def set(self, key, value, expire=None, read=False, tag=None, retry=False): """Set `key` and `value` item in cache. @@ -2746,18 +2982,32 @@ def volume(self): return total_size def close(self): - """Close database connection.""" - con = getattr(self._local, 'con', None) + """Close database connection. - if con is None: + Closes the connection opened by the calling thread and the + connections of threads that have exited. Connections that other + live threads use stay open. The cache reconnects on next use. + + """ + if self._cons.pid != os.getpid(): + self._discard_connections() return - con.close() + local = self._local + con = getattr(local, 'con', None) + cons = self._cons.pop(threading.get_ident()) - try: - delattr(self._local, 'con') - except AttributeError: - pass + if con is not None and con not in cons: + cons.append(con) + + for name in ('con', 'sentinel'): + try: + delattr(local, name) + except AttributeError: + pass + + for con in cons: + con.close() def __enter__(self): # Create connection in thread. diff --git a/docs/tutorial.rst b/docs/tutorial.rst index 00f2d31..0ec5bad 100644 --- a/docs/tutorial.rst +++ b/docs/tutorial.rst @@ -81,12 +81,24 @@ be shared between threads. Two Cache objects may also reference the same directory from separate threads or processes. In this way, they are also process-safe and support cross-process communication. -Cache objects open and maintain one or more file handles. But unlike files, all -Cache operations are atomic and Cache objects support process-forking and may -be serialized using Pickle. Each thread that accesses a cache should also call -:meth:`close <.Cache.close>` on the cache. Cache objects can be used -in a `with` statement to safeguard calling :meth:`close -`. +Cache objects open and maintain one or more file handles: one SQLite +connection per thread that uses the cache. But unlike files, all Cache +operations are atomic and Cache objects may be serialized using Pickle. +Calling :meth:`close <.Cache.close>` closes the calling thread's connection +and the connections of threads that have exited; connections that other +running threads use stay open until those threads call ``close`` or the Cache +object is garbage collected. Cache objects can be used in a `with` statement +to safeguard calling :meth:`close `. + +Cache objects also support process forking. Before a fork, the forking +thread's connections are closed in the parent, where they reopen on the next +use, and the child process opens connections of its own; connections that +other threads hold at the time of the fork are left untouched in the child. +Forking inside a :meth:`transact <.Cache.transact>` block is not supported. +On macOS, a Python that links the system SQLite library cannot open any +database in a forked child once the parent has used SQLite; that limitation +is not specific to DiskCache, and the ``spawn`` start method of +:mod:`multiprocessing`, the default on macOS, avoids it. >>> cache.close() >>> with Cache(cache.directory) as reference: diff --git a/tests/test_core.py b/tests/test_core.py index 6f078da..12ff1dc 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -1,6 +1,8 @@ """Test diskcache.core.Cache.""" +import _thread import errno +import gc import hashlib import io import os @@ -8,6 +10,7 @@ import pathlib import pickle import shutil +import signal import sqlite3 import stat import subprocess as sp @@ -1751,3 +1754,333 @@ def test_existing_directory_keeps_mode(): finally: os.umask(umask) shutil.rmtree(directory, ignore_errors=True) + + +def resource_warnings(caught): + return [w for w in caught if issubclass(w.category, ResourceWarning)] + + +def wait_for_child(pid, timeout=30): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + rpid, status = os.waitpid(pid, os.WNOHANG) + if rpid: + return os.waitstatus_to_exitcode(status) + time.sleep(0.01) + os.kill(pid, signal.SIGKILL) + os.waitpid(pid, 0) + pytest.fail('child process did not exit within %s seconds' % timeout) + + +def forget_inherited(con): + # Undo an in-process simulation of a fork child: here it is safe to + # close the connection, because this process did open it. + dc.core._INHERITED_CONNECTIONS.remove(con) + con.close() + + +def test_close_closes_connections_of_exited_threads(cache): + cache.set('main', 0) + threads = [ + threading.Thread(target=cache.set, args=(num, num)) for num in range(3) + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + done = threading.Event() + + def foreign(): + # A thread with no threading.Thread object, like a Qt thread. + cache.set('foreign', 1) + done.set() + + _thread.start_new_thread(foreign, ()) + assert done.wait(10) + items = cache._cons._items + main = threading.get_ident() + deadline = time.monotonic() + 10 + while any(ref() for ident, (ref, _) in items.items() if ident != main): + assert time.monotonic() < deadline, 'thread-local state not released' + time.sleep(0.01) + # Opening a connection closes those of exited threads, so at most the + # main thread's and the last thread's connections are left. + assert 1 <= len(items) <= 2 + gc.collect() + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + cache.close() + gc.collect() + + assert not resource_warnings(caught) + assert cache._cons._items == {} + assert cache.get('main') == 0 # Reconnects on use. + assert cache.get('foreign') == 1 + + +def test_close_leaves_connections_of_live_threads(cache): + ready = threading.Event() + release = threading.Event() + results = [] + + def work(): + cache.set('thread', 1) + ready.set() + release.wait(10) + results.append(cache.get('thread')) + + thread = threading.Thread(target=work) + thread.start() + assert ready.wait(10) + cache.close() + release.set() + thread.join(10) + assert results == [1] + + +def test_connection_of_lost_thread_state_is_closed(cache): + cache.set('a', 1) + first = cache._local.con + # The interpreter drops a thread's local state when the thread exits or + # when a foreign thread gets a new thread state; simulate that here. + del cache._local.con + del cache._local.sentinel + cache.set('b', 2) + assert cache._local.con is not first + assert len(cache._cons._items) == 1 + with pytest.raises(sqlite3.ProgrammingError): + first.execute('SELECT 1') + + +def test_garbage_collected_cache_closes_connections(): + cache = dc.Cache() + directory = cache.directory + cache.set('a', 1) + gc.collect() + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + del cache + gc.collect() + + assert not resource_warnings(caught) + shutil.rmtree(directory, ignore_errors=True) + + +def test_after_fork_in_child_discards_connections(cache): + # Run the at-fork child handler in this process. + cache.set('a', 1) + con = cache._local.con + assert con not in dc.core._INHERITED_CONNECTIONS + dc.core._after_fork_in_child() + assert con in dc.core._INHERITED_CONNECTIONS + assert getattr(cache._local, 'con', None) is None + assert cache._cons._items == {} + assert cache.get('a') == 1 + assert cache._local.con is not con + forget_inherited(con) + + +def test_before_fork_closes_idle_connection(cache): + # Run the at-fork parent handlers in this process. + cache.set('a', 1) + con = cache._local.con + + with cache.transact(): + dc.core._before_fork() + dc.core._FORK_GATE.after_fork_parent() + assert cache._local.con is con # Not closed inside a transaction. + + dc.core._before_fork() + dc.core._FORK_GATE.after_fork_parent() + assert getattr(cache._local, 'con', None) is None + assert con not in dc.core._INHERITED_CONNECTIONS + with pytest.raises(sqlite3.ProgrammingError): + con.execute('SELECT 1') + assert cache.get('a') == 1 # Reconnects on use. + + +def test_pid_change_discards_connections(cache): + # The fallback for a fork the at-fork handler did not see. + cache.set('a', 1) + con = cache._local.con + cache._cons.pid = -1 + assert cache.get('a') == 1 + assert cache._local.con is not con + assert cache._cons.pid == os.getpid() + assert con in dc.core._INHERITED_CONNECTIONS + forget_inherited(con) + + +def test_close_after_pid_change_discards_connections(cache): + cache.set('a', 1) + con = cache._local.con + cache._cons.pid = -1 + cache.close() + assert con in dc.core._INHERITED_CONNECTIONS + assert cache._cons.pid == os.getpid() + assert cache.get('a') == 1 + forget_inherited(con) + + +def test_close_all_after_pid_change_discards_connections(): + cons = dc.core._Connections() + con = sqlite3.connect(':memory:', check_same_thread=False) + sentinel = dc.core._ThreadSentinel() + assert cons.add(threading.get_ident(), sentinel, con) == [] + cons.pid = -1 + cons.close_all() + assert con in dc.core._INHERITED_CONNECTIONS + forget_inherited(con) + + +def test_connections_add_replaces_same_thread(cache): + cons = dc.core._Connections() + sentinel = dc.core._ThreadSentinel() + first = sqlite3.connect(':memory:', check_same_thread=False) + second = sqlite3.connect(':memory:', check_same_thread=False) + ident = threading.get_ident() + assert cons.add(ident, sentinel, first) == [] + assert cons.add(ident, sentinel, second) == [first] + assert cons.pop(ident) == [second] + assert cons.pop(ident) == [] + first.close() + second.close() + + +def test_fork_gate_timeout(): + gate = dc.core._ForkGate() + gate.timeout = 0.05 + entered = threading.Event() + release = threading.Event() + + def hold(): + with gate: + entered.set() + release.wait(10) + + thread = threading.Thread(target=hold) + thread.start() + assert entered.wait(10) + + with pytest.warns(RuntimeWarning, match='forking while'): + gate.before_fork() + + gate.after_fork_parent() + release.set() + thread.join(10) + + +def test_fork_gate_nesting_and_child_reset(): + gate = dc.core._ForkGate() + + with gate: + with gate: + assert gate._active == 1 + gate.before_fork() # The forking thread's own transaction. + gate.after_fork_parent() + assert gate._active == 1 + + assert gate._active == 0 + + with gate: + gate.after_fork_child() + + assert gate._active == 0 + + +def test_fork_gate_blocks_new_transactions_while_forking(): + gate = dc.core._ForkGate() + gate.before_fork() + entered = threading.Event() + + def enter(): + with gate: + entered.set() + + thread = threading.Thread(target=enter) + thread.start() + assert not entered.wait(0.2) + gate.after_fork_parent() + assert entered.wait(10) + thread.join(10) + + +@pytest.mark.skipif(not hasattr(os, 'fork'), reason='requires os.fork') +def test_fork_child_opens_its_own_connections(cache): + # Two caches on one directory, then fork: upstream issue 325. + other = dc.Cache(cache.directory) + cache.set('parent', 1) + other.set('other', 2) + inherited = len(dc.core._INHERITED_CONNECTIONS) + read_fd, write_fd = os.pipe() + pid = os.fork() + + if pid == 0: # pragma: no cover + try: + checks = ( + cache.get('parent') == 1, + other.get('other') == 2, + # Both connections were closed before the fork. + len(dc.core._INHERITED_CONNECTIONS) == inherited, + ) + other.set('child', 3) + os.write(write_fd, bytes(checks)) + except BaseException: + os.write(write_fd, b'\xff') + finally: + os._exit(0) + + os.close(write_fd) + assert wait_for_child(pid) == 0 + assert os.read(read_fd, 3) == b'\x01\x01\x01' + os.close(read_fd) + assert other.get('child') == 3 + other.close() + + +@pytest.mark.skipif(not hasattr(os, 'fork'), reason='requires os.fork') +def test_fork_waits_for_transactions_in_other_threads(cache): + entered = threading.Event() + release = threading.Event() + + def hold(): + with cache.transact(): + cache.set('held', 1) + entered.set() + release.wait(10) + + holder = threading.Thread(target=hold) + holder.start() + assert entered.wait(10) + forked = threading.Event() + pids = [] + read_fd, write_fd = os.pipe() + + def forker(): + pid = os.fork() + if pid == 0: # pragma: no cover + try: + cache.set('child', 2) + os.write(write_fd, b'\x01') + except BaseException: + os.write(write_fd, b'\xff') + finally: + os._exit(0) + pids.append(pid) + forked.set() + + thread = threading.Thread(target=forker) + thread.start() + assert not forked.wait(0.5) # The fork waits for the transaction. + release.set() + holder.join(10) + assert forked.wait(10) + thread.join(10) + os.close(write_fd) + assert wait_for_child(pids[0]) == 0 + assert os.read(read_fd, 1) == b'\x01' + os.close(read_fd) + assert cache.get('held') == 1 + assert cache.get('child') == 2