retry.java

package retry;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.SQLException;
import java.sql.Savepoint;
import java.sql.Statement;

public class RetryTest
{
    private RetryTest() {}

    private static final String COCKROACH_VERSION = "22.1";
    private static final String COCKROACH_CONTAINER = "cockroachdb/cockroach:latest-v%s".formatted(COCKROACH_VERSION);

    private static final int MAX_RETRIES = 5;

    private static final String SAVEPOINT_NAME = "cockroach_restart";

    public static void main(String[] args)
            throws Exception
    {
        ExtendedCockroachContainer extendedCockroachContainer = new ExtendedCockroachContainer();
        extendedCockroachContainer.start();

        try (Connection connection = DriverManager.getConnection(extendedCockroachContainer.getJdbcUrl(), extendedCockroachContainer.getUsername(), extendedCockroachContainer.getPassword())) {
            try (Statement statement = connection.createStatement()) {
                statement.execute("CREATE TABLE test (id int NOT NULL PRIMARY KEY, name STRING NOT NULL)");
            }
        }


        for (int i = 0; i < 100; ++i ) {
            int index = i;
            try (Connection connection = DriverManager.getConnection(extendedCockroachContainer.getJdbcUrl(), extendedCockroachContainer.getUsername(), extendedCockroachContainer.getPassword())) {
                System.out.println("Trying " + i);
                retry(connection, statement -> {
                    statement.execute("SELECT * FROM test");
                    statement.execute("INSERT INTO test (id, name) VALUES (%s, 'hey')".formatted(index));
                });
            }
        }
    }

    private interface Handler
    {
        void call(Statement statement)
                throws SQLException;
    }

    private static void retry(Connection connection, Handler handler)
            throws Exception
    {
        connection.setAutoCommit(false);
        try (Statement statement = connection.createStatement()) {
            statement.execute("SET inject_retry_errors_enabled = 'true'");
        }

        Savepoint savepoint = connection.setSavepoint(SAVEPOINT_NAME);
        for (int attempt = 0; attempt < MAX_RETRIES; ++attempt) {
            try {
                try (Statement statement = connection.createStatement()) {
                    handler.call(statement);
                }

                if (attempt == 0) {
                    throw new RuntimeException("Didn't retry");
                }

                return;
            }
            catch (Exception e) {
                connection.rollback(savepoint);
            }
        }

        throw new RuntimeException("Retries expired");
    }
}