Repository: rasterio/rasterio

Issue:
Access config options using thread-safe methods
A follow-up to #993. Also guard against registering drivers in any but the main thread. Main thread can be identified like

`isinstance(threading.current_thread(), threading._MainThread)` or

`assert threading.current_thread() == threading.main_thread()` for Python >= 3.4.


Discussion:

Resolving patch:
diff --git a/rasterio/env.py b/rasterio/env.py
--- a/rasterio/env.py
+++ b/rasterio/env.py
@@ -1,46 +1,49 @@
 """Rasterio's GDAL/AWS environment"""
 
 from functools import wraps
-import itertools as it
 import logging
+import threading
 
 from rasterio._env import (
     GDALEnv, del_gdal_config, get_gdal_config, set_gdal_config)
+from rasterio.compat import string_types
 from rasterio.dtypes import check_dtype
 from rasterio.errors import EnvError
-from rasterio.compat import string_types
 from rasterio.transform import guard_transform
 from rasterio.vfs import parse_path, vsi_path
 
 
-# The currently active GDAL/AWS environment is a private attribute.
-_env = None
-
-
-# When the outermost 'rasterio.Env()' executes '__enter__' it probes the
-# GDAL environment to see if any of the supplied config options already
-# exist, the assumption being that they were set with
-# 'osgeo.gdal.SetConfigOption()' or possibly 'rasterio.env.set_gdal_config()'.
-# The discovered options are reinstated when the outermost Rasterio environment
-# exits.  Without this check any environment options that are present in the
-# GDAL environment and are also passed to 'rasterio.Env()' will be unset
-# when 'rasterio.Env()' tears down, regardless of their value.  For example:
-#
-#   from osgeo import gdal
-#   import rasterio
-#
-#   gdal.SetConfigOption('key', 'value')
-#   with rasterio.Env(key='something'):
-#       pass
-#
-# The config option 'key' would be unset when 'Env()' exits.  A more
-# comprehensive solution would also leverage https://trac.osgeo.org/gdal/changeset/37273
-# but this gets Rasterio + older versions of GDAL halfway there.  One major
-# assumption is that environment variables are not set directly with
-# 'osgeo.gdal.SetConfigOption()' OR 'rasterio.env.set_gdal_config()' inside
-# of a 'rasterio.Env()'.
-_discovered_options = None
-
+class ThreadEnv(threading.local):
+    def __init__(self):
+        self._env = None  # Initialises in each thread
+
+        # When the outermost 'rasterio.Env()' executes '__enter__' it
+        # probes the GDAL environment to see if any of the supplied
+        # config options already exist, the assumption being that they
+        # were set with 'osgeo.gdal.SetConfigOption()' or possibly
+        # 'rasterio.env.set_gdal_config()'.  The discovered options are
+        # reinstated when the outermost Rasterio environment exits.
+        # Without this check any environment options that are present in
+        # the GDAL environment and are also passed to 'rasterio.Env()'
+        # will be unset when 'rasterio.Env()' tears down, regardless of
+        # their value.  For example:
+        #
+        #   from osgeo import gdal import rasterio
+        #
+        #   gdal.SetConfigOption('key', 'value') with
+        #   rasterio.Env(key='something'): pass
+        #
+        # The config option 'key' would be unset when 'Env()' exits.
+        # A more comprehensive solution would also leverage
+        # https://trac.osgeo.org/gdal/changeset/37273 but this gets
+        # Rasterio + older versions of GDAL halfway there.  One major
+        # assumption is that environment variables are not set directly
+        # with 'osgeo.gdal.SetConfigOption()' OR
+        # 'rasterio.env.set_gdal_config()' inside of a 'rasterio.Env()'.
+        self._discovered_options = None
+
+
+local = ThreadEnv()
 
 log = logging.getLogger(__name__)
 
@@ -48,7 +51,7 @@
 default_options = {
     'CHECK_WITH_INVERT_PROJ': True,
     'GTIFF_IMPLICIT_JPEG_OVR': False,
-    "I'M_ON_RASTERIO": True
+    "RASTERIO_ENV": True
 }
 
 class Env(object):
@@ -159,34 +162,30 @@ def get_aws_credentials(self):
             options.update(aws_region=self.aws_session.region_name)
 
         # Pass these credentials to the GDAL environment.
-        global _env
-        _env.update_config_options(**options)
+        local._env.update_config_options(**options)
 
     def drivers(self):
         """Return a mapping of registered drivers."""
