fix(importer): batch puts together (#4428)

The sequential puts is holding up the importer. Try batching them to see
if it improves things.

Author: michaelkedar

deployment/clouddeploy/gke-workers/base/importer.yaml

@@ -27,10 +27,10 @@ spec:
             resources:
               requests:
                 cpu: "1"
-                memory: "1G"
+                memory: "4G"
               limits:
                 cpu: "1"
-                memory: "2G"
+                memory: "8G"
           nodeSelector:
             cloud.google.com/gke-nodepool: importer-pool
           restartPolicy: Never

gcp/workers/importer/importer.py

@@ -56,6 +56,7 @@
 _BUCKET_THREAD_COUNT = 20
 _HTTP_LAST_MODIFIED_FORMAT = '%a, %d %b %Y %H:%M:%S %Z'
 _TIMEOUT_SECONDS = 60
+_NDB_PUT_BATCH_SIZE = 500
 
 _client_store = threading.local()
 
@@ -610,6 +611,7 @@ def _process_updates_git(self, source_repo: osv.SourceRepository):
             changed_entries[rel_path] = None
 
     import_failure_logs = []
+    changed_entries_to_process = []
     # Create tasks for changed files.
     for changed_entry, ts in changed_entries.items():
       path = os.path.join(osv.repo_path(repo), changed_entry)
@@ -641,10 +643,17 @@ def _process_updates_git(self, source_repo: osv.SourceRepository):
 
       logging.info('Re-analysis triggered for %s', changed_entry)
       original_sha256 = osv.sha256(path)
-      put_if_newer(vuln, source_repo.name, path)
-
-      self._request_analysis_external(
-          source_repo, original_sha256, changed_entry, source_timestamp=ts)
+      # Collect for batch processing
+      changed_entries_to_process.append(
+          (vuln, path, ts, original_sha256, changed_entry))
+
+    if changed_entries_to_process:
+      put_if_newer_batch(
+          [(v, p) for v, p, _, _, _ in changed_entries_to_process],
+          source_repo.name)
+      for vuln, path, ts, orig_sha256, entry in changed_entries_to_process:
+        self._request_analysis_external(
+            source_repo, orig_sha256, entry, source_timestamp=ts)
 
     replace_importer_log(storage.Client(), source_repo.name,
                          self._public_log_bucket, import_failure_logs)
@@ -733,8 +742,7 @@ def _process_updates_bucket(self, source_repo: osv.SourceRepository):
         logging.info('Requesting analysis of bucket entry: %s/%s',
                      source_repo.bucket, blob_name)
 
-        for vuln in vulns:
-          put_if_newer(vuln, source_repo.name, blob_name)
+        put_if_newer_batch([(v, blob_name) for v in vulns], source_repo.name)
 
         self._request_analysis_external(
             source_repo,
@@ -952,6 +960,7 @@ def _process_updates_rest(self, source_repo: osv.SourceRepository):
     vulns_last_modified = last_update_date
     logging.info('%d records to consider', len(vulns))
     # Create tasks for changed files.
+    vulns_to_process = []
     for vuln in vulns:
       import_failure_logs = []
       vuln_modified = vuln.modified.ToDatetime(datetime.UTC)
@@ -981,15 +990,9 @@ def _process_updates_rest(self, source_repo: osv.SourceRepository):
           self._record_quality_finding(source_repo.name, bug_id)
           continue
 
-        put_if_newer(v, source_repo.name, v.id + source_repo.extension)
-        logging.info('Requesting analysis of REST record: %s',
-                     vuln.id + source_repo.extension)
         ts = None if ignore_last_import else vuln_modified
-        self._request_analysis_external(
-            source_repo,
-            osv.sha256_bytes(single_vuln.text.encode()),
-            vuln.id + source_repo.extension,
-            source_timestamp=ts)
+        vulns_to_process.append((v, vuln.id + source_repo.extension, ts,
+                                 osv.sha256_bytes(single_vuln.text.encode())))
       except osv.sources.KeyPathError:
         # Key path doesn't exist in the vulnerability.
         # No need to log a full error, as this is expected result.
@@ -1001,6 +1004,13 @@ def _process_updates_rest(self, source_repo: osv.SourceRepository):
         import_failure_logs.append(f'Failed to parse vulnerability "{vuln.id}"')
         continue
 
+    if vulns_to_process:
+      put_if_newer_batch([(v, p) for v, p, _, _ in vulns_to_process],
+                         source_repo.name)
+      for v, path, ts, sha256 in vulns_to_process:
+        self._request_analysis_external(
+            source_repo, sha256, path, source_timestamp=ts)
+
     replace_importer_log(storage.Client(), source_repo.name,
                          self._public_log_bucket, import_failure_logs)
 
@@ -1235,25 +1245,59 @@ def put_if_newer(vuln: vulnerability_pb2.Vulnerability, source: str, path: str):
   unchanged. Does not write if vuln's modified date is older than what's already
   in datastore.
   """
-  preprocess_vuln(vuln)
-  bug = osv.Bug.get_by_id(vuln.id)
-  if bug is None:
-    bug = new_bug_from_vuln(vuln, source, path)
-    bug.put()
-    log_update_latency(bug)
-    return
+  put_if_newer_batch([(vuln, path)], source)
 
-  # Only update if the incoming vulnerability is newer.
-  orig_modified = vuln.modified.ToDatetime(datetime.UTC)
-  if bug.import_last_modified and orig_modified <= bug.import_last_modified:
-    logging.info(
-        'Skipping update for %s because incoming modification time'
-        ' (%s) is not newer than existing record (%s)', vuln.id, orig_modified,
-        bug.import_last_modified)
+
+def put_if_newer_batch(
+    vulns_and_paths: list[tuple[vulnerability_pb2.Vulnerability,
+                                str]], source: str):
+  """Try to write vulnerabilities to datastore in batch, keeping enumerated
+  versions if unchanged. Does not write if vuln's modified date is older than
+  what's already in datastore.
+  """
+  if not vulns_and_paths:
     return
-  update_bug_from_vuln(bug, vuln, source, path)
-  bug.put()
-  log_update_latency(bug)
+
+  # Deduplicate by vuln.id, keeping the last one.
+  unique_vulns_and_paths = {}
+  for vuln, path in vulns_and_paths:
+    unique_vulns_and_paths[vuln.id] = (vuln, path)
+  vulns_and_paths = list(unique_vulns_and_paths.values())
+
+  for vuln, _ in vulns_and_paths:
+    preprocess_vuln(vuln)
+
+  keys = [ndb.Key(osv.Bug, v.id) for v, _ in vulns_and_paths]
+  existing_bugs = []
+  for i in range(0, len(keys), _NDB_PUT_BATCH_SIZE):
+    batch_keys = keys[i:i + _NDB_PUT_BATCH_SIZE]
+    existing_bugs.extend(ndb.get_multi(batch_keys))
+
+  bugs_to_put = []
+  for i, (vuln, path) in enumerate(vulns_and_paths):
+    bug = existing_bugs[i]
+    if bug is None:
+      bug = new_bug_from_vuln(vuln, source, path)
+      bugs_to_put.append(bug)
+      continue
+
+    # Only update if the incoming vulnerability is newer.
+    orig_modified = vuln.modified.ToDatetime(datetime.UTC)
+    if bug.import_last_modified and orig_modified <= bug.import_last_modified:
+      logging.info(
+          'Skipping update for %s because incoming modification time'
+          ' (%s) is not newer than existing record (%s)', vuln.id,
+          orig_modified, bug.import_last_modified)
+      continue
+    update_bug_from_vuln(bug, vuln, source, path)
+    bugs_to_put.append(bug)
+
+  if bugs_to_put:
+    for i in range(0, len(bugs_to_put), _NDB_PUT_BATCH_SIZE):
+      batch = bugs_to_put[i:i + _NDB_PUT_BATCH_SIZE]
+      ndb.put_multi(batch)
+      for bug in batch:
+        log_update_latency(bug)
 
 
 def main():