fix: Add context manager for clearer, safer locks

ldclient/__init__.py

@@ -37,17 +37,16 @@ def set_config(config: Config):
     global __config
     global __client
     global __lock
-    try:
-        __lock.lock()
-        if __client:
-            log.info("Reinitializing LaunchDarkly Client " + VERSION + " with new config")
-            new_client = LDClient(config=config, start_wait=start_wait)
-            old_client = __client
-            __client = new_client
-            old_client.close()
-    finally:
-        __config = config
-        __lock.unlock()
+    with __lock.write():
+        try:
+            if __client:
+                log.info("Reinitializing LaunchDarkly Client " + VERSION + " with new config")
+                new_client = LDClient(config=config, start_wait=start_wait)
+                old_client = __client
+                __client = new_client
+                old_client.close()
+        finally:
+            __config = config
 
 
 def get() -> LDClient:
@@ -63,35 +62,27 @@ def get() -> LDClient:
     global __config
     global __client
     global __lock
-    try:
-        __lock.rlock()
+    with __lock.read():
         if __client:
             return __client
         if __config is None:
             raise Exception("set_config was not called")
-    finally:
-        __lock.runlock()
 
-    try:
-        __lock.lock()
+    with __lock.write():
         if not __client:
             log.info("Initializing LaunchDarkly Client " + VERSION)
             __client = LDClient(config=__config, start_wait=start_wait)
         return __client
-    finally:
-        __lock.unlock()
 
 
 # for testing only
 def _reset_client():
     global __client
     global __lock
-    try:
-        __lock.lock()
+    c = None
+    with __lock.write():
         c = __client
         __client = None
-    finally:
-        __lock.unlock()
     if c:
         c.close()
 

ldclient/client.py

@@ -111,13 +111,10 @@ def __wrapper(self, fn: Callable):
             raise
 
     def __update_availability(self, available: bool):
-        try:
-            self.__lock.lock()
+        with self.__lock.write():
             if available == self.__last_available:
                 return
             self.__last_available = available
-        finally:
-            self.__lock.unlock()
 
         status = DataStoreStatus(available, False)
 
@@ -127,23 +124,19 @@ def __update_availability(self, available: bool):
         self.__store_update_sink.update_status(status)
 
         if available:
-            try:
-                self.__lock.lock()
+            with self.__lock.write():
                 if self.__poller is not None:
                     self.__poller.stop()
                     self.__poller = None
-            finally:
-                self.__lock.unlock()
 
             return
 
         log.warn("Detected persistent store unavailability; updates will be cached until it recovers")
         task = RepeatingTask("ldclient.check-availability", 0.5, 0, self.__check_availability)
 
-        self.__lock.lock()
-        self.__poller = task
-        self.__poller.start()
-        self.__lock.unlock()
+        with self.__lock.write():
+            self.__poller = task
+            self.__poller.start()
 
     def __check_availability(self):
         try:
@@ -717,9 +710,8 @@ def add_hook(self, hook: Hook):
         if not isinstance(hook, Hook):
             return
 
-        self.__hooks_lock.lock()
-        self.__hooks.append(hook)
-        self.__hooks_lock.unlock()
+        with self.__hooks_lock.write():
+            self.__hooks.append(hook)
 
     def __evaluate_with_hooks(self, key: str, context: Context, default_value: Any, method: str, block: Callable[[], _EvaluationWithHookResult]) -> _EvaluationWithHookResult:
         """
@@ -733,15 +725,11 @@ def __evaluate_with_hooks(self, key: str, context: Context, default_value: Any,
         # :return:
         """
         hooks = []  # type: List[Hook]
-        try:
-            self.__hooks_lock.rlock()
-
+        with self.__hooks_lock.read():
             if len(self.__hooks) == 0:
                 return block()
 
             hooks = self.__hooks.copy()
-        finally:
-            self.__hooks_lock.runlock()
 
         series_context = EvaluationSeriesContext(key=key, context=context, default_value=default_value, method=method)
         hook_data = self.__execute_before_evaluation(hooks, series_context)

ldclient/feature_store.py

@@ -77,8 +77,7 @@ def is_available(self) -> bool:
 
     def get(self, kind: VersionedDataKind, key: str, callback: Callable[[Any], Any] = lambda x: x) -> Any:
         """ """
-        try:
-            self._lock.rlock()
+        with self._lock.read():
             itemsOfKind = self._items[kind]
             item = itemsOfKind.get(key)
             if item is None:
