Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 15 additions & 14 deletions src/snowflake/connector/crl_cache.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@
import platform
import threading
from abc import ABC, abstractmethod
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from pathlib import Path
Expand Down Expand Up @@ -426,7 +425,7 @@ class CRLCacheFactory:
_instance_lock = threading.RLock()

# Cleanup management
_cleanup_executor: ThreadPoolExecutor | None = None
_cleanup_thread: threading.Thread | None = None
_cleanup_shutdown: threading.Event = threading.Event()
_cleanup_interval: timedelta | None = None
_atexit_registered: bool = False
Expand Down Expand Up @@ -504,17 +503,19 @@ def start_periodic_cleanup(cls, cleanup_interval: timedelta) -> None:
cls.stop_periodic_cleanup()

cls._cleanup_interval = cleanup_interval
cls._cleanup_executor = ThreadPoolExecutor(
max_workers=1, thread_name_prefix="crl-cache-cleanup"
cls._cleanup_thread = threading.Thread(
target=cls._cleanup_loop,
name="crl-cache-cleanup",
daemon=True, # Make it a daemon thread so it doesn't block program exit
)

# Register atexit handler for graceful shutdown (only once)
if not cls._atexit_registered:
atexit.register(cls._atexit_cleanup_handler)
cls._atexit_registered = True

# Submit the cleanup task
cls._cleanup_executor.submit(cls._cleanup_loop)
# Start the cleanup thread
cls._cleanup_thread.start()

logger.debug(
f"Scheduled CRL cache cleanup task to run every {cleanup_interval.total_seconds()} seconds."
Expand All @@ -523,29 +524,29 @@ def start_periodic_cleanup(cls, cleanup_interval: timedelta) -> None:
@classmethod
def stop_periodic_cleanup(cls) -> None:
"""Stop the periodic cleanup task."""
executor_to_shutdown = None
thread_to_join = None

with cls._instance_lock:
if cls._cleanup_executor is None or cls._cleanup_shutdown.is_set():
if cls._cleanup_thread is None or cls._cleanup_shutdown.is_set():
return

cls._cleanup_shutdown.set()
executor_to_shutdown = cls._cleanup_executor
thread_to_join = cls._cleanup_thread

# Shutdown outside of lock to avoid deadlock
if executor_to_shutdown is not None:
executor_to_shutdown.shutdown(wait=True)
# Join thread outside of lock to avoid deadlock
if thread_to_join is not None and thread_to_join.is_alive():
thread_to_join.join(timeout=5.0)

with cls._instance_lock:
cls._cleanup_shutdown.clear()
cls._cleanup_executor = None
cls._cleanup_thread = None
cls._cleanup_interval = None

@classmethod
def is_periodic_cleanup_running(cls) -> bool:
"""Check if periodic cleanup task is running."""
with cls._instance_lock:
return cls._cleanup_executor is not None
return cls._cleanup_thread is not None and cls._cleanup_thread.is_alive()

@classmethod
def _cleanup_loop(cls) -> None:
Expand Down
2 changes: 1 addition & 1 deletion test/unit/test_crl_cache.py
Original file line number Diff line number Diff line change
Expand Up @@ -436,7 +436,7 @@ def test_cleanup_loop_double_stop_is_safe(cache_factory):
def test_cleanup_loop_double_start_is_safe_and_restarts(
cache_factory, mem_cache_mock, disk_cache_mock
):
"""Test that calling start_periodic_cleanup multiple times creates new executors"""
"""Test that calling start_periodic_cleanup multiple times creates new threads"""
# Set up singleton instances to be cleaned
cache_factory._memory_cache_instance = mem_cache_mock
cache_factory._file_cache_instance = disk_cache_mock
Expand Down
Loading