-        global _env
-        return _env.drivers()
+        return local._env.drivers()
 
     def __enter__(self):
-        global _env
-        global _discovered_options
         log.debug("Entering env context: %r", self)
 
         # No parent Rasterio environment exists.
-        if _env is None:
+        if local._env is None:
             logging.debug("Starting outermost env")
             self._has_parent_env = False
 
             # See note directly above where _discovered_options is globally
             # defined.  This MUST happen before calling 'defenv()'.
-            _discovered_options = {}
-            # Don't want to reinstate the "I'M_ON_RASTERIO" option.
-            probe_env = {k for k in default_options if k != "I'M_ON_RASTERIO"}
+            local._discovered_options = {}
+            # Don't want to reinstate the "RASTERIO_ENV" option.
+            probe_env = {k for k in default_options if k != "RASTERIO_ENV"}
             probe_env |= set(self.options.keys())
             for key in probe_env:
                 val = get_gdal_config(key, normalize=False)
                 if val is not None:
-                    _discovered_options[key] = val
+                    local._discovered_options[key] = val
                     logging.debug("Discovered option: %s=%s", key, val)
 
             defenv()
@@ -199,8 +198,6 @@ def __enter__(self):
         return self
 
     def __exit__(self, exc_type=None, exc_val=None, exc_tb=None):
-        global _env
-        global _discovered_options
         log.debug("Exiting env context: %r", self)
         delenv()
         if self._has_parent_env:
@@ -210,66 +207,65 @@ def __exit__(self, exc_type=None, exc_val=None, exc_tb=None):
             logging.debug("Exiting outermost env")
             # See note directly above where _discovered_options is globally
             # defined.
-            while _discovered_options:
-                key, val = _discovered_options.popitem()
+            while local._discovered_options:
+                key, val = local._discovered_options.popitem()
                 set_gdal_config(key, val, normalize=False)
                 logging.debug(
                     "Set discovered option back to: '%s=%s", key, val)
-            _discovered_options = None
+            local._discovered_options = None
         log.debug("Exited env context: %r", self)
 
 
 def defenv():
     """Create a default environment if necessary."""
-    global _env
-    if _env:
-        log.debug("GDAL environment exists: %r", _env)
+    if local._env:
+        log.debug("GDAL environment exists: %r", local._env)
     else:
         log.debug("No GDAL environment exists")
-        _env = GDALEnv()
-        _env.update_config_options(**default_options)
+        local._env = GDALEnv()
+        local._env.update_config_options(**default_options)
         log.debug(
-            "New GDAL environment %r created", _env)
-    _env.start()
+            "New GDAL environment %r created", local._env)
+    local._env.start()
 
 
 def getenv():
     """Get a mapping of current options."""
-    global _env
-    if not _env:
+    if not local._env:
         raise EnvError("No GDAL environment exists")
     else:
-        log.debug("Got a copy of environment %r options", _env)
-        return _env.options.copy()
+        log.debug("Got a copy of environment %r options", local._env)
+        return local._env.options.copy()
 
 
 def setenv(**options):
     """Set options in the existing environment."""
-    global _env
-    if not _env:
+    if not local._env:
         raise EnvError("No GDAL environment exists")
     else:
-        _env.update_config_options(**options)
-        log.debug("Updated existing %r with options %r", _env, options)
+        local._env.update_config_options(**options)
+        log.debug("Updated existing %r with options %r", local._env, options)
 
 
 def delenv():
     """Delete options in the existing environment."""
-    global _env
-    if not _env:
+    if not local._env:
         raise EnvError("No GDAL environment exists")
     else:
-        _env.clear_config_options()
-        log.debug("Cleared existing %r options", _env)
-    _env.stop()
-    _env = None
+        local._env.clear_config_options()
+        log.debug("Cleared existing %r options", local._env)
+    local._env.stop()
+    local._env = None
 
 
 def ensure_env(f):
     """A decorator that ensures an env exists before a function
     calls any GDAL C functions."""
-    @wraps(f)
-    def wrapper(*args, **kwds):
-        with Env(WITH_RASTERIO_ENV=True):
-            return f(*args, **kwds)
-    return wrapper
+    if local._env:
+        return f
+    else:
+        @wraps(f)
+        def wrapper(*args, **kwds):
+            with Env():
+                return f(*args, **kwds)
+        return wrapper
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -310,6 +310,10 @@ def run(self):
         'pytest>=2.8.2', 'pytest-cov>=2.2.0', 'boto3>=1.2.4', 'packaging'],
     'docs': ['ghp-import', 'numpydoc', 'sphinx', 'sphinx-rtd-theme']}
 
