Ensure use of cachetools is thread-safe (#2178)

* Ensure use of cachetools is thread-safe

* Minor cleanup

* fix formatting

Author: rayokota

src/confluent_kafka/schema_registry/_async/json_schema.py

@@ -14,6 +14,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import asyncio as _locks
 import io
 import logging
 from typing import Any, Callable, Coroutine, Optional, Tuple, Union, cast
@@ -214,6 +215,7 @@ class AsyncJSONSerializer(AsyncBaseSerializer):
         '_to_dict',
         '_parsed_schemas',
         '_validators',
+        '_validators_lock',
         '_validate',
         '_json_encode',
     ]
@@ -255,6 +257,7 @@ async def __init_impl(
         self._known_subjects: set[str] = set()
         self._parsed_schemas = ParsedSchemaCache()
         self._validators: LRUCache[Schema, Validator] = LRUCache(1000)
+        self._validators_lock = _locks.Lock()
 
         if to_dict is not None and not callable(to_dict):
             raise ValueError(
@@ -404,7 +407,7 @@ def field_transformer(rule_ctx, field_transform, msg):
 
         if self._validate and schema is not None and parsed_schema is not None and ref_registry is not None:
             try:
-                validator = self._get_validator(schema, parsed_schema, ref_registry)
+                validator = await self._get_validator(schema, parsed_schema, ref_registry)
                 validator.validate(value)
             except ValidationError as ve:
                 raise SerializationError(ve.message)
@@ -441,16 +444,18 @@ async def _get_parsed_schema(self, schema: Optional[Schema]) -> Tuple[Optional[J
         self._parsed_schemas.set(schema, (parsed_schema, ref_registry))
         return parsed_schema, ref_registry
 
-    def _get_validator(self, schema: Schema, parsed_schema: JsonSchema, registry: Registry) -> Validator:
-        validator = self._validators.get(schema, None)
-        if validator is not None:
-            return validator
+    async def _get_validator(self, schema: Schema, parsed_schema: JsonSchema, registry: Registry) -> Validator:
+        async with self._validators_lock:
+            validator = self._validators.get(schema, None)
+            if validator is not None:
+                return validator
 
         cls = validator_for(parsed_schema)
         cls.check_schema(parsed_schema)
         validator = cls(parsed_schema, registry=registry)
 
-        self._validators[schema] = validator
+        async with self._validators_lock:
+            self._validators[schema] = validator
         return validator
 
 
@@ -518,6 +523,7 @@ class AsyncJSONDeserializer(AsyncBaseDeserializer):
         '_schema',
         '_parsed_schemas',
         '_validators',
+        '_validators_lock',
         '_validate',
         '_json_decode',
     ]
@@ -562,6 +568,7 @@ async def __init_impl(
         self._rule_registry = rule_registry if rule_registry else RuleRegistry.get_global_instance()
         self._parsed_schemas = ParsedSchemaCache()
         self._validators: LRUCache[Schema, Validator] = LRUCache(1000)
+        self._validators_lock = _locks.Lock()
         self._json_decode = json_decode or orjson.loads
         self._use_schema_id = None
 
@@ -708,7 +715,7 @@ def field_transformer(rule_ctx, field_transform, message):
         if self._validate:
             if reader_schema_raw is not None and reader_schema is not None and reader_ref_registry is not None:
                 try:
-                    validator = self._get_validator(reader_schema_raw, reader_schema, reader_ref_registry)
+                    validator = await self._get_validator(reader_schema_raw, reader_schema, reader_ref_registry)
                     validator.validate(obj_dict)
                 except ValidationError as ve:
                     raise SerializationError(ve.message)
@@ -745,14 +752,16 @@ async def _get_parsed_schema(self, schema: Schema) -> Tuple[Optional[JsonSchema]
         self._parsed_schemas.set(schema, (parsed_schema, ref_registry))
         return parsed_schema, ref_registry
 
-    def _get_validator(self, schema: Schema, parsed_schema: JsonSchema, registry: Registry) -> Validator:
-        validator = self._validators.get(schema, None)
-        if validator is not None:
-            return validator
+    async def _get_validator(self, schema: Schema, parsed_schema: JsonSchema, registry: Registry) -> Validator:
+        async with self._validators_lock:
+            validator = self._validators.get(schema, None)
+            if validator is not None:
+                return validator
 
         cls = validator_for(parsed_schema)
         cls.check_schema(parsed_schema)
         validator = cls(parsed_schema, registry=registry)
 
-        self._validators[schema] = validator
+        async with self._validators_lock:
+            self._validators[schema] = validator
         return validator

src/confluent_kafka/schema_registry/_async/schema_registry_client.py

@@ -16,6 +16,7 @@
 # limitations under the License.
 #
 import asyncio
+import asyncio as _locks
 import json
 import logging
 import os
@@ -642,6 +643,7 @@ def __init__(self, conf: dict):
         self._conf = conf
         self._rest_client = _AsyncRestClient(conf)
         self._cache = _SchemaCache()
+        self._latest_lock = _locks.Lock()
         cache_capacity = self._rest_client.cache_capacity
         cache_ttl = self._rest_client.cache_latest_ttl_sec
         self._latest_version_cache: Cache[Any, Any]
@@ -1045,7 +1047,8 @@ async def get_latest_version(self, subject_name: str, fmt: Optional[str] = None)
             `GET Subject Version API Reference <https://docs.confluent.io/current/schema-registry/develop/api.html#get--subjects-(string-%20subject)-versions-(versionId-%20version)>`_
         """  # noqa: E501
 
-        registered_schema = self._latest_version_cache.get(subject_name, None)
+        async with self._latest_lock:
+            registered_schema = self._latest_version_cache.get(subject_name, None)
         if registered_schema is not None:
             return registered_schema
 
@@ -1056,7 +1059,8 @@ async def get_latest_version(self, subject_name: str, fmt: Optional[str] = None)
 
         registered_schema = RegisteredSchema.from_dict(response)
 
-        self._latest_version_cache[subject_name] = registered_schema
+        async with self._latest_lock:
+            self._latest_version_cache[subject_name] = registered_schema
 
         return registered_schema
 
@@ -1080,7 +1084,8 @@ async def get_latest_with_metadata(
         """  # noqa: E501
 
         cache_key = (subject_name, frozenset(metadata.items()), deleted)
-        registered_schema = self._latest_with_metadata_cache.get(cache_key, None)
+        async with self._latest_lock:
+            registered_schema = self._latest_with_metadata_cache.get(cache_key, None)
         if registered_schema is not None:
             return registered_schema
 
@@ -1096,7 +1101,8 @@ async def get_latest_with_metadata(
 
         registered_schema = RegisteredSchema.from_dict(response)
 
-        self._latest_with_metadata_cache[cache_key] = registered_schema
+        async with self._latest_lock:
+            self._latest_with_metadata_cache[cache_key] = registered_schema
 
         return registered_schema
 
@@ -1548,13 +1554,15 @@ async def get_contexts(self, offset: int = 0, limit: int = -1) -> List[str]:
         result = await self._rest_client.get('contexts', query={'offset': offset, 'limit': limit})
         return result
 
-    def clear_latest_caches(self):
-        self._latest_version_cache.clear()
-        self._latest_with_metadata_cache.clear()
+    async def clear_latest_caches(self):
+        async with self._latest_lock:
+            self._latest_version_cache.clear()
+            self._latest_with_metadata_cache.clear()
 
-    def clear_caches(self):
-        self._latest_version_cache.clear()
-        self._latest_with_metadata_cache.clear()
+    async def clear_caches(self):
+        async with self._latest_lock:
+            self._latest_version_cache.clear()
+            self._latest_with_metadata_cache.clear()
         self._cache.clear()
 
     @staticmethod

src/confluent_kafka/schema_registry/_sync/json_schema.py

@@ -16,6 +16,7 @@
 # limitations under the License.
 import io
 import logging
+import threading as _locks
 from typing import Any, Callable, Optional, Tuple, Union, cast
 
 import orjson
@@ -212,6 +213,7 @@ class JSONSerializer(BaseSerializer):
         '_to_dict',
         '_parsed_schemas',
         '_validators',
+        '_validators_lock',
         '_validate',
         '_json_encode',
     ]
@@ -253,6 +255,7 @@ def __init_impl(
         self._known_subjects: set[str] = set()
         self._parsed_schemas = ParsedSchemaCache()
         self._validators: LRUCache[Schema, Validator] = LRUCache(1000)
+        self._validators_lock = _locks.Lock()
 
         if to_dict is not None and not callable(to_dict):
             raise ValueError(
@@ -440,15 +443,17 @@ def _get_parsed_schema(self, schema: Optional[Schema]) -> Tuple[Optional[JsonSch
         return parsed_schema, ref_registry
 
     def _get_validator(self, schema: Schema, parsed_schema: JsonSchema, registry: Registry) -> Validator:
-        validator = self._validators.get(schema, None)
-        if validator is not None:
-            return validator
+        with self._validators_lock:
+            validator = self._validators.get(schema, None)
+            if validator is not None:
+                return validator
 
         cls = validator_for(parsed_schema)
         cls.check_schema(parsed_schema)
         validator = cls(parsed_schema, registry=registry)
 
-        self._validators[schema] = validator
+        with self._validators_lock:
+            self._validators[schema] = validator
         return validator
 
 
@@ -515,6 +520,7 @@ class JSONDeserializer(BaseDeserializer):
         '_schema',
         '_parsed_schemas',
         '_validators',
+        '_validators_lock',
         '_validate',
         '_json_decode',
     ]
@@ -559,6 +565,7 @@ def __init_impl(
         self._rule_registry = rule_registry if rule_registry else RuleRegistry.get_global_instance()
         self._parsed_schemas = ParsedSchemaCache()
         self._validators: LRUCache[Schema, Validator] = LRUCache(1000)
+        self._validators_lock = _locks.Lock()
         self._json_decode = json_decode or orjson.loads
         self._use_schema_id = None
 
@@ -741,13 +748,15 @@ def _get_parsed_schema(self, schema: Schema) -> Tuple[Optional[JsonSchema], Opti
         return parsed_schema, ref_registry
 
     def _get_validator(self, schema: Schema, parsed_schema: JsonSchema, registry: Registry) -> Validator:
-        validator = self._validators.get(schema, None)
-        if validator is not None:
-            return validator
+        with self._validators_lock:
+            validator = self._validators.get(schema, None)
+            if validator is not None:
+                return validator
 
         cls = validator_for(parsed_schema)
         cls.check_schema(parsed_schema)
         validator = cls(parsed_schema, registry=registry)
 
-        self._validators[schema] = validator
+        with self._validators_lock:
+            self._validators[schema] = validator
         return validator

src/confluent_kafka/schema_registry/_sync/schema_registry_client.py

@@ -20,6 +20,7 @@
 import logging
 import os
 import ssl
+import threading as _locks
 import time
 import urllib
 from typing import Any, Callable, Dict, List, Literal, Optional, Union
@@ -640,6 +641,7 @@ def __init__(self, conf: dict):
         self._conf = conf
         self._rest_client = _RestClient(conf)
         self._cache = _SchemaCache()
+        self._latest_lock = _locks.Lock()
         cache_capacity = self._rest_client.cache_capacity
         cache_ttl = self._rest_client.cache_latest_ttl_sec
         self._latest_version_cache: Cache[Any, Any]
@@ -1041,7 +1043,8 @@ def get_latest_version(self, subject_name: str, fmt: Optional[str] = None) -> 'R
             `GET Subject Version API Reference <https://docs.confluent.io/current/schema-registry/develop/api.html#get--subjects-(string-%20subject)-versions-(versionId-%20version)>`_
         """  # noqa: E501
 
-        registered_schema = self._latest_version_cache.get(subject_name, None)
+        with self._latest_lock:
+            registered_schema = self._latest_version_cache.get(subject_name, None)
         if registered_schema is not None:
             return registered_schema
 
@@ -1050,7 +1053,8 @@ def get_latest_version(self, subject_name: str, fmt: Optional[str] = None) -> 'R
 
         registered_schema = RegisteredSchema.from_dict(response)
 
-        self._latest_version_cache[subject_name] = registered_schema
+        with self._latest_lock:
+            self._latest_version_cache[subject_name] = registered_schema
 
         return registered_schema
 
@@ -1074,7 +1078,8 @@ def get_latest_with_metadata(
         """  # noqa: E501
 
         cache_key = (subject_name, frozenset(metadata.items()), deleted)
-        registered_schema = self._latest_with_metadata_cache.get(cache_key, None)
+        with self._latest_lock:
+            registered_schema = self._latest_with_metadata_cache.get(cache_key, None)
         if registered_schema is not None:
             return registered_schema
 
@@ -1090,7 +1095,8 @@ def get_latest_with_metadata(
 
         registered_schema = RegisteredSchema.from_dict(response)
 
-        self._latest_with_metadata_cache[cache_key] = registered_schema
+        with self._latest_lock:
+            self._latest_with_metadata_cache[cache_key] = registered_schema
 
         return registered_schema
 
@@ -1537,12 +1543,14 @@ def get_contexts(self, offset: int = 0, limit: int = -1) -> List[str]:
         return result
 
     def clear_latest_caches(self):
-        self._latest_version_cache.clear()
-        self._latest_with_metadata_cache.clear()
+        with self._latest_lock:
+            self._latest_version_cache.clear()
+            self._latest_with_metadata_cache.clear()
 
     def clear_caches(self):
-        self._latest_version_cache.clear()
-        self._latest_with_metadata_cache.clear()
+        with self._latest_lock:
+            self._latest_version_cache.clear()
+            self._latest_with_metadata_cache.clear()
         self._cache.clear()
 
     @staticmethod

tools/unasync.py

@@ -33,7 +33,9 @@
 SUBS = [
     ('from confluent_kafka.schema_registry.common import asyncinit', ''),
     ('@asyncinit', ''),
+    ('import asyncio as _locks', 'import threading as _locks'),
     ('import asyncio', ''),
+    ('async with self.', 'with self.'),
     ('asyncio.sleep', 'time.sleep'),
     ('._async.', '._sync.'),
     ('Async([A-Z][A-Za-z0-9_]*)', r'\2'),