diff --git a/storage/inkless/src/main/java/io/aiven/inkless/control_plane/postgres/JobUtils.java b/storage/inkless/src/main/java/io/aiven/inkless/control_plane/postgres/JobUtils.java index e2d22210938..76a359f739f 100644 --- a/storage/inkless/src/main/java/io/aiven/inkless/control_plane/postgres/JobUtils.java +++ b/storage/inkless/src/main/java/io/aiven/inkless/control_plane/postgres/JobUtils.java @@ -17,8 +17,16 @@ */ package io.aiven.inkless.control_plane.postgres; +import org.apache.kafka.common.utils.ExponentialBackoff; import org.apache.kafka.common.utils.Time; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.sql.SQLException; +import java.util.Collections; +import java.util.IdentityHashMap; +import java.util.Set; import java.util.concurrent.Callable; import java.util.function.Consumer; @@ -26,24 +34,59 @@ import io.aiven.inkless.control_plane.ControlPlaneException; public class JobUtils { - public static void run(final Runnable runnable) { - try { + private static final Logger LOGGER = LoggerFactory.getLogger(JobUtils.class); + + /** + * Total number of attempts (initial try + retries) for a transient database failure. + */ + static final int MAX_ATTEMPTS = 3; + static final long INITIAL_BACKOFF_MS = 50; + static final long MAX_BACKOFF_MS = 1_000; + static final int BACKOFF_MULTIPLIER = 2; + static final double BACKOFF_JITTER = 0.2; + + /** + * Backoff between retries. Uses jitter (a randomized factor in {@code [1 - jitter, 1 + jitter]}) + * so that many jobs failing against the same struggling standby at once do not retry in lockstep + * (thundering herd) and re-collide on the same recovery-conflict window. + */ + private static final ExponentialBackoff BACKOFF = + new ExponentialBackoff(INITIAL_BACKOFF_MS, BACKOFF_MULTIPLIER, MAX_BACKOFF_MS, BACKOFF_JITTER); + + /** + * PostgreSQL {@code SQLState}s that are safe to retry because they guarantee the transaction did + * not commit (the server rolled it back), so replaying the whole transaction cannot + * double-apply a write. + * + * + * + *

We deliberately do not retry generic connection-loss states (class {@code 08*}, + * {@code 57P01}) here: a connection dropped during commit leaves the transaction + * in-doubt, so blindly replaying a write job could apply it twice. The recovery-conflict family + * that motivated this retry logic is fully covered by the rollback-guaranteed states above. + */ + private static final Set RETRIABLE_SQL_STATES = Set.of("40001", "40P01"); + + public static void run(final Runnable runnable, final Time time, final Consumer durationCallback) { + run(() -> { runnable.run(); - } catch (final Exception e) { - // TODO add retry with backoff - if (e instanceof ControlPlaneException) { - throw (ControlPlaneException) e; - } else { - throw new RuntimeException(e); - } - } + return null; + }, time, durationCallback); } - public static void run(final Runnable runnable, final Time time, final Consumer durationCallback) { + public static T run(final Callable callable, final Time time, final Consumer durationCallback) { try { - TimeUtils.measureDurationMs(time, runnable, durationCallback); + return runWithRetry(callable, time, durationCallback); } catch (final Exception e) { - // TODO add retry with backoff if (e instanceof ControlPlaneException) { throw (ControlPlaneException) e; } else { @@ -52,29 +95,82 @@ public static void run(final Runnable runnable, final Time time, final Consumer< } } - public static T run(final Callable callable) { - try { - return callable.call(); - } catch (final Exception e) { - // TODO add retry with backoff - if (e instanceof ControlPlaneException) { - throw (ControlPlaneException) e; - } else { - throw new RuntimeException(e); + /** + * Execute {@code callable}, retrying with jittered exponential backoff on transient PostgreSQL + * failures (see {@link #RETRIABLE_SQL_STATES}). On exhaustion, or for any non-retriable failure, + * the original exception is rethrown so callers observe the unchanged error contract. + * + *

Only the decisive attempt (the one that succeeds, or the final failure) is reported + * to {@code durationCallback}. Intermediate retried attempts are timed internally but not + * forwarded, so failed-attempt durations do not inflate the latency/rate metrics — exactly when + * an operator relies on them during a conflict storm. + */ + private static T runWithRetry(final Callable callable, final Time time, final Consumer durationCallback) throws Exception { + // measureDurationMs records the just-finished attempt's duration into this holder (via its + // finally block, so it fires on failure too); we forward it to durationCallback only once, + // when we stop retrying. + final long[] lastAttemptDurationMs = {0L}; + final Consumer captureDuration = d -> lastAttemptDurationMs[0] = d; + + for (int attempt = 1; ; attempt++) { + try { + final T result = TimeUtils.measureDurationMs(time, callable, captureDuration); + durationCallback.accept(lastAttemptDurationMs[0]); + return result; + } catch (final Exception e) { + if (attempt >= MAX_ATTEMPTS || !isRetriable(e)) { + // Decisive failure: record its duration once, then propagate unchanged. + durationCallback.accept(lastAttemptDurationMs[0]); + if (attempt > 1) { + LOGGER.warn("Giving up after {} attempts on transient database error", attempt, e); + } + throw e; + } + final long backoffMs = BACKOFF.backoff(attempt - 1); + LOGGER.warn("Transient database error on attempt {}/{}, retrying in {} ms", + attempt, MAX_ATTEMPTS, backoffMs, e); + time.sleep(backoffMs); + if (Thread.currentThread().isInterrupted()) { + // Interrupted during backoff (e.g. broker shutdown). Utils.sleep restores the + // interrupt flag but returns normally, so we must check it explicitly: stop + // retrying, keep the flag set, and surface the last error rather than spinning + // through the remaining attempts and delaying shutdown. + durationCallback.accept(lastAttemptDurationMs[0]); + throw e; + } } } } - public static T run(final Callable callable, final Time time, final Consumer durationCallback) { - try { - return TimeUtils.measureDurationMs(time, callable, durationCallback); - } catch (final Exception e) { - // TODO add retry with backoff - if (e instanceof ControlPlaneException) { - throw (ControlPlaneException) e; - } else { - throw new RuntimeException(e); + /** + * Returns {@code true} if any throwable in the cause chain (including suppressed throwables) is a + * {@link SQLException} whose {@code SQLState} is in {@link #RETRIABLE_SQL_STATES}. The retriable + * cause is typically wrapped several layers deep (e.g. {@code ControlPlaneException} -> + * {@code DataAccessException} -> {@code PSQLException}), so the whole tree is scanned. + */ + static boolean isRetriable(final Throwable throwable) { + // Identity-based visited set guards against self-referential cause chains. + return hasRetriableCause(throwable, Collections.newSetFromMap(new IdentityHashMap<>())); + } + + private static boolean hasRetriableCause(final Throwable throwable, final Set seen) { + if (throwable == null || !seen.add(throwable)) { + return false; + } + if (throwable instanceof SQLException) { + final String sqlState = ((SQLException) throwable).getSQLState(); + if (sqlState != null && RETRIABLE_SQL_STATES.contains(sqlState)) { + return true; + } + } + if (hasRetriableCause(throwable.getCause(), seen)) { + return true; + } + for (final Throwable suppressed : throwable.getSuppressed()) { + if (hasRetriableCause(suppressed, seen)) { + return true; } } + return false; } } diff --git a/storage/inkless/src/test/java/io/aiven/inkless/control_plane/postgres/JobUtilsTest.java b/storage/inkless/src/test/java/io/aiven/inkless/control_plane/postgres/JobUtilsTest.java new file mode 100644 index 00000000000..66f8c2e039b --- /dev/null +++ b/storage/inkless/src/test/java/io/aiven/inkless/control_plane/postgres/JobUtilsTest.java @@ -0,0 +1,248 @@ +/* + * Inkless + * Copyright (C) 2024 - 2025 Aiven OY + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as published by + * the Free Software Foundation, either version 3 of the License, or + * (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU Affero General Public License for more details. + * + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see . + */ +package io.aiven.inkless.control_plane.postgres; + +import org.apache.kafka.common.utils.MockTime; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +import java.sql.SQLException; +import java.util.concurrent.Callable; +import java.util.concurrent.atomic.AtomicInteger; + +import io.aiven.inkless.control_plane.ControlPlaneException; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Unit tests for {@link JobUtils} retry behavior. These use {@link MockTime}, so the backoff + * {@code sleep} advances virtual time only and the tests never block on wall-clock. + */ +@Timeout(10) +class JobUtilsTest { + + /** Serialization failure / recovery conflict — PostgreSQL uses 40001 for both. */ + private static SQLException recoveryConflict() { + return new SQLException("canceling statement due to conflict with recovery", "40001"); + } + + private static SQLException deadlock() { + return new SQLException("deadlock detected", "40P01"); + } + + @Test + void succeedsWithoutRetryWhenNoError() { + final MockTime time = new MockTime(); + final AtomicInteger attempts = new AtomicInteger(); + final AtomicInteger durationCallbacks = new AtomicInteger(); + final long startMs = time.milliseconds(); + + final String result = JobUtils.run(() -> { + attempts.incrementAndGet(); + return "ok"; + }, time, ignored -> durationCallbacks.incrementAndGet()); + + assertThat(result).isEqualTo("ok"); + assertThat(attempts).hasValue(1); + // Exactly one metric sample, no backoff on the happy path. + assertThat(durationCallbacks).hasValue(1); + assertThat(time.milliseconds()).isEqualTo(startMs); + } + + @Test + void retriesRecoveryConflictThenSucceeds() { + final MockTime time = new MockTime(); + final AtomicInteger attempts = new AtomicInteger(); + final AtomicInteger durationCallbacks = new AtomicInteger(); + + // Reproduces the production incident: the read replica cancels the first attempts with a + // recovery conflict (SQLState 40001), wrapped exactly as the control plane wraps it, then a + // later attempt on a fresh pooled connection succeeds. + final Callable flaky = () -> { + if (attempts.incrementAndGet() < 3) { + throw new ControlPlaneException("Error finding batches", + new RuntimeException("Cannot commit transaction", recoveryConflict())); + } + return "recovered"; + }; + + final String result = JobUtils.run(flaky, time, ignored -> durationCallbacks.incrementAndGet()); + + assertThat(result).isEqualTo("recovered"); + assertThat(attempts).hasValue(3); + // Guards against metric double-counting (finding #1): even though the job ran 3 times, only + // the decisive attempt is reported to the duration callback. + assertThat(durationCallbacks).hasValue(1); + } + + @Test + void backsOffWithExponentialJitterBetweenRetries() { + final MockTime time = new MockTime(); + final AtomicInteger attempts = new AtomicInteger(); + final long startMs = time.milliseconds(); + + // Fails on the first two attempts, succeeds on the third, so two backoff sleeps happen. + JobUtils.run(() -> { + if (attempts.incrementAndGet() < 3) { + throw new RuntimeException(recoveryConflict()); + } + return "ok"; + }, time, ignored -> { }); + + // Jittered exponential backoff: attempt 1 sleeps ~50ms, attempt 2 sleeps ~100ms, each + // scaled by a factor in [1 - JITTER, 1 + JITTER]. Assert the total elapsed virtual time + // falls within those combined bounds (this also proves a sleep actually happened). + final long elapsed = time.milliseconds() - startMs; + final long nominal = JobUtils.INITIAL_BACKOFF_MS + + JobUtils.INITIAL_BACKOFF_MS * JobUtils.BACKOFF_MULTIPLIER; + final long minExpected = (long) (nominal * (1 - JobUtils.BACKOFF_JITTER)); + final long maxExpected = (long) (nominal * (1 + JobUtils.BACKOFF_JITTER)); + assertThat(elapsed).isBetween(minExpected, maxExpected); + } + + @Test + void retriesDeadlock() { + final MockTime time = new MockTime(); + final AtomicInteger attempts = new AtomicInteger(); + + final String result = JobUtils.run(() -> { + if (attempts.incrementAndGet() < 2) { + throw new RuntimeException(deadlock()); + } + return "ok"; + }, time, ignored -> { }); + + assertThat(result).isEqualTo("ok"); + assertThat(attempts).hasValue(2); + } + + @Test + void stopsAfterMaxAttemptsAndRethrowsOriginal() { + final MockTime time = new MockTime(); + final AtomicInteger attempts = new AtomicInteger(); + final AtomicInteger durationCallbacks = new AtomicInteger(); + + // Always fails with a retriable error: the caller must observe the unchanged + // ControlPlaneException contract after retries are exhausted. + assertThatThrownBy(() -> JobUtils.run((Callable) () -> { + attempts.incrementAndGet(); + throw new ControlPlaneException("Error finding batches", + new RuntimeException(recoveryConflict())); + }, time, ignored -> durationCallbacks.incrementAndGet())) + .isInstanceOf(ControlPlaneException.class) + .hasMessage("Error finding batches"); + + assertThat(attempts).hasValue(JobUtils.MAX_ATTEMPTS); + // Even on exhaustion the decisive (final) failure is reported exactly once. + assertThat(durationCallbacks).hasValue(1); + } + + @Test + void doesNotRetryNonRetriableSqlState() { + final MockTime time = new MockTime(); + final AtomicInteger attempts = new AtomicInteger(); + final long startMs = time.milliseconds(); + + // 23505 (unique_violation) is a deterministic error; retrying would just fail again. + assertThatThrownBy(() -> JobUtils.run((Callable) () -> { + attempts.incrementAndGet(); + throw new RuntimeException(new SQLException("duplicate key", "23505")); + }, time, ignored -> { })) + .isInstanceOf(RuntimeException.class); + + assertThat(attempts).hasValue(1); + // A deterministic failure must fail fast: no backoff sleep before giving up. + assertThat(time.milliseconds()).isEqualTo(startMs); + } + + @Test + void doesNotRetryPlainRuntimeException() { + final MockTime time = new MockTime(); + final AtomicInteger attempts = new AtomicInteger(); + final long startMs = time.milliseconds(); + + assertThatThrownBy(() -> JobUtils.run((Callable) () -> { + attempts.incrementAndGet(); + throw new IllegalStateException("bug, not transient"); + }, time, ignored -> { })) + .isInstanceOf(RuntimeException.class); + + assertThat(attempts).hasValue(1); + assertThat(time.milliseconds()).isEqualTo(startMs); + } + + @Test + void abandonsRetriesWhenInterruptedDuringBackoff() { + final MockTime time = new MockTime(); + final AtomicInteger attempts = new AtomicInteger(); + + // Set the interrupt flag so that after the first backoff the loop observes an interrupted + // thread (mirrors a broker shutdown interrupting the retry) and stops immediately instead + // of exhausting all attempts. + Thread.currentThread().interrupt(); + try { + assertThatThrownBy(() -> JobUtils.run((Callable) () -> { + attempts.incrementAndGet(); + throw new RuntimeException(recoveryConflict()); + }, time, ignored -> { })) + .isInstanceOf(RuntimeException.class); + + // Only the first attempt ran; the interrupt aborted the retry loop before a second try. + assertThat(attempts).hasValue(1); + // The interrupt status is preserved for cooperative shutdown downstream. + assertThat(Thread.currentThread().isInterrupted()).isTrue(); + } finally { + // Clear the flag so it does not leak into other tests. + Thread.interrupted(); + } + } + + @Test + void wrapsNonControlPlaneExceptionInRuntimeException() { + final MockTime time = new MockTime(); + + assertThatThrownBy(() -> JobUtils.run((Callable) () -> { + throw new SQLException("non-retriable", "08006"); + }, time, ignored -> { })) + .isInstanceOf(RuntimeException.class) + .hasCauseInstanceOf(SQLException.class); + } + + @Test + void isRetriableScansSuppressedThrowables() { + // The connection-terminated form surfaces the retriable cause as a suppressed throwable + // (the failed rollback), not on the main cause chain. + final RuntimeException top = new RuntimeException("Cannot commit transaction"); + top.addSuppressed(new SQLException("terminating connection due to conflict with recovery", "40001")); + + assertThat(JobUtils.isRetriable(top)).isTrue(); + } + + @Test + void isRetriableHandlesSelfReferentialCauseChain() { + final SQLException a = new SQLException("a", "23505"); + final SQLException b = new SQLException("b", "23505"); + a.initCause(b); + b.initCause(a); + + // Must terminate, not StackOverflow, and report no retriable state. + assertThat(JobUtils.isRetriable(a)).isFalse(); + } +}