+# Add futures to 'test' for Python < 3.2.
+if sys.version_info < (3, 2):
+    extra_reqs['test'].append('futures')
+
 # Add all extra requirements
 extra_reqs['all'] = list(set(itertools.chain(*extra_reqs.values())))
 


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -247,9 +247,9 @@ def gdalenv(request):
     import rasterio.env
 
     def fin():
-        if rasterio.env._env:
+        if rasterio.env.local._env:
             rasterio.env.delenv()
-            rasterio.env._env = None
+            rasterio.env.local._env = None
     request.addfinalizer(fin)
 
 
diff --git a/tests/test_concurrency.py b/tests/test_concurrency.py
new file mode 100644
--- /dev/null
+++ b/tests/test_concurrency.py
@@ -0,0 +1,39 @@
+"""Tests for correct behavior of Rasterio's GDALEnv in concurrent programs"""
+
+from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
+
+import rasterio as rio
+
+
+def get_data(path):
+    """Return all raster bands as an ndarray"""
+    with rio.open(path) as src:
+        return src.read()
+
+
+def test_threads_main_env():
+    """Get raster data using ThreadPoolExecutor with main thread Env"""
+    with rio.Env(), ThreadPoolExecutor(4) as pool:
+        for res in pool.map(get_data, ['tests/data/RGB.byte.tif'] * 10):
+            assert res.any()
+
+
+def test_threads_no_main_env():
+    """Get raster data using ThreadPoolExecutor with no main thread Env"""
+    with ThreadPoolExecutor(4) as pool:
+        for res in pool.map(get_data, ['tests/data/RGB.byte.tif'] * 10):
+            assert res.any()
+
+
+def testget_data_main_env():
+    """Get raster data using ProcessPoolExecutor with main thread Env"""
+    with rio.Env(), ProcessPoolExecutor() as pool:
+        for res in pool.map(get_data, ['tests/data/RGB.byte.tif'] * 10):
+            assert res.any()
+
+
+def testget_data_no_main_env():
+    """Get raster data using ProcessPoolExecutor with main thread Env"""
+    with ProcessPoolExecutor() as pool:
+        for res in pool.map(get_data, ['tests/data/RGB.byte.tif'] * 10):
+            assert res.any()
diff --git a/tests/test_env.py b/tests/test_env.py
--- a/tests/test_env.py
+++ b/tests/test_env.py
@@ -85,7 +85,7 @@ def test_env_accessors(gdalenv):
     setenv(foo='1', bar='2')
     expected = default_options.copy()
     expected.update({'foo': '1', 'bar': '2'})
-    assert getenv() == rasterio.env._env.options
+    assert getenv() == rasterio.env.local._env.options
     assert getenv() == expected
     assert get_gdal_config('foo') == '1'
     assert get_gdal_config('bar') == '2'
@@ -109,7 +109,7 @@ def test_env_accessors_no_env():
 def test_ensure_env_decorator(gdalenv):
     @ensure_env
     def f():
-        return getenv()['WITH_RASTERIO_ENV']
+        return getenv()['RASTERIO_ENV']
     assert f() is True
 
 
@@ -129,7 +129,7 @@ def test_env_defaults(gdalenv):
     with env:
         assert get_gdal_config('CHECK_WITH_INVERT_PROJ') is True
         assert get_gdal_config('GTIFF_IMPLICIT_JPEG_OVR') is False
-        assert get_gdal_config("I'M_ON_RASTERIO") is True
+        assert get_gdal_config("RASTERIO_ENV") is True
 
 
 def test_aws_session(gdalenv):
@@ -163,12 +163,12 @@ def test_with_aws_session_credentials(gdalenv):
             aws_access_key_id='id', aws_secret_access_key='key',
             aws_session_token='token', region_name='null-island-1') as s:
         expected = default_options.copy()
-        assert getenv() == rasterio.env._env.options == expected
+        assert getenv() == rasterio.env.local._env.options == expected
         s.get_aws_credentials()
         expected.update({
             'aws_access_key_id': 'id', 'aws_region': 'null-island-1',
             'aws_secret_access_key': 'key', 'aws_session_token': 'token'})
-        assert getenv() == rasterio.env._env.options == expected
+        assert getenv() == rasterio.env.local._env.options == expected
 
 
 def test_session_env_lazy(monkeypatch, gdalenv):
