Fix ComponentsTest failures with CockroachDB

- Component addition/edition/removal is now executed in a retriable transaction.

Closes #13209

server-spi-private/src/main/java/org/keycloak/models/utils/KeycloakModelUtils.java

@@ -39,6 +39,7 @@ import org.keycloak.models.IdentityProviderModel;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.KeycloakSessionFactory;
 import org.keycloak.models.KeycloakSessionTask;
+import org.keycloak.models.KeycloakSessionTaskWithResult;
 import org.keycloak.models.KeycloakTransaction;
 import org.keycloak.models.RealmModel;
 import org.keycloak.models.RealmProvider;
@@ -59,11 +60,13 @@ import java.security.KeyPair;
 import java.security.PrivateKey;
 import java.security.PublicKey;
 import java.security.cert.X509Certificate;
+import java.sql.SQLException;
 import java.util.Collection;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.Random;
 import java.util.Set;
 import java.util.UUID;
 import java.util.stream.Collectors;
@@ -273,6 +276,83 @@ public final class KeycloakModelUtils {
         }
     }
 
+    /**
+     * Creates a new {@link KeycloakSession} and runs the specified callable in a new transaction. If the transaction fails
+     * with a SQL retriable error, the method re-executes the specified callable until it either succeeds or the maximum number
+     * of attempts is reached, leaving some increasing random delay milliseconds between the invocations. It uses the exponential
+     * backoff + jitter algorithm to compute the delay, which is limited to {@code attemptsCount * retryIntervalMillis}.
+     * More details https://aws.amazon.com/blogs/architecture/exponential-backoff-and-jitter/
+     *
+     * @param factory a reference to the {@link KeycloakSessionFactory}.
+     * @param callable a reference to the {@link KeycloakSessionTaskWithResult} that will be executed in a retriable way.
+     * @param attemptsCount the maximum number of attempts to execute the callable.
+     * @param retryIntervalMillis the base interval value in millis used to compute the delay.
+     * @param <V> the type returned by the callable.
+     * @return the value computed by the callable.
+     */
+    public static <V> V runJobInRetriableTransaction(final KeycloakSessionFactory factory, final KeycloakSessionTaskWithResult<V> callable,
+                                                     final int attemptsCount, final int retryIntervalMillis) {
+        int retryCount = 0;
+        Random rand = new Random();
+        V result;
+        while (true) {
+            KeycloakSession session = factory.create();
+            KeycloakTransaction tx = session.getTransactionManager();
+            try {
+                tx.begin();
+                result = callable.run(session);
+                if (tx.isActive()) {
+                    if (tx.getRollbackOnly()) {
+                        tx.rollback();
+                    } else {
+                        tx.commit();
+                    }
+                }
+                break;
+            } catch (RuntimeException re) {
+                if (tx.isActive()) {
+                    tx.rollback();
+                }
+                if (isExceptionRetriable(re) && ++retryCount < attemptsCount) {
+                    int delay = Math.min(retryIntervalMillis * attemptsCount, (1 << retryCount) * retryIntervalMillis)
+                            + rand.nextInt(retryIntervalMillis);
+                    try {
+                        Thread.sleep(delay);
+                    } catch (InterruptedException ie) {
+                        ie.addSuppressed(re);
+                        throw new RuntimeException(ie);
+                    }
+                } else {
+                    throw re;
+                }
+            } finally {
+                session.close();
+            }
+        }
+        return result;
+    }
+
+    /**
+     * Checks if the specified exception is retriable or not. A retriable exception must be an instance of {@code SQLException}
+     * and must have a 40001 SQL retriable state. This is a standard SQL state as defined in SQL standard, and across the
+     * implementations its meaning boils down to "deadlock" (applies to Postgres, MSSQL, Oracle, MySQL, and others).
+     *
+     * @param exception the exception to be checked.
+     * @return {@code true} if the exception is retriable; {@code false} otherwise.
+     */
+    public static boolean isExceptionRetriable(final Exception exception) {
+        Objects.requireNonNull(exception);
+        // first find the root cause and check if it is a SQLException
+        Throwable rootCause = exception;
+        while (rootCause.getCause() != null && rootCause.getCause() != rootCause) {
+            rootCause = rootCause.getCause();
+        }
+        if (rootCause instanceof SQLException) {
+            // check if the exception state is a recoverable one (40001)
+            return "40001".equals(((SQLException) rootCause).getSQLState());
+        }
+        return false;
+    }
 
     /**
      * Wrap given runnable job into KeycloakTransaction. Set custom timeout for the JTA transaction (in case we're in the environment with JTA enabled)

server-spi-private/src/main/java/org/keycloak/utils/LockObjectsForModification.java

@@ -18,6 +18,7 @@
 package org.keycloak.utils;
 
 import org.keycloak.models.KeycloakSession;
+import org.keycloak.models.RealmModel;
 import org.keycloak.models.UserSessionModel;
 
 import java.util.HashSet;
@@ -65,6 +66,10 @@ public class LockObjectsForModification {
         return lockObjectsForModification(session, UserSessionModel.class, callable);
     }
 
+    public static <V> V lockRealmsForModification(KeycloakSession session, CallableWithoutThrowingAnException<V> callable) {
+        return lockObjectsForModification(session, RealmModel.class, callable);
+    }
+
     private static <V> V lockObjectsForModification(KeycloakSession session, Class<?> model, CallableWithoutThrowingAnException<V> callable) {
         if (LockObjectsForModification.isEnabled(session, model)) {
             // If someone nests the call, and it would already be locked, don't try to lock it a second time.

server-spi/src/main/java/org/keycloak/models/KeycloakSessionTaskWithResult.java

@@ -0,0 +1,35 @@
+/*
+ * Copyright 2022 Red Hat, Inc. and/or its affiliates
+ * and other contributors as indicated by the @author tags.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.keycloak.models;
+
+/**
+ * Interface for tasks that compute a result and need access to the {@link KeycloakSession}.
+ *
+ * @param <V> the type of the computed result.
+ * @author <a href="mailto:sguilhen@redhat.com">Stefan Guilhen</a>
+ */
+@FunctionalInterface
+public interface KeycloakSessionTaskWithResult<V> {
+
+    /**
+     * Computes a result.
+     *
+     * @param session a reference to the {@link KeycloakSession}.
+     * @return the computed result.
+     */
+    V run(final KeycloakSession session);
+}

