Test for memory leaks

posthog/hogql/test/_test_parser.py

@@ -6,11 +6,18 @@
 from posthog.hogql.errors import HogQLException, SyntaxException
 from posthog.hogql.parser import parse_expr, parse_order_expr, parse_select
 from posthog.hogql.visitor import clear_locations
-from posthog.test.base import BaseTest
+from posthog.test.base import BaseTest, MemoryLeakTestMixin
 
 
 def parser_test_factory(backend: Literal["python", "cpp"]):
-    class TestParser(BaseTest):
+    base_classes = (MemoryLeakTestMixin, BaseTest) if backend == "cpp" else (BaseTest,)
+
+    class TestParser(*base_classes):
+        MEMORY_INCREASE_PER_PARSE_LIMIT_B = 10_000
+        MEMORY_INCREASE_INCREMENTAL_FACTOR_LIMIT = 0.1
+        MEMORY_PRIMING_RUNS_N = 2
+        MEMORY_LEAK_CHECK_RUNS_N = 100
+
         maxDiff = None
 
         def _expr(self, expr: str, placeholders: Optional[Dict[str, ast.Expr]] = None) -> ast.Expr:

posthog/test/base.py

@@ -1,11 +1,12 @@
 import datetime as dt
 import inspect
 import re
+import resource
 import threading
 import uuid
 from contextlib import contextmanager
 from functools import wraps
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 from unittest.mock import patch
 
 import freezegun
@@ -133,7 +134,7 @@ def validation_error_response(
         return {"type": "validation_error", "code": code, "detail": message, "attr": attr}
 
 
-class TestMixin:
+class TestMixin(TestCase if TYPE_CHECKING else object):
     CONFIG_ORGANIZATION_NAME: str = "Test"
     CONFIG_EMAIL: Optional[str] = "[email protected]"
     CONFIG_PASSWORD: Optional[str] = "testpassword12345"
@@ -194,6 +195,41 @@ def validate_basic_html(self, html_message, site_url, preheader=None):
             self.assertIn(preheader, html_message)  # type: ignore
 
 
+class MemoryLeakTestMixin(TestCase if TYPE_CHECKING else object):
+    MEMORY_INCREASE_PER_PARSE_LIMIT_B: int
+    """Parsing more than once can never increase memory by this much (on average)"""
+    MEMORY_INCREASE_INCREMENTAL_FACTOR_LIMIT: float
+    """Parsing more than once can never increase memory by more than this factor * first run's increase (on average)"""
+    MEMORY_PRIMING_RUNS_N: int
+    """How many times to run every test method to prime the heap"""
+    MEMORY_LEAK_CHECK_RUNS_N: int
+    """How many times to run every test method to check for memory leaks"""
+
+    def _callTestMethod(self, method):
+        mem_original_b = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+        for _ in range(self.MEMORY_PRIMING_RUNS_N):  # Priming runs
+            method()
+        mem_primed_b = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+        for _ in range(self.MEMORY_LEAK_CHECK_RUNS_N):  # Memory leak check runs
+            method()
+        mem_tested_b = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+        avg_memory_priming_increase_b = (mem_primed_b - mem_original_b) / self.MEMORY_PRIMING_RUNS_N
+        avg_memory_test_increase_b = (mem_tested_b - mem_primed_b) / self.MEMORY_LEAK_CHECK_RUNS_N
+        avg_memory_increase_factor = (
+            avg_memory_test_increase_b / avg_memory_priming_increase_b if avg_memory_priming_increase_b else 0
+        )
+        self.assertLessEqual(
+            avg_memory_test_increase_b,
+            self.MEMORY_INCREASE_PER_PARSE_LIMIT_B,
+            f"Possible memory leak - exceeded {self.MEMORY_INCREASE_PER_PARSE_LIMIT_B}-byte limit of incremental memory per parse",
+        )
+        self.assertLessEqual(
+            avg_memory_increase_factor,
+            self.MEMORY_INCREASE_INCREMENTAL_FACTOR_LIMIT,
+            f"Possible memory leak - exceeded {self.MEMORY_INCREASE_INCREMENTAL_FACTOR_LIMIT*100:.2f}% limit of incremental memory per parse",
+        )
+
+
 class BaseTest(TestMixin, ErrorResponsesMixin, TestCase):
     """
     Base class for performing Postgres-based backend unit tests on.