chore(batch-exports): Add memory cap to buffer queue

PostHog · Oct 16, 2024 · 20bba96 · 20bba96
1 parent c30c1df
commit 20bba96
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/posthog/settings/temporal.py b/posthog/settings/temporal.py
@@ -17,6 +17,7 @@
 BATCH_EXPORT_BIGQUERY_UPLOAD_CHUNK_SIZE_BYTES: int = 1024 * 1024 * 100  # 100MB
 BATCH_EXPORT_HTTP_UPLOAD_CHUNK_SIZE_BYTES: int = 1024 * 1024 * 50  # 50MB
 BATCH_EXPORT_HTTP_BATCH_SIZE: int = 5000
+BATCH_EXPORT_BUFFER_QUEUE_MAX_SIZE_BYTES: int = 1024 * 1024 * 300  # 300MB
 
 UNCONSTRAINED_TIMESTAMP_TEAM_IDS: list[str] = get_list(os.getenv("UNCONSTRAINED_TIMESTAMP_TEAM_IDS", ""))
 ASYNC_ARROW_STREAMING_TEAM_IDS: list[str] = get_list(os.getenv("ASYNC_ARROW_STREAMING_TEAM_IDS", ""))

diff --git a/posthog/temporal/batch_exports/batch_exports.py b/posthog/temporal/batch_exports/batch_exports.py
@@ -370,7 +370,7 @@ def start_produce_batch_export_record_batches(
     extra_query_parameters = parameters.pop("extra_query_parameters") or {}
     parameters = {**parameters, **extra_query_parameters}
 
-    queue = RecordBatchQueue()
+    queue = RecordBatchQueue(max_size_bytes=settings.BATCH_EXPORT_BUFFER_QUEUE_MAX_SIZE_BYTES)
     query_id = uuid.uuid4()
     done_event = asyncio.Event()
     produce_task = asyncio.create_task(