fix(subscriptions): cap png export height to prevent OOM (#40754)

Author: andyzzhao

ee/tasks/subscriptions/subscription_utils.py

@@ -23,6 +23,9 @@
 
 UTM_TAGS_BASE = "utm_source=posthog&utm_campaign=subscription_report"
 DEFAULT_MAX_ASSET_COUNT = 6
+# Maximum height for screenshots in pixels. This prevents Chrome from consuming excessive memory
+# when rendering very tall pages (e.g., tables with thousands of rows).
+MAX_SCREENSHOT_HEIGHT_PIXELS = 5000
 
 
 def _get_failed_asset_info(assets: list[ExportedAsset], resource: Union[Subscription, SharingConfiguration]) -> dict:
@@ -187,7 +190,9 @@ async def export_single_asset(asset: ExportedAsset) -> None:
                     subscription_id=getattr(resource, "id", None),
                     team_id=resource.team_id,
                 )
-                await database_sync_to_async(exporter.export_asset_direct, thread_sensitive=False)(asset)
+                await database_sync_to_async(exporter.export_asset_direct, thread_sensitive=False)(
+                    asset, max_height_pixels=MAX_SCREENSHOT_HEIGHT_PIXELS
+                )
                 logger.info(
                     "generate_assets_async.asset_exported",
                     asset_id=asset.id,

ee/tasks/test/subscriptions/test_generate_assets_async.py

@@ -221,7 +221,7 @@ async def test_async_foreign_key_access_with_real_subscription(team, user, dashb
 
 @patch("ee.tasks.subscriptions.subscription_utils.exporter.export_asset_direct")
 async def test_async_generate_assets_basic(mock_export: MagicMock, team, user) -> None:
-    def export_success(asset: ExportedAsset) -> None:
+    def export_success(asset: ExportedAsset, **kwargs) -> None:
         asset.content = b"fake image data"
         asset.save()
 
@@ -283,7 +283,7 @@ async def test_async_generate_assets_timeout_continues_with_partial_results(
 async def test_async_generate_assets_partial_success(mock_export: MagicMock, team, user) -> None:
     call_count = 0
 
-    def export_with_partial_success(asset: ExportedAsset) -> None:
+    def export_with_partial_success(asset: ExportedAsset, **kwargs) -> None:
         nonlocal call_count
         call_count += 1
         if call_count <= 2:

posthog/api/test/test_exports.py

@@ -230,7 +230,7 @@ def test_can_create_new_valid_export_insight(self, mock_exporter_task, mock_expo
 
             # Should warm up the cache
             export_image(exported_asset)
-            mock_export_to_png.assert_called_once_with(exported_asset)
+            mock_export_to_png.assert_called_once_with(exported_asset, max_height_pixels=None)
 
             mock_process_query_dict.assert_called_once()
 

posthog/tasks/exporter.py

@@ -62,17 +62,25 @@
     max_retries=3,
 )
 @transaction.atomic
-def export_asset(exported_asset_id: int, limit: Optional[int] = None) -> None:
+def export_asset(
+    exported_asset_id: int,
+    limit: Optional[int] = None,  # For CSV/XLSX: max row count
+    max_height_pixels: Optional[int] = None,  # For images: max screenshot height in pixels
+) -> None:
     # if Celery is lagging then you can end up with an exported asset that has had a TTL added
     # and that TTL has passed, in the exporter we don't care about that.
     # the TTL is for later cleanup.
     exported_asset: ExportedAsset = ExportedAsset.objects_including_ttl_deleted.select_related(
         "created_by", "team", "team__organization"
     ).get(pk=exported_asset_id)
-    export_asset_direct(exported_asset, limit)
+    export_asset_direct(exported_asset, limit=limit, max_height_pixels=max_height_pixels)
 
 
-def export_asset_direct(exported_asset: ExportedAsset, limit: Optional[int] = None) -> None:
+def export_asset_direct(
+    exported_asset: ExportedAsset,
+    limit: Optional[int] = None,  # For CSV/XLSX: max row count
+    max_height_pixels: Optional[int] = None,  # For images: max screenshot height in pixels
+) -> None:
     from posthog.tasks.exports import csv_exporter, image_exporter
 
     start_time = perf_counter()
@@ -107,7 +115,7 @@ def export_asset_direct(exported_asset: ExportedAsset, limit: Optional[int] = No
             csv_exporter.export_tabular(exported_asset, limit=limit)
             EXPORT_QUEUED_COUNTER.labels(type="csv").inc()
         else:
-            image_exporter.export_image(exported_asset)
+            image_exporter.export_image(exported_asset, max_height_pixels=max_height_pixels)
             EXPORT_QUEUED_COUNTER.labels(type="image").inc()
 
         logger.info(

posthog/tasks/exports/image_exporter.py

@@ -74,7 +74,7 @@ def get_driver() -> webdriver.Chrome:
     )
 
 
-def _export_to_png(exported_asset: ExportedAsset) -> None:
+def _export_to_png(exported_asset: ExportedAsset, max_height_pixels: Optional[int] = None) -> None:
     """
     Exporting an Insight means:
     1. Loading the Insight from the web app in a dedicated rendering mode
@@ -150,7 +150,9 @@ def _export_to_png(exported_asset: ExportedAsset) -> None:
 
         logger.info("exporting_asset", asset_id=exported_asset.id, render_url=url_to_render)
 
-        _screenshot_asset(image_path, url_to_render, screenshot_width, wait_for_css_selector, screenshot_height)
+        _screenshot_asset(
+            image_path, url_to_render, screenshot_width, wait_for_css_selector, screenshot_height, max_height_pixels
+        )
 
         with open(image_path, "rb") as image_file:
             image_data = image_file.read()
@@ -181,6 +183,7 @@ def _screenshot_asset(
     screenshot_width: ScreenWidth,
     wait_for_css_selector: CSSSelector,
     screenshot_height: int = 600,
+    max_height_pixels: Optional[int] = None,
 ) -> None:
     driver: Optional[webdriver.Chrome] = None
     try:
@@ -238,6 +241,15 @@ def _screenshot_asset(
         """
         )
 
+        if max_height_pixels and height > max_height_pixels:
+            logger.warning(
+                "screenshot_height_capped",
+                original_height=height,
+                capped_height=max_height_pixels,
+                url=url_to_render,
+            )
+            height = max_height_pixels
+
         # For example funnels use a table that can get very wide, so try to get its width
         # For replay players, check for player width
         width = driver.execute_script(
@@ -280,6 +292,15 @@ def _screenshot_asset(
         """
         )
 
+        if max_height_pixels and final_height > max_height_pixels:
+            logger.warning(
+                "screenshot_final_height_capped",
+                original_final_height=final_height,
+                capped_height=max_height_pixels,
+                url=url_to_render,
+            )
+            final_height = max_height_pixels
+
         # Set final window size
         driver.set_window_size(width, final_height + HEIGHT_OFFSET)
         driver.save_screenshot(image_path)
@@ -302,7 +323,7 @@ def _screenshot_asset(
             driver.quit()
 
 
-def export_image(exported_asset: ExportedAsset) -> None:
+def export_image(exported_asset: ExportedAsset, max_height_pixels: Optional[int] = None) -> None:
     with posthoganalytics.new_context():
         posthoganalytics.tag("team_id", exported_asset.team_id if exported_asset else "unknown")
         posthoganalytics.tag("asset_id", exported_asset.id if exported_asset else "unknown")
@@ -324,7 +345,7 @@ def export_image(exported_asset: ExportedAsset) -> None:
 
             if exported_asset.export_format == "image/png":
                 with EXPORT_TIMER.labels(type="image").time():
-                    _export_to_png(exported_asset)
+                    _export_to_png(exported_asset, max_height_pixels=max_height_pixels)
                 EXPORT_SUCCEEDED_COUNTER.labels(type="image").inc()
             else:
                 raise NotImplementedError(