Repository: corteva/rioxarray

Issue:
overview not read when specifying chunks or lock
#### Code Sample, a copy-pastable example if possible

```python
import rasterio
import rioxarray

# from https://openaerialmap.org/
cog_url = (
    "https://oin-hotosm.s3.amazonaws.com/"
    "5d7dad0becaf880008a9bc88/0/5d7dad0becaf880008a9bc89.tif"
)
```


```python
ds = rioxarray.open_rasterio(cog_url, lock=False, chunks=(4, "auto", -1))
ds.shape
```




    (3, 9984, 22016)




```python
dsLockChunksOverview = rioxarray.open_rasterio(cog_url, lock=False, chunks=(4, "auto", -1), overview_level=1)
dsLockChunksOverview.shape
```




    (3, 9984, 22016)




```python
dsChunksOverview = rioxarray.open_rasterio(cog_url, chunks=(4, "auto", -1), overview_level=1)
dsChunksOverview.shape
```




    (3, 9984, 22016)




```python
dsLockOverview = rioxarray.open_rasterio(cog_url, lock=False, overview_level=1)
dsLockOverview.shape
```




    (3, 9984, 22016)




```python
dsOverview = rioxarray.open_rasterio(cog_url, overview_level=1)
dsOverview.shape
```




    (3, 2496, 5504)




#### Problem description

