[worker] Worker performance increase (#11889)

Author: gileri

opencti-worker/src/message_queue_consumer.py

@@ -15,12 +15,13 @@ class MessageQueueConsumer:  # pylint: disable=too-many-instance-attributes
     pika_parameters: pika.ConnectionParameters
     execution_pool: ThreadPoolExecutor
     handle_message: Callable[[str], Literal["ack", "nack", "requeue"]]
+    queue_concurrency_enabled: bool
     should_stop: bool = field(default=False, init=False)
 
     def __post_init__(self) -> None:
         self.pika_connection = pika.BlockingConnection(self.pika_parameters)
         self.channel = self.pika_connection.channel()
-        self.channel.basic_qos(prefetch_count=1)
+        self.channel.basic_qos(prefetch_count=(self.execution_pool._max_workers + 1) if self.queue_concurrency_enabled else 1)
         self.thread = Thread(target=self.consume_queue, name=self.queue_name)
         self.thread.start()
 
@@ -83,9 +84,10 @@ def consume_queue(self) -> None:
                     method.delivery_tag,
                     body,
                 )
-                while task_future.running():  # Loop while the thread is processing
-                    self.pika_connection.sleep(0.05)
-                self.logger.info("Message processed, thread terminated")
+                if not self.queue_concurrency_enabled:
+                    while task_future.running():  # Loop while the thread is processing
+                        self.pika_connection.sleep(0.05)
+                    self.logger.info("Message processed, thread terminated")
         except Exception as e:
             self.logger.error("Unhandled exception", {"exception": e})
         finally:

opencti-worker/src/push_handler.py

@@ -1,6 +1,7 @@
 import base64
 import datetime
 import json
+import threading
 from dataclasses import dataclass
 from typing import Any, Dict, Union, Literal
 
@@ -29,13 +30,19 @@ class PushHandler:  # pylint: disable=too-many-instance-attributes
     objects_max_refs: int
 
     def __post_init__(self) -> None:
-        self.api = OpenCTIApiClient(
-            url=self.opencti_url,
-            token=self.opencti_token,
-            log_level=self.log_level,
-            json_logging=self.json_logging,
-            ssl_verify=self.ssl_verify,
+        self.local_api = threading.local()
+
+    # OpenCTIClient is not thread safe, use a thread local to ensure to work on a dedicated client when creating and sending a request
+    def get_api_client(self) -> OpenCTIApiClient:
+        if not hasattr(self.local_api, "client"):
+            self.local_api.client = OpenCTIApiClient(
+                url=self.opencti_url,
+                token=self.opencti_token,
+                log_level=self.log_level,
+                json_logging=self.json_logging,
+                ssl_verify=self.ssl_verify,
         )
+        return self.local_api.client
 
     def send_bundle_to_specific_queue(
         self,
@@ -76,17 +83,18 @@ def handle_message(
         imported_items = []
         start_processing = datetime.datetime.now()
         try:
+            api = self.get_api_client()
             # Set the API headers
-            self.api.set_applicant_id_header(data.get("applicant_id"))
-            self.api.set_playbook_id_header(data.get("playbook_id"))
-            self.api.set_event_id(data.get("event_id"))
-            self.api.set_draft_id(data.get("draft_id"))
-            self.api.set_synchronized_upsert_header(data.get("synchronized", False))
-            self.api.set_previous_standard_header(data.get("previous_standard"))
+            api.set_applicant_id_header(data.get("applicant_id"))
+            api.set_playbook_id_header(data.get("playbook_id"))
+            api.set_event_id(data.get("event_id"))
+            api.set_draft_id(data.get("draft_id"))
+            api.set_synchronized_upsert_header(data.get("synchronized", False))
+            api.set_previous_standard_header(data.get("previous_standard"))
             work_id = data.get("work_id")
             # Check if work is still valid
             if work_id is not None:
-                is_work_alive = self.api.work.get_is_work_alive(work_id)
+                is_work_alive = api.work.get_is_work_alive(work_id)
                 # If work no longer exists, bundle can be acked without doing anything
                 if not is_work_alive:
                     return "ack"
@@ -107,7 +115,7 @@ def handle_message(
                 if len(content["objects"]) == 1 or data.get("no_split", False):
                     update = data.get("update", False)
                     imported_items, too_large_items_bundles = (
-                        self.api.stix2.import_bundle_from_json(
+                        api.stix2.import_bundle_from_json(
                             raw_content, update, types, work_id, self.objects_max_refs
                         )
                     )
@@ -159,7 +167,7 @@ def handle_message(
                             )
                             # Add expectations to the work
                             if work_id is not None:
-                                self.api.work.add_expectations(work_id, expectations)
+                                api.work.add_expectations(work_id, expectations)
                             # For each split bundle, send it to the same queue
                             for bundle in bundles:
                                 self.send_bundle_to_specific_queue(
@@ -179,7 +187,7 @@ def handle_message(
                             "type": "bundle",
                             "objects": [content["data"]],
                         }
-                        imported_items = self.api.stix2.import_bundle(
+                        imported_items = api.stix2.import_bundle(
                             bundle, True, types, work_id
                         )
                     # Specific knowledge merge
@@ -200,7 +208,7 @@ def handle_message(
                             "type": "bundle",
                             "objects": [merge_object],
                         }
-                        imported_items = self.api.stix2.import_bundle(
+                        imported_items = api.stix2.import_bundle(
                             bundle, True, types, work_id
                         )
                     # All standard operations
@@ -223,7 +231,7 @@ def handle_message(
                             "type": "bundle",
                             "objects": [data_object],
                         }
-                        imported_items = self.api.stix2.import_bundle(
+                        imported_items = api.stix2.import_bundle(
                             bundle, True, types, work_id
                         )
                     case _:

opencti-worker/src/worker.py

@@ -157,6 +157,14 @@ def __post_init__(self) -> None:
             True,
             0,
         )
+        self.worker_queue_concurrency_enabled = get_config_variable(
+            "WORKER_QUEUE_CONCURRENCY_ENABLED",
+            ["worker", "queue_concurrency_enabled"],
+            config,
+            False,
+            False,
+        )
+
         # Telemetry
         if self.telemetry_enabled:
             self.prom_httpd, self.prom_t = start_http_server(
@@ -275,6 +283,7 @@ def start(self) -> None:
                             pika_parameters,
                             execution_pool,
                             push_handler.handle_message,
+                            self.worker_queue_concurrency_enabled,
                         )
 
                     # Listen for webhook message
@@ -299,6 +308,7 @@ def start(self) -> None:
                                 self.build_pika_parameters(connector_config),
                                 listen_execution_pool,
                                 listen_handler.handle_message,
+                                self.worker_queue_concurrency_enabled,
                             )
 
                 # Stop consumers whose queues no longer exist