Skip to content

perf(mongodb): batch finding writes (bulk_write) instead of one DB call per item #1176

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ocervell wants to merge 1 commit into
base: main
Choose a base branch
from
Open

perf(mongodb): batch finding writes (bulk_write) instead of one DB call per item #1176

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
86 changes: 58 additions & 28 deletions secator/hooks/mongodb.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
from secator.hooks._dedup import compute_duplicate_updates
from secator.output_types import OUTPUT_TYPES
from secator.runners import Scan, Task, Workflow
from secator.utils import debug, escape_mongodb_url
from secator.utils import debug, escape_mongodb_url, should_update

# import gevent.monkey
# gevent.monkey.patch_all()
Expand All @@ -19,6 +19,10 @@
MONGODB_CONNECT_TIMEOUT = CONFIG.addons.mongodb.server_selection_timeout_ms
MONGODB_MAX_POOL_SIZE = CONFIG.addons.mongodb.max_pool_size

# Max buffered finding upserts before a forced bulk flush (bounds worker memory).
# Time-based flushing is throttled by CONFIG.runners.backend_update_frequency.
MONGODB_FLUSH_SIZE = 1000

logger = logging.getLogger(__name__)

_mongodb_client = None
Expand Down Expand Up @@ -103,32 +107,58 @@ def update_runner(self):


def update_finding(self, item):
"""Buffer a finding upsert; the write is batched (see flush_findings).

Findings used to be written one-by-one (one update_one/insert_one per item),
which is millions of round-trips on a large crawl. We now mint the Mongo _id
client-side (ObjectId() needs no round-trip) so the item has a stable id
immediately, buffer an upsert, and flush in bulk on a size cap, on_interval
(throttled by backend_update_frequency) and on_end.
"""
if type(item) not in OUTPUT_TYPES:
return item
if not ObjectId.is_valid(str(item._uuid)):
item._uuid = str(ObjectId())
buffer = getattr(self, '_mongodb_findings_buffer', None)
if buffer is None:
buffer = self._mongodb_findings_buffer = []
buffer.append(
pymongo.UpdateOne({'_id': ObjectId(item._uuid)}, {'$set': item.toDict()}, upsert=True)
)
Comment on lines +122 to +127

@coderabbitai coderabbitai Bot Jun 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

🧩 Analysis chain

🌐 Web query:

In PyMongo and MongoDB bulk_write with ordered=False, are operations guaranteed to execute in the order provided, especially multiple UpdateOne upserts targeting the same _id?

💡 Result:

When using PyMongo or any MongoDB driver with bulk_write and ordered=False, the operations are not guaranteed to execute in the order provided [1][2][3]. In unordered mode, MongoDB may execute operations in parallel and reorder them for performance optimization [1][2][3]. Consequently, applications should not rely on the execution order of these operations [2][3]. Regarding multiple UpdateOne upserts targeting the same _id in an unordered bulk write: because there is no guarantee of serial execution, the final state of the document is non-deterministic if multiple operations in the same batch target the same _id [1][2][3]. You cannot rely on the operations "winning" based on their index in the input array. If your application logic requires a specific final state for a document, you must either use ordered=True (the default) or ensure that your bulk operations do not have overlapping targets [1][2][3].

Citations:

🏁 Script executed:

cd secator && wc -l hooks/mongodb.py

Repository: freelabz/secator

Length of output: 81

🏁 Script executed:

cd secator && sed -n '115,150p' hooks/mongodb.py

Repository: freelabz/secator

Length of output: 1379

Coalesce repeated finding updates before unordered bulk writes.

The buffer appends every UpdateOne, then Line 142 flushes with ordered=False. With unordered bulk writes, operations are not guaranteed to execute in order, making the final document state non-deterministic when the same finding is buffered multiple times before a flush—an older snapshot can be applied after the newer one. Keep only the latest operation per item._uuid by using a dict keyed by UUID, or switch to ordered writes.

Proposed coalescing fix 
-	buffer = getattr(self, '_mongodb_findings_buffer', None)
-	if buffer is None:
-		buffer = self._mongodb_findings_buffer = []
-	buffer.append(
-		pymongo.UpdateOne({'_id': ObjectId(item._uuid)}, {'$set': item.toDict()}, upsert=True)
-	)
+	buffer = getattr(self, '_mongodb_findings_buffer', None)
+	if buffer is None:
+		buffer = self._mongodb_findings_buffer = {}
+	buffer[item._uuid] = pymongo.UpdateOne(
+		{'_id': ObjectId(item._uuid)},
+		{'$set': item.toDict()},
+		upsert=True,
+	)
 	if len(buffer) >= MONGODB_FLUSH_SIZE:
 		flush_findings_buffer(self)
 	return item
@@
-	count = len(buffer)
-	db.findings.bulk_write(buffer, ordered=False)
-	self._mongodb_findings_buffer = []
+	operations = list(buffer.values())
+	count = len(operations)
+	db.findings.bulk_write(operations, ordered=False)
+	self._mongodb_findings_buffer = {}

