Fix MongoReplicaSetClient race conditions PYTHON-467

* RSState, Member, and MovingAverage are now immutable
* In refresh(), try up members before down ones
* A test_ha fixup (clear process-list after killing them in each test, so we don't re-kill previous tests' processes)

Author: A. Jesse Jiryu Davis

pymongo/mongo_replica_set_client.py

@@ -15,8 +15,6 @@
 """Utilities for choosing which member of a replica set to read from."""
 
 import random
-import threading
-from collections import deque
 
 from pymongo.errors import ConfigurationError
 
@@ -86,7 +84,6 @@ def mongos_mode(mode):
 def mongos_enum(enum):
     return _mongos_modes.index(enum)
 
-
 def select_primary(members):
     for member in members:
         if member.is_primary:
@@ -151,6 +148,7 @@ def select_member(
         return select_primary(members)
 
     elif mode == PRIMARY_PREFERRED:
+        # Recurse.
         candidate_primary = select_member(members, PRIMARY, [{}], latency)
         if candidate_primary:
             return candidate_primary
@@ -166,6 +164,7 @@ def select_member(
         return None
 
     elif mode == SECONDARY_PREFERRED:
+        # Recurse.
         candidate_secondary = select_member(
             members, SECONDARY, tag_sets, latency)
         if candidate_secondary:
@@ -196,32 +195,16 @@ def select_member(
 
 
 class MovingAverage(object):
-    """Tracks a moving average.
-    """
-    def __init__(self, window_sz):
-        self.window_sz = window_sz
-        self.samples = deque()
-        self.total = 0
-        self.lock = threading.Lock()
-
-    def update(self, sample):
-        # One reason we synchronize MovingAverage is that Jython's
-        # popleft isn't safe: http://bugs.jython.org/issue2001
-        self.lock.acquire()
-        try:
-            self.samples.append(sample)
-            self.total += sample
-            if len(self.samples) > self.window_sz:
-                self.total -= self.samples.popleft()
-        finally:
-            self.lock.release()
+    def __init__(self, samples):
+        """Immutable structure to track a 5-sample moving average.
+        """
+        self.samples = samples[-5:]
+        assert self.samples
+        self.average = sum(self.samples) / float(len(self.samples))
+
+    def clone_with(self, sample):
+        """Get a copy of this instance plus a new sample"""
+        return MovingAverage(self.samples + [sample])
 
     def get(self):
-        self.lock.acquire()
-        try:
-            if self.samples:
-                return self.total / float(len(self.samples))
-            else:
-                return None
-        finally:
-            self.lock.release()
+        return self.average

pymongo/read_preferences.py

@@ -20,14 +20,17 @@
 import pymongo
 from pymongo.errors import ConnectionFailure
 
-host = os.environ.get("DB_IP", 'localhost')
+# hostnames retrieved by MongoReplicaSetClient from isMaster will be of unicode
+# type in Python 2, so ensure these hostnames are unicodes, too. It makes tests
+# like `test_repr` predictable.
+host = unicode(os.environ.get("DB_IP", 'localhost'))
 port = int(os.environ.get("DB_PORT", 27017))
 pair = '%s:%d' % (host, port)
 
-host2 = os.environ.get("DB_IP2", 'localhost')
+host2 = unicode(os.environ.get("DB_IP2", 'localhost'))
 port2 = int(os.environ.get("DB_PORT2", 27018))
 
-host3 = os.environ.get("DB_IP3", 'localhost')
+host3 = unicode(os.environ.get("DB_IP3", 'localhost'))
 port3 = int(os.environ.get("DB_PORT3", 27019))
 
 

test/__init__.py

@@ -27,14 +27,11 @@
 from ha_tools import use_greenlets
 
 
-from pymongo import (MongoReplicaSetClient,
-                     ReadPreference)
-from pymongo.mongo_replica_set_client import (
-    Member, Monitor, MongoReplicaSetClient)
-from pymongo.mongo_client import _partition_node
-from pymongo.mongo_client import MongoClient
 from pymongo.errors import AutoReconnect, OperationFailure, ConnectionFailure
-from pymongo.read_preferences import modes
+from pymongo.mongo_replica_set_client import Member, Monitor
+from pymongo.mongo_replica_set_client import MongoReplicaSetClient
+from pymongo.mongo_client import MongoClient, _partition_node
+from pymongo.read_preferences import ReadPreference, modes
 
 from test import utils
 from test.utils import one
@@ -52,7 +49,15 @@
 NEAREST = ReadPreference.NEAREST
 
 
-class TestDirectConnection(unittest.TestCase):
+class HATestCase(unittest.TestCase):
+    """A test case for connections to replica sets or mongos."""
+
+    def tearDown(self):
+        ha_tools.kill_all_members()
+        ha_tools.nodes.clear()
+
+
+class TestDirectConnection(HATestCase):
 
     def setUp(self):
         members = [{}, {}, {'arbiterOnly': True}]
@@ -145,10 +150,10 @@ def test_secondary_connection(self):
 
     def tearDown(self):
         self.c.close()
-        ha_tools.kill_all_members()
+        super(TestDirectConnection, self).tearDown()
 
 
-class TestPassiveAndHidden(unittest.TestCase):
+class TestPassiveAndHidden(HATestCase):
 
     def setUp(self):
         members = [{},
@@ -177,10 +182,10 @@ def test_passive_and_hidden(self):
 
     def tearDown(self):
         self.c.close()
-        ha_tools.kill_all_members()
+        super(TestPassiveAndHidden, self).tearDown()
 
 
-class TestMonitorRemovesRecoveringMember(unittest.TestCase):
+class TestMonitorRemovesRecoveringMember(HATestCase):
     # Members in STARTUP2 or RECOVERING states are shown in the primary's
     # isMaster response, but aren't secondaries and shouldn't be read from.
     # Verify that if a secondary goes into RECOVERING mode, the Monitor removes
@@ -211,10 +216,10 @@ def test_monitor_removes_recovering_member(self):
 
     def tearDown(self):
         self.c.close()
-        ha_tools.kill_all_members()
+        super(TestMonitorRemovesRecoveringMember, self).tearDown()
 
 
-class TestTriggeredRefresh(unittest.TestCase):
+class TestTriggeredRefresh(HATestCase):
     # Verify that if a secondary goes into RECOVERING mode or if the primary
     # changes, the next exception triggers an immediate refresh.
 
@@ -282,14 +287,14 @@ def test_stepdown_triggers_refresh(self):
         # We've detected the stepdown
         self.assertTrue(
             not c_find_one.primary
-            or primary != _partition_node(c_find_one.primary))
+            or _partition_node(primary) != c_find_one.primary)
 
     def tearDown(self):
         Monitor._refresh_interval = MONITOR_INTERVAL
-        ha_tools.kill_all_members()
+        super(TestTriggeredRefresh, self).tearDown()
 
 
-class TestHealthMonitor(unittest.TestCase):
+class TestHealthMonitor(HATestCase):
 
     def setUp(self):
         res = ha_tools.start_replica_set([{}, {}, {}])
@@ -357,17 +362,10 @@ def primary_changed():
 
         ha_tools.stepdown_primary()
         self.assertTrue(primary_changed())
-
-        # There can be a delay between finding the primary and updating
-        # secondaries
-        sleep(5)
         self.assertNotEqual(secondaries, c.secondaries)
 
-    def tearDown(self):
-        ha_tools.kill_all_members()
-
 
-class TestWritesWithFailover(unittest.TestCase):
+class TestWritesWithFailover(HATestCase):
 
     def setUp(self):
         res = ha_tools.start_replica_set([{}, {}, {}])
@@ -398,11 +396,8 @@ def try_write():
         self.assertTrue(primary != c.primary)
         self.assertEqual('baz', db.test.find_one({'bar': 'baz'})['bar'])
 
-    def tearDown(self):
-        ha_tools.kill_all_members()
-
 
-class TestReadWithFailover(unittest.TestCase):
+class TestReadWithFailover(HATestCase):
 
     def setUp(self):
         res = ha_tools.start_replica_set([{}, {}, {}])
@@ -435,11 +430,8 @@ def iter_cursor(cursor):
         self.assertTrue(iter_cursor(cursor))
         self.assertEqual(10, cursor._Cursor__retrieved)
 
-    def tearDown(self):
-        ha_tools.kill_all_members()
 
-
-class TestReadPreference(unittest.TestCase):
+class TestReadPreference(HATestCase):
     def setUp(self):
         members = [
             # primary
@@ -784,11 +776,10 @@ def test_pinning(self):
 
     def tearDown(self):
         self.c.close()
-        ha_tools.kill_all_members()
-        self.clear_ping_times()
+        super(TestReadPreference, self).tearDown()
 
 
-class TestReplicaSetAuth(unittest.TestCase):
+class TestReplicaSetAuth(HATestCase):
     def setUp(self):
         members = [
             {},
@@ -829,10 +820,10 @@ def test_auth_during_failover(self):
 
     def tearDown(self):
         self.c.close()
-        ha_tools.kill_all_members()
+        super(TestReplicaSetAuth, self).tearDown()
 
 
-class TestAlive(unittest.TestCase):
+class TestAlive(HATestCase):
     def setUp(self):
         members = [{}, {}]
         self.seed, self.name = ha_tools.start_replica_set(members)
@@ -866,11 +857,8 @@ def test_alive(self):
         finally:
             rsc.close()
 
-    def tearDown(self):
-        ha_tools.kill_all_members()
 
-        
-class TestMongosHighAvailability(unittest.TestCase):
+class TestMongosHighAvailability(HATestCase):
     def setUp(self):
         seed_list = ha_tools.create_sharded_cluster()
         self.dbname = 'pymongo_mongos_ha'
@@ -910,10 +898,10 @@ def test_mongos_ha(self):
 
     def tearDown(self):
         self.client.drop_database(self.dbname)
-        ha_tools.kill_all_members()
+        super(TestMongosHighAvailability, self).tearDown()
 
 
-class TestReplicaSetRequest(unittest.TestCase):
+class TestReplicaSetRequest(HATestCase):
     def setUp(self):
         members = [{}, {}, {'arbiterOnly': True}]
         res = ha_tools.start_replica_set(members)
@@ -928,8 +916,9 @@ def test_request_during_failover(self):
         self.assertTrue(self.c.auto_start_request)
         self.assertTrue(self.c.in_request())
 
-        primary_pool = self.c._MongoReplicaSetClient__members[primary].pool
-        secondary_pool = self.c._MongoReplicaSetClient__members[secondary].pool
+        rs_state = self.c._MongoReplicaSetClient__rs_state
+        primary_pool = rs_state.get(primary).pool
+        secondary_pool = rs_state.get(secondary).pool
 
         # Trigger start_request on primary pool
         utils.assertReadFrom(self, self.c, primary, PRIMARY)
@@ -967,7 +956,7 @@ def test_request_during_failover(self):
 
     def tearDown(self):
         self.c.close()
-        ha_tools.kill_all_members()
+        super(TestReplicaSetRequest, self).tearDown()
 
 
 if __name__ == '__main__':

test/high_availability/test_ha.py

@@ -100,7 +100,9 @@ def test_host_w_port(self):
             ConnectionFailure, MongoClient, "%s:1234567" % (host,), port)
 
     def test_repr(self):
-        self.assertEqual(repr(MongoClient(host, port)),
+        # Making host a str avoids the 'u' prefix in Python 2, so the repr is
+        # the same in Python 2 and 3.
+        self.assertEqual(repr(MongoClient(str(host), port)),
                          "MongoClient('%s', %d)" % (host, port))
 
     def test_getters(self):

test/test_client.py

@@ -180,7 +180,7 @@ def test_nearest(self):
         not_used = data_members.difference(used)
         latencies = ', '.join(
             '%s: %dms' % (member.host, member.ping_time.get())
-            for member in c._MongoReplicaSetClient__members.values())
+            for member in c._MongoReplicaSetClient__rs_state.members)
 
         self.assertFalse(not_used,
             "Expected to use primary and all secondaries for mode NEAREST,"
@@ -426,51 +426,24 @@ def test_aggregate(self):
 
 
 class TestMovingAverage(unittest.TestCase):
-    def test_empty_moving_average(self):
-        avg = MovingAverage(0)
-        self.assertEqual(None, avg.get())
-        avg.update(10)
-        self.assertEqual(None, avg.get())
-
-    def test_trivial_moving_average(self):
-        avg = MovingAverage(1)
-        self.assertEqual(None, avg.get())
-        avg.update(10)
-        self.assertEqual(10, avg.get())
-        avg.update(20)
-        self.assertEqual(20, avg.get())
-        avg.update(0)
-        self.assertEqual(0, avg.get())
-
-    def test_2_sample_moving_average(self):
-        avg = MovingAverage(2)
-        self.assertEqual(None, avg.get())
-        avg.update(10)
-        self.assertEqual(10, avg.get())
-        avg.update(20)
-        self.assertEqual(15, avg.get())
-        avg.update(30)
-        self.assertEqual(25, avg.get())
-        avg.update(-100)
-        self.assertEqual(-35, avg.get())
-
-    def test_5_sample_moving_average(self):
-        avg = MovingAverage(5)
-        self.assertEqual(None, avg.get())
-        avg.update(10)
+    def test_empty_init(self):
+        self.assertRaises(AssertionError, MovingAverage, [])
+
+    def test_moving_average(self):
+        avg = MovingAverage([10])
         self.assertEqual(10, avg.get())
-        avg.update(20)
-        self.assertEqual(15, avg.get())
-        avg.update(30)
-        self.assertEqual(20, avg.get())
-        avg.update(-100)
-        self.assertEqual((10 + 20 + 30 - 100) / 4, avg.get())
-        avg.update(17)
-        self.assertEqual((10 + 20 + 30 - 100 + 17) / 5., avg.get())
-        avg.update(43)
-        self.assertEqual((20 + 30 - 100 + 17 + 43) / 5., avg.get())
-        avg.update(-1111)
-        self.assertEqual((30 - 100 + 17 + 43 - 1111) / 5., avg.get())
+        avg2 = avg.clone_with(20)
+        self.assertEqual(15, avg2.get())
+        avg3 = avg2.clone_with(30)
+        self.assertEqual(20, avg3.get())
+        avg4 = avg3.clone_with(-100)
+        self.assertEqual((10 + 20 + 30 - 100) / 4., avg4.get())
+        avg5 = avg4.clone_with(17)
+        self.assertEqual((10 + 20 + 30 - 100 + 17) / 5., avg5.get())
+        avg6 = avg5.clone_with(43)
+        self.assertEqual((20 + 30 - 100 + 17 + 43) / 5., avg6.get())
+        avg7 = avg6.clone_with(-1111)
+        self.assertEqual((30 - 100 + 17 + 43 - 1111) / 5., avg7.get())
 
 
 class TestMongosConnection(unittest.TestCase):