Implement TLS session cache for faster reconnections

Co-authored-by: mykaul <[email protected]>

Author: Copilot

cassandra/cluster.py

@@ -875,6 +875,37 @@ def default_retry_policy(self, policy):
     .. versionadded:: 3.17.0
     """
 
+    tls_session_cache_enabled = True
+    """
+    Enable TLS session caching for faster reconnections. When enabled, TLS sessions
+    are cached per endpoint and reused for subsequent connections to the same server.
+    This reduces handshake latency and CPU usage during reconnections.
+    
+    Defaults to True when SSL/TLS is enabled. Set to False to disable session caching.
+    
+    .. versionadded:: 3.30.0
+    """
+
+    tls_session_cache_size = 100
+    """
+    Maximum number of TLS sessions to cache. When the cache is full, the least
+    recently used session is evicted.
+    
+    Defaults to 100.
+    
+    .. versionadded:: 3.30.0
+    """
+
+    tls_session_cache_ttl = 3600
+    """
+    Time-to-live for cached TLS sessions in seconds. Sessions older than this
+    are not reused and are removed from the cache.
+    
+    Defaults to 3600 seconds (1 hour).
+    
+    .. versionadded:: 3.30.0
+    """
+
     sockopts = None
     """
     An optional list of tuples which will be used as arguments to
