From 529022ece2bae9dc55c458a76d3876767973f0f8 Mon Sep 17 00:00:00 2001
From: Ben Timby <btimby@gmail.com>
Date: Fri, 2 Aug 2024 13:21:42 -0400
Subject: [PATCH] Alternative implementation.

Signed-off-by: Ben Timby <btimby@gmail.com>
---
 prometheus_client/metrics.py  |  9 +++++----
 prometheus_client/registry.py |  4 ++--
 prometheus_client/utils.py    |  6 ------
 prometheus_client/values.py   |  6 +++---
 tests/test_core.py            | 10 +++++++++-
 5 files changed, 19 insertions(+), 16 deletions(-)

diff --git a/prometheus_client/metrics.py b/prometheus_client/metrics.py
index 684cf1ae..3bda92c4 100644
--- a/prometheus_client/metrics.py
+++ b/prometheus_client/metrics.py
@@ -1,4 +1,5 @@
 import os
+from threading import RLock
 import time
 import types
 from typing import (
@@ -15,7 +16,7 @@
 )
 from .registry import Collector, CollectorRegistry, REGISTRY
 from .samples import Exemplar, Sample
-from .utils import floatToGoString, INF, Lock
+from .utils import floatToGoString, INF
 
 T = TypeVar('T', bound='MetricWrapperBase')
 F = TypeVar("F", bound=Callable[..., Any])
@@ -143,7 +144,7 @@ def __init__(self: T,
 
         if self._is_parent():
             # Prepare the fields needed for child metrics.
-            self._lock = Lock()
+            self._lock = RLock()
             self._metrics: Dict[Sequence[str], T] = {}
 
         if self._is_observable():
@@ -696,7 +697,7 @@ class Info(MetricWrapperBase):
 
     def _metric_init(self):
         self._labelname_set = set(self._labelnames)
-        self._lock = Lock()
+        self._lock = RLock()
         self._value = {}
 
     def info(self, val: Dict[str, str]) -> None:
@@ -758,7 +759,7 @@ def __init__(self,
 
     def _metric_init(self) -> None:
         self._value = 0
-        self._lock = Lock()
+        self._lock = RLock()
 
     def state(self, state: str) -> None:
         """Set enum metric state."""
diff --git a/prometheus_client/registry.py b/prometheus_client/registry.py
index ab002445..4326b39a 100644
--- a/prometheus_client/registry.py
+++ b/prometheus_client/registry.py
@@ -1,9 +1,9 @@
 from abc import ABC, abstractmethod
 import copy
+from threading import RLock
 from typing import Dict, Iterable, List, Optional
 
 from .metrics_core import Metric
-from .utils import Lock
 
 
 # Ideally this would be a Protocol, but Protocols are only available in Python >= 3.8.
@@ -30,7 +30,7 @@ def __init__(self, auto_describe: bool = False, target_info: Optional[Dict[str,
         self._collector_to_names: Dict[Collector, List[str]] = {}
         self._names_to_collectors: Dict[str, Collector] = {}
         self._auto_describe = auto_describe
-        self._lock = Lock()
+        self._lock = RLock()
         self._target_info: Optional[Dict[str, str]] = {}
         self.set_target_info(target_info)
 
diff --git a/prometheus_client/utils.py b/prometheus_client/utils.py
index 51dea014..0d2b0948 100644
--- a/prometheus_client/utils.py
+++ b/prometheus_client/utils.py
@@ -1,5 +1,4 @@
 import math
-from threading import _PyRLock  # type: ignore[attr-defined]
 
 INF = float("inf")
 MINUS_INF = float("-inf")
@@ -23,8 +22,3 @@ def floatToGoString(d):
             mantissa = f'{s[0]}.{s[1:dot]}{s[dot + 1:]}'.rstrip('0.')
             return f'{mantissa}e+0{dot - 1}'
         return s
-
-
-class Lock(_PyRLock):
-    def locked(self):
-        return bool(self._count)
diff --git a/prometheus_client/values.py b/prometheus_client/values.py
index 0dfe8621..05331f82 100644
--- a/prometheus_client/values.py
+++ b/prometheus_client/values.py
@@ -1,8 +1,8 @@
 import os
+from threading import RLock
 import warnings
 
 from .mmap_dict import mmap_key, MmapedDict
-from .utils import Lock
 
 
 class MutexValue:
@@ -13,7 +13,7 @@ class MutexValue:
     def __init__(self, typ, metric_name, name, labelnames, labelvalues, help_text, **kwargs):
         self._value = 0.0
         self._exemplar = None
-        self._lock = Lock()
+        self._lock = RLock()
 
     def inc(self, amount):
         with self._lock:
@@ -50,7 +50,7 @@ def MultiProcessValue(process_identifier=os.getpid):
     # Use a single global lock when in multi-processing mode
     # as we presume this means there is no threading going on.
     # This avoids the need to also have mutexes in __MmapDict.
-    lock = Lock()
+    lock = RLock()
 
     class MmapedValue:
         """A float protected by a mutex backed by a per-process mmaped file."""
diff --git a/tests/test_core.py b/tests/test_core.py
index 8a54a02d..f80fb882 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -16,6 +16,14 @@
 from prometheus_client.metrics import _get_use_created
 
 
+def is_locked(lock):
+    "Tries to obtain a lock, returns True on success, False on failure."
+    locked = lock.acquire(blocking=False)
+    if locked:
+        lock.release()
+    return not locked
+
+
 def assert_not_observable(fn, *args, **kwargs):
     """
     Assert that a function call falls with a ValueError exception containing
@@ -963,7 +971,7 @@ def test_restricted_registry_does_not_yield_while_locked(self):
         m = Metric('target', 'Target metadata', 'info')
         m.samples = [Sample('target_info', {'foo': 'bar'}, 1)]
         for _ in registry.restricted_registry(['target_info', 's_sum']).collect():
-            self.assertFalse(registry._lock.locked())
+            self.assertFalse(is_locked(registry._lock))
 
 
 if __name__ == '__main__':