@@ -178,7 +178,7 @@ def test_session_env_lazy(monkeypatch, gdalenv):
     monkeypatch.setenv('AWS_SESSION_TOKEN', 'token')
     with rasterio.Env() as s:
         s.get_aws_credentials()
-        assert getenv() == rasterio.env._env.options
+        assert getenv() == rasterio.env.local._env.options
         expected = {
             'aws_access_key_id': 'id',
             'aws_secret_access_key': 'key',
@@ -283,7 +283,7 @@ def _check_defaults():
         pass
 
     _check_defaults()
-    assert rasterio.env._env is None
+    assert rasterio.env.local._env is None
 
 
 @pytest.mark.parametrize("key,val", [
@@ -299,7 +299,7 @@ def test_env_discovery(key, val):
     environment default.
     """
 
-    assert rasterio.env._discovered_options is None, \
+    assert rasterio.env.local._discovered_options is None, \
         "Something has gone horribly wrong."
 
     try:
@@ -309,20 +309,20 @@ def test_env_discovery(key, val):
         # Start an environment and overwrite the value that should persist
         with rasterio.Env(**{key: True}):
             assert get_gdal_config(key) is True
-            assert rasterio.env._discovered_options == {key: val}
+            assert rasterio.env.local._discovered_options == {key: val}
 
             # Start another nested environment, again overwriting the value
             # that should persist
             with rasterio.Env(**{key: False}):
-                assert rasterio.env._discovered_options == {key: val}
+                assert rasterio.env.local._discovered_options == {key: val}
                 assert get_gdal_config(key) is False
 
             # Ensure the outer state is restored.
-            assert rasterio.env._discovered_options == {key: val}
+            assert rasterio.env.local._discovered_options == {key: val}
             assert get_gdal_config(key) is True
 
         # Ensure the discovered value remains unchanged.
-        assert rasterio.env._discovered_options is None
+        assert rasterio.env.local._discovered_options is None
         assert get_gdal_config(key, normalize=False) == val
 
     # Leaving this option in the GDAL environment could cause a problem
diff --git a/tests/test_threading.py b/tests/test_threading.py
new file mode 100644
--- /dev/null
+++ b/tests/test_threading.py
@@ -0,0 +1,82 @@
+from threading import Thread
+import time
+import unittest
+
+import rasterio as rio
+from rasterio.env import get_gdal_config
+
+
+class TestThreading(unittest.TestCase):
+    def test_multiopen(self):
+        """
+        Open a file from different threads.
+
+        Regression test for issue #986
+        """
+
+        def func(delay):
+            try:
+                with rio.open('tests/data/RGB.byte.tif'):
+                    time.sleep(delay)
+            except Exception as err:
+                global exceptions
+                exceptions.append(err)
+
+        global exceptions
+        exceptions = []
+
+        t1 = Thread(target=func, args=(0.1,))
+        t2 = Thread(target=func, args=(0,))
+
+        with rio.Env():
+            t1.start()
+            t2.start()  # potential error if Env manages globals unsafely
+
+            t1.join()
+            t2.join()
+
+        assert not exceptions
+
+    def test_reliability(self):
+        """Allow for nondeterminism of race condition"""
+        for i in range(3):
+            self.test_multiopen()
+
+
+def test_child_thread_inherits_env():
+    """A new thread inherit's the main thread's env"""
+    def func():
+        with rio.Env(lol='wut'):
+            assert get_gdal_config('lol') == 'wut'
+            # The next config option will have been set in the main thread.
+            assert get_gdal_config('FROM_MAIN') is True
+
+    t1 = Thread(target=func)
+
+    with rio.Env(FROM_MAIN=True):
+        t1.start()
+        assert get_gdal_config('FROM_MAIN') is True
+        assert get_gdal_config('lol') is None
+        t1.join()
+
+
+def test_child_thread_isolation():
+    """Child threads have isolated environments"""
+    def func(key, value, other_key):
+        env = {key: value}
+        with rio.Env(**env):
+            assert get_gdal_config(key) == value
+            # The other key is one set in another child thread.
+            assert get_gdal_config(other_key) is None
+
+    t1 = Thread(target=func, args=('is_t1', True, 'is_t2'))
+    t2 = Thread(target=func, args=('is_t2', True, 'is_t1'))
+
+    t1.start()
+    t2.start()
+    t1.join()
+    t2.join()
+
+
+if __name__ == '__main__':
+    unittest.main()