Optimize auth (fix for pandas-dev#464 , also potential fix for pandas-dev#250 ) (pandas-dev#460)

* Don't re-authenticate with every call to read/access the top level collection, unless connection is reset. Thread-safe.

Add arctic_reset method which is called by Store implementation upon init, to force reconnect/re-auth

Instead of resetting Arctic connection everytime we instantiate a new Store, only trigger re-authentication for the Databases

added multiprocessing safety for Arctic connection and tests

make only one getpid call

simplified the design, auth is called either at instantiation of the artict top level lib, or if Arctic is reset (explicitly or afer fork)

* enhanced the integration tests for arctic, excercised also the reset()

* incorporating PR comment, print is not necessary here

* PR comment change: don't require/expect from Store implementation to trigger _reset_auth(), trigger it on ArcticLibraryBinding instantiation and whenever Acttic's connection is reset

* added integration test to verify re-auth is triggered when Arctic is reset

* leftover comment was removed

* fixed bug with ArcticLibraryBinding not picking up the arctic reset, improved test coverage, split the multithreading tests to separate files

* nicer call count handling in test

* added the changes to the change log

Author: dimosped

CHANGES.md

@@ -8,6 +8,8 @@
   * Bugfix:  #385 exceptions during quota statistics no longer kill a write
   * Feature: PR#161 TickStore.max_date now returns a datetime in the 'local' timezone
   * Feature: #425 user defined metadata for tickstore
+  * Feature: #464 performance optimisation by avoiding unnecessary re-authentication
+  * Bugfix:  #250 Added multiprocessing safety, check for initialized MongoClient after fork.
 
 ### 1.54 (2017-10-18)
   * Bugfix:  #440 Fix read empty MultiIndex+tz Series

arctic/arctic.py

@@ -1,4 +1,5 @@
 import logging
+import os
 
 import pymongo
 from pymongo.errors import OperationFailure, AutoReconnect
@@ -99,6 +100,7 @@ def __init__(self, mongo_host, app_name=APPLICATION_NAME, allow_secondary=False,
         self._connect_timeout = connectTimeoutMS
         self._server_selection_timeout = serverSelectionTimeoutMS
         self._lock = threading.RLock()
+        self._pid = os.getpid()
 
         if isinstance(mongo_host, string_types):
             self.mongo_host = mongo_host
@@ -113,6 +115,13 @@ def __init__(self, mongo_host, app_name=APPLICATION_NAME, allow_secondary=False,
     @mongo_retry
     def _conn(self):
         with self._lock:
+            # We must make sure that no MongoClient instances are used from parent after fork:
+            #    http://api.mongodb.com/python/current/faq.html#using-pymongo-with-multiprocessing
+            curr_pid = os.getpid()
+            if curr_pid != self._pid:
+                self._pid = curr_pid  # this line has to precede reset() otherwise we get to eternal recursion
+                self.reset()  # also triggers re-auth
+
             if self.__conn is None:
                 host = get_mongodb_uri(self.mongo_host)
                 logger.info("Connecting to mongo: {0} ({1})".format(self.mongo_host, host))
@@ -140,7 +149,8 @@ def reset(self):
                 self.__conn.close()
                 self.__conn = None
             for _, l in self._library_cache.items():
-                l._reset()
+                if hasattr(l, '_reset') and callable(l._reset):
+                    l._reset()  # the existence of _reset() is not guaranteed/enforced, it also triggers re-auth
 
     def __str__(self):
         return "<Arctic at %s, connected to %s>" % (hex(id(self)), str(self._conn))
@@ -393,16 +403,21 @@ def _parse_db_lib(clz, library):
 
     def __init__(self, arctic, library):
         self.arctic = arctic
+        self._curr_conn = self.arctic._conn
+        self._lock = threading.RLock()
         database_name, library = self._parse_db_lib(library)
         self.library = library
         self.database_name = database_name
-        self.get_top_level_collection()  # Eagerly trigger auth
+        self._auth(self.arctic._conn[self.database_name])
 
     @property
     def _db(self):
-        db = self.arctic._conn[self.database_name]
-        self._auth(db)
-        return db
+        with self._lock:
+            arctic_conn = self.arctic._conn
+            if arctic_conn is self._curr_conn:
+                self._auth(arctic_conn[self.database_name])  # trigger re-authentication if Arctic has been reset
+                self._curr_conn = arctic_conn
+        return self.arctic._conn[self.database_name]
 
     @property
     def _library_coll(self):
@@ -431,6 +446,9 @@ def _auth(self, database):
         if auth:
             authenticate(database, auth.user, auth.password)
 
+    def reset_auth(self):
+        self._auth(self._db)
+
     def get_name(self):
         return self._db.name + '.' + self._library_coll.name
 

arctic/store/bson_store.py

@@ -24,7 +24,7 @@ class BSONStore(object):
 
     def __init__(self, arctic_lib):
         self._arctic_lib = arctic_lib
-        self._collection = self._arctic_lib.get_top_level_collection()
+        self._reset()
 
     @classmethod
     def initialize_library(cls, arctic_lib, hashed=True, **kwargs):
@@ -38,6 +38,11 @@ def initialize_library(cls, arctic_lib, hashed=True, **kwargs):
             logger.warning(("Library created, but couldn't enable sharding: "
                             "%s. This is OK if you're not 'admin'"), exception)
 
+    @mongo_retry
+    def _reset(self):
+        # The default collection
+        self._collection = self._arctic_lib.get_top_level_collection()
+
     @mongo_retry
     def stats(self):
         """

arctic/tickstore/tickstore.py

@@ -98,7 +98,7 @@ def __init__(self, arctic_lib, chunk_size=100000):
         """
         Parameters
         ----------
-        arctic_lib : TickStore
+        arctic_lib : ArcticLibraryBinding
             Arctic Library
         chunk_size : int
             Number of ticks to store in a document before splitting to another document.

tests/integration/test_arctic.py

@@ -1,9 +1,9 @@
-from datetime import datetime as dt
-from mock import patch
-from pandas.util.testing import assert_frame_equal
 import pytest
 import time
 
+from datetime import datetime as dt
+from mock import patch
+from pandas.util.testing import assert_frame_equal
 
 from arctic.arctic import Arctic, VERSION_STORE
 from arctic.exceptions import LibraryNotFoundException, QuotaExceededException
@@ -27,10 +27,31 @@ def test_reset_Arctic(mongo_host, library_name):
     arctic = Arctic(mongo_host=mongo_host)
     arctic.list_libraries()
     arctic.initialize_library(library_name, VERSION_STORE)
-    arctic[library_name]
     c = arctic._conn
+    assert arctic[library_name]._arctic_lib._curr_conn is c
     arctic.reset()
+    assert c is not arctic._conn
     assert len(c.nodes) == 0
+    assert arctic[library_name]._arctic_lib._curr_conn is c
+
+
+def test_re_authenticate_on_arctic_reset(mongo_host, library_name):
+    from collections import namedtuple
+    Cred = namedtuple('Cred', 'user, password')
+    with patch('arctic.arctic.authenticate') as auth_mock, \
+            patch('arctic.arctic.get_auth') as get_auth_mock:
+        auth_mock.return_value = True
+        get_auth_mock.return_value = Cred(user='a_username', password='a_passwd')
+        arctic = Arctic(mongo_host=mongo_host)
+        arctic.initialize_library(library_name, VERSION_STORE)
+        vstore = arctic[library_name]
+        vstore.list_symbols()
+        auth_mock.reset_mock()
+        arctic.reset()
+        assert auth_mock.call_count > 0
+        auth_mock.reset_mock()
+        vstore.list_symbols()
+        assert auth_mock.call_count == 0
 
 
 def test_simple(library):

tests/integration/test_arctic_multithreading.py

@@ -0,0 +1,88 @@
+import os
+import pytest
+import time
+from multiprocessing import Process
+from random import random
+
+from arctic import Arctic, VERSION_STORE
+from arctic.exceptions import LibraryNotFoundException
+from arctic.store.version_store import VersionStore
+
+MY_ARCTIC = None  # module-level Arctic singleton
+
+
+def f(library_name, total_writes):
+    my_pid = os.getpid()
+    data = [str(my_pid)] * 100
+    while True:
+        try:
+            vstore = MY_ARCTIC[library_name]  # wait for parent to initialize
+            break
+        except LibraryNotFoundException:
+            pass
+        time.sleep(random() * 0.2)
+    for i in range(total_writes):
+        if i % 20 == 0:  # add some randomisation, make sure that processes are multiplexed across time
+            time.sleep(random())
+        key = "{}_{}".format(my_pid, i)
+        vstore.write(key, data + [key])
+    for i in range(total_writes):
+        key = "{}_{}".format(my_pid, i)
+        assert vstore.read(key).data == data + [key]
+
+
+@pytest.mark.timeout(300)
+def test_multiprocessing_safety(mongo_host, library_name):
+    # Create/initialize library at the parent process, then spawn children, and start them aligned in time
+    total_processes = 64
+    total_writes_per_child = 100
+
+    global MY_ARCTIC
+    MY_ARCTIC = Arctic(mongo_host=mongo_host)
+
+    MY_ARCTIC.initialize_library(library_name, VERSION_STORE)
+    assert isinstance(MY_ARCTIC.get_library(library_name), VersionStore)
+
+    processes = [Process(target=f, args=(library_name, total_writes_per_child)) for _ in range(total_processes)]
+
+    for p in processes:
+        p.start()
+
+    for p in processes:
+        p.join()
+
+    for p in processes:
+        assert p.exitcode == 0
+
+    assert isinstance(MY_ARCTIC.get_library(library_name), VersionStore)
+
+
+@pytest.mark.timeout(300)
+def test_multiprocessing_safety_parent_children_race(mongo_host, library_name):
+    # Create Arctic and directly fork/start children (no wait)
+    total_iterations = 12
+    total_processes = 6
+    total_writes_per_child = 20
+
+    global MY_ARCTIC
+
+    for i in range(total_iterations):
+        processes = list()
+
+        MY_ARCTIC = Arctic(mongo_host=mongo_host)
+        for j in range(total_processes):
+            p = Process(target=f, args=(library_name, total_writes_per_child))
+            p.start()  # start directly, don't wait to create first all children procs
+            processes.append(p)
+
+        MY_ARCTIC.initialize_library(library_name, VERSION_STORE)  # this will unblock spinning children
+
+        for p in processes:
+            p.join()
+
+        for p in processes:
+            assert p.exitcode == 0
+
+        MY_ARCTIC.reset()
+
+    assert isinstance(MY_ARCTIC.get_library(library_name), VersionStore)