From 7b6561d8bb83b51d6407e8c9c4413da46818f4df Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 07:58:17 +0000 Subject: [PATCH 1/2] perf(telemetry): bound how long interpreter exit waits for telemetry The final telemetry flush on close is submitted to a ThreadPoolExecutor, whose workers are joined at interpreter exit. When /telemetry-ext is slow or returns 429 (retried with 10 s sleeps), the process could not exit until the retries finished: 10-50 s on a rate-limited endpoint, up to the 900 s request timeout on a hung one. Run telemetry sends on daemon worker threads instead, and at exit wait at most EXIT_FLUSH_TIMEOUT_SECONDS (3 s) for in-flight sends, so the normal final flush is still delivered before the process exits. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- CHANGELOG.md | 1 + .../sql/telemetry/telemetry_client.py | 118 +++++++++++++++++- tests/unit/test_telemetry.py | 105 ++++++++++++++++ 3 files changed, 218 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d9c58ebf3..e04f35f5e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,7 @@ # Unreleased - Fix inline parameters (`use_inline_params=True`) so a sequence directly after `IN` renders as a list again: `x IN %(ids)s` renders as `x IN (1,2,3)` instead of `x IN ARRAY(1,2,3)`, which Spark rejects. Inside the list, tuples render as rows for a multi-column `IN`, and a `set` or `frozenset` is accepted again. Sequences anywhere else still render as `ARRAY(...)` ([#629](https://github.com/databricks/databricks-sql-python/issues/629)) - Speed up Thrift response handling: read replies in 1 MiB chunks and decode them with Thrift's C-accelerated protocol, and stop stringifying every request/response for the default no-op debug logger. Wide result schemas, metadata calls with many rows, and inline Arrow results benefit most. +- Stop a slow or rate-limited (HTTP 429) telemetry endpoint from delaying interpreter exit after connections close: exit now waits at most 3 seconds for in-flight telemetry instead of the tens of seconds spent on retries. # 4.9.0 (2026-11-04) - Upgrade Databricks SQL Kernel to 1.4.0, which keeps Lakehouse Real-Time session affinity and flushes shared telemetry only after the final session closes. diff --git a/src/databricks/sql/telemetry/telemetry_client.py b/src/databricks/sql/telemetry/telemetry_client.py index fed2001cb..00792457b 100644 --- a/src/databricks/sql/telemetry/telemetry_client.py +++ b/src/databricks/sql/telemetry/telemetry_client.py @@ -1,9 +1,10 @@ +import atexit import threading import time import logging import json -from queue import Queue, Full -from concurrent.futures import ThreadPoolExecutor +import weakref +from queue import Queue, Full, SimpleQueue from concurrent.futures import Future from datetime import datetime, timezone from typing import List, Dict, Any, Optional, TYPE_CHECKING @@ -462,6 +463,89 @@ def __del__(self): pass +class _DaemonThreadPoolExecutor: + """Minimal ThreadPoolExecutor stand-in whose workers are daemon threads. + + ThreadPoolExecutor workers are joined at interpreter exit, so a telemetry + request that is slow or retrying a 429 used to hold the process open for + tens of seconds. Exit instead waits a bounded time for pending sends (see + TelemetryClientFactory._wait_for_pending_sends). + """ + + def __init__(self, max_workers: int, thread_name_prefix: str): + self._max_workers = max_workers + self._thread_name_prefix = thread_name_prefix + self._work_queue: SimpleQueue = SimpleQueue() + self._threads: List[threading.Thread] = [] + self._unfinished = 0 + self._lock = threading.Lock() + self._idle = threading.Condition(self._lock) + self._shutdown = False + + def submit(self, fn, *args, **kwargs) -> Future: + with self._lock: + if self._shutdown: + raise RuntimeError("cannot schedule new futures after shutdown") + future: Future = Future() + self._work_queue.put((future, fn, args, kwargs)) + self._unfinished += 1 + # Add a worker only when all live ones are busy (dead ones: forked child). + self._threads = [t for t in self._threads if t.is_alive()] + if len(self._threads) < min(self._unfinished, self._max_workers): + thread = threading.Thread( + target=self._worker, + name=f"{self._thread_name_prefix}_{len(self._threads)}", + daemon=True, + ) + thread.start() + self._threads.append(thread) + return future + + @staticmethod + def _run(future: Future, fn, args, kwargs): + if not future.set_running_or_notify_cancel(): + return + try: + result = fn(*args, **kwargs) + except BaseException as exc: + future.set_exception(exc) + else: + future.set_result(result) + + def _worker(self): + while True: + item = self._work_queue.get() + if item is None: + self._work_queue.put(None) # wake the next worker + return + self._run(*item) + # Drop references (e.g. to the TelemetryClient) before blocking again. + del item + with self._lock: + self._unfinished -= 1 + if self._unfinished == 0: + self._idle.notify_all() + + def shutdown(self, wait: bool = True): + with self._lock: + if not self._shutdown: + self._shutdown = True + self._work_queue.put(None) + threads = list(self._threads) + if wait: + for thread in threads: + thread.join() + + def wait_until_idle(self, timeout: float) -> bool: + """Wait up to ``timeout`` seconds for submitted work to finish.""" + with self._lock: + return self._idle.wait_for( + lambda: self._unfinished == 0 + or not any(t.is_alive() for t in self._threads), + timeout, + ) + + class _TelemetryClientHolder: """ Holds a telemetry client with reference counting. @@ -495,12 +579,18 @@ class TelemetryClientFactory: _clients: Dict[str, _TelemetryClientHolder] = ( {} ) # Map of host_url -> TelemetryClientHolder - _executor: Optional[ThreadPoolExecutor] = None + _executor: Optional[_DaemonThreadPoolExecutor] = None + # Every executor that may still run sends, including ones already shut down. + _executors: weakref.WeakSet = weakref.WeakSet() _initialized: bool = False _lock = threading.RLock() # Thread safety for factory operations # used RLock instead of Lock to avoid deadlocks when garbage collection is triggered _original_excepthook = None _excepthook_installed = False + _exit_hook_installed = False + + # Upper bound on how long interpreter exit waits for in-flight telemetry sends. + EXIT_FLUSH_TIMEOUT_SECONDS = 3.0 # Shared flush thread for all clients _flush_thread = None @@ -531,10 +621,12 @@ def _initialize(cls): if not cls._initialized: cls._clients = {} - cls._executor = ThreadPoolExecutor( - max_workers=10 - ) # Thread pool for async operations + cls._executor = _DaemonThreadPoolExecutor( + max_workers=10, thread_name_prefix="telemetry-sender" + ) + cls._executors.add(cls._executor) cls._install_exception_hook() + cls._install_exit_hook() cls._start_flush_thread() cls._initialized = True logger.debug( @@ -577,6 +669,20 @@ def _install_exception_hook(cls): cls._excepthook_installed = True logger.debug("Global exception handler installed for telemetry") + @classmethod + def _install_exit_hook(cls): + if not cls._exit_hook_installed: + atexit.register(cls._wait_for_pending_sends) + cls._exit_hook_installed = True + + @classmethod + def _wait_for_pending_sends(cls): + """At interpreter exit, give in-flight sends a bounded time to finish.""" + deadline = time.monotonic() + cls.EXIT_FLUSH_TIMEOUT_SECONDS + for executor in list(cls._executors): + if not executor.wait_until_idle(max(0.0, deadline - time.monotonic())): + logger.debug("Dropping telemetry still in flight at interpreter exit") + @classmethod def _handle_unhandled_exception(cls, exc_type, exc_value, exc_traceback): """Handle unhandled exceptions by sending telemetry and flushing thread pool""" diff --git a/tests/unit/test_telemetry.py b/tests/unit/test_telemetry.py index f000eaa43..89f35228d 100644 --- a/tests/unit/test_telemetry.py +++ b/tests/unit/test_telemetry.py @@ -1,4 +1,10 @@ +import subprocess +import sys +import textwrap +import threading +import time import uuid +import weakref import pytest from unittest.mock import patch, MagicMock import json @@ -9,6 +15,7 @@ NoopTelemetryClient, TelemetryClientFactory, TelemetryHelper, + _DaemonThreadPoolExecutor, ) from databricks.sql.common.feature_flag import ( FeatureFlagsContextFactory, @@ -531,6 +538,104 @@ def test_connection_failure_does_not_send_telemetry_for_kernel( mock_export_failure_log.assert_not_called() +class TestDaemonThreadPoolExecutor: + """Telemetry sends run on daemon threads so they cannot block interpreter exit.""" + + def test_runs_work_on_bounded_daemon_threads(self): + executor = _DaemonThreadPoolExecutor(max_workers=2, thread_name_prefix="t") + release = threading.Event() + try: + futures = [executor.submit(release.wait, 5) for _ in range(5)] + assert len(executor._threads) == 2 + assert all(t.daemon for t in executor._threads) + release.set() + assert [f.result(timeout=5) for f in futures] == [True] * 5 + finally: + release.set() + executor.shutdown(wait=True) + + def test_result_and_exception_propagate(self): + executor = _DaemonThreadPoolExecutor(max_workers=1, thread_name_prefix="t") + try: + assert executor.submit(lambda x, y: x + y, 1, y=2).result(timeout=5) == 3 + with pytest.raises(ValueError): + executor.submit(int, "not-an-int").result(timeout=5) + finally: + executor.shutdown(wait=True) + + def test_shutdown_drains_queued_work_and_rejects_new_work(self): + executor = _DaemonThreadPoolExecutor(max_workers=1, thread_name_prefix="t") + futures = [executor.submit(time.sleep, 0.01) for _ in range(3)] + executor.shutdown(wait=True) + assert all(f.done() for f in futures) + with pytest.raises(RuntimeError): + executor.submit(time.sleep, 0) + + def test_wait_until_idle(self): + executor = _DaemonThreadPoolExecutor(max_workers=1, thread_name_prefix="t") + release = threading.Event() + try: + executor.submit(release.wait, 5) + assert executor.wait_until_idle(0.1) is False + release.set() + assert executor.wait_until_idle(5) is True + finally: + release.set() + executor.shutdown(wait=True) + + +class TestTelemetryExitWait: + def test_exit_wait_is_bounded_and_covers_shut_down_executors(self, monkeypatch): + monkeypatch.setattr(TelemetryClientFactory, "_executors", weakref.WeakSet()) + monkeypatch.setattr(TelemetryClientFactory, "EXIT_FLUSH_TIMEOUT_SECONDS", 0.3) + quick = _DaemonThreadPoolExecutor(max_workers=1, thread_name_prefix="t") + stuck = _DaemonThreadPoolExecutor(max_workers=1, thread_name_prefix="t") + TelemetryClientFactory._executors.add(quick) + TelemetryClientFactory._executors.add(stuck) + release = threading.Event() + try: + sent = quick.submit(time.sleep, 0.1) + quick.shutdown(wait=False) # as when the last connection closes + stuck.submit(release.wait, 30) + + start = time.monotonic() + TelemetryClientFactory._wait_for_pending_sends() + + assert sent.done() + assert time.monotonic() - start < 2 + finally: + release.set() + stuck.shutdown(wait=True) + + def test_interpreter_exit_waits_for_quick_send_but_not_stuck_one(self): + script = textwrap.dedent( + """ + import time + from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory as F + + F.EXIT_FLUSH_TIMEOUT_SECONDS = 1.0 + with F._lock: + F._initialize() + + def send(delay, label): + time.sleep(delay) + print(label, flush=True) + + F._executor.submit(send, 0.2, "quick-sent") + F._executor.submit(send, 60, "stuck-sent") + F._executor.shutdown(wait=False) + """ + ) + start = time.monotonic() + proc = subprocess.run( + [sys.executable, "-c", script], capture_output=True, text=True, timeout=30 + ) + assert proc.returncode == 0, proc.stderr + assert "quick-sent" in proc.stdout + assert "stuck-sent" not in proc.stdout + assert time.monotonic() - start < 15 + + @patch("databricks.sql.client.Session") class TestTelemetryFeatureFlag: """Tests the interaction between the telemetry feature flag and connection parameters.""" From 26efb872b8d7992b92093c4d5b5bfc177c980664 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 08:10:02 +0000 Subject: [PATCH 2/2] fix(telemetry): reset the daemon executor in forked children Sends in flight in the parent at fork time stayed counted in the child, so a child that sent its own telemetry always waited the full exit bound. Reset each executor's queue, workers, count and lock after fork (the parent still sends the work it queued), and start a worker before queueing so the count only covers work some worker will run. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- .../sql/telemetry/telemetry_client.py | 50 ++++++++++++------- tests/unit/test_telemetry.py | 31 ++++++++++-- 2 files changed, 61 insertions(+), 20 deletions(-) diff --git a/src/databricks/sql/telemetry/telemetry_client.py b/src/databricks/sql/telemetry/telemetry_client.py index 00792457b..a4615fd55 100644 --- a/src/databricks/sql/telemetry/telemetry_client.py +++ b/src/databricks/sql/telemetry/telemetry_client.py @@ -1,4 +1,5 @@ import atexit +import os import threading import time import logging @@ -472,33 +473,38 @@ class _DaemonThreadPoolExecutor: TelemetryClientFactory._wait_for_pending_sends). """ + # Every executor that may still run work, including ones already shut down. + _instances: weakref.WeakSet = weakref.WeakSet() + def __init__(self, max_workers: int, thread_name_prefix: str): self._max_workers = max_workers self._thread_name_prefix = thread_name_prefix self._work_queue: SimpleQueue = SimpleQueue() self._threads: List[threading.Thread] = [] - self._unfinished = 0 + self._unfinished = 0 # queued + running self._lock = threading.Lock() self._idle = threading.Condition(self._lock) self._shutdown = False + _DaemonThreadPoolExecutor._instances.add(self) def submit(self, fn, *args, **kwargs) -> Future: with self._lock: if self._shutdown: raise RuntimeError("cannot schedule new futures after shutdown") - future: Future = Future() - self._work_queue.put((future, fn, args, kwargs)) - self._unfinished += 1 - # Add a worker only when all live ones are busy (dead ones: forked child). - self._threads = [t for t in self._threads if t.is_alive()] - if len(self._threads) < min(self._unfinished, self._max_workers): + # All workers are busy: add one (before queueing, so a failed start + # leaves nothing counted that no worker will run). + num_threads = len(self._threads) + if self._unfinished >= num_threads and num_threads < self._max_workers: thread = threading.Thread( target=self._worker, - name=f"{self._thread_name_prefix}_{len(self._threads)}", + name=f"{self._thread_name_prefix}_{num_threads}", daemon=True, ) thread.start() self._threads.append(thread) + future: Future = Future() + self._work_queue.put((future, fn, args, kwargs)) + self._unfinished += 1 return future @staticmethod @@ -539,11 +545,24 @@ def shutdown(self, wait: bool = True): def wait_until_idle(self, timeout: float) -> bool: """Wait up to ``timeout`` seconds for submitted work to finish.""" with self._lock: - return self._idle.wait_for( - lambda: self._unfinished == 0 - or not any(t.is_alive() for t in self._threads), - timeout, - ) + return self._idle.wait_for(lambda: self._unfinished == 0, timeout) + + def _reset_after_fork(self): + # Workers don't survive fork, and the parent still sends the work it queued. + self._work_queue = SimpleQueue() + self._threads = [] + self._unfinished = 0 + self._lock = threading.Lock() + self._idle = threading.Condition(self._lock) + + @classmethod + def _reset_all_after_fork(cls): + for executor in list(cls._instances): + executor._reset_after_fork() + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_DaemonThreadPoolExecutor._reset_all_after_fork) class _TelemetryClientHolder: @@ -580,8 +599,6 @@ class TelemetryClientFactory: {} ) # Map of host_url -> TelemetryClientHolder _executor: Optional[_DaemonThreadPoolExecutor] = None - # Every executor that may still run sends, including ones already shut down. - _executors: weakref.WeakSet = weakref.WeakSet() _initialized: bool = False _lock = threading.RLock() # Thread safety for factory operations # used RLock instead of Lock to avoid deadlocks when garbage collection is triggered @@ -624,7 +641,6 @@ def _initialize(cls): cls._executor = _DaemonThreadPoolExecutor( max_workers=10, thread_name_prefix="telemetry-sender" ) - cls._executors.add(cls._executor) cls._install_exception_hook() cls._install_exit_hook() cls._start_flush_thread() @@ -679,7 +695,7 @@ def _install_exit_hook(cls): def _wait_for_pending_sends(cls): """At interpreter exit, give in-flight sends a bounded time to finish.""" deadline = time.monotonic() + cls.EXIT_FLUSH_TIMEOUT_SECONDS - for executor in list(cls._executors): + for executor in list(_DaemonThreadPoolExecutor._instances): if not executor.wait_until_idle(max(0.0, deadline - time.monotonic())): logger.debug("Dropping telemetry still in flight at interpreter exit") diff --git a/tests/unit/test_telemetry.py b/tests/unit/test_telemetry.py index 89f35228d..53cdae356 100644 --- a/tests/unit/test_telemetry.py +++ b/tests/unit/test_telemetry.py @@ -1,3 +1,4 @@ +import os import subprocess import sys import textwrap @@ -586,12 +587,10 @@ def test_wait_until_idle(self): class TestTelemetryExitWait: def test_exit_wait_is_bounded_and_covers_shut_down_executors(self, monkeypatch): - monkeypatch.setattr(TelemetryClientFactory, "_executors", weakref.WeakSet()) + monkeypatch.setattr(_DaemonThreadPoolExecutor, "_instances", weakref.WeakSet()) monkeypatch.setattr(TelemetryClientFactory, "EXIT_FLUSH_TIMEOUT_SECONDS", 0.3) quick = _DaemonThreadPoolExecutor(max_workers=1, thread_name_prefix="t") stuck = _DaemonThreadPoolExecutor(max_workers=1, thread_name_prefix="t") - TelemetryClientFactory._executors.add(quick) - TelemetryClientFactory._executors.add(stuck) release = threading.Event() try: sent = quick.submit(time.sleep, 0.1) @@ -635,6 +634,32 @@ def send(delay, label): assert "stuck-sent" not in proc.stdout assert time.monotonic() - start < 15 + @pytest.mark.skipif(not hasattr(os, "fork"), reason="requires os.fork") + def test_forked_child_does_not_wait_for_parent_sends(self): + script = textwrap.dedent( + """ + import os, sys, threading, time + from databricks.sql.telemetry.telemetry_client import _DaemonThreadPoolExecutor + + executor = _DaemonThreadPoolExecutor(max_workers=2, thread_name_prefix="t") + release = threading.Event() + executor.submit(release.wait, 30) # in flight in the parent at fork time + pid = os.fork() + if pid == 0: + sent = executor.submit(time.sleep, 0.05) + start = time.monotonic() + idle = executor.wait_until_idle(5) + os._exit(0 if idle and sent.done() and time.monotonic() - start < 2 else 1) + _, status = os.waitpid(pid, 0) + release.set() + sys.exit(os.waitstatus_to_exitcode(status)) + """ + ) + proc = subprocess.run( + [sys.executable, "-c", script], capture_output=True, text=True, timeout=30 + ) + assert proc.returncode == 0, proc.stderr + @patch("databricks.sql.client.Session") class TestTelemetryFeatureFlag: