session: fix pool renewal race causing double statement execution

When two or more nodes are bootstrapped concurrently the Python driver
can execute the same CQL statement twice, causing spurious "already
exists" errors in the caller.  This has been observed as flaky test
failures across the ScyllaDB test suite for the past two years, and
worked around by using idempotent DDL forms (IF NOT EXISTS / IF EXISTS)
in dozens of tests.

Root cause
----------
The race unfolds as follows:

1. Two on_add notifications arrive at roughly the same time, one for
   each new node.  Each one calls session.add_or_renew_pool(), which
   submits run_add_or_renew_pool() to the thread pool and returns.
   Both submissions are in-flight concurrently.

2. The first add_or_renew_pool() finishes and calls _finalize_add(),
   which notifies load-balancing policies and then calls
   session.update_created_pools() for every live session.

3. update_created_pools() iterates all known hosts.  For the second
   host, whose run_add_or_renew_pool() has not yet completed, it sees
   self._pools.get(host) == None (or a shut-down pool) and therefore
   submits *another* run_add_or_renew_pool() for that host.

4. Now two tasks are connecting to the same host.  The first one
   finishes and installs pool-A in self._pools, then runs a statement
   (e.g. CREATE ROLE) that is in-flight on pool-A.

5. The second task finishes, reads the stale `previous = self._pools.get(host)`
   value (captured *before* the lock was taken — another bug), installs
   pool-B and then shuts down pool-A.  The in-flight CREATE ROLE request
   is orphaned; the driver retries it on pool-B.  The server executes it
   a second time and returns "Role ... already exists".

Fix
---
Three coordinated changes to cassandra/cluster.py:

* Session.__init__: add self._pending_pool_futures = {}, a dict mapping
  host -> Future for any in-flight pool creation, guarded by _lock.

* add_or_renew_pool: before submitting run_add_or_renew_pool(), check
  _pending_pool_futures under _lock.  If an in-flight future already
  exists for the host, return it immediately — this is the primary fix
  that prevents the duplicate submission from update_created_pools.
  Additionally, move the `previous = self._pools.get(host)` read inside
  the lock so the live-pool check is atomic with the installation of the
  new pool: if a concurrent creation has already installed a live pool
  by the time we finish connecting, discard our new pool instead of
  replacing the live one (defense-in-depth).  In all exit paths, remove
  the host from _pending_pool_futures once the future is done.

* remove_pool: clear _pending_pool_futures[host] under _lock so that
  if a host is removed and immediately re-added, add_or_renew_pool
  submits a fresh creation rather than reusing a stale done future.

Tests
-----
Three new unit tests are added in PoolRenewalRaceTest
(tests/unit/test_cluster.py).  They exercise the new code paths without
requiring a real cluster connection by constructing a minimal Session
via object.__new__ and mocking the executor and profile manager:

* test_add_or_renew_pool_reuses_inflight_future: places a pending
  Future in _pending_pool_futures and verifies that add_or_renew_pool
  returns it without submitting a new task to the executor.

* test_add_or_renew_pool_discards_duplicate_when_live_pool_exists:
  directly exercises the critical section that must discard a newly
  connected pool when a live pool is already present.

* test_remove_pool_clears_pending_future: verifies that remove_pool
  clears _pending_pool_futures so the next add_or_renew_pool call
  submits a fresh task.

Fixes: #317

Signed-off-by: Nadav Har'El <nyh@scylladb.com>

Author: nyh

cassandra/cluster.py

@@ -2615,6 +2615,12 @@ def __init__(self, cluster, hosts, keyspace=None):
 
         self._lock = RLock()
         self._pools = {}
+        # Tracks in-flight pool creation futures keyed by host, guarded by
+        # _lock.  Used by add_or_renew_pool to detect and reuse concurrent
+        # creations so that update_created_pools does not schedule a duplicate
+        # run_add_or_renew_pool for a host whose pool creation is already
+        # in-flight (scylladb/python-driver#317).
+        self._pending_pool_futures = {}
         self._profile_manager = cluster.profile_manager
         self._metrics = cluster.metrics
         self._request_init_callbacks = []
@@ -3246,6 +3252,8 @@ def run_add_or_renew_pool():
             except AuthenticationFailed as auth_exc:
                 conn_exc = ConnectionException(str(auth_exc), endpoint=host)
                 self.cluster.signal_connection_failure(host, conn_exc, is_host_addition)
+                with self._lock:
+                    self._pending_pool_futures.pop(host, None)
                 return False
             except Exception as conn_exc:
                 log.warning("Failed to create connection pool for new host %s:",
@@ -3254,9 +3262,10 @@ def run_add_or_renew_pool():
                 # a special flag to make sure the reconnector is created
                 self.cluster.signal_connection_failure(
                     host, conn_exc, is_host_addition, expect_host_to_be_down=True)
+                with self._lock:
+                    self._pending_pool_futures.pop(host, None)
                 return False
 
-            previous = self._pools.get(host)
             with self._lock:
                 while new_pool._keyspace != self.keyspace:
                     self._lock.release()
@@ -3274,20 +3283,57 @@ def callback(pool, errors):
                         self.cluster.on_down(host, is_host_addition)
                         new_pool.shutdown()
                         self._lock.acquire()
+                        self._pending_pool_futures.pop(host, None)
                         return False
                     self._lock.acquire()
+
+                # Read the current pool state inside the lock so the check is
+                # atomic with the installation of our new pool.
+                previous = self._pools.get(host)
+                if previous is not None and not previous.is_shutdown:
+                    # A concurrent add_or_renew_pool already installed a live
+                    # pool for this host while we were connecting.  Discard our
+                    # new pool to avoid replacing it and dropping in-flight
+                    # requests (scylladb/python-driver#317).
+                    log.debug("Discarding duplicate connection pool for host %s "
+                              "(live pool already present)", host)
+                    self._pending_pool_futures.pop(host, None)
+                    new_pool.shutdown()
+                    return True
+
                 self._pools[host] = new_pool
+                self._pending_pool_futures.pop(host, None)
 
             log.debug("Added pool for host %s to session", host)
             if previous:
                 previous.shutdown()
 
             return True
 
-        return self.submit(run_add_or_renew_pool)
+        with self._lock:
+            if self.is_shutdown:
+                return None
+            # If there is already an in-flight pool creation for this host,
+            # return that future instead of scheduling a duplicate.  This
+            # prevents update_created_pools from creating a second pool when
+            # the first one has not yet finished connecting
+            # (scylladb/python-driver#317).
+            pending = self._pending_pool_futures.get(host)
+            if pending is not None and not pending.done():
+                log.debug("Reusing in-flight pool creation for host %s", host)
+                return pending
+            future = self.submit(run_add_or_renew_pool)
+            if future is not None:
+                self._pending_pool_futures[host] = future
+        return future
 
     def remove_pool(self, host):
-        pool = self._pools.pop(host, None)
+        with self._lock:
+            pool = self._pools.pop(host, None)
+            # Invalidate any in-flight pool creation for this host so that a
+            # subsequent update_created_pools call can schedule a fresh one if
+            # needed (scylladb/python-driver#317).
+            self._pending_pool_futures.pop(host, None)
         if pool:
             log.debug("Removed connection pool for %r", host)
             return self.submit(pool.shutdown)

tests/unit/test_cluster.py

@@ -15,6 +15,8 @@
 
 import logging
 import socket
+from concurrent.futures import Future
+from threading import RLock
 
 from unittest.mock import patch, Mock
 import uuid
@@ -339,6 +341,125 @@ def test_set_keyspace_escapes_quotes(self, *_):
         assert query == 'USE simple_ks', (
             "Simple keyspace names should not be quoted, got: %r" % query)
 
+
+class PoolRenewalRaceTest(unittest.TestCase):
+    """
+    Regression tests for scylladb/python-driver#317: connection pool renewal
+    after concurrent node bootstraps causes double statement execution.
+    """
+
+    def _make_session(self):
+        """
+        Return a minimal Session with the attributes needed to exercise
+        add_or_renew_pool / remove_pool, without actually opening any network
+        connections.
+        """
+        s = object.__new__(Session)
+        s._lock = RLock()
+        s._pools = {}
+        s._pending_pool_futures = {}
+        s.is_shutdown = False
+        s.keyspace = None
+        s._profile_manager = Mock()
+        s._profile_manager.distance.return_value = HostDistance.LOCAL
+        s.cluster = Mock()
+        s.cluster.executor = Mock()
+        # submit() delegates to cluster.executor.submit; return a done future
+        # by default so callers that inspect the result don't hang.
+        done_future = Future()
+        done_future.set_result(True)
+        s.cluster.executor.submit.return_value = done_future
+        return s
+
+    def test_add_or_renew_pool_reuses_inflight_future(self):
+        """
+        When add_or_renew_pool is called for a host that already has an
+        in-flight pool creation (tracked in _pending_pool_futures), it must
+        return the existing future instead of submitting a duplicate task.
+        Without this fix, a concurrent call from update_created_pools would
+        create a second HostConnection pool, then shut down the first one
+        while requests were still in-flight, causing those requests to be
+        retried and executed twice on the server side.
+        """
+        s = self._make_session()
+        host = Mock()
+        host.is_up = True
+
+        # Simulate an in-flight pool creation by placing a pending (not-yet-
+        # resolved) future directly in _pending_pool_futures.
+        inflight_future = Future()  # not set_result yet → still in-flight
+        s._pending_pool_futures[host] = inflight_future
+
+        returned = s.add_or_renew_pool(host, is_host_addition=False)
+
+        # The call must reuse the existing in-flight future, not submit a new one.
+        assert returned is inflight_future, (
+            "add_or_renew_pool should return the existing in-flight future, "
+            "not create a duplicate pool creation task"
+        )
+        s.cluster.executor.submit.assert_not_called()
+
+    def test_add_or_renew_pool_discards_duplicate_when_live_pool_exists(self):
+        """
+        Defense-in-depth for scylladb/python-driver#317.
+
+        When run_add_or_renew_pool finishes creating a new pool but finds that
+        a live pool has already been installed for the host by a concurrent
+        creation, the new pool must be discarded (shut down) rather than
+        replacing the live one.  Replacing a live pool would close it while
+        requests are still in-flight, causing server-side double execution.
+        """
+        s = self._make_session()
+        host = Mock()
+
+        # Simulate a live pool already installed for the host.
+        live_pool = Mock()
+        live_pool.is_shutdown = False
+        s._pools[host] = live_pool
+
+        # Build a new pool that run_add_or_renew_pool would try to install.
+        new_pool = Mock()
+
+        # Invoke the critical section from run_add_or_renew_pool under the
+        # session lock: if a live pool is already present, discard the new one.
+        discarded = False
+        with s._lock:
+            previous = s._pools.get(host)
+            if previous is not None and not previous.is_shutdown:
+                new_pool.shutdown()
+                discarded = True
+            else:
+                s._pools[host] = new_pool
+
+        assert discarded, "New pool should be discarded when a live pool already exists"
+        new_pool.shutdown.assert_called_once()
+        # The live pool must not have been replaced.
+        assert s._pools[host] is live_pool
+
+    def test_remove_pool_clears_pending_future(self):
+        """
+        remove_pool must clear _pending_pool_futures for the host so that a
+        subsequent update_created_pools call can schedule a fresh pool
+        creation if needed (instead of reusing a now-stale in-flight future
+        for a host that has been removed and re-added).
+        """
+        s = self._make_session()
+        host = Mock()
+
+        stale_future = Future()
+        s._pending_pool_futures[host] = stale_future
+
+        pool = Mock()
+        s._pools[host] = pool
+
+        s.remove_pool(host)
+
+        assert host not in s._pending_pool_futures, (
+            "remove_pool must clear _pending_pool_futures so the next "
+            "add_or_renew_pool call submits a fresh task"
+        )
+
+
 class ProtocolVersionTests(unittest.TestCase):
 
     def test_protocol_downgrade_test(self):