Avoid querying with secondary columns which might fetch and lock additional rows (#20474)

* Accessing UserSession by primary key

This resolves problematic locking queries databases running on SERIALIZABLE isolation level like CockroachDB

Closes #16977

* Avoid querying with expiring column

This resolves problematic locking queries databases running on SERIALIZABLE isolation level like CockroachDB

Closes #16977

model/map-jpa/src/main/java/org/keycloak/models/map/storage/jpa/JpaMapStorage.java

@@ -23,10 +23,10 @@ import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.UUID;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import jakarta.persistence.EntityManager;
 import jakarta.persistence.LockModeType;
-import jakarta.persistence.Parameter;
 import jakarta.persistence.PersistenceException;
 import jakarta.persistence.TypedQuery;
 import jakarta.persistence.criteria.CriteriaBuilder;
@@ -181,7 +181,10 @@ public abstract class JpaMapStorage<RE extends JpaRootEntity, E extends Abstract
         }
 
         JpaPredicateFunction<RE> predicateFunc = mcb.getPredicateFunc();
-        if (this.isExpirableEntity) {
+        if (this.isExpirableEntity && (queryParameters.getLimit() != null || queryParameters.getOffset() != null)) {
+            // only when using pagination exclude expired entities in the query directly
+            // for all other queries, remove the expired results later as those additional predicates might confuse the database
+            // to use a bad index (see: CockroachDB), and we assume that expired entities are cleaned from the DB regularly
             predicateFunc = predicateFunc != null ? predicateFunc.andThen(predicate -> cb.and(predicate, notExpired(cb, query::subquery, root)))
                                                   : this::notExpired;
         }
@@ -197,6 +200,10 @@ public abstract class JpaMapStorage<RE extends JpaRootEntity, E extends Abstract
             // In order to cache the result, the full result needs to be retrieved.
             // There is also no difference to that in Hibernate, as Hibernate will first retrieve all elements from the ResultSet.
             List<RE> resultList = emQuery.getResultList();
+            if (isExpirableEntity) {
+                // remove expired entities when those haven't been excluded by a predicate
+                resultList = resultList.stream().filter(e -> !isExpired((ExpirableEntity) e, true)).collect(Collectors.toList());
+            }
             cache.put(queryCacheKey, resultList);
 
             return closing(resultList.stream()).map(this::mapToEntityDelegateUnique);

model/map/src/main/java/org/keycloak/models/map/singleUseObject/MapSingleUseObjectProvider.java

@@ -144,14 +144,15 @@ public class MapSingleUseObjectProvider implements SingleUseObjectProvider {
         DefaultModelCriteria<SingleUseObjectValueModel> mcb = criteria();
         mcb = mcb.compare(SingleUseObjectValueModel.SearchableFields.OBJECT_KEY, ModelCriteriaBuilder.Operator.EQ, key);
 
-        MapSingleUseObjectEntity singleUseEntity = singleUseObjectTx.read(withCriteria(mcb)).findFirst().orElse(null);
+        return singleUseObjectTx.read(withCriteria(mcb))
-        if (singleUseEntity != null) {
+                .filter(entity -> {
-            if (isExpired(singleUseEntity, false)) {
+                    if (isExpired(entity, false)) {
-                singleUseObjectTx.delete(singleUseEntity.getId());
+                        singleUseObjectTx.delete(entity.getId());
-            } else {
+                        return false;
-                return singleUseEntity;
+                    } else {
-            }
+                        return true;
-        }
+                    }
-        return null;
+                })
+                .findFirst().orElse(null);
     }
  }

model/map/src/main/java/org/keycloak/models/map/userSession/MapUserSessionProvider.java

@@ -187,13 +187,19 @@ public class MapUserSessionProvider implements UserSessionProvider {
             return userEntityToAdapterFunc(realm).apply(userSessionEntity);
         }
 
-        DefaultModelCriteria<UserSessionModel> mcb = realmAndOfflineCriteriaBuilder(realm, false)
+        // This is an exceptional case where not to use the criteria query:
-                .compare(UserSessionModel.SearchableFields.ID, Operator.EQ, id);
+        // As the ID is already known, and we expect in almost all cases to have exactly one row being returned,
+        // the provider fetches the instance by ID and does the filtering in the Java code afterward instead
+        // of using the criteria query. When using a criteria query in earlier versions, the store (CockroachDB) would pick
+        // a wrong optimization path and lock too many DB rows which would result in transaction-not-serializable exceptions
+        // on concurrent transactions. This change has been done in the assumption that all stores would be faster
+        // to evaluate the fetch-by-id than a criteria query.
+        userSessionEntity = storeWithRealm(realm).read(id);
+        if (userSessionEntity != null && Objects.equals(userSessionEntity.getRealmId(), realm.getId()) && !userSessionEntity.isOffline()) {
+            return userEntityToAdapterFunc(realm).apply(userSessionEntity);
+        }
 
-        return storeWithRealm(realm).read(withCriteria(mcb))
+        return null;
-                .findFirst()
-                .map(userEntityToAdapterFunc(realm))
-                .orElse(null);
     }
 
     @Override
@@ -314,7 +320,16 @@ public class MapUserSessionProvider implements UserSessionProvider {
 
         LOG.tracef("removeUserSession(%s, %s)%s", realm, session, getShortStackTrace());
 
-        storeWithRealm(realm).delete(withCriteria(mcb));
+        // This is an exceptional case where not to use the criteria query:
+        // As the ID is already known, the provider does the filtering in the Java code and uses delete-by-id
+        // instead of using the criteria query to delete rows.
+        // When using a criteria query in earlier versions, the store (CockroachDB) would pick
+        // a wrong optimization path and lock too many DB rows which would result in transaction-not-serializable exceptions
+        // on concurrent transactions. This change has been done in the assumption that delete-by-id would be faster than
+        // delete-by-criteria for all stores.
+        if (Objects.equals(session.getRealm(), realm) && !session.isOffline()) {
+            storeWithRealm(realm).delete(session.getId());
+        }
     }
 
     @Override