PYTHON-1419 Connection failure to SNI endpoint when first host is unavailable

cassandra/connection.py

@@ -245,9 +245,9 @@ def create(self, row):
 class SniEndPoint(EndPoint):
     """SNI Proxy EndPoint implementation."""
 
-    def __init__(self, proxy_address, server_name, port=9042):
+    def __init__(self, proxy_address, server_name, port=9042, init_index=0):
         self._proxy_address = proxy_address
-        self._index = 0
+        self._index = init_index
         self._resolved_address = None  # resolved address
         self._port = port
         self._server_name = server_name
@@ -267,8 +267,7 @@ def ssl_options(self):
 
     def resolve(self):
         try:
-            resolved_addresses = socket.getaddrinfo(self._proxy_address, self._port,
-                                                    socket.AF_UNSPEC, socket.SOCK_STREAM)
+            resolved_addresses = self._resolve_proxy_addresses()
         except socket.gaierror:
             log.debug('Could not resolve sni proxy hostname "%s" '
                       'with port %d' % (self._proxy_address, self._port))
@@ -280,6 +279,10 @@ def resolve(self):
 
         return self._resolved_address, self._port
 
+    def _resolve_proxy_addresses(self):
+        return socket.getaddrinfo(self._proxy_address, self._port,
+                                  socket.AF_UNSPEC, socket.SOCK_STREAM)
+
     def __eq__(self, other):
         return (isinstance(other, SniEndPoint) and
                 self.address == other.address and self.port == other.port and
@@ -305,16 +308,24 @@ class SniEndPointFactory(EndPointFactory):
     def __init__(self, proxy_address, port):
         self._proxy_address = proxy_address
         self._port = port
+        # Initial lookup index to prevent all SNI endpoints to be resolved
+        # into the same starting IP address (which might not be available currently).
+        # If SNI resolves to 3 IPs, first endpoint will connect to first
+        # IP address, and subsequent resolutions to next IPs in round-robin
+        # fusion.
+        self._init_index = -1
 
     def create(self, row):
         host_id = row.get("host_id")
         if host_id is None:
             raise ValueError("No host_id to create the SniEndPoint")
 
-        return SniEndPoint(self._proxy_address, str(host_id), self._port)
+        self._init_index += 1
+        return SniEndPoint(self._proxy_address, str(host_id), self._port, self._init_index)
 
     def create_from_sni(self, sni):
-        return SniEndPoint(self._proxy_address, sni, self._port)
+        self._init_index += 1
+        return SniEndPoint(self._proxy_address, sni, self._port, self._init_index)
 
 
 @total_ordering

tests/integration/long/test_sni_connection.py

@@ -0,0 +1,94 @@
+# Copyright DataStax, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import socket
+import ssl
+import unittest
+from unittest.mock import patch
+
+from cassandra import ProtocolVersion
+from cassandra.cluster import Cluster, ControlConnection, ProfileManager
+from cassandra.connection import SniEndPoint
+from cassandra.datastax.cloud import CloudConfig
+from cassandra.policies import ConstantReconnectionPolicy
+from tests.integration import (
+    get_cluster, remove_cluster, CASSANDRA_IP
+)
+from tests.integration.long.test_ssl import setup_cluster_ssl, CLIENT_CA_CERTS, DRIVER_CERTFILE, DRIVER_KEYFILE
+
+
+class SniConnectionTests(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        setup_cluster_ssl(client_auth=True)
+
+    @classmethod
+    def tearDownClass(cls):
+        ccm_cluster = get_cluster()
+        ccm_cluster.stop()
+        remove_cluster()
+
+    def _mocked_cloud_config(self, cloud_config, create_pyopenssl_context):
+        config = CloudConfig.from_dict({})
+        config.sni_host = 'proxy.datastax.com'
+        config.sni_port = 9042
+        config.host_ids = ['8c4b6ed7-f505-4226-b7a4-41f322520c1f', '2e25021d-8d72-41a7-a247-3da85c5d92d2']
+
+        ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS)
+        ssl_context.load_verify_locations(CLIENT_CA_CERTS)
+        ssl_context.verify_mode = ssl.CERT_REQUIRED
+        ssl_context.load_cert_chain(certfile=DRIVER_CERTFILE, keyfile=DRIVER_KEYFILE)
+        config.ssl_context = ssl_context
+
+        return config
+
+    def _mocked_proxy_dns_resolution(self):
+        return [
+            # return wrong IP at first position, so that we make sure all SNI endpoints
+            # do not start with first IP only
+            (socket.AF_UNIX, socket.SOCK_STREAM, 0, None, ('100.101.102.103', 9042)),
+            (socket.AF_UNIX, socket.SOCK_STREAM, 0, None, (CASSANDRA_IP, 9042))
+        ]
+
+    def _mocked_refresh_node_list_and_token_map(self, connection, preloaded_results=None,
+                                                force_token_rebuild=False):
+        return
+
+    def _mocked_refresh_schema(self, connection, preloaded_results=None, schema_agreement_wait=None,
+                               force=False, **kwargs):
+        return
+
+    def _mocked_check_supported(self):
+        return
+
+    # Tests verifies that driver can connect to SNI endpoint even when one IP
+    # returned by the DNS resolution of SNI does not respond. Mocked SNI resolution method
+    # returns two IPs where only one corresponds to online C* cluster started with CCM.
+    def test_round_robin_dns_resolution(self):
+        with patch('cassandra.datastax.cloud.get_cloud_config', self._mocked_cloud_config):
+            with patch.object(SniEndPoint, '_resolve_proxy_addresses', self._mocked_proxy_dns_resolution):
+                # Mock below three functions, because host ID returned from proxy will not match ID present in C*
+                # Network connection should be already made, so we can consider our test successful
+                with patch.object(ControlConnection, '_refresh_node_list_and_token_map',
+                                  self._mocked_refresh_node_list_and_token_map):
+                    with patch.object(ControlConnection, '_refresh_schema',
+                                      self._mocked_refresh_schema):
+                        with patch.object(ProfileManager, 'check_supported', self._mocked_check_supported):
+                            cloud_config = {
+                                'secure_connect_bundle': '/path/to/secure-connect-dbname.zip'
+                            }
+                            cluster = Cluster(cloud=cloud_config, protocol_version=ProtocolVersion.V4, reconnection_policy=ConstantReconnectionPolicy(10))
+                            session = cluster.connect()
+                            session.shutdown()
+                            cluster.shutdown()

tests/unit/test_endpoints.py

@@ -65,3 +65,15 @@ def test_endpoint_resolve(self):
         for i in range(10):
             (address, _) = endpoint.resolve()
             self.assertEqual(address, next(it))
+
+    def test_sni_resolution_start_index(self):
+        factory = SniEndPointFactory("proxy.datastax.com", 9999)
+        initial_index = factory._init_index
+
+        endpoint1 = factory.create_from_sni('sni1')
+        self.assertEqual(factory._init_index, initial_index + 1)
+        self.assertEqual(endpoint1._index, factory._init_index)
+
+        endpoint2 = factory.create_from_sni('sni2')
+        self.assertEqual(factory._init_index, initial_index + 2)
+        self.assertEqual(endpoint2._index, factory._init_index)