fix(thrift): recover from transient Mongo timeouts

Author: rondinelisaad

articlemeta/controller.py

@@ -1,4 +1,5 @@
 # coding: utf-8
+import os
 from urllib.parse import urlparse
 import warnings
 import json
@@ -173,6 +174,24 @@ def get_date_range_filter(from_date=None, until_date=None):
 def get_dbconn(db_dsn):
     """Connects to the MongoDB server and returns a database handler."""
 
+    def _mongo_timeouts():
+        return {
+            "serverSelectionTimeoutMS": int(os.environ.get("MONGO_SERVER_SELECTION_TIMEOUT_MS", 5000)),
+            "connectTimeoutMS": int(os.environ.get("MONGO_CONNECT_TIMEOUT_MS", 5000)),
+            "socketTimeoutMS": int(os.environ.get("MONGO_SOCKET_TIMEOUT_MS", 15000)),
+            "waitQueueTimeoutMS": int(os.environ.get("MONGO_WAIT_QUEUE_TIMEOUT_MS", 5000)),
+            "maxPoolSize": int(os.environ.get("MONGO_MAX_POOL_SIZE", 100)),
+            "retryWrites": False,
+        }
+
+    def _normalize_dsn(raw_dsn):
+        # Accepts legacy forms like "host:27017/db", "//host:27017/db", or full mongodb URI.
+        if raw_dsn.startswith("mongodb://") or raw_dsn.startswith("mongodb+srv://"):
+            return raw_dsn
+        if raw_dsn.startswith("//"):
+            return "mongodb:%s" % raw_dsn
+        return "mongodb://%s" % raw_dsn
+
     def _create_indexes(db):
         """
         Ensures that an index exists on specified collections.
@@ -254,10 +273,23 @@ def _create_indexes(db):
                     db[collection].create_index(index[0], **index[1])
 
     print('End Creation index')
-    db_url = urlparse(db_dsn)
-    conn = pymongo.MongoClient('mongodb://%s' % db_url.netloc)
-    db = conn[db_url.path[1:]]
-    _create_indexes(db)
+    mongo_uri = _normalize_dsn(db_dsn)
+    db_url = urlparse(mongo_uri)
+    db_name = db_url.path[1:] or "articlemeta"
+    conn = pymongo.MongoClient(mongo_uri, **_mongo_timeouts())
+    db = conn[db_name]
+    try:
+        _create_indexes(db)
+    except (
+        pymongo.errors.AutoReconnect,
+        pymongo.errors.NetworkTimeout,
+        pymongo.errors.ServerSelectionTimeoutError,
+        pymongo.errors.ConnectionFailure,
+    ) as exc:
+        warnings.warn(
+            "MongoDB unavailable during index creation; continuing startup "
+            "without blocking. Error: %s" % exc
+        )
     return db
 
 

articlemeta/thrift/server.py

@@ -2,10 +2,12 @@
 import json
 import logging
 import os
+import threading
 import uuid
 
 import thriftpywrap
 import thriftpy2
+from pymongo import errors as pymongo_errors
 
 from articlemeta.controller import DataBroker, get_dbconn
 from articlemeta import utils
@@ -64,6 +66,44 @@
 articlemeta_thrift = thriftpy2.load(
     os.path.join(os.path.dirname(__file__), 'articlemeta.thrift'))
 
+TRANSIENT_MONGO_ERRORS = (
+    pymongo_errors.AutoReconnect,
+    pymongo_errors.NetworkTimeout,
+    pymongo_errors.ServerSelectionTimeoutError,
+    pymongo_errors.ConnectionFailure,
+)
+
+
+class ResilientDataBroker(object):
+    """Wraps DataBroker calls with one reconnect/retry on transient Mongo errors."""
+
+    def __init__(self, broker_factory):
+        self._broker_factory = broker_factory
+        self._lock = threading.Lock()
+        self._broker = broker_factory()
+
+    def _reset(self):
+        with self._lock:
+            self._broker = self._broker_factory()
+
+    def __getattr__(self, attr_name):
+        attr = getattr(self._broker, attr_name)
+        if not callable(attr):
+            return attr
+
+        def _wrapped(*args, **kwargs):
+            try:
+                return getattr(self._broker, attr_name)(*args, **kwargs)
+            except TRANSIENT_MONGO_ERRORS:
+                logger.exception(
+                    "Transient MongoDB error on '%s', reconnecting and retrying once",
+                    attr_name,
+                )
+                self._reset()
+                return getattr(self._broker, attr_name)(*args, **kwargs)
+
+        return _wrapped
+
 
 class Dispatcher(object):
     def __init__(self):
@@ -73,9 +113,9 @@ def __init__(self):
         db_dsn = os.environ.get('MONGODB_HOST', settings.get('mongo_uri', '127.0.0.1:27017'))
 
         self._admintoken = os.environ.get('ADMIN_TOKEN', None) or settings['app:main'].get('admintoken', uuid.uuid4().hex)
-
-        db_client = get_dbconn(db_dsn)
-        self._databroker = DataBroker(db_client)
+        self._databroker = ResilientDataBroker(
+            lambda: DataBroker(get_dbconn(db_dsn))
+        )
 
     def getInterfaceVersion(self):
         return articlemeta_thrift.VERSION