BUG: Close cached raster files before interpreter shutdown

docs/history.rst

@@ -3,6 +3,7 @@ History
 
 Latest
 ------
+- BUG: Close cached raster files before interpreter shutdown (issue #929)
 
 0.23.0
 ------

rioxarray/_io.py

@@ -6,12 +6,14 @@
 Source file: https://github.com/pydata/xarray/blob/1d7bcbdc75b6d556c04e2c7d7a042e4379e15303/xarray/backends/rasterio_.py # noqa
 """
 # pylint: disable=too-many-lines
+import atexit
 import contextlib
 import importlib.metadata
 import os
 import re
 import threading
 import warnings
+import weakref
 from collections import defaultdict
 from collections.abc import Hashable, Iterable
 from typing import Any, Optional, Union
@@ -25,7 +27,7 @@
 from xarray import Dataset, IndexVariable
 from xarray.backends.common import BackendArray
 from xarray.backends.file_manager import CachingFileManager, FileManager
-from xarray.backends.locks import SerializableLock
+from xarray.backends.locks import SerializableLock, acquire
 from xarray.coding import times, variables
 from xarray.core import indexing
 from xarray.core.dataarray import DataArray
@@ -43,6 +45,38 @@
 # TODO: should this be GDAL_LOCK instead?
 RASTERIO_LOCK = SerializableLock()
 NO_LOCK = contextlib.nullcontext()
+# xarray disables CachingFileManager.__del__ during interpreter shutdown.
+# Track the open files instead of their managers because concurrent manager
+# finalizers can leave files in xarray's cache after the managers are gone.
+_LIVE_RASTERIO_FILES: weakref.WeakKeyDictionary[Any, Any] = weakref.WeakKeyDictionary()
+_LIVE_RASTERIO_FILES_LOCK = threading.Lock()
+
+
+@atexit.register
+def _close_live_rasterio_files():
+    if not _LIVE_RASTERIO_FILES_LOCK.acquire(  # pylint: disable=consider-using-with
+        blocking=False
+    ):
+        return
+    try:
+        live_files = list(_LIVE_RASTERIO_FILES.items())
+    finally:
+        _LIVE_RASTERIO_FILES_LOCK.release()
+
+    for riods, lock in live_files:
+        if acquire(lock, blocking=False):
+            try:
+                riods.close()
+            finally:
+                lock.release()
+
+
+def _acquire_rasterio(manager, lock, needs_lock=True):
+    riods = manager.acquire(needs_lock=needs_lock)
+    if isinstance(manager, CachingFileManager):
+        with _LIVE_RASTERIO_FILES_LOCK:
+            _LIVE_RASTERIO_FILES[riods] = lock
+    return riods
 
 
 def _ensure_warped_vrt(riods, vrt_params):
@@ -312,7 +346,7 @@ def __init__(
         self.mask_and_scale = mask_and_scale
 
         # cannot save riods as an attribute: this would break pickleability
-        riods = _ensure_warped_vrt(manager.acquire(), vrt_params)
+        riods = _ensure_warped_vrt(_acquire_rasterio(manager, lock), vrt_params)
         self.vrt_params = vrt_params
         self._shape = (riods.count, riods.height, riods.width)
         self._dtype = None
@@ -430,7 +464,8 @@ def _getitem(self, key):
         else:
             with self.lock:
                 riods = _ensure_warped_vrt(
-                    self.manager.acquire(needs_lock=False), self.vrt_params
+                    _acquire_rasterio(self.manager, self.lock, needs_lock=False),
+                    self.vrt_params,
                 )
                 out = riods.read(band_key, window=window, masked=self.masked)
                 if self._unsigned_dtype is not None:
@@ -1137,7 +1172,7 @@ def open_rasterio(
             )
         else:
             manager = URIManager(file_opener, filename, mode="r", kwargs=open_kwargs)
-        riods = manager.acquire()
+        riods = _acquire_rasterio(manager, lock)
         captured_warnings = rio_warnings.copy()
 
     # raise the NotGeoreferencedWarning if applicable

test/integration/test_integration__io.py

@@ -5,6 +5,7 @@
 import os
 import pickle
 import shutil
+import subprocess
 import sys
 import tempfile
 import warnings
@@ -54,6 +55,93 @@ def _assert_tmmx_source(source):
         )
 
 
+def _assert_clean_rasterio_shutdown(script):
+    raster_path = os.path.abspath(os.path.join(TEST_INPUT_DATA_DIR, "2d_test.tif"))
+    result = subprocess.run(
+        [sys.executable, "-c", script, raster_path],
+        capture_output=True,
+        check=True,
+        text=True,
+        timeout=15,
+    )
+
+    assert result.stderr == ""
+
+
+def test_open_rasterio_closes_files_at_interpreter_shutdown():
+    _assert_clean_rasterio_shutdown(
+        """
+import rioxarray
+import sys
+
+for _ in range(5):
+    rioxarray.open_rasterio(sys.argv[1])
+"""
+    )
+
+
+def test_failed_open_rasterio_closes_files_at_interpreter_shutdown():
+    _assert_clean_rasterio_shutdown(
+        """
+import rioxarray
+import rioxarray._io
+import sys
+
+def raise_dtype_error(_):
+    raise RuntimeError("metadata failure")
+
+rioxarray._io._rasterio_to_numpy_dtype = raise_dtype_error
+try:
+    rioxarray.open_rasterio(sys.argv[1])
+except RuntimeError:
+    saved_exception = sys.exc_info()
+
+assert saved_exception[0] is RuntimeError
+"""
+    )
+
+
+def test_open_rasterio_shutdown_does_not_wait_for_locked_files():
+    _assert_clean_rasterio_shutdown(
+        """
+import rioxarray
+import rioxarray._io
+import sys
+import threading
+
+raster = rioxarray.open_rasterio(sys.argv[1])
+lock_acquired = threading.Event()
+
+def hold_file_lock():
+    with rioxarray._io.RASTERIO_LOCK:
+        lock_acquired.set()
+        threading.Event().wait()
+
+threading.Thread(target=hold_file_lock, daemon=True).start()
+assert lock_acquired.wait(timeout=5)
+"""
+    )
+
+
+def test_open_rasterio_closes_orphaned_cached_files_at_shutdown():
+    _assert_clean_rasterio_shutdown(
+        """
+from concurrent.futures import ThreadPoolExecutor
+import gc
+import rioxarray
+import sys
+
+def open_and_read(_):
+    raster = rioxarray.open_rasterio(sys.argv[1])
+    return raster.isel(x=0, y=0).values.tolist()
+
+with ThreadPoolExecutor(max_workers=16) as pool:
+    list(pool.map(open_and_read, range(100)))
+gc.collect()
+"""
+    )
+
+
 @pytest.mark.parametrize(
     "subdataset, variable, group, match",
     [