AlertMerger query optimizations

streamalert/alert_merger/main.py

@@ -60,6 +60,10 @@ class AlertMerger:
     # Set the max payload size to slightly under that to account for the rest of the message.
     MAX_LAMBDA_PAYLOAD_SIZE = 126000
 
+    # The maximum number of alerts that are loaded into memory for a single rule, during a single
+    # loop in the Alert Merger.
+    ALERT_GENERATOR_DEFAULT_LIMIT = 5000
+
     @classmethod
     def get_instance(cls):
         """Get an instance of the AlertMerger, using a cached version if possible."""
@@ -73,17 +77,32 @@ def __init__(self):
         self.alert_proc_timeout = int(os.environ['ALERT_PROCESSOR_TIMEOUT_SEC'])
         self.lambda_client = boto3.client('lambda')
 
-    def _get_alerts(self, rule_name):
-        """Build a list of Alert instances triggered from the given rule name."""
-        alerts = []
+        # FIXME (derek.wang) Maybe make this configurable in the future
+        self._alert_generator_limit = self.ALERT_GENERATOR_DEFAULT_LIMIT
+
+    def _alert_generator(self, rule_name):
+        """
+        Returns a generator that yields Alert instances triggered from the given rule name.
 
+        To limit memory consumption, the generator yields a maximum number of alerts, defined
+        by self._alert_generator_limit.
+        """
+        alert_count = 0
         for record in self.table.get_alert_records(rule_name, self.alert_proc_timeout):
             try:
-                alerts.append(Alert.create_from_dynamo_record(record))
+                yield Alert.create_from_dynamo_record(record)
+                alert_count += 1
             except AlertCreationError:
                 LOGGER.exception('Invalid alert record %s', record)
+                continue
 
-        return alerts
+            if alert_count >= self._alert_generator_limit:
+                LOGGER.warning(
+                    'Alert Merger reached alert limit of %d for rule "%s"',
+                    self._alert_generator_limit,
+                    rule_name
+                )
+                return
 
     @staticmethod
     def _merge_groups(alerts):
@@ -144,14 +163,9 @@ def dispatch(self):
         merged_alerts = []  # List of newly created merge alerts
         alerts_to_delete = []  # List of alerts which can be deleted
 
-        # TODO: Find a way to avoid a full table scan just to get rule names
-        for rule_name in self.table.rule_names():
-            alerts = self._get_alerts(rule_name)
-            if not alerts:
-                continue
-
+        for rule_name in self.table.rule_names_generator():
             merge_enabled_alerts = []
-            for alert in alerts:
+            for alert in self._alert_generator(rule_name):
                 if alert.remaining_outputs:
                     # If an alert still has pending outputs, it needs to be sent immediately.
                     # For example, all alerts are sent to the default firehose now even if they will

streamalert/shared/alert_table.py

@@ -59,6 +59,10 @@ def _paginate(func, func_kwargs):
     def rule_names(self):
         """Find all of the distinct rule names in the table.
 
+        @deprecated
+          This method is deprecated due to requiring a full table scan prior to returning any
+          records.
+
         Returns:
             set: Set of string rule names
         """
@@ -68,6 +72,36 @@ def rule_names(self):
         }
         return set(item['RuleName'] for item in self._paginate(self._table.scan, kwargs))
 
+    def rule_names_generator(self):
+        """Returns a generator that yields unique rule_names of items on the table
+
+        Each unique name is yielded only once. Additionally, because the names materialized over
+        several paginated calls, it is not 100% guaranteed to return every possible rule_name on
+        the Alert Table; if there are other operations that are writing items to the DynamoDB table,
+        it is possible for certain names to get skipped.
+
+        Returns:
+            generator
+        """
+        kwargs = {
+            'ProjectionExpression': 'RuleName',
+            'Select': 'SPECIFIC_ATTRIBUTES',
+
+            # It is acceptable to use inconsistent reads here to reduce read capacity units
+            # consumed, as there is already no guarantee of consistent in the rule names due to
+            # pagination.
+            'ConsistentRead': False,
+        }
+
+        rule_names = set()
+        for item in self._paginate(self._table.scan, kwargs):
+            name = item['RuleName']
+            if name in rule_names:
+                continue
+
+            rule_names.add(name)
+            yield name
+
     def get_alert_records(self, rule_name, alert_proc_timeout_sec):
         """Find all alerts for the given rule which need to be dispatched to the alert processor.
 

tests/unit/streamalert/alert_merger/test_main.py

@@ -78,15 +78,15 @@ def teardown(self):
         self.lambda_mock.stop()
 
     @patch.object(main, 'LOGGER')
-    def test_get_alerts(self, mock_logger):
+    def test_alert_generator(self, mock_logger):
         """Alert Merger - Sorted Alerts - Invalid Alerts are Logged"""
         records = [
             Alert('test_rule', {}, {'output'}).dynamo_record(),
             {'Nonsense': 'Record'}
         ]
 
         with patch.object(self.merger.table, 'get_alert_records', return_value=records):
-            result = self.merger._get_alerts('test_rule')
+            result = list(self.merger._alert_generator('test_rule'))
             # Valid record is returned
             assert_equal(1, len(result))
             assert_equal(records[0]['AlertID'], result[0].alert_id)

tests/unit/streamalert/shared/test_alert_table.py

@@ -74,6 +74,10 @@ def test_rule_names(self):
         """Alert Table - Rule Names From Table Scan"""
         assert_equal({'even', 'odd'}, self.alert_table.rule_names())
 
+    def test_rule_names_generator(self):
+        """Alert Table - Rule Names Generator From Table Scan"""
+        assert_equal({'even', 'odd'}, set(self.alert_table.rule_names_generator()))
+
     def test_get_alert_records(self):
         """Alert Table - Pending Alerts From Table Query"""
         result = list(self.alert_table.get_alert_records('odd', _ALERT_PROCESSOR_TIMEOUT_SEC))