Repository: snowflakedb/snowflake-connector-python

Issue:
SNOW-856569: Segmentation Fault in cache.py pickle dump
### Python version

Python 3.8.12 (default, Nov 17 2021, 08:36:07)  [Clang 7.1.0 (tags/RELEASE_710/final)]

### Operating system and processor architecture

Linux _x86_64_

### Installed packages

```python
dbt-core==1.5.2
dbt-snowflake==1.5.2
snowflake-connector-python[secure-local-storage] == 3.0.3
```


### What did you do?

Running DBT on our linux machines with more than 1 thread triggers a segmentation fault in snowflake-connector-python for versions >= 3.x.x
We're not seeing this issue locally on MacOS.


### What did you expect to see?

The offending line is [here](https://github.com/snowflakedb/snowflake-connector-python/blob/v3.0.0/src/snowflake/connector/cache.py#L511C37-L511C37) and the truncated stack trace:
```
Fatal Python error: Segmentation fault
Stack (most recent call first):
File "pypi_snowflake_connector_python/site-packages/snowflake/connector/cache.py", line 511 in _save
```


Discussion:

Patch:
diff --git a/DESCRIPTION.md b/DESCRIPTION.md
index 9eb579c1..a7614f44 100644
--- a/DESCRIPTION.md
+++ b/DESCRIPTION.md
@@ -41,6 +41,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
   - Improved error handling of connection reset error.
   - Fixed a bug about deleting the temporary files happened when running PUT command.
   - Allowed to pass `type_mapper` to `fetch_pandas_batches()` and `fetch_pandas_all()`.
+  - Fixed a bug where pickle.dump segfaults during cache serialization in multi-threaded scenarios.
 
 - v3.0.4(May 23,2023)
   - Fixed a bug in which `cursor.execute()` could modify the argument statement_params dictionary object when executing a multistatement query.
diff --git a/src/snowflake/connector/cache.py b/src/snowflake/connector/cache.py
index 25c698be..99287edd 100644
--- a/src/snowflake/connector/cache.py
+++ b/src/snowflake/connector/cache.py
@@ -20,6 +20,7 @@ from filelock import FileLock, Timeout
 from typing_extensions import NamedTuple, Self
 
 from . import constants
+from .constants import ENV_VAR_TEST_MODE
 
 now = datetime.datetime.now
 getmtime = os.path.getmtime
@@ -28,6 +29,8 @@ T = TypeVar("T")
 
 logger = logging.getLogger(__name__)
 
+test_mode = os.getenv(ENV_VAR_TEST_MODE, "").lower() == "true"
+
 
 class CacheEntry(NamedTuple, Generic[T]):
     expiry: datetime.datetime
@@ -75,8 +78,7 @@ class SFDictCache(Generic[K, V]):
         the dictionary provided.
         """
         cache = cls(**kw)
-        for k, v in _dict.items():
-            cache[k] = v
+        cache.update(_dict)
         return cache
 
     def _getitem(
@@ -90,6 +92,10 @@ class SFDictCache(Generic[K, V]):
         This should only be used by internal functions when already
         holding self._lock.
         """
+        if test_mode:
+            assert (
+                self._lock.locked()
+            ), "The mutex self._lock should be locked by this thread"
         try:
             t, v = self._cache[k]
         except KeyError:
@@ -101,14 +107,8 @@ class SFDictCache(Generic[K, V]):
             self._hit(k)
         return v
 
-    def _getitem_non_locking(
-        self,
-        k: K,
-        *,
-        should_record_hits: bool = True,
-    ) -> V:
-        # aliasing _getitem to unify the api with SFDictFileCache
-        return self._getitem(k, should_record_hits=should_record_hits)
+    # aliasing _getitem to unify the api with SFDictFileCache
+    _getitem_non_locking = _getitem
 
     def _setitem(self, k: K, v: V) -> None:
         """Non-locking version of __setitem__.
@@ -116,6 +116,10 @@ class SFDictCache(Generic[K, V]):
         This should only be used by internal functions when already
         holding self._lock.
         """
+        if test_mode:
+            assert (
+                self._lock.locked()
+            ), "The mutex self._lock should be locked by this thread"
         self._cache[k] = CacheEntry(
             expiry=now() + self._entry_lifetime,
             entry=v,
@@ -182,6 +186,10 @@ class SFDictCache(Generic[K, V]):
         This should only be used by internal functions when already
         holding self._lock.
         """
+        if test_mode:
+            assert (
+                self._lock.locked()
+            ), "The mutex self._lock should be locked by this thread"
         del self._cache[key]
         self._add_or_remove()
 
@@ -209,6 +217,15 @@ class SFDictCache(Generic[K, V]):
         other: dict[K, V] | list[tuple[K, V]] | SFDictCache[K, V],
         update_newer_only: bool = False,
     ) -> bool:
+        """Non-locking version of update.
+
+        This should only be used by internal functions when already
+        holding self._lock and other._lock.
+        """
+        if test_mode:
+            assert (
+                self._lock.locked()
+            ), "The mutex self._lock should be locked by this thread"
         to_insert: dict[K, CacheEntry[V]]
         self._clear_expired_entries()
         if isinstance(other, (list, dict)):
@@ -219,7 +236,7 @@ class SFDictCache(Generic[K, V]):
                 g = iter(other.items())
             to_insert = {k: CacheEntry(expiry=expiry, entry=v) for k, v in g}
         elif isinstance(other, SFDictCache):
-            other._clear_expired_entries()
+            other.clear_expired_entries()
             others_items = list(other._cache.items())
             # Only accept values from another cache if their key is not in self,
             #  or if expiry is later the self known one
@@ -240,11 +257,13 @@ class SFDictCache(Generic[K, V]):
         self._cache.update(to_insert)
         if to_insert:
             self._add_or_remove()
+        # TODO: this should really save_if_should
         return len(to_insert) > 0
 
     def update(
         self,
         other: dict[K, V] | list[tuple[K, V]] | SFDictCache[K, V],
+        update_newer_only: bool = False,
     ) -> bool:
         """Insert multiple values at the same time, if self could learn from the other.
 
@@ -262,7 +281,7 @@ class SFDictCache(Generic[K, V]):
         into the other caches.
         """
         with self._lock:
-            return self._update(other)
+            return self._update(other, update_newer_only)
 
     def update_newer(
         self,
@@ -276,6 +295,10 @@ class SFDictCache(Generic[K, V]):
             )
 
     def _clear_expired_entries(self) -> None:
+        if test_mode:
+            assert (
+                self._lock.locked()
+            ), "The mutex self._lock should be locked by this thread"
         cache_updated = False
         for k in list(self._cache.keys()):
             try:
@@ -423,7 +446,8 @@ class SFDictFileCache(SFDictCache):
         self._file_lock = FileLock(self._file_lock_path, timeout=self.file_timeout)
         self.last_loaded: datetime.datetime | None = None
         if os.path.exists(self.file_path):
-            self._load()
+            with self._lock:
+                self._load()
         # indicate whether the cache is modified or not, this variable is for
         # SFDictFileCache to determine whether to dump cache to file when _save is called
         self._cache_modified = False
@@ -463,41 +487,8 @@ class SFDictFileCache(SFDictCache):
 
     def __getitem__(self, k: K) -> V:
         """Returns an element if it hasn't expired yet in a thread-safe way."""
-        self._lock.acquire()
-        # TODO: This variable could be replaced by a wrapper class that keeps track
-        #  of whether the lock is locked, but unless this function gets extended I
-        #  feel like it's an overkill. Make sure to change the bool right after
-        #  self._lock.acquire() and self._lock.release().
-        currently_holding = True
-        try:
-            if k not in self._cache:
-                self._lock.release()
-                currently_holding = False
-                loaded = self._load_if_should()
-                self._lock.acquire()
-                currently_holding = True
-                if (not loaded) or k not in self._cache:
-                    self._miss(k)
-                    raise KeyError
-            t, v = self._cache[k]
-            if is_expired(t):
-                self._lock.release()
-                currently_holding = False
-                loaded = self._load_if_should()
-                self._lock.acquire()
-                currently_holding = True
-                expire_item = True
-                if loaded:
-                    t, v = self._cache[k]
-                    expire_item = is_expired(t)
-                if expire_item:
-                    # Raises KeyError
-                    self._expire(k)
-            self._hit(k)
-            return v
-        finally:
-            if currently_holding:
-                self._lock.release()
+        with self._lock:
+            return self._getitem_non_locking(k)
 
     def _setitem(self, k: K, v: V) -> None:
         super()._setitem(k, v)
@@ -507,19 +498,46 @@ class SFDictFileCache(SFDictCache):
         """Load cache from disk if possible, returns whether it was able to load."""
         try:
             with open(self.file_path, "rb") as r_file:
-                other = pickle.load(r_file)
-            self._update(
-                other,
+                other: SFDictFileCache = pickle.load(r_file)
+            # Since we want to know whether we are dirty after loading
+            #  we have to know whether the file could learn anything from self
+            #  so instead of calling self.update we call other.update and swap
+            #  the 2 underlying caches after.
+            self._lock.release()
+            cache_file_learnt = other.update(
+                self,
                 update_newer_only=True,
             )
+            self._lock.acquire()
+            self._cache, other._cache = other._cache, self._cache
+            self.telemetry["size"] = other.telemetry["size"]
+            self._cache_modified = cache_file_learnt
             self.last_loaded = now()
             return True
+        except (AssertionError, RuntimeError):
+            raise
         except Exception as e:
             logger.debug("Fail to read cache from disk due to error: %s", e)
             return False
 
+    def load(self) -> bool:
+        """Load cache from disk if possible, returns whether it was able to load.
+
+        This is the public version of _load, it makes sure that all the
+        necessary locks are acquired.
+        """
+        with self._lock:
+            return self._load()
+
     def _save(self, load_first: bool = True, force_flush: bool = False) -> bool:
-        """Save cache to disk if possible, returns whether it was able to save."""
+        """Save cache to disk if possible, returns whether it was able to save.
+
+        This function is non-locking when it comes to self._lock.
+        """
+        if test_mode:
+            assert (
+                self._lock.locked()
+            ), "The mutex self._lock should be locked by this thread"
         self._clear_expired_entries()
         if not self._cache_modified and not force_flush:
             # cache is not updated, so there is no need to dump cache to file, we just return
@@ -571,10 +589,21 @@ class SFDictFileCache(SFDictCache):
             logger.debug(
                 f"acquiring {self._file_lock_path} timed out, skipping saving..."
             )
+        except (AssertionError, RuntimeError):
+            raise
         except Exception as e:
             logger.debug("Fail to write cache to disk due to error: %s", e)
         return False
 
+    def save(self, load_first: bool = True) -> bool:
+        """Save cache to disk if possible, returns whether it was able to save.
+
+        This is the public version of _save, it makes sure that all the
+        necessary locks are acquired.
+        """
+        with self._lock:
+            return self._save(load_first)
+
     def _save_if_should(self) -> bool:
         """Saves file to disk if necessary and returns whether it saved.
 
@@ -624,17 +653,15 @@ class SFDictFileCache(SFDictCache):
             # At tear-down time builtins module might be already gone, ignore every error
             pass
 
-    def clear_expired_entries(self) -> None:
-        super().clear_expired_entries()
-        self._save_if_should()
-
     def clear(self) -> None:
         super().clear()
         # This unlink prevents us from loading just before saving
         with self._file_lock:
             if os.path.exists(self.file_path) and os.path.isfile(self.file_path):
                 os.unlink(self.file_path)
-        self._save(load_first=False, force_flush=True)
+        # TODO: is this necessary?
+        with self._lock:
+            self._save(load_first=False, force_flush=True)
 
     # Custom pickling implementation
 
diff --git a/src/snowflake/connector/constants.py b/src/snowflake/connector/constants.py
index a0029b13..a530424b 100644
--- a/src/snowflake/connector/constants.py
+++ b/src/snowflake/connector/constants.py
@@ -323,3 +323,4 @@ DAY_IN_SECONDS = 60 * 60 * 24
 
 # TODO: all env variables definitions should be here
 ENV_VAR_PARTNER = "SF_PARTNER"
+ENV_VAR_TEST_MODE = "SNOWFLAKE_TEST_MODE"
diff --git a/src/snowflake/connector/ocsp_snowflake.py b/src/snowflake/connector/ocsp_snowflake.py
index d73616df..c01d9c3c 100644
--- a/src/snowflake/connector/ocsp_snowflake.py
+++ b/src/snowflake/connector/ocsp_snowflake.py
@@ -587,7 +587,7 @@ class OCSPCache:
         """
         if OCSPCache.CACHE_UPDATED:
             if isinstance(OCSP_RESPONSE_VALIDATION_CACHE, SFDictFileCache):
-                OCSP_RESPONSE_VALIDATION_CACHE._save()
+                OCSP_RESPONSE_VALIDATION_CACHE.save()
             OCSPCache.update_ocsp_response_cache_file(
                 ocsp, OCSPCache.OCSP_RESPONSE_CACHE_URI
             )
@@ -1618,7 +1618,7 @@ class SnowflakeOCSP:
                             self, cert_id, cache_key=cache_key, lock_cache=False
                         )
             if new_cache_dict:
-                OCSP_RESPONSE_VALIDATION_CACHE._update(new_cache_dict)
+                OCSP_RESPONSE_VALIDATION_CACHE.update(new_cache_dict)
                 OCSPCache.CACHE_UPDATED = True
         except Exception as ex:
             logger.debug("Caught here - %s", ex)
diff --git a/tox.ini b/tox.ini
index 31323300..0a4bdc22 100644
--- a/tox.ini
+++ b/tox.ini
@@ -50,6 +50,7 @@ setenv =
     SNOWFLAKE_PYTEST_COV_LOCATION = {env:JUNIT_REPORT_DIR:{toxworkdir}}/junit.{envname}-{env:cloud_provider:dev}.xml
     SNOWFLAKE_PYTEST_COV_CMD = --cov snowflake.connector --junitxml {env:SNOWFLAKE_PYTEST_COV_LOCATION} --cov-report=
     SNOWFLAKE_PYTEST_CMD = pytest {env:SNOWFLAKE_PYTEST_OPTS:} {env:SNOWFLAKE_PYTEST_COV_CMD}
+    SNOWFLAKE_TEST_MODE = true
 passenv =
     AWS_ACCESS_KEY_ID
     AWS_SECRET_ACCESS_KEY


Test patch:
diff --git a/test/unit/test_cache.py b/test/unit/test_cache.py
index b9be2345..d490e4c0 100644
--- a/test/unit/test_cache.py
+++ b/test/unit/test_cache.py
@@ -127,7 +127,8 @@ class TestSFDictCache:
         c._entry_lifetime = NO_LIFETIME
         c["d"] = 4
         assert c.get("d") is None
-        assert c._getitem_non_locking(1) == "a"
+        with c._lock:
+            assert c._getitem_non_locking(1) == "a"
 
     def test_items(self):
         c = cache.SFDictCache()
@@ -255,8 +256,11 @@ if cache is not None:
 
 class TestSFDictFileCache:
     # Since the __getitem__ is overwritten copied over tests to test it
-    def test_simple_usage(self):
-        c = cache.SFDictCache.from_dict({1: "a", 2: "b"})
+    def test_simple_usage(self, tmpdir):
+        c = cache.SFDictFileCache.from_dict(
+            {1: "a", 2: "b"},
+            file_path=os.path.join(tmpdir, "c.txt"),
+        )
         assert 1 in c and 2 in c
         assert c[1] == "a"
         assert c[2] == "b"
@@ -267,9 +271,10 @@ class TestSFDictFileCache:
             "size": 2,
         }
 
-    def test_miss(self):
-        c = cache.SFDictCache.from_dict(
+    def test_miss(self, tmpdir):
+        c = cache.SFDictFileCache.from_dict(
             {"a": 1},
+            file_path=os.path.join(tmpdir, "c.txt"),
         )
         with pytest.raises(KeyError):
             c["b"]
@@ -280,10 +285,12 @@ class TestSFDictFileCache:
             "size": 1,
         }
 
-    def test_expiration(self):
-        c = cache.SFDictCache.from_dict(
+    def test_expiration(self, tmpdir):
+        tmp_file = os.path.join(tmpdir, "c.txt")
+        c = NeverSaveSFDictFileCache.from_dict(
             {"a": 1},
             entry_lifetime=0,
+            file_path=tmp_file,
         )
         with pytest.raises(KeyError):
             c["a"]
@@ -293,6 +300,11 @@ class TestSFDictFileCache:
             "expiration": 1,
             "size": 0,
         }
+        c2 = AlwaysSaveSFDictFileCache(file_path=tmp_file)
+        c["a"] = 1
+        c2["a"] = 2
+        assert c2["a"] == 2
+        assert c["a"] == 2
 
     # The rest of the tests test the file cache portion
 
@@ -335,7 +347,8 @@ class TestSFDictFileCache:
                 # c will dump cache to file every time an item is set as it's AlwaysSaveSFDictFileCache
                 # so the program exeucted at this point, c._cache_modified is false as it has the latest cache
                 # changes, we set force_flush to True
-                assert c._save(force_flush=True) is False
+                with c._lock:
+                    assert not c._save(force_flush=True)
             assert caplog.record_tuples == [
                 (
                     "snowflake.connector.cache",
@@ -366,14 +379,14 @@ class TestSFDictFileCache:
         c2 = pickle.loads(pickle.dumps(c1))
         c1["a"] = 1
         c1["b"] = 1
-        assert c1._save()
+        assert c1.save()
         assert c2.telemetry == {
             "hit": 0,
             "miss": 0,
             "expiration": 0,
             "size": 0,
         }
-        assert c2._load()
+        assert c2.load()
         assert c2.telemetry == {
             "hit": 0,
             "miss": 0,
@@ -388,8 +401,8 @@ class TestSFDictFileCache:
         c1._entry_lifetime = NO_LIFETIME
         c1["b"] = 1
         c1["c"] = 1
-        assert c1._save()  # Save calls self._clear_expired_entries()
-        assert c2._load()
+        assert c1.save()  # Save calls self._clear_expired_entries()
+        assert c2.load()
         assert c2.keys() == ["a"]
 
     def test_path_choosing(self, tmpdir):
@@ -450,7 +463,7 @@ class TestSFDictFileCache:
         cache_path = os.path.join(tmpdir, "cache.txt")
         c1 = cache.SFDictFileCache(file_path=cache_path)
         c1["a"] = 1
-        c1._save()
+        c1.save()
 
         # load cache
         c2 = cache.SFDictFileCache(
@@ -461,66 +474,65 @@ class TestSFDictFileCache:
         # cache will expire immediately due to the NO_LIFETIME setting
         # when loading again, clearing cache logic will be triggered
         # load will trigger clear expired entries, and then further call _getitem
-        c2._load()
+        c2.load()
 
         assert len(c2) == 1 and c2["a"] == 1 and c2._getitem_non_locking("a") == 1
 
     def test_broken_pickle_objects(self, tmpdir):
         cache_path = os.path.join(tmpdir, "cache.txt")
-        c1 = cache.SFDictFileCache(file_path=cache_path)
+        c1 = NeverSaveSFDictFileCache(file_path=cache_path)
         c1["key"] = BrokenReadingPickleObject()
-        # __setitem__ might have already saved the item, so cache is unchanged after saving, thus we force flush here
-        assert c1._save(force_flush=True)
+        assert c1.save()
         assert os.path.exists(cache_path) and os.path.isfile(cache_path)
 
-        c2 = cache.SFDictFileCache(file_path=cache_path)
+        c2 = NeverSaveSFDictFileCache(file_path=cache_path)
         assert not c2._load()  # load should return false due to ModuleNotFound error
 
         cache_path = os.path.join(tmpdir, "cache2.txt")
-        c1 = cache.SFDictFileCache(file_path=cache_path)
+        c1 = NeverSaveSFDictFileCache(file_path=cache_path)
         c1["key"] = BrokenWritingPickleObject()
-        assert not c1._save()
+        assert not c1.save()
         assert not os.path.exists(cache_path)
 
-        c2 = cache.SFDictFileCache(file_path=cache_path)
-        assert not c2._load()  # load should return false due to no file
+        c2 = NeverSaveSFDictFileCache(file_path=cache_path)
+        assert not c2.load()  # load should return false due to no file
 
 
 def test_file_is_not_updated(tmpdir):
     tmp_cache_file = os.path.join(tmpdir, "tmp_cache")
-    sfcache = cache.SFDictFileCache(file_path=tmp_cache_file, entry_lifetime=1)
+    sfcache = NeverSaveSFDictFileCache(file_path=tmp_cache_file, entry_lifetime=1)
     sfcache.update({"key": "value"})
     assert sfcache._cache_modified
-    sfcache._save()  # this save call will dump cache to file because item was added
+    sfcache.save()  # this save call will dump cache to file because item was added
     assert not sfcache._cache_modified
     updated_time = os.path.getmtime(tmp_cache_file)
-    sfcache._save()  # this save call will be a no-op since there is no update
+    sfcache.save()  # this save call will be a no-op since there is no update
     assert os.path.getmtime(tmp_cache_file) == updated_time
     sfcache.update({"key": "value2"})
     assert sfcache._cache_modified
     time.sleep(0.1)  # sleep 0.1 to avoid flushing too fast
-    sfcache._save()  # this save call will dump cache to file
+    sfcache.save()  # this save call will dump cache to file
     assert not sfcache._cache_modified
     second_updated_time = os.path.getmtime(tmp_cache_file)
     assert second_updated_time > updated_time
     assert sfcache["key"] == "value2"
     assert not sfcache._cache_modified
-    sfcache._save()  # this save call will be a no-op since there is no update
+    sfcache.save()  # this save call will be a no-op since there is no update
     assert not sfcache._cache_modified
     assert os.path.getmtime(tmp_cache_file) == second_updated_time
     time.sleep(1)
-    sfcache._save()  # this save call will dump cache because cache item is expired
+    sfcache.save()  # this save call will dump cache because cache item is expired
     assert not sfcache._cache_modified
     assert os.path.getmtime(tmp_cache_file) > second_updated_time
 
-    cache3 = cache.SFDictFileCache(file_path=tmp_cache_file, entry_lifetime=1)
+    cache3 = NeverSaveSFDictFileCache(file_path=tmp_cache_file, entry_lifetime=1)
     assert len(cache3) == 0
     os.unlink(tmp_cache_file)
 
 
 def test_cache_do_not_write_while_set_item(tmpdir):
     tmp_cache_file = os.path.join(tmpdir, "tmp_cache")
-    sfcache = cache.SFDictFileCache(tmp_cache_file)
+    sfcache = NeverSaveSFDictFileCache(tmp_cache_file)
     to_update_dict = {}
     for i in range(1000):
         to_update_dict[i] = i
@@ -528,14 +540,14 @@ def test_cache_do_not_write_while_set_item(tmpdir):
     # there should be no file created as setting item won't trigger creation
     assert sfcache._cache_modified
     assert not os.path.exists(tmp_cache_file)
-    sfcache._save()
+    sfcache.save()
     assert not sfcache._cache_modified
     assert os.path.exists(tmp_cache_file)
     file_modified_time = os.path.getmtime(tmp_cache_file)
-    sfcache2 = cache.SFDictFileCache(tmp_cache_file)
+    sfcache2 = NeverSaveSFDictFileCache(tmp_cache_file)
     time.sleep(0.01)
     for i in range(1000):
         assert sfcache2[i] == i
     assert not sfcache2._cache_modified
-    sfcache2._save()  # this save should be a no-op since there is no change
+    sfcache2.save()  # this save should be a no-op since there is no change
     assert os.path.getmtime(tmp_cache_file) == file_modified_time