debug on TelemetryClientFactory lock

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>

Author: saishreeeee

.github/workflows/code-quality-checks.yml

@@ -112,7 +112,7 @@ jobs:
       #              run test suite
       #----------------------------------------------
       -   name: Run tests
-          run: poetry run python -m pytest tests/unit
+          run: poetry run python -m pytest tests/unit -v -s
   check-linting:
     runs-on: ubuntu-latest
     strategy:

src/databricks/sql/telemetry/telemetry_client.py

@@ -31,6 +31,76 @@
 logger = logging.getLogger(__name__)
 
 
+class DebugLock:
+    """A wrapper around threading.Lock that provides detailed debugging for lock acquisition/release"""
+    def __init__(self, name: str = "DebugLock"):
+        self._lock = threading.Lock()
+        self._name = name
+        self._owner = None
+        self._waiters = []
+        self._debug_logger = logging.getLogger(f"{__name__}.{name}")
+        # Ensure debug logging is visible
+        if not self._debug_logger.handlers:
+            handler = logging.StreamHandler()
+            formatter = logging.Formatter(
+                ":lock: %(asctime)s [%(threadName)s-%(thread)d] LOCK-%(name)s: %(message)s"
+            )
+            handler.setFormatter(formatter)
+            self._debug_logger.addHandler(handler)
+            self._debug_logger.setLevel(logging.DEBUG)
+    def acquire(self, blocking=True, timeout=-1):
+        current = threading.current_thread()
+        thread_info = f"{current.name}-{current.ident}"
+        if self._owner:
+            self._debug_logger.warning(
+                f":rotating_light: WAITING: {thread_info} waiting for lock held by {self._owner}"
+            )
+            self._waiters.append(thread_info)
+        else:
+            self._debug_logger.debug(
+                f":large_green_circle: TRYING: {thread_info} attempting to acquire lock"
+            )
+        # Try to acquire the lock
+        acquired = self._lock.acquire(blocking, timeout)
+        if acquired:
+            self._owner = thread_info
+            self._debug_logger.info(f":white_check_mark: ACQUIRED: {thread_info} got the lock")
+            if self._waiters:
+                self._debug_logger.info(
+                    f":clipboard: WAITERS: {len(self._waiters)} threads waiting: {self._waiters}"
+                )
+        else:
+            self._debug_logger.error(
+                f":x: FAILED: {thread_info} failed to acquire lock (timeout)"
+            )
+            if thread_info in self._waiters:
+                self._waiters.remove(thread_info)
+        return acquired
+    def release(self):
+        current = threading.current_thread()
+        thread_info = f"{current.name}-{current.ident}"
+        if self._owner != thread_info:
+            self._debug_logger.error(
+                f":rotating_light: ERROR: {thread_info} trying to release lock owned by {self._owner}"
+            )
+        else:
+            self._debug_logger.info(f":unlock: RELEASED: {thread_info} released the lock")
+            self._owner = None
+            # Remove from waiters if present
+            if thread_info in self._waiters:
+                self._waiters.remove(thread_info)
+            if self._waiters:
+                self._debug_logger.info(
+                    f":loudspeaker: NEXT: {len(self._waiters)} threads still waiting: {self._waiters}"
+                )
+        self._lock.release()
+    def __enter__(self):
+        self.acquire()
+        return self
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.release()
+
+
 class TelemetryHelper:
     """Helper class for getting telemetry related information."""
 
@@ -355,7 +425,8 @@ class TelemetryClientFactory:
     ] = {}  # Map of session_id_hex -> BaseTelemetryClient
     _executor: Optional[ThreadPoolExecutor] = None
     _initialized: bool = False
-    _lock = threading.Lock()  # Thread safety for factory operations
+    # _lock = threading.Lock()  # Thread safety for factory operations
+    _lock = DebugLock("TelemetryClientFactory")  # Thread safety for factory operations with debugging
     _original_excepthook = None
     _excepthook_installed = False