services/src/main/java/org/keycloak/keys/JavaKeystoreKeyProviderFactory.java

@@ -75,8 +75,7 @@ public class JavaKeystoreKeyProviderFactory extends AbstractRsaKeyProviderFactor
                 .checkSingle(KEY_PASSWORD_PROPERTY, true);
 
         try {
-            new JavaKeystoreKeyProvider(session.getContext().getRealm(), model)
+            new JavaKeystoreKeyProvider(realm, model).loadKey(realm, model);
-                    .loadKey(session.getContext().getRealm(), model);
         } catch (Throwable t) {
             logger.error("Failed to load keys.", t);
             throw new ComponentValidationException("Failed to load keys. " + t.getMessage(), t);

services/src/main/java/org/keycloak/services/resources/admin/ComponentResource.java

@@ -28,6 +28,7 @@ import org.keycloak.events.admin.OperationType;
 import org.keycloak.events.admin.ResourceType;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.RealmModel;
+import org.keycloak.models.utils.KeycloakModelUtils;
 import org.keycloak.models.utils.ModelToRepresentation;
 import org.keycloak.models.utils.RepresentationToModel;
 import org.keycloak.models.utils.StripSecretsUtils;
@@ -39,6 +40,7 @@ import org.keycloak.representations.idm.ComponentTypeRepresentation;
 import org.keycloak.representations.idm.ConfigPropertyRepresentation;
 import org.keycloak.services.ErrorResponse;
 import org.keycloak.services.resources.admin.permissions.AdminPermissionEvaluator;
+import org.keycloak.utils.LockObjectsForModification;
 
 import javax.ws.rs.BadRequestException;
 import javax.ws.rs.Consumes;
@@ -126,19 +128,22 @@ public class ComponentResource {
     @Consumes(MediaType.APPLICATION_JSON)
     public Response create(ComponentRepresentation rep) {
         auth.realm().requireManageRealm();
-        try {
+        return KeycloakModelUtils.runJobInRetriableTransaction(session.getKeycloakSessionFactory(), kcSession -> {
-            ComponentModel model = RepresentationToModel.toModel(session, rep);
+            RealmModel realmModel = LockObjectsForModification.lockRealmsForModification(kcSession, () -> kcSession.realms().getRealm(realm.getId()));
-            if (model.getParentId() == null) model.setParentId(realm.getId());
+            try {
+                ComponentModel model = RepresentationToModel.toModel(kcSession, rep);
+                if (model.getParentId() == null) model.setParentId(realmModel.getId());
 
-            model = realm.addComponentModel(model);
+                model = realmModel.addComponentModel(model);
 
-            adminEvent.operation(OperationType.CREATE).resourcePath(session.getContext().getUri(), model.getId()).representation(StripSecretsUtils.strip(session, rep)).success();
+                adminEvent.operation(OperationType.CREATE).resourcePath(kcSession.getContext().getUri(), model.getId()).representation(StripSecretsUtils.strip(kcSession, rep)).success();
-            return Response.created(session.getContext().getUri().getAbsolutePathBuilder().path(model.getId()).build()).build();
+                return Response.created(kcSession.getContext().getUri().getAbsolutePathBuilder().path(model.getId()).build()).build();
-        } catch (ComponentValidationException e) {
+            } catch (ComponentValidationException e) {
-            return localizedErrorResponse(e);
+                return localizedErrorResponse(e);
-        } catch (IllegalArgumentException e) {
+            } catch (IllegalArgumentException e) {
-            throw new BadRequestException(e);
+                throw new BadRequestException(e);
-        }
+            }
+        }, 10, 100);
     }
 
     @GET
@@ -160,32 +165,39 @@ public class ComponentResource {
     @Consumes(MediaType.APPLICATION_JSON)
     public Response updateComponent(@PathParam("id") String id, ComponentRepresentation rep) {
         auth.realm().requireManageRealm();
-        try {
+        return KeycloakModelUtils.runJobInRetriableTransaction(session.getKeycloakSessionFactory(), kcSession -> {
-            ComponentModel model = realm.getComponent(id);
+            RealmModel realmModel = LockObjectsForModification.lockRealmsForModification(kcSession, () -> kcSession.realms().getRealm(realm.getId()));
-            if (model == null) {
+            try {
-                throw new NotFoundException("Could not find component");
+                ComponentModel model = realmModel.getComponent(id);
+                if (model == null) {
+                    throw new NotFoundException("Could not find component");
+                }
+                RepresentationToModel.updateComponent(kcSession, rep, model, false);
+                adminEvent.operation(OperationType.UPDATE).resourcePath(kcSession.getContext().getUri()).representation(StripSecretsUtils.strip(kcSession, rep)).success();
+                realmModel.updateComponent(model);
+                return Response.noContent().build();
+            } catch (ComponentValidationException e) {
+                return localizedErrorResponse(e);
+            } catch (IllegalArgumentException e) {
+                throw new BadRequestException();
             }
-            RepresentationToModel.updateComponent(session, rep, model, false);
+        }, 10, 100);
-            adminEvent.operation(OperationType.UPDATE).resourcePath(session.getContext().getUri()).representation(StripSecretsUtils.strip(session, rep)).success();
-            realm.updateComponent(model);
-            return Response.noContent().build();
-        } catch (ComponentValidationException e) {
-            return localizedErrorResponse(e);
-        } catch (IllegalArgumentException e) {
-            throw new BadRequestException();
-        }
     }
     @DELETE
     @Path("{id}")
     public void removeComponent(@PathParam("id") String id) {
         auth.realm().requireManageRealm();
-        ComponentModel model = realm.getComponent(id);
+        KeycloakModelUtils.runJobInRetriableTransaction(session.getKeycloakSessionFactory(), kcSession -> {
-        if (model == null) {
+            RealmModel realmModel = LockObjectsForModification.lockRealmsForModification(kcSession, () -> kcSession.realms().getRealm(realm.getId()));
-            throw new NotFoundException("Could not find component");
-        }
-        adminEvent.operation(OperationType.DELETE).resourcePath(session.getContext().getUri()).success();
-        realm.removeComponent(model);
 
+            ComponentModel model = realmModel.getComponent(id);
+            if (model == null) {
+                throw new NotFoundException("Could not find component");
+            }
+            adminEvent.operation(OperationType.DELETE).resourcePath(kcSession.getContext().getUri()).success();
+            realmModel.removeComponent(model);
+            return null;
+        }, 10 , 100);
     }
 
     private Response localizedErrorResponse(ComponentValidationException cve) {