connection: tolerate clean startup close

Author: dkropachev

cassandra/connection.py

@@ -986,8 +986,6 @@ def factory(cls, endpoint, timeout, host_conn = None, *args, **kwargs):
             conn.close()
             raise OperationTimedOut("Timed out creating connection (%s seconds)" % timeout,
                                     timeout=timeout)
-        elif conn.is_closed:
-            raise ConnectionShutdown("Connection to %s was closed by server" % conn.endpoint)
         else:
             return conn
 

tests/unit/test_connection.py

@@ -12,10 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import itertools
+import socket
 import unittest
 from io import BytesIO
 import time
-from threading import Lock
+from threading import Event, Lock, Thread
 from unittest.mock import Mock, ANY, call, patch
 
 from cassandra import OperationTimedOut
@@ -383,6 +384,102 @@ def test_wait_for_responses_shutdown_includes_last_error(self):
         assert "already closed" in error_message
         assert "Bad file descriptor" in error_message
 
+    def test_factory_returns_maintenance_mode_clean_startup_close(self):
+        """
+        Maintenance mode accepts regular CQL sockets and closes them during
+        startup. The factory must leave that clean close for the owning
+        pool/control connection instead of turning it into an immediate failure.
+        """
+
+        class MaintenanceModeCqlServer(object):
+            def __init__(self):
+                self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+                self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+                self._sock.bind(('127.0.0.1', 0))
+                self._sock.listen(1)
+                self._sock.settimeout(2)
+                self.port = self._sock.getsockname()[1]
+                self.first_frame = b''
+                self.ready = Event()
+                self.received_frame = Event()
+                self.error = None
+                self.thread = Thread(target=self._run)
+                self.thread.daemon = True
+                self.thread.start()
+
+            def _run(self):
+                self.ready.set()
+                try:
+                    client, _ = self._sock.accept()
+                    with client:
+                        client.settimeout(2)
+                        while len(self.first_frame) < 9:
+                            chunk = client.recv(9 - len(self.first_frame))
+                            if not chunk:
+                                break
+                            self.first_frame += chunk
+                except Exception as exc:
+                    self.error = exc
+                finally:
+                    self.received_frame.set()
+
+            def close(self):
+                self._sock.close()
+                self.thread.join(2)
+
+        class MaintenanceModeConnection(Connection):
+            def __init__(self, *args, **kwargs):
+                super(MaintenanceModeConnection, self).__init__(*args, **kwargs)
+                self._reader = None
+                self._connect_socket()
+                self._send_options_message()
+                self._reader = Thread(target=self._read_until_server_closes)
+                self._reader.daemon = True
+                self._reader.start()
+
+            def push(self, data):
+                self._socket.sendall(data)
+
+            def close(self):
+                with self.lock:
+                    if self.is_closed:
+                        return
+                    self.is_closed = True
+
+                if self._socket:
+                    self._socket.close()
+
+                if not self.is_defunct:
+                    self.error_all_requests(
+                        ConnectionShutdown("Connection to %s was closed" % self.endpoint))
+                    self.connected_event.set()
+
+            def _read_until_server_closes(self):
+                try:
+                    while True:
+                        data = self._socket.recv(self.in_buffer_size)
+                        if not data:
+                            self.close()
+                            return
+                        self._iobuf.write(data)
+                        self.process_io_buffer()
+                except socket.error as exc:
+                    if not self.is_closed:
+                        self.defunct(exc)
+
+        server = MaintenanceModeCqlServer()
+        try:
+            assert server.ready.wait(2)
+            conn = MaintenanceModeConnection.factory(
+                DefaultEndPoint('127.0.0.1', server.port), timeout=2)
+
+            assert conn.is_closed
+            assert server.received_frame.wait(2)
+            assert server.error is None
+            assert server.first_frame[4] == 0x05  # OPTIONS
+        finally:
+            server.close()
+
 
 @patch('cassandra.connection.ConnectionHeartbeat._raise_if_stopped')
 class ConnectionHeartbeatTest(unittest.TestCase):