fix(experimentation): retry experiment computes on transient ClickHouse errors

api/experimentation/services.py

@@ -107,6 +107,7 @@
 
 CLICKHOUSE_CONNECT_TIMEOUT_SECONDS = 5
 CLICKHOUSE_QUERY_TIMEOUT_SECONDS = 30
+CLICKHOUSE_BACKGROUND_QUERY_TIMEOUT_SECONDS = 120
 CLICKHOUSE_VERIFY_TIMEOUT_SECONDS = 5
 CLICKHOUSE_EVENT_NAMES_TIMEOUT_SECONDS = 15
 CUSTOMER_EVENT_STATS_CACHE_SECONDS = 60
@@ -151,17 +152,19 @@ def is_experiment_feature_enabled(organisation: Organisation) -> bool:
     )
 
 
-@lru_cache(maxsize=1)
-def _get_clickhouse_client() -> Client:
+@lru_cache(maxsize=2)
+def _get_clickhouse_client(
+    send_receive_timeout: int = CLICKHOUSE_QUERY_TIMEOUT_SECONDS,
+) -> Client:
     """Build a clickhouse-driver client for the experimentation event store.
 
     The database is taken from the DSN path, so queries can reference the
     `events` table unqualified. Connect and query timeouts are bounded unless the
-    DSN overrides them.
+    DSN overrides them. One client is cached per requested timeout.
     """
     host, kwargs = parse_url(settings.EXPERIMENTATION_CLICKHOUSE_URL)
     kwargs.setdefault("connect_timeout", CLICKHOUSE_CONNECT_TIMEOUT_SECONDS)
-    kwargs.setdefault("send_receive_timeout", CLICKHOUSE_QUERY_TIMEOUT_SECONDS)
+    kwargs.setdefault("send_receive_timeout", send_receive_timeout)
     kwargs.setdefault("client_name", settings.CLICKHOUSE_CONNECTION_CLIENT_NAME)
     return Client(host, **kwargs)
 
@@ -347,7 +350,9 @@ def get_exposure_buckets(
     window_end: datetime,
     granularity: ExposureGranularity,
 ) -> list[ExposureBucket]:
-    rows = _get_clickhouse_client().execute(
+    rows = _get_clickhouse_client(
+        send_receive_timeout=CLICKHOUSE_BACKGROUND_QUERY_TIMEOUT_SECONDS,
+    ).execute(
         EXPOSURE_BUCKETS_QUERY.format(
             bucket_function=_EXPOSURE_BUCKET_FUNCTIONS[granularity]
         ),
@@ -390,9 +395,9 @@ def get_metric_variant_stats(
     }
     builder.add_metric_params(params)
 
-    rows, columns = _get_clickhouse_client().execute(
-        builder.build_query(), params, with_column_types=True
-    )
+    rows, columns = _get_clickhouse_client(
+        send_receive_timeout=CLICKHOUSE_BACKGROUND_QUERY_TIMEOUT_SECONDS,
+    ).execute(builder.build_query(), params, with_column_types=True)
     exposure_counts, metric_stats = builder.decode_rows(
         rows, [name for name, _type in columns]
     )

api/experimentation/tasks.py

@@ -6,6 +6,7 @@
     register_recurring_task,
     register_task_handler,
 )
+from task_processor.exceptions import TaskBackoffError
 
 from environments.models import Environment, EnvironmentAPIKey
 from experimentation import ingestion_sync_service
@@ -28,6 +29,8 @@
     deliver_warehouse_events,
 )
 
+COMPUTE_TASK_TIMEOUT = timedelta(minutes=3)
+
 logger = structlog.get_logger("experimentation")
 
 
@@ -162,7 +165,7 @@ def clean_up_old_warehouse_delivery_logs() -> None:
     ).delete()
 
 
