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..a4615fd55 100644 --- a/src/databricks/sql/telemetry/telemetry_client.py +++ b/src/databricks/sql/telemetry/telemetry_client.py @@ -1,9 +1,11 @@ +import atexit +import os 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 +464,107 @@ 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). + """ + + # 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 # 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") + # 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}_{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 + 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, 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: """ Holds a telemetry client with reference counting. @@ -495,12 +598,16 @@ class TelemetryClientFactory: _clients: Dict[str, _TelemetryClientHolder] = ( {} ) # Map of host_url -> TelemetryClientHolder - _executor: Optional[ThreadPoolExecutor] = None + _executor: Optional[_DaemonThreadPoolExecutor] = None _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 +638,11 @@ 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._install_exception_hook() + cls._install_exit_hook() cls._start_flush_thread() cls._initialized = True logger.debug( @@ -577,6 +685,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(_DaemonThreadPoolExecutor._instances): + 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..53cdae356 100644 --- a/tests/unit/test_telemetry.py +++ b/tests/unit/test_telemetry.py @@ -1,4 +1,11 @@ +import os +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 +16,7 @@ NoopTelemetryClient, TelemetryClientFactory, TelemetryHelper, + _DaemonThreadPoolExecutor, ) from databricks.sql.common.feature_flag import ( FeatureFlagsContextFactory, @@ -531,6 +539,128 @@ 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(_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") + 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 + + @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: """Tests the interaction between the telemetry feature flag and connection parameters."""