Large COGs consequently have large overviews. It may be beneficial to read the overview in parallel without locks, as documented in [Reading COGs in Parallel](https://corteva.github.io/rioxarray/stable/examples/read-locks.html). Perhaps this behavior is documented, but I could not find it.

#### Expected Output

The correct shape of the overview `(3, 2496, 5504)` when specifying `overview_level` AND `chunks` or `lock`.

#### Environment Information

rioxarray (0.4.0) deps:
  rasterio: 1.2.3
    xarray: 0.18.2
      GDAL: 3.2.2

Other python deps:
     scipy: 1.6.3
    pyproj: 3.0.1

System:
    python: 3.7.3 (default, Jan 22 2021, 20:04:44)  [GCC 8.3.0]
executable: /usr/bin/python3
   machine: Linux-5.4.104+-x86_64-with-debian-10.9

#### Installation method
pypi

Comments:

Patch:
diff --git a/.flake8 b/.flake8
index 06ffc9b..fdcf6e3 100644
--- a/.flake8
+++ b/.flake8
@@ -1,6 +1,7 @@
 [flake8]
 max_line_length = 88
 ignore =
+    E501 # line too long - let black handle that
     C408 # Unnecessary dict/list/tuple call - rewrite as a literal
     E203 # whitespace before ':' - doesn't work well with black
     E225 # missing whitespace around operator - let black worry about that
diff --git a/docs/history.rst b/docs/history.rst
index 217dc86..fcd17ec 100644
--- a/docs/history.rst
+++ b/docs/history.rst
@@ -4,6 +4,7 @@ History
 Latest
 ------
 - BUG: pass kwargs with lock=False (issue #344)
+- BUG: Close file handle with lock=False (pull #346)
 
 0.4.0
 ------
diff --git a/rioxarray/_io.py b/rioxarray/_io.py
index aeda3b7..c243f38 100644
--- a/rioxarray/_io.py
+++ b/rioxarray/_io.py
@@ -9,6 +9,7 @@ Source file: https://github.com/pydata/xarray/blob/1d7bcbdc75b6d556c04e2c7d7a042
 import contextlib
 import os
 import re
+import threading
 import warnings
 from distutils.version import LooseVersion
 
@@ -35,6 +36,63 @@ RASTERIO_LOCK = SerializableLock()
 NO_LOCK = contextlib.nullcontext()
 
 
+class FileHandleLocal(threading.local):
+    """
+    This contains the thread local ThreadURIManager
+    """
+
+    def __init__(self):  # pylint: disable=super-init-not-called
+        self.thread_manager = None  # Initialises in each thread
+
+
+class ThreadURIManager:
+    """
+    This handles opening & closing file handles in each thread.
+    """
+
+    def __init__(
+        self,
+        opener,
+        *args,
+        mode="r",
+        kwargs=None,
+    ):
+        self._opener = opener
+        self._args = args
+        self._mode = mode
+        self._kwargs = {} if kwargs is None else dict(kwargs)
+        self._file_handle = None
+
+    @property
+    def file_handle(self):
+        """
+        File handle returned by the opener.
+        """
+        if self._file_handle is not None:
+            return self._file_handle
+        print("OPEN")
+        self._file_handle = self._opener(*self._args, mode=self._mode, **self._kwargs)
+        return self._file_handle
+
+    def close(self):
+        """
+        Close file handle.
+        """
+        if self._file_handle is not None:
+            print("CLOSE")
+            self._file_handle.close()
+            self._file_handle = None
+
+    def __del__(self):
+        self.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type_, value, traceback):
+        self.close()
+
+
 class URIManager(FileManager):
     """
     The URI manager is used for lockless reading
@@ -51,15 +109,39 @@ class URIManager(FileManager):
         self._args = args
         self._mode = mode
         self._kwargs = {} if kwargs is None else dict(kwargs)
+        self._local = FileHandleLocal()
 
     def acquire(self, needs_lock=True):
-        return self._opener(*self._args, mode=self._mode, **self._kwargs)
+        if self._local.thread_manager is None:
+            self._local.thread_manager = ThreadURIManager(
+                self._opener, *self._args, mode=self._mode, kwargs=self._kwargs
+            )
+        return self._local.thread_manager.file_handle
 
+    @contextlib.contextmanager
     def acquire_context(self, needs_lock=True):
-        yield self.acquire(needs_lock=needs_lock)
+        try:
+            yield self.acquire(needs_lock=needs_lock)
+        except Exception:
+            self.close(needs_lock=needs_lock)
+            raise
 
     def close(self, needs_lock=True):
-        pass
+        if self._local.thread_manager is not None:
+            self._local.thread_manager.close()
+            self._local.thread_manager = None
+
+    def __del__(self):
+        self.close(needs_lock=False)
+
+    def __getstate__(self):
+        """State for pickling."""
+        return (self._opener, self._args, self._mode, self._kwargs)
+
+    def __setstate__(self, state):
+        """Restore from a pickle."""
+        opener, args, mode, kwargs = state
+        self.__init__(opener, *args, mode=mode, kwargs=kwargs)
 
 
 class RasterioArrayWrapper(BackendArray):


Test patch:
diff --git a/test/integration/test_integration__io.py b/test/integration/test_integration__io.py
index d471375..339767a 100644
--- a/test/integration/test_integration__io.py
+++ b/test/integration/test_integration__io.py
@@ -1009,14 +1009,14 @@ def test_nc_attr_loading__disable_decode_times(open_rasterio):
 
 def test_lockless():
     with rioxarray.open_rasterio(
-        os.path.join(TEST_INPUT_DATA_DIR, "PLANET_SCOPE_3D.nc"), lock=False, chunks=True
+        os.path.join(TEST_INPUT_DATA_DIR, "cog.tif"), lock=False, chunks=True
     ) as rds:
         rds.mean().compute()
 
 
 def test_lock_true():
     with rioxarray.open_rasterio(
-        os.path.join(TEST_INPUT_DATA_DIR, "PLANET_SCOPE_3D.nc"), lock=True, chunks=True
+        os.path.join(TEST_INPUT_DATA_DIR, "cog.tif"), lock=True, chunks=True
     ) as rds:
         rds.mean().compute()
 
diff --git a/test/integration/test_integration__io_uri_manager.py b/test/integration/test_integration__io_uri_manager.py
new file mode 100644
index 0000000..e3251ae
--- /dev/null
+++ b/test/integration/test_integration__io_uri_manager.py
@@ -0,0 +1,139 @@
+"""
+Tests based on: https://github.com/pydata/xarray/blob/071da2a900702d65c47d265192bc7e424bb57932/xarray/tests/test_backends_file_manager.py
+"""
+import concurrent.futures
+import gc
+import pickle
+from unittest import mock
+
+import pytest
+
+from rioxarray._io import URIManager
+
+
+def test_uri_manager_mock_write():
+    mock_file = mock.Mock()
+    opener = mock.Mock(spec=open, return_value=mock_file)
+
+    manager = URIManager(opener, "filename")
+    f = manager.acquire()
+    f.write("contents")
+    manager.close()
+
+    opener.assert_called_once_with("filename", mode="r")
+    mock_file.write.assert_called_once_with("contents")
+    mock_file.close.assert_called_once_with()
+
+
+def test_uri_manager_mock_write__threaded():
+    mock_file = mock.Mock()
+    opener = mock.Mock(spec=open, return_value=mock_file)
+
+    manager = URIManager(opener, "filename")
+
+    def write(iter):
+        nonlocal manager
+        fh = manager.acquire()
+        fh.write("contents")
+        manager._local.thread_manager = None
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
+        for result in executor.map(write, range(5)):
+            pass
+
+    gc.collect()
+
+    opener.assert_has_calls([mock.call("filename", mode="r") for _ in range(5)])
+    mock_file.write.assert_has_calls([mock.call("contents") for _ in range(5)])
+    mock_file.close.assert_has_calls([mock.call() for _ in range(5)])
+
+
+@pytest.mark.parametrize("expected_warning", [None, RuntimeWarning])
+def test_uri_manager_autoclose(expected_warning):
+    mock_file = mock.Mock()
+    opener = mock.Mock(return_value=mock_file)
+
+    manager = URIManager(opener, "filename")
+    manager.acquire()
+
+    del manager
+    gc.collect()
+
+    mock_file.close.assert_called_once_with()
+
+
+def test_uri_manager_write_concurrent(tmpdir):
+    path = str(tmpdir.join("testing.txt"))
+    manager = URIManager(open, path, mode="w")
+    f1 = manager.acquire()
+    f2 = manager.acquire()
+    f3 = manager.acquire()
+    assert f1 is f2
+    assert f2 is f3
+    f1.write("foo")
+    f1.flush()
+    f2.write("bar")
+    f2.flush()
+    f3.write("baz")
+    f3.flush()
+
+    del manager
+    gc.collect()
+
+    with open(path) as f:
+        assert f.read() == "foobarbaz"
+
+
+def test_uri_manager_write_pickle(tmpdir):
+    path = str(tmpdir.join("testing.txt"))
+    manager = URIManager(open, path, mode="a")
+    f = manager.acquire()
+    f.write("foo")
+    f.flush()
+    manager2 = pickle.loads(pickle.dumps(manager))
+    f2 = manager2.acquire()
+    f2.write("bar")
+    del manager
+    del manager2
+    gc.collect()
+
+    with open(path) as f:
+        assert f.read() == "foobar"
+
+
+def test_uri_manager_read(tmpdir):
+    path = str(tmpdir.join("testing.txt"))
+
+    with open(path, "w") as f:
+        f.write("foobar")
+
+    manager = URIManager(open, path)
+    f = manager.acquire()
+    assert f.read() == "foobar"
+    manager.close()
+
+
+def test_uri_manager_acquire_context(tmpdir):
+    path = str(tmpdir.join("testing.txt"))
+
+    with open(path, "w") as f:
+        f.write("foobar")
+
+    class AcquisitionError(Exception):
+        pass
+
+    manager = URIManager(open, path)
+    with pytest.raises(AcquisitionError):
+        with manager.acquire_context() as f:
+            assert f.read() == "foobar"
+            raise AcquisitionError
+
+    with manager.acquire_context() as f:
+        assert f.read() == "foobar"
+
+    with pytest.raises(AcquisitionError):
+        with manager.acquire_context() as f:
+            f.seek(0)
+            assert f.read() == "foobar"
+            raise AcquisitionError
+    manager.close()