From acaf1724dde6fd0b2d3071c8743d7bf73fc77769 Mon Sep 17 00:00:00 2001 From: Stefan Guilhen Date: Mon, 29 Aug 2022 16:43:38 -0300 Subject: [PATCH] Fix ComponentsTest failures with CockroachDB - Component addition/edition/removal is now executed in a retriable transaction. Closes #13209 --- .../models/utils/KeycloakModelUtils.java | 80 +++++++++++++++++++ .../utils/LockObjectsForModification.java | 5 ++ .../models/KeycloakSessionTaskWithResult.java | 35 ++++++++ .../keys/JavaKeystoreKeyProviderFactory.java | 3 +- .../resources/admin/ComponentResource.java | 72 ++++++++++------- 5 files changed, 163 insertions(+), 32 deletions(-) create mode 100644 server-spi/src/main/java/org/keycloak/models/KeycloakSessionTaskWithResult.java diff --git a/server-spi-private/src/main/java/org/keycloak/models/utils/KeycloakModelUtils.java b/server-spi-private/src/main/java/org/keycloak/models/utils/KeycloakModelUtils.java index bf27f768f4..f3c070ef1e 100755 --- a/server-spi-private/src/main/java/org/keycloak/models/utils/KeycloakModelUtils.java +++ b/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 the type returned by the callable. + * @return the value computed by the callable. + */ + public static V runJobInRetriableTransaction(final KeycloakSessionFactory factory, final KeycloakSessionTaskWithResult 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) diff --git a/server-spi-private/src/main/java/org/keycloak/utils/LockObjectsForModification.java b/server-spi-private/src/main/java/org/keycloak/utils/LockObjectsForModification.java index 61c63b0701..986dfd3d9b 100644 --- a/server-spi-private/src/main/java/org/keycloak/utils/LockObjectsForModification.java +++ b/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 lockRealmsForModification(KeycloakSession session, CallableWithoutThrowingAnException callable) { + return lockObjectsForModification(session, RealmModel.class, callable); + } + private static V lockObjectsForModification(KeycloakSession session, Class model, CallableWithoutThrowingAnException 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. diff --git a/server-spi/src/main/java/org/keycloak/models/KeycloakSessionTaskWithResult.java b/server-spi/src/main/java/org/keycloak/models/KeycloakSessionTaskWithResult.java new file mode 100644 index 0000000000..e21a261096 --- /dev/null +++ b/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 the type of the computed result. + * @author Stefan Guilhen + */ +@FunctionalInterface +public interface KeycloakSessionTaskWithResult { + + /** + * Computes a result. + * + * @param session a reference to the {@link KeycloakSession}. + * @return the computed result. + */ + V run(final KeycloakSession session); +} diff --git a/services/src/main/java/org/keycloak/keys/JavaKeystoreKeyProviderFactory.java b/services/src/main/java/org/keycloak/keys/JavaKeystoreKeyProviderFactory.java index d8918f6857..ff39ab9448 100644 --- a/services/src/main/java/org/keycloak/keys/JavaKeystoreKeyProviderFactory.java +++ b/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) - .loadKey(session.getContext().getRealm(), model); + new JavaKeystoreKeyProvider(realm, model).loadKey(realm, model); } catch (Throwable t) { logger.error("Failed to load keys.", t); throw new ComponentValidationException("Failed to load keys. " + t.getMessage(), t); diff --git a/services/src/main/java/org/keycloak/services/resources/admin/ComponentResource.java b/services/src/main/java/org/keycloak/services/resources/admin/ComponentResource.java index 459edb423d..2311aa889d 100644 --- a/services/src/main/java/org/keycloak/services/resources/admin/ComponentResource.java +++ b/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 { - ComponentModel model = RepresentationToModel.toModel(session, rep); - if (model.getParentId() == null) model.setParentId(realm.getId()); + return KeycloakModelUtils.runJobInRetriableTransaction(session.getKeycloakSessionFactory(), kcSession -> { + RealmModel realmModel = LockObjectsForModification.lockRealmsForModification(kcSession, () -> kcSession.realms().getRealm(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(); - return Response.created(session.getContext().getUri().getAbsolutePathBuilder().path(model.getId()).build()).build(); - } catch (ComponentValidationException e) { - return localizedErrorResponse(e); - } catch (IllegalArgumentException e) { - throw new BadRequestException(e); - } + adminEvent.operation(OperationType.CREATE).resourcePath(kcSession.getContext().getUri(), model.getId()).representation(StripSecretsUtils.strip(kcSession, rep)).success(); + return Response.created(kcSession.getContext().getUri().getAbsolutePathBuilder().path(model.getId()).build()).build(); + } catch (ComponentValidationException e) { + return localizedErrorResponse(e); + } catch (IllegalArgumentException 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 { - ComponentModel model = realm.getComponent(id); - if (model == null) { - throw new NotFoundException("Could not find component"); + return KeycloakModelUtils.runJobInRetriableTransaction(session.getKeycloakSessionFactory(), kcSession -> { + RealmModel realmModel = LockObjectsForModification.lockRealmsForModification(kcSession, () -> kcSession.realms().getRealm(realm.getId())); + try { + 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); - 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(); - } + }, 10, 100); } @DELETE @Path("{id}") public void removeComponent(@PathParam("id") String id) { auth.realm().requireManageRealm(); - ComponentModel model = realm.getComponent(id); - if (model == null) { - throw new NotFoundException("Could not find component"); - } - adminEvent.operation(OperationType.DELETE).resourcePath(session.getContext().getUri()).success(); - realm.removeComponent(model); + KeycloakModelUtils.runJobInRetriableTransaction(session.getKeycloakSessionFactory(), kcSession -> { + RealmModel realmModel = LockObjectsForModification.lockRealmsForModification(kcSession, () -> kcSession.realms().getRealm(realm.getId())); + 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) {