@@ -88,17 +87,12 @@ def get(self, kind: VersionedDataKind, key: str, callback: Callable[[Any], Any]
                 log.debug("Attempted to get deleted key %s in '%s', returning None", key, kind.namespace)
                 return callback(None)
             return callback(item)
-        finally:
-            self._lock.runlock()
 
     def all(self, kind, callback):
         """ """
-        try:
-            self._lock.rlock()
+        with self._lock.read():
             itemsOfKind = self._items[kind]
             return callback(dict((k, i) for k, i in itemsOfKind.items() if ('deleted' not in i) or not i['deleted']))
-        finally:
-            self._lock.runlock()
 
     def init(self, all_data):
         """ """
@@ -108,51 +102,39 @@ def init(self, all_data):
             for key, item in items.items():
                 items_decoded[key] = kind.decode(item)
             all_decoded[kind] = items_decoded
-        try:
-            self._lock.rlock()
+        with self._lock.write():
             self._items.clear()
             self._items.update(all_decoded)
             self._initialized = True
             for k in all_data:
                 log.debug("Initialized '%s' store with %d items", k.namespace, len(all_data[k]))
-        finally:
-            self._lock.runlock()
 
     # noinspection PyShadowingNames
     def delete(self, kind, key: str, version: int):
         """ """
-        try:
-            self._lock.rlock()
+        with self._lock.write():
             itemsOfKind = self._items[kind]
             i = itemsOfKind.get(key)
             if i is None or i['version'] < version:
                 i = {'deleted': True, 'version': version}
                 itemsOfKind[key] = i
-        finally:
-            self._lock.runlock()
 
     def upsert(self, kind, item):
         """ """
         decoded_item = kind.decode(item)
         key = item['key']
-        try:
-            self._lock.rlock()
+        with self._lock.write():
             itemsOfKind = self._items[kind]
             i = itemsOfKind.get(key)
             if i is None or i['version'] < item['version']:
                 itemsOfKind[key] = decoded_item
                 log.debug("Updated %s in '%s' to version %d", key, kind.namespace, item['version'])
-        finally:
-            self._lock.runlock()
 
     @property
     def initialized(self) -> bool:
         """ """
-        try:
-            self._lock.rlock()
+        with self._lock.read():
             return self._initialized
-        finally:
-            self._lock.runlock()
 
     def describe_configuration(self, config):
         return 'memory'

ldclient/impl/datasource/status.py

@@ -29,11 +29,8 @@ def __init__(self, store: FeatureStore, status_listeners: Listeners, flag_change
 
     @property
     def status(self) -> DataSourceStatus:
-        try:
-            self.__lock.rlock()
+        with self.__lock.read():
             return self.__status
-        finally:
-            self.__lock.runlock()
 
     def init(self, all_data: Mapping[VersionedDataKind, Mapping[str, dict]]):
         old_data = None
@@ -70,8 +67,7 @@ def delete(self, kind: VersionedDataKind, key: str, version: int):
     def update_status(self, new_state: DataSourceState, new_error: Optional[DataSourceErrorInfo]):
         status_to_broadcast = None
 
-        try:
-            self.__lock.lock()
+        with self.__lock.write():
             old_status = self.__status
 
             if new_state == DataSourceState.INTERRUPTED and old_status.state == DataSourceState.INITIALIZING:
@@ -83,8 +79,6 @@ def update_status(self, new_state: DataSourceState, new_error: Optional[DataSour
             self.__status = DataSourceStatus(new_state, self.__status.since if new_state == self.__status.state else time.time(), self.__status.error if new_error is None else new_error)
 
             status_to_broadcast = self.__status
-        finally:
-            self.__lock.unlock()
 
         if status_to_broadcast is not None:
             self.__status_listeners.notify(status_to_broadcast)

ldclient/impl/datastore/status.py

@@ -27,16 +27,12 @@ def listeners(self) -> Listeners:
         return self.__listeners
 
     def status(self) -> DataStoreStatus:
-        self.__lock.rlock()
-        status = copy(self.__status)
-        self.__lock.runlock()
-
-        return status
+        with self.__lock.read():
+            return copy(self.__status)
 
     def update_status(self, status: DataStoreStatus):
-        self.__lock.lock()
-        old_value, self.__status = self.__status, status
-        self.__lock.unlock()
+        with self.__lock.write():
+            old_value, self.__status = self.__status, status
 
         if old_value != status:
             self.__listeners.notify(status)