From 26243f84e61e7de95190f6446f67995399746e6e Mon Sep 17 00:00:00 2001 From: Konstantinos Stefanidis Vozikis Date: Thu, 13 Aug 2026 14:27:36 +0200 Subject: [PATCH] fix: add bounded jitter to commit retries --- .github/workflows/test-python.yml | 5 +- src/tower/_tables.py | 156 ++++++++-------- tests/tower/test_table_retries.py | 284 ++++++++++++++++++++++++++++++ 3 files changed, 366 insertions(+), 79 deletions(-) create mode 100644 tests/tower/test_table_retries.py diff --git a/.github/workflows/test-python.yml b/.github/workflows/test-python.yml index bbb1c01e..18ad47b8 100644 --- a/.github/workflows/test-python.yml +++ b/.github/workflows/test-python.yml @@ -95,4 +95,7 @@ jobs: PY - name: Run Iceberg compatibility tests - run: uv run --no-sync pytest tests/tower/test_tables.py tests/tower/test_storage.py + run: >- + uv run --no-sync pytest + tests/tower/test_table*.py + tests/tower/test_storage.py diff --git a/src/tower/_tables.py b/src/tower/_tables.py index 2b0d8401..b477fe79 100644 --- a/src/tower/_tables.py +++ b/src/tower/_tables.py @@ -1,15 +1,16 @@ from __future__ import annotations +import math import os +import random +import time from dataclasses import dataclass -from typing import List, Optional, TypeVar, Union +from typing import Callable, List, Optional, TypeVar, Union from pyiceberg.exceptions import CommitFailedException, NoSuchTableError TTable = TypeVar("TTable", bound="Table") - -import random -import time +TRetryResult = TypeVar("TRetryResult") import polars as pl import pyarrow as pa @@ -37,6 +38,8 @@ namespace_or_default, ) +_MAX_COMMIT_RETRY_DELAY_SECONDS = 30.0 + @dataclass class RowsAffectedInformation: @@ -249,8 +252,34 @@ def rows_affected(self) -> RowsAffectedInformation: def _validate_retry_args(max_retries: int, retry_delay_seconds: float) -> None: if max_retries < 0: raise ValueError("max_retries must be >= 0") - if retry_delay_seconds < 0: - raise ValueError("retry_delay_seconds must be >= 0") + if not math.isfinite(retry_delay_seconds) or retry_delay_seconds < 0: + raise ValueError("retry_delay_seconds must be finite and >= 0") + + def _commit_with_retry( + self, + operation: Callable[[], TRetryResult], + max_retries: int, + initial_retry_ceiling_seconds: float, + ) -> TRetryResult: + retry_ceiling_seconds = min( + initial_retry_ceiling_seconds, _MAX_COMMIT_RETRY_DELAY_SECONDS + ) + + for attempt in range(max_retries + 1): + try: + return operation() + except CommitFailedException: + if attempt == max_retries: + raise + + delay_seconds = random.uniform(0.0, retry_ceiling_seconds) + time.sleep(delay_seconds) + self._table.refresh() + retry_ceiling_seconds = min( + retry_ceiling_seconds * 2, _MAX_COMMIT_RETRY_DELAY_SECONDS + ) + + raise AssertionError("unreachable") def insert( self, @@ -270,8 +299,9 @@ def insert( must match the schema of the target table. max_retries (int): Maximum number of retry attempts on commit conflicts. Defaults to 5. - retry_delay_seconds (float): Wait time in seconds between retries. - Defaults to 0.5 seconds. + retry_delay_seconds (float): Maximum randomized wait before the first retry, + in seconds. The maximum doubles after each conflict but never exceeds + 30 seconds; values above 30 are treated as 30. Defaults to 0.5 seconds. Returns: TTable: The table instance with the newly inserted rows, allowing for method chaining. @@ -296,23 +326,11 @@ def insert( self._validate_retry_args(max_retries, retry_delay_seconds) self._ensure_read_write_table() - last_exception = None - - for attempt in range(max_retries + 1): - try: - if attempt > 0: - self._table.refresh() - - self._table.append(data) - self._stats.inserts += data.num_rows - return self - - except CommitFailedException as e: - last_exception = e - if attempt < max_retries: - time.sleep(retry_delay_seconds) - - raise last_exception + self._commit_with_retry( + lambda: self._table.append(data), max_retries, retry_delay_seconds + ) + self._stats.inserts += data.num_rows + return self def upsert( self, @@ -337,8 +355,9 @@ def upsert( If not provided, all columns will be used for matching. max_retries (int): Maximum number of retry attempts on commit conflicts. Defaults to 5. - retry_delay_seconds (float): Wait time in seconds between retries. - Defaults to 0.5 seconds. + retry_delay_seconds (float): Maximum randomized wait before the first retry, + in seconds. The maximum doubles after each conflict but never exceeds + 30 seconds; values above 30 are treated as 30. Defaults to 0.5 seconds. Returns: TTable: The table instance with the upserted rows, allowing for method chaining. @@ -370,34 +389,24 @@ def upsert( self._validate_retry_args(max_retries, retry_delay_seconds) self._ensure_read_write_table() - last_exception = None - - for attempt in range(max_retries + 1): - try: - if attempt > 0: - self._table.refresh() - - res = self._table.upsert( - data, - join_cols=join_cols, - # All upserts will always be case sensitive. Perhaps we'll add this - # as a parameter in the future? - case_sensitive=True, - # These are the defaults, but we're including them to be complete. - when_matched_update_all=True, - when_not_matched_insert_all=True, - ) - - self._stats.updates += res.rows_updated - self._stats.inserts += res.rows_inserted - return self - - except CommitFailedException as e: - last_exception = e - if attempt < max_retries: - time.sleep(retry_delay_seconds) + res = self._commit_with_retry( + lambda: self._table.upsert( + data, + join_cols=join_cols, + # All upserts will always be case sensitive. Perhaps we'll add this + # as a parameter in the future? + case_sensitive=True, + # These are the defaults, but we're including them to be complete. + when_matched_update_all=True, + when_not_matched_insert_all=True, + ), + max_retries, + retry_delay_seconds, + ) - raise last_exception + self._stats.updates += res.rows_updated + self._stats.inserts += res.rows_inserted + return self def delete( self, @@ -421,8 +430,9 @@ def delete( - A string expression max_retries (int): Maximum number of retry attempts on commit conflicts. Defaults to 5. - retry_delay_seconds (float): Wait time in seconds between retries. - Defaults to 0.5 seconds. + retry_delay_seconds (float): Maximum randomized wait before the first retry, + in seconds. The maximum doubles after each conflict but never exceeds + 30 seconds; values above 30 are treated as 30. Defaults to 0.5 seconds. Returns: TTable: The table instance with the deleted rows, allowing for method chaining. @@ -455,30 +465,20 @@ def delete( next_filters = convert_pyarrow_expressions(filters) filters = next_filters - last_exception = None - - for attempt in range(max_retries + 1): - try: - if attempt > 0: - self._table.refresh() - - self._table.delete( - delete_filter=filters, - # We want this to always be the case. Not sure why you wouldn't? - case_sensitive=True, - ) - - # NOTE: There is, unfortunately, no way to get the number of rows - # deleted besides comparing the two snapshots that were created. - - return self + self._commit_with_retry( + lambda: self._table.delete( + delete_filter=filters, + # We want this to always be the case. Not sure why you wouldn't? + case_sensitive=True, + ), + max_retries, + retry_delay_seconds, + ) - except CommitFailedException as e: - last_exception = e - if attempt < max_retries: - time.sleep(retry_delay_seconds) + # NOTE: There is, unfortunately, no way to get the number of rows + # deleted besides comparing the two snapshots that were created. - raise last_exception + return self def schema(self) -> pa.Schema: """ diff --git a/tests/tower/test_table_retries.py b/tests/tower/test_table_retries.py new file mode 100644 index 00000000..6de9a76e --- /dev/null +++ b/tests/tower/test_table_retries.py @@ -0,0 +1,284 @@ +from types import SimpleNamespace + +import httpx +import pyarrow as pa +import pytest +from pyiceberg.exceptions import ( + AuthorizationExpiredError, + BadRequestError, + CommitFailedException, + CommitStateUnknownException, + ForbiddenError, + NoSuchTableError, + ServerError, + ServiceUnavailableError, + UnauthorizedError, + WaitingForLockException, +) + +import tower._tables as tables_module +from tower._context import TowerContext + + +class FakeMutationTable: + def __init__(self, failures=(), refresh_error=None): + self.failures = list(failures) + self.refresh_error = refresh_error + self.mutations = [] + self.refresh_calls = 0 + self.events = [] + + def _mutate(self, operation): + self.mutations.append(operation) + self.events.append(("mutation", operation)) + if self.failures: + raise self.failures.pop(0) + return SimpleNamespace(rows_inserted=1, rows_updated=2) + + def append(self, data): + return self._mutate("insert") + + def upsert(self, data, **kwargs): + return self._mutate("upsert") + + def delete(self, **kwargs): + return self._mutate("delete") + + def refresh(self): + self.refresh_calls += 1 + self.events.append(("refresh",)) + if self.refresh_error is not None: + raise self.refresh_error + + +def make_table(iceberg_table): + context = TowerContext( + tower_url="https://api.example.com", + environment="production", + ) + return tables_module.Table(context, iceberg_table) + + +def run_mutation(table, operation, max_retries, retry_delay_seconds): + if operation == "insert": + return table.insert( + pa.table({"id": [1, 2, 3]}), + max_retries=max_retries, + retry_delay_seconds=retry_delay_seconds, + ) + if operation == "upsert": + return table.upsert( + pa.table({"id": [1, 2, 3]}), + join_cols=["id"], + max_retries=max_retries, + retry_delay_seconds=retry_delay_seconds, + ) + if operation == "delete": + return table.delete( + "id = 1", + max_retries=max_retries, + retry_delay_seconds=retry_delay_seconds, + ) + raise AssertionError(f"unknown operation: {operation}") + + +@pytest.mark.parametrize( + ("operation", "expected_inserts", "expected_updates"), + [ + ("insert", 3, 0), + ("upsert", 1, 2), + ("delete", 0, 0), + ], +) +def test_mutations_retry_commit_conflicts_with_exponential_full_jitter( + monkeypatch, operation, expected_inserts, expected_updates +): + iceberg_table = FakeMutationTable( + [CommitFailedException("conflict 1"), CommitFailedException("conflict 2")] + ) + table = make_table(iceberg_table) + uniform_calls = [] + sleep_calls = [] + + def uniform(low, high): + uniform_calls.append((low, high)) + iceberg_table.events.append(("jitter", low, high)) + return high / 2 + + def sleep(delay): + sleep_calls.append(delay) + iceberg_table.events.append(("sleep", delay)) + + monkeypatch.setattr(tables_module.random, "uniform", uniform) + monkeypatch.setattr(tables_module.time, "sleep", sleep) + + result = run_mutation(table, operation, max_retries=2, retry_delay_seconds=0.5) + + assert result is table + assert iceberg_table.mutations == [operation, operation, operation] + assert iceberg_table.refresh_calls == 2 + assert uniform_calls == [(0.0, 0.5), (0.0, 1.0)] + assert sleep_calls == [0.25, 0.5] + assert iceberg_table.events == [ + ("mutation", operation), + ("jitter", 0.0, 0.5), + ("sleep", 0.25), + ("refresh",), + ("mutation", operation), + ("jitter", 0.0, 1.0), + ("sleep", 0.5), + ("refresh",), + ("mutation", operation), + ] + assert table.rows_affected() == tables_module.RowsAffectedInformation( + inserts=expected_inserts, + updates=expected_updates, + ) + + +def test_commit_retry_backoff_is_capped(monkeypatch): + iceberg_table = FakeMutationTable( + [CommitFailedException(f"conflict {attempt}") for attempt in range(5)] + ) + table = make_table(iceberg_table) + uniform_calls = [] + + def uniform(low, high): + uniform_calls.append((low, high)) + return 0.0 + + monkeypatch.setattr(tables_module.random, "uniform", uniform) + monkeypatch.setattr(tables_module.time, "sleep", lambda delay: None) + + table.insert(pa.table({"id": [1]}), max_retries=5, retry_delay_seconds=10.0) + + assert uniform_calls == [ + (0.0, 10.0), + (0.0, 20.0), + (0.0, 30.0), + (0.0, 30.0), + (0.0, 30.0), + ] + + +def test_commit_retry_initial_ceiling_is_clamped(monkeypatch): + iceberg_table = FakeMutationTable([CommitFailedException("conflict")]) + table = make_table(iceberg_table) + uniform_calls = [] + + def uniform(low, high): + uniform_calls.append((low, high)) + return 0.0 + + monkeypatch.setattr(tables_module.random, "uniform", uniform) + monkeypatch.setattr(tables_module.time, "sleep", lambda delay: None) + + table.insert(pa.table({"id": [1]}), max_retries=1, retry_delay_seconds=300.0) + + assert uniform_calls == [(0.0, 30.0)] + + +@pytest.mark.parametrize("max_retries", [0, 2]) +def test_commit_retry_exhaustion_preserves_final_exception(monkeypatch, max_retries): + failures = [ + CommitFailedException(f"conflict {attempt}") + for attempt in range(max_retries + 1) + ] + iceberg_table = FakeMutationTable(failures) + table = make_table(iceberg_table) + sleep_calls = [] + + monkeypatch.setattr(tables_module.random, "uniform", lambda low, high: 0.0) + monkeypatch.setattr(tables_module.time, "sleep", sleep_calls.append) + + with pytest.raises(CommitFailedException) as exc_info: + table.insert( + pa.table({"id": [1]}), + max_retries=max_retries, + retry_delay_seconds=0.5, + ) + + assert exc_info.value is failures[-1] + assert iceberg_table.mutations == ["insert"] * (max_retries + 1) + assert iceberg_table.refresh_calls == max_retries + assert len(sleep_calls) == max_retries + assert table.rows_affected().inserts == 0 + + +@pytest.mark.parametrize( + "exception_type", + [ + CommitStateUnknownException, + ServiceUnavailableError, + AuthorizationExpiredError, + UnauthorizedError, + ForbiddenError, + NoSuchTableError, + ServerError, + BadRequestError, + WaitingForLockException, + httpx.TimeoutException, + httpx.ConnectError, + ValueError, + ], +) +def test_mutation_errors_other_than_commit_conflicts_are_not_retried( + monkeypatch, exception_type +): + failure = exception_type("not retryable") + iceberg_table = FakeMutationTable([failure]) + table = make_table(iceberg_table) + + def unexpected_call(*args, **kwargs): + raise AssertionError("non-retryable errors must not back off") + + monkeypatch.setattr(tables_module.random, "uniform", unexpected_call) + monkeypatch.setattr(tables_module.time, "sleep", unexpected_call) + + with pytest.raises(exception_type) as exc_info: + table.insert( + pa.table({"id": [1]}), + max_retries=5, + retry_delay_seconds=0.5, + ) + + assert exc_info.value is failure + assert iceberg_table.mutations == ["insert"] + assert iceberg_table.refresh_calls == 0 + assert table.rows_affected().inserts == 0 + + +def test_refresh_failure_is_not_retried(monkeypatch): + refresh_failure = RuntimeError("refresh failed") + iceberg_table = FakeMutationTable( + [CommitFailedException("conflict")], refresh_error=refresh_failure + ) + table = make_table(iceberg_table) + + monkeypatch.setattr(tables_module.random, "uniform", lambda low, high: 0.0) + monkeypatch.setattr(tables_module.time, "sleep", lambda delay: None) + + with pytest.raises(RuntimeError) as exc_info: + table.insert( + pa.table({"id": [1]}), + max_retries=5, + retry_delay_seconds=0.5, + ) + + assert exc_info.value is refresh_failure + assert iceberg_table.mutations == ["insert"] + assert iceberg_table.refresh_calls == 1 + assert table.rows_affected().inserts == 0 + + +@pytest.mark.parametrize( + "retry_delay_seconds", [float("nan"), float("inf"), float("-inf")] +) +def test_commit_retry_rejects_non_finite_delay(retry_delay_seconds): + iceberg_table = FakeMutationTable() + table = make_table(iceberg_table) + + with pytest.raises(ValueError, match="must be finite and >= 0"): + table.insert(pa.table({"id": [1]}), retry_delay_seconds=retry_delay_seconds) + + assert iceberg_table.mutations == []