feat: add async pymongo support with new doc feeder and client methods

Author: DylanWelzel

biothings/hub/databuild/backend.py

@@ -122,7 +122,7 @@ def validate_sources(self, sources=None):
 
     def get_src_master_docs(self):
         if self.src_masterdocs is None:
-            self.src_masterdocs = dict([(src["_id"], src) for src in list(self.master.find())])
+            self.src_masterdocs = {src["_id"]: src for src in mongo.doc_feeder(self.master)}
         return self.src_masterdocs
 
     def get_src_metadata(self):
@@ -145,10 +145,10 @@ def get_src_metadata(self):
                 srcs.append(doc["_id"])
             srcs = list(set(srcs))
         else:
-            srcs = [d["_id"] for d in self.dump.find()]
+            srcs = [d["_id"] for d in mongo.doc_feeder(self.dump, fields={"_id": 1})]
         # we need to return main_source named, but if accessed, it's been through sub-source names
         # query is different in that case
-        for src in self.dump.find({"_id": {"$in": srcs}}):
+        for src in mongo.doc_feeder(self.dump, query={"_id": {"$in": srcs}}):
             # now merge other extra information from src_master (src_meta key). src_master _id
             # are sub-source names, not main source so we have to deal with src_dump as well
             # in order to resolve/map main/sub source name

biothings/hub/databuild/builder.py

@@ -20,6 +20,15 @@
 
 from biothings import config as btconfig
 from biothings.hub import BUILDER_CATEGORY, UPLOADER_CATEGORY
+from biothings.hub.databuild.backend import (
+    LinkTargetDocMongoBackend,
+    SourceDocMongoBackend,
+    TargetDocMongoBackend,
+    create_backend,
+)
+from biothings.hub.databuild.buildconfig import AutoBuildConfig
+from biothings.hub.databuild.mapper import TransparentMapper
+from biothings.hub.dataload.uploader import ResourceNotReady
 from biothings.hub.manager import BaseManager
 from biothings.utils import mongo
 from biothings.utils.backend import DocMongoBackend
@@ -36,24 +45,14 @@
     get_source_fullname,
     get_src_build,
     get_src_build_config,
+    get_src_db,
     get_src_dump,
     get_src_master,
-    get_src_db,
 )
 from biothings.utils.loggers import get_logger
 from biothings.utils.manager import JobManager
 from biothings.utils.mongo import doc_feeder, id_feeder
 
-from biothings.hub.databuild.backend import (
-    LinkTargetDocMongoBackend,
-    SourceDocMongoBackend,
-    TargetDocMongoBackend,
-    create_backend,
-)
-from biothings.hub.databuild.buildconfig import AutoBuildConfig
-from biothings.hub.databuild.mapper import TransparentMapper
-from biothings.hub.dataload.uploader import ResourceNotReady
-
 logging = btconfig.logger
 
 
@@ -1141,7 +1140,7 @@ def __init__(
 
     def clean_stale_status(self):
         src_build = get_src_build()
-        for build in src_build.find():
+        for build in doc_feeder(src_build):
             dirty = False
             for job in build.get("jobs", []):
                 if job.get("status") == "building":

biothings/hub/databuild/differ.py

@@ -19,7 +19,7 @@
 from biothings.utils.hub_db import get_src_build
 from biothings.utils.jsondiff import make as jsondiff
 from biothings.utils.loggers import get_logger
-from biothings.utils.mongo import get_previous_collection, get_target_db, id_feeder
+from biothings.utils.mongo import doc_feeder, get_previous_collection, get_target_db, id_feeder
 from biothings.utils.serializer import to_json_file
 
 from .backend import create_backend, merge_src_build_metadata
@@ -1004,7 +1004,7 @@ def __init__(self, poll_schedule=None, *args, **kwargs):
 
     def clean_stale_status(self):
         src_build = get_src_build()
-        for build in src_build.find():
+        for build in doc_feeder(src_build):
             dirty = False
             for job in build.get("jobs", []):
                 if job.get("status") == "diffing":

biothings/hub/databuild/syncer.py

@@ -836,7 +836,7 @@ def __init__(self, *args, **kwargs):
 
     def clean_stale_status(self):
         src_build = get_src_build()
-        for build in src_build.find():
+        for build in doc_feeder(src_build):
             dirty = False
             for job in build.get("jobs", []):
                 if job.get("status") == "syncing":

biothings/hub/default_config.py

@@ -83,12 +83,8 @@
 import biothings.utils.jsondiff
 
 # set_default_folder is needed for evaluating some default values below
-from biothings.utils.configuration import (
-    ConfigurationDefault,
-    ConfigurationError,
-    ConfigurationValue,
-    set_default_folder,  # pylint: disable=unused-import      # noqa
-)
+from biothings.utils.configuration import set_default_folder  # pylint: disable=unused-import      # noqa
+from biothings.utils.configuration import ConfigurationDefault, ConfigurationError, ConfigurationValue
 
 # * 1. General *#
 # Hub name/icon url/version, for display purpose
@@ -174,6 +170,31 @@
     desc="Define path to folder which will contain cache files, set to None to disable",
 )
 
+# Experimental: use async pymongo-backed feeder in `biothings.utils.mongo.doc_feeder`
+# Enabled path is currently opt-in and only applies to Mongo source/target/hubdb collections.
+USE_ASYNC_DOC_FEEDER = ConfigurationDefault(
+    default=False,
+    desc="Use async pymongo doc feeder (experimental, defaults to synchronous behavior)",
+)
+
+# Size of the async producer -> sync consumer prefetch queue used by doc_feeder bridge.
+ASYNC_DOC_FEEDER_BUFFER_SIZE = ConfigurationDefault(
+    default=4,
+    desc="Prefetch queue size for async doc feeder bridge (higher may improve throughput at the cost of memory)",
+)
+
+# Emit comparable feeder benchmark lines with counts and elapsed time.
+DOC_FEEDER_BENCHMARK = ConfigurationDefault(
+    default=False,
+    desc="Enable doc_feeder benchmark summary logs (DOC_FEEDER_BENCH)",
+)
+
+# Optional benchmark label to group A/B runs in logs.
+DOC_FEEDER_BENCHMARK_TAG = ConfigurationDefault(
+    default="",
+    desc="Optional label attached to DOC_FEEDER_BENCH logs (for example baseline/async/module-swap)",
+)
+
 # Path to a folder to store all 3rd party parsers, dumpers, etc...
 DATA_PLUGIN_FOLDER = ConfigurationDefault(
     default="./plugins", desc="Define path to folder which will contain all 3rd party parsers, dumpers, etc..."