Also applies to: 141-143

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@secator/hooks/mongodb.py` around lines 122 - 127, The
_mongodb_findings_buffer currently uses a list to append every UpdateOne
operation, but when flushed with unordered bulk writes, operations for the same
finding (same item._uuid) can execute out of order, leaving the document in a
non-deterministic state. Replace the list-based buffer with a dict keyed by
item._uuid so that repeated updates for the same finding are coalesced and only
the latest UpdateOne operation per UUID is retained. This ensures that when the
buffer is flushed at the bulk_write call, each finding is updated only once with
its most recent state, regardless of the unordered execution.

if len(buffer) >= MONGODB_FLUSH_SIZE:
flush_findings_buffer(self)
return item


def flush_findings_buffer(self):
"""Write all buffered finding upserts to MongoDB in a single bulk_write."""
buffer = getattr(self, '_mongodb_findings_buffer', None)
if not buffer:
return
start_time = time.time()
client = get_mongodb_client()
db = client.main
update = item.toDict()
_type = item._type
_id = ObjectId(item._uuid) if ObjectId.is_valid(item._uuid) else None
if _id:
finding = db['findings'].update_one({'_id': _id}, {'$set': update})
status = 'UPDATED'
else:
finding = db['findings'].insert_one(update)
item._uuid = str(finding.inserted_id)
status = 'CREATED'
end_time = time.time()
elapsed = end_time - start_time
debug_obj = {
_type: status,
'type': 'finding',
'class': self.__class__.__name__,
'caller': self.config.name,
**self.context
}
debug(f'in {elapsed:.4f}s', sub='hooks.mongodb', id=str(item._uuid), obj=debug_obj, obj_after=False) # noqa: E501
return item
count = len(buffer)
db.findings.bulk_write(buffer, ordered=False)
self._mongodb_findings_buffer = []
self._last_findings_flush = start_time
debug(f'flushed {count} findings in {time.time() - start_time:.4f}s', sub='hooks.mongodb', obj_after=False)


def flush_findings(self):
"""on_interval hook: flush buffered findings, throttled by backend_update_frequency."""
if should_update(CONFIG.runners.backend_update_frequency, getattr(self, '_last_findings_flush', None)):
flush_findings_buffer(self)


def flush_findings_final(self):
"""on_end hook: always flush remaining findings.

This is required for correctness, not just throughput: the next runner in a
scan re-hydrates these findings from the DB (get_results), so they must be
persisted before this runner finishes.
"""
flush_findings_buffer(self)


def find_duplicates(self):
Expand Down Expand Up @@ -220,23 +250,23 @@ def tag_duplicates(ws_id: str = None, full_scan: bool = False, exclude_types=[],
Scan: {
'on_init': [update_runner],
'on_start': [update_runner],
'on_interval': [update_runner],
'on_interval': [update_runner, flush_findings],
'on_duplicate': [update_finding],
'on_end': [update_runner],
'on_end': [update_runner, flush_findings_final],
Comment on lines +253 to +255

@coderabbitai coderabbitai Bot Jun 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Run findings flush hooks before update_runner.

For on_interval, secator/runners/_base.py:948-983 re-checks self.last_updated_db before each hook; update_runner sets it on Line 96, so placing flush_findings after update_runner causes the interval flush to be skipped. For on_end, persisting the runner first can mark completion before the mandatory final findings flush. Move flush hooks first.

Proposed hook ordering fix 
 	Scan: {
 		'on_init': [update_runner],
 		'on_start': [update_runner],
-		'on_interval': [update_runner, flush_findings],
+		'on_interval': [flush_findings, update_runner],
 		'on_duplicate': [update_finding],
-		'on_end': [update_runner, flush_findings_final],
+		'on_end': [flush_findings_final, update_runner],
 	},
 	Workflow: {
 		'on_init': [update_runner],
 		'on_start': [update_runner],
-		'on_interval': [update_runner, flush_findings],
+		'on_interval': [flush_findings, update_runner],
 		'on_duplicate': [update_finding],
-		'on_end': [update_runner, flush_findings_final],
+		'on_end': [flush_findings_final, update_runner],
 	},
 	Task: {
 		'on_init': [update_runner],
 		'on_start': [update_runner],
 		'on_item': [update_finding],
 		'on_duplicate': [update_finding],
-		'on_interval': [update_runner, flush_findings],
-		'on_end': [update_runner, flush_findings_final]
+		'on_interval': [flush_findings, update_runner],
+		'on_end': [flush_findings_final, update_runner]
 	}

Also applies to: 260-262, 269-270

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@secator/hooks/mongodb.py` around lines 253 - 255, The hook ordering in the
MongoDB hooks file needs to be adjusted so that flush operations execute before
update_runner, because update_runner sets self.last_updated_db which the runner
uses to determine whether to skip interval-based flushes. In the on_interval
list (lines 253-255), move flush_findings before update_runner in the list. In
the on_end list (lines 269-270), move flush_findings_final before update_runner
so that findings are persisted before the runner is marked as complete. For the
on_duplicate hooks at lines 260-262, apply the same principle if flush
operations are present there.

},
Workflow: {
'on_init': [update_runner],
'on_start': [update_runner],
'on_interval': [update_runner],
'on_interval': [update_runner, flush_findings],
'on_duplicate': [update_finding],
'on_end': [update_runner],
'on_end': [update_runner, flush_findings_final],
},
Task: {
'on_init': [update_runner],
'on_start': [update_runner],
'on_item': [update_finding],
'on_duplicate': [update_finding],
'on_interval': [update_runner],
'on_end': [update_runner]
'on_interval': [update_runner, flush_findings],
'on_end': [update_runner, flush_findings_final]
}
}
Loading