-@register_task_handler()
+@register_task_handler(timeout=COMPUTE_TASK_TIMEOUT)
 def compute_experiment_exposures(experiment_id: int) -> None:
     experiment = (
         Experiment.objects.select_related("environment__project", "feature")
@@ -194,12 +197,14 @@ def compute_experiment_exposures(experiment_id: int) -> None:
             environment__id=experiment.environment_id,
             organisation__id=experiment.environment.project.organisation_id,
         )
+        if isinstance(exc, OSError):
+            raise TaskBackoffError() from exc
         return
 
     exposures.record_refresh(summary, as_of)
 
 
-@register_task_handler()
+@register_task_handler(timeout=COMPUTE_TASK_TIMEOUT)
 def compute_experiment_results(experiment_id: int) -> None:
     experiment = (
         Experiment.objects.select_related("environment__project", "feature")
@@ -229,6 +234,8 @@ def compute_experiment_results(experiment_id: int) -> None:
             environment__id=experiment.environment_id,
             organisation__id=experiment.environment.project.organisation_id,
         )
+        if isinstance(exc, OSError):
+            raise TaskBackoffError() from exc
         return
 
     results.record_refresh(summary, as_of)

api/tests/unit/experimentation/test_services.py

@@ -109,6 +109,36 @@ def test_get_clickhouse_client__dsn_timeouts__are_preserved(
     services._get_clickhouse_client.cache_clear()
 
 
+def test_get_clickhouse_client__per_timeout__caches_distinct_clients(
+    mocker: MockerFixture,
+    settings: SettingsWrapper,
+) -> None:
+    # Given
+    settings.EXPERIMENTATION_CLICKHOUSE_URL = "clickhouse://ch.example.com/db"
+    mock_client_cls = mocker.patch(
+        "experimentation.services.Client",
+        side_effect=lambda *args, **kwargs: mocker.Mock(),
+    )
+    services._get_clickhouse_client.cache_clear()
+
+    # When
+    client = services._get_clickhouse_client()
+    same_client = services._get_clickhouse_client()
+    background_client = services._get_clickhouse_client(
+        send_receive_timeout=services.CLICKHOUSE_BACKGROUND_QUERY_TIMEOUT_SECONDS,
+    )
+
+    # Then
+    assert client is same_client
+    assert background_client is not client
+    assert mock_client_cls.call_count == 2
+    assert (
+        mock_client_cls.call_args_list[1].kwargs["send_receive_timeout"]
+        == services.CLICKHOUSE_BACKGROUND_QUERY_TIMEOUT_SECONDS
+    )
+    services._get_clickhouse_client.cache_clear()
+
+
 @pytest.mark.parametrize(
     "rows, expected",
     [
@@ -327,7 +357,7 @@ def test_get_exposure_buckets__day_granularity__queries_and_maps_rows(
     ]
     mock_client = mocker.Mock()
     mock_client.execute.return_value = rows
-    mocker.patch(
+    mock_get_client = mocker.patch(
         "experimentation.services._get_clickhouse_client",
         return_value=mock_client,
     )
@@ -378,6 +408,9 @@ def test_get_exposure_buckets__day_granularity__queries_and_maps_rows(
         "window_start": window_start,
         "window_end": window_end,
     }
+    mock_get_client.assert_called_once_with(
+        send_receive_timeout=services.CLICKHOUSE_BACKGROUND_QUERY_TIMEOUT_SECONDS,
+    )
 
 
 def test_get_exposure_buckets__hour_granularity__buckets_by_hour(
@@ -795,7 +828,7 @@ def test_get_metric_variant_stats__metrics__queries_and_maps_rows(
     ]
     mock_client = mocker.Mock()
     mock_client.execute.return_value = (rows, _result_columns(4))
-    mocker.patch(
+    mock_get_client = mocker.patch(
         "experimentation.services._get_clickhouse_client",
         return_value=mock_client,
     )
@@ -862,6 +895,9 @@ def test_get_metric_variant_stats__metrics__queries_and_maps_rows(
     assert params["metric_2_event"] == "page_view"
     assert params["metric_3_event"] == "session"
     assert params["window_end"] == window_end
+    mock_get_client.assert_called_once_with(
+        send_receive_timeout=services.CLICKHOUSE_BACKGROUND_QUERY_TIMEOUT_SECONDS,
+    )
 
 
 def test_get_metric_variant_stats__three_variants__maps_all_variants(

api/tests/unit/experimentation/test_tasks.py

@@ -14,6 +14,7 @@
 from prometheus_client import REGISTRY
 from pytest_mock import MockerFixture
 from pytest_structlog import StructuredLogCapture
+from task_processor.exceptions import TaskBackoffError
 
 from environments.models import Environment, EnvironmentAPIKey
 from experimentation import warehouse_delivery_service
@@ -39,6 +40,7 @@
     WarehouseDeliveryOutcome,
     WarehouseType,
 )
+from experimentation.services import CLICKHOUSE_BACKGROUND_QUERY_TIMEOUT_SECONDS
 from experimentation.stats import VariantStats
 from experimentation.tasks import (
     clean_up_old_warehouse_delivery_logs,
@@ -355,6 +357,30 @@ def test_compute_experiment_exposures__warehouse_error__records_failure(
     )
 
 
+def test_compute_experiment_exposures__transient_warehouse_error__records_failure_and_backs_off(
+    experiment: Experiment,
+    mocker: MockerFixture,
+    log: StructuredLogCapture,
+) -> None:
+    # Given
+    experiment.status = ExperimentStatus.RUNNING
+    experiment.started_at = datetime(2026, 6, 10, tzinfo=dt_timezone.utc)
+    experiment.save()
+    mocker.patch(
+        "experimentation.tasks.compute_exposures_summary",
+        side_effect=TimeoutError("The read operation timed out"),
+    )
+
+    # When
+    with pytest.raises(TaskBackoffError):
+        compute_experiment_exposures(experiment_id=experiment.id)
+
+    # Then
+    exposures = ExperimentExposures.objects.get(experiment=experiment)
+    assert exposures.last_error_at is not None
+    assert log.has("exposures.compute_failed", level="error")
+
+
 def test_compute_experiment_exposures__not_started_experiment__skips(
     experiment: Experiment,
     mocker: MockerFixture,
@@ -531,6 +557,55 @@ def test_compute_experiment_results__warehouse_error__records_failure(
     ]
 
 
+@pytest.mark.parametrize(
+    "exc",
+    [
+        TimeoutError("The read operation timed out"),
+        ConnectionResetError("Connection reset by peer"),
+    ],
+    ids=["timeout", "reset"],
+)
+def test_compute_experiment_results__transient_warehouse_error__records_failure_and_backs_off(
+    experiment: Experiment,
+    mocker: MockerFixture,
+    log: StructuredLogCapture,
+    exc: Exception,
+) -> None:
+    # Given
+    experiment.status = ExperimentStatus.RUNNING
+    experiment.started_at = datetime(2026, 6, 10, tzinfo=dt_timezone.utc)
+    experiment.save()
+    mocker.patch(
+        "experimentation.tasks.compute_results_summary",
+        side_effect=exc,
+    )
+
+    # When
+    with pytest.raises(TaskBackoffError):
+        compute_experiment_results(experiment_id=experiment.id)
+
+    # Then
+    results = ExperimentResults.objects.get(experiment=experiment)
+    assert results.last_error_at is not None
+    assert log.has("results.compute_failed", level="error")
+
+
+@pytest.mark.parametrize(
+    "task_handler",
+    [compute_experiment_exposures, compute_experiment_results],
+    ids=["exposures", "results"],
+)
+def test_compute_experiment_task_handlers__task_timeout__exceeds_background_query_timeout(
+    task_handler: Any,
+) -> None:
+    # Given
+    task_timeout = task_handler.timeout
+
+    # When / Then
+    assert task_timeout is not None
+    assert task_timeout.total_seconds() > CLICKHOUSE_BACKGROUND_QUERY_TIMEOUT_SECONDS
+
+
 def test_compute_experiment_results__not_started_experiment__skips(
     experiment: Experiment,
     mocker: MockerFixture,