perf(telemetry): don't let slow or rate-limited telemetry delay interpreter exit

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.

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"""

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."""