@@ -1204,6 +1235,9 @@ def __init__(self,
                  idle_heartbeat_timeout=30,
                  no_compact=False,
                  ssl_context=None,
+                 tls_session_cache_enabled=True,
+                 tls_session_cache_size=100,
+                 tls_session_cache_ttl=3600,
                  endpoint_factory=None,
                  application_name=None,
                  application_version=None,
@@ -1420,6 +1454,19 @@ def __init__(self,
 
         self.ssl_options = ssl_options
         self.ssl_context = ssl_context
+        self.tls_session_cache_enabled = tls_session_cache_enabled
+        self.tls_session_cache_size = tls_session_cache_size
+        self.tls_session_cache_ttl = tls_session_cache_ttl
+        
+        # Initialize TLS session cache if SSL is enabled
+        self._tls_session_cache = None
+        if (ssl_context or ssl_options) and tls_session_cache_enabled:
+            from cassandra.connection import TLSSessionCache
+            self._tls_session_cache = TLSSessionCache(
+                max_size=tls_session_cache_size,
+                ttl=tls_session_cache_ttl
+            )
+        
         self.sockopts = sockopts
         self.cql_version = cql_version
         self.max_schema_agreement_wait = max_schema_agreement_wait
@@ -1661,6 +1708,7 @@ def _make_connection_kwargs(self, endpoint, kwargs_dict):
         kwargs_dict.setdefault('sockopts', self.sockopts)
         kwargs_dict.setdefault('ssl_options', self.ssl_options)
         kwargs_dict.setdefault('ssl_context', self.ssl_context)
+        kwargs_dict.setdefault('tls_session_cache', self._tls_session_cache)
         kwargs_dict.setdefault('cql_version', self.cql_version)
         kwargs_dict.setdefault('protocol_version', self.protocol_version)
         kwargs_dict.setdefault('user_type_map', self._user_types)

cassandra/connection.py

@@ -128,6 +128,108 @@ def decompress(byts):
 frame_header_v3 = struct.Struct('>BhBi')
 
 
+class TLSSessionCache:
+    """
+    Thread-safe cache for TLS sessions to enable session resumption.
+    
+    This cache stores TLS sessions per endpoint (host:port) to allow
+    quick TLS renegotiation when reconnecting to the same server.
+    Sessions are automatically expired after a TTL and the cache has
+    a maximum size with LRU eviction.
+    """
+    
+    def __init__(self, max_size=100, ttl=3600):
+        """
+        Initialize the TLS session cache.
+        
+        Args:
+            max_size: Maximum number of sessions to cache (default: 100)
+            ttl: Time-to-live for cached sessions in seconds (default: 3600)
+        """
+        self._sessions = {}  # {endpoint_key: (session, timestamp, access_time)}
+        self._lock = RLock()
+        self._max_size = max_size
+        self._ttl = ttl
+    
+    def _make_key(self, host, port):
+        """Create a cache key from host and port."""
+        return (host, port)
+    
+    def get_session(self, host, port):
+        """
+        Get a cached TLS session for the given endpoint.
+        
+        Args:
+            host: The hostname or IP address
+            port: The port number
+            
+        Returns:
+            ssl.SSLSession object if a valid cached session exists, None otherwise
+        """
+        key = self._make_key(host, port)
+        with self._lock:
+            if key not in self._sessions:
+                return None
+            
+            session, timestamp, _ = self._sessions[key]
+            
+            # Check if session has expired
+            if time.time() - timestamp > self._ttl:
+                del self._sessions[key]
+                return None
+            
+            # Update access time for LRU
+            self._sessions[key] = (session, timestamp, time.time())
+            return session
+    
+    def set_session(self, host, port, session):
+        """
+        Store a TLS session for the given endpoint.
+        
+        Args:
+            host: The hostname or IP address
+            port: The port number
+            session: The ssl.SSLSession object to cache
+        """
+        if session is None:
+            return
+        
+        key = self._make_key(host, port)
+        current_time = time.time()
+        
+        with self._lock:
+            # If cache is at max size, remove least recently used entry
+            if len(self._sessions) >= self._max_size and key not in self._sessions:
+                # Find entry with oldest access time
+                oldest_key = min(self._sessions.keys(), 
+                               key=lambda k: self._sessions[k][2])
+                del self._sessions[oldest_key]
+            
+            # Store session with creation time and access time
+            self._sessions[key] = (session, current_time, current_time)
+    
+    def clear_expired(self):
+        """Remove all expired sessions from the cache."""
+        current_time = time.time()
+        with self._lock:
+            expired_keys = [
+                key for key, (_, timestamp, _) in self._sessions.items()
+                if current_time - timestamp > self._ttl
+            ]
+            for key in expired_keys:
+                del self._sessions[key]
+    
+    def clear(self):
+        """Clear all sessions from the cache."""
+        with self._lock:
+            self._sessions.clear()
+    
+    def size(self):
+        """Return the current number of cached sessions."""
+        with self._lock:
+            return len(self._sessions)
+
+
 class EndPoint(object):
     """
     Represents the information to connect to a cassandra node.
@@ -687,6 +789,8 @@ class Connection(object):
     endpoint = None
     ssl_options = None
     ssl_context = None
+    tls_session_cache = None
+    session_reused = False
     last_error = None
 
     # The current number of operations that are in flight. More precisely,
@@ -763,14 +867,16 @@ def __init__(self, host='127.0.0.1', port=9042, authenticator=None,
                  ssl_options=None, sockopts=None, compression: Union[bool, str] = True,
                  cql_version=None, protocol_version=ProtocolVersion.MAX_SUPPORTED, is_control_connection=False,
                  user_type_map=None, connect_timeout=None, allow_beta_protocol_version=False, no_compact=False,
-                 ssl_context=None, owning_pool=None, shard_id=None, total_shards=None,
+                 ssl_context=None, tls_session_cache=None, owning_pool=None, shard_id=None, total_shards=None,
                  on_orphaned_stream_released=None, application_info: Optional[ApplicationInfoBase] = None):
         # TODO next major rename host to endpoint and remove port kwarg.
         self.endpoint = host if isinstance(host, EndPoint) else DefaultEndPoint(host, port)
 
         self.authenticator = authenticator
         self.ssl_options = ssl_options.copy() if ssl_options else {}
         self.ssl_context = ssl_context
+        self.tls_session_cache = tls_session_cache
+        self.session_reused = False
         self.sockopts = sockopts
         self.compression = compression
         self.cql_version = cql_version
@@ -913,7 +1019,28 @@ def _wrap_socket_from_context(self):
             server_hostname = self.endpoint.address
             opts['server_hostname'] = server_hostname
 
-        return self.ssl_context.wrap_socket(self._socket, **opts)
+        # Try to get a cached TLS session for resumption
+        if self.tls_session_cache:
+            cached_session = self.tls_session_cache.get_session(
+                self.endpoint.address, self.endpoint.port)
+            if cached_session:
+                opts['session'] = cached_session
+                log.debug("Using cached TLS session for %s:%s", 
+                         self.endpoint.address, self.endpoint.port)
+
+        ssl_socket = self.ssl_context.wrap_socket(self._socket, **opts)
+        
+        # Store the session for future reuse
+        if self.tls_session_cache and ssl_socket.session:
+            self.tls_session_cache.set_session(
+                self.endpoint.address, self.endpoint.port, ssl_socket.session)
+            # Track if the session was reused
+            self.session_reused = ssl_socket.session_reused
+            if self.session_reused:
+                log.debug("TLS session was reused for %s:%s", 
+                         self.endpoint.address, self.endpoint.port)
+        
+        return ssl_socket
 
     def _initiate_connection(self, sockaddr):
         if self.features.shard_id is not None: