2021-08-17 06:30:52 -05:00
/ *
* Copyright 2021 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.map.storage.hotRod ;
2023-02-28 06:42:08 -06:00
import org.infinispan.client.hotrod.MetadataValue ;
2021-08-17 06:30:52 -05:00
import org.infinispan.client.hotrod.RemoteCache ;
import org.infinispan.client.hotrod.Search ;
import org.infinispan.commons.util.CloseableIterator ;
import org.infinispan.query.dsl.Query ;
import org.infinispan.query.dsl.QueryFactory ;
import org.jboss.logging.Logger ;
2022-06-20 04:48:57 -05:00
import org.keycloak.common.util.Time ;
2023-01-25 08:33:26 -06:00
import org.keycloak.models.AbstractKeycloakTransaction ;
2021-08-17 06:30:52 -05:00
import org.keycloak.models.KeycloakSession ;
2023-02-28 06:42:08 -06:00
import org.keycloak.models.UserSessionModel ;
2021-08-17 06:30:52 -05:00
import org.keycloak.models.map.common.AbstractEntity ;
import org.keycloak.models.map.common.DeepCloner ;
2022-06-20 04:48:57 -05:00
import org.keycloak.models.map.common.ExpirableEntity ;
2022-03-08 09:02:14 -06:00
import org.keycloak.models.map.common.StringKeyConverter ;
2021-08-17 06:30:52 -05:00
import org.keycloak.models.map.storage.MapKeycloakTransaction ;
import org.keycloak.models.map.storage.MapStorage ;
2023-01-25 08:33:26 -06:00
import org.keycloak.models.map.storage.ModelEntityUtil ;
2021-08-17 06:30:52 -05:00
import org.keycloak.models.map.storage.QueryParameters ;
import org.keycloak.models.map.storage.chm.ConcurrentHashMapCrudOperations ;
import org.keycloak.models.map.storage.chm.ConcurrentHashMapKeycloakTransaction ;
import org.keycloak.models.map.storage.chm.MapFieldPredicates ;
import org.keycloak.models.map.storage.chm.MapModelCriteriaBuilder ;
2023-01-25 08:33:26 -06:00
import org.keycloak.models.map.storage.criteria.DefaultModelCriteria ;
import org.keycloak.models.map.storage.hotRod.common.AbstractHotRodEntity ;
import org.keycloak.models.map.storage.hotRod.common.HotRodEntityDelegate ;
import org.keycloak.models.map.storage.hotRod.common.HotRodEntityDescriptor ;
import org.keycloak.models.map.storage.hotRod.connections.DefaultHotRodConnectionProviderFactory ;
import org.keycloak.models.map.storage.hotRod.connections.HotRodConnectionProvider ;
import org.keycloak.models.map.storage.hotRod.locking.HotRodLocksUtils ;
2023-02-01 09:24:07 -06:00
import org.keycloak.models.map.storage.hotRod.transaction.AllAreasHotRodTransactionsWrapper ;
2023-01-25 08:33:26 -06:00
import org.keycloak.models.map.storage.hotRod.transaction.NoActionHotRodTransactionWrapper ;
2021-08-17 06:30:52 -05:00
import org.keycloak.storage.SearchableModelField ;
2023-01-25 08:33:26 -06:00
import org.keycloak.utils.LockObjectsForModification ;
2021-08-17 06:30:52 -05:00
2023-02-28 06:42:08 -06:00
import javax.persistence.OptimisticLockException ;
2023-01-25 08:33:26 -06:00
import java.time.Duration ;
2023-02-28 06:42:08 -06:00
import java.util.HashMap ;
2021-08-17 06:30:52 -05:00
import java.util.Map ;
import java.util.Objects ;
import java.util.Spliterators ;
2022-08-29 09:12:28 -05:00
import java.util.concurrent.TimeUnit ;
2021-11-23 09:46:17 -06:00
import java.util.function.Function ;
2021-08-17 06:30:52 -05:00
import java.util.stream.Collectors ;
import java.util.stream.Stream ;
import java.util.stream.StreamSupport ;
2023-01-25 08:33:26 -06:00
import static org.keycloak.common.util.StackUtil.getShortStackTrace ;
2021-11-29 09:50:35 -06:00
import static org.keycloak.models.map.storage.hotRod.common.HotRodUtils.paginateQuery ;
2021-08-17 06:30:52 -05:00
import static org.keycloak.utils.StreamsUtil.closing ;
2022-07-11 01:54:39 -05:00
public class HotRodMapStorage < K , E extends AbstractHotRodEntity , V extends AbstractEntity & HotRodEntityDelegate < E > , M > implements MapStorage < V , M > , ConcurrentHashMapCrudOperations < V , M > {
2021-08-17 06:30:52 -05:00
private static final Logger LOG = Logger . getLogger ( HotRodMapStorage . class ) ;
2023-01-25 08:33:26 -06:00
private final KeycloakSession session ;
2021-11-23 09:46:17 -06:00
private final RemoteCache < K , E > remoteCache ;
2022-05-25 03:43:47 -05:00
protected final StringKeyConverter < K > keyConverter ;
protected final HotRodEntityDescriptor < E , V > storedEntityDescriptor ;
2021-11-23 09:46:17 -06:00
private final Function < E , V > delegateProducer ;
2022-05-25 03:43:47 -05:00
protected final DeepCloner cloner ;
2022-06-20 04:48:57 -05:00
protected boolean isExpirableEntity ;
2023-02-01 09:24:07 -06:00
private final AllAreasHotRodTransactionsWrapper txWrapper ;
2023-01-25 08:33:26 -06:00
private final Map < SearchableModelField < ? super M > , MapModelCriteriaBuilder . UpdatePredicatesFunc < K , V , M > > fieldPredicates ;
private final Long lockTimeout ;
private final RemoteCache < String , String > locksCache ;
2023-02-28 06:42:08 -06:00
private final Map < K , Long > entityVersionCache = new HashMap < > ( ) ;
2021-08-17 06:30:52 -05:00
2023-01-25 08:33:26 -06:00
public HotRodMapStorage ( KeycloakSession session , RemoteCache < K , E > remoteCache , StringKeyConverter < K > keyConverter , HotRodEntityDescriptor < E , V > storedEntityDescriptor , DeepCloner cloner , AllAreasHotRodTransactionsWrapper txWrapper , Long lockTimeout ) {
this . session = session ;
2021-08-17 06:30:52 -05:00
this . remoteCache = remoteCache ;
2022-03-08 09:02:14 -06:00
this . keyConverter = keyConverter ;
2021-08-17 06:30:52 -05:00
this . storedEntityDescriptor = storedEntityDescriptor ;
this . cloner = cloner ;
2021-11-23 09:46:17 -06:00
this . delegateProducer = storedEntityDescriptor . getHotRodDelegateProvider ( ) ;
2022-06-20 04:48:57 -05:00
this . isExpirableEntity = ExpirableEntity . class . isAssignableFrom ( ModelEntityUtil . getEntityType ( storedEntityDescriptor . getModelTypeClass ( ) ) ) ;
2023-02-01 09:24:07 -06:00
this . txWrapper = txWrapper ;
2023-01-25 08:33:26 -06:00
this . fieldPredicates = MapFieldPredicates . getPredicates ( ( Class < M > ) storedEntityDescriptor . getModelTypeClass ( ) ) ;
this . lockTimeout = lockTimeout ;
HotRodConnectionProvider cacheProvider = session . getProvider ( HotRodConnectionProvider . class ) ;
this . locksCache = cacheProvider . getRemoteCache ( DefaultHotRodConnectionProviderFactory . HOT_ROD_LOCKS_CACHE_NAME ) ;
2021-08-17 06:30:52 -05:00
}
@Override
public V create ( V value ) {
2022-03-08 09:02:14 -06:00
K key = keyConverter . fromStringSafe ( value . getId ( ) ) ;
2021-08-17 06:30:52 -05:00
if ( key = = null ) {
2022-03-08 09:02:14 -06:00
key = keyConverter . yieldNewUniqueKey ( ) ;
value = cloner . from ( keyConverter . keyToString ( key ) , value ) ;
2021-08-17 06:30:52 -05:00
}
2022-08-29 09:12:28 -05:00
if ( isExpirableEntity ) {
Long lifespan = getLifespan ( value ) ;
if ( lifespan ! = null ) {
if ( lifespan > 0 ) {
remoteCache . putIfAbsent ( key , value . getHotRodEntity ( ) , lifespan , TimeUnit . MILLISECONDS ) ;
} else {
LOG . warnf ( " Skipped creation of entity %s in storage due to negative/zero lifespan. " , key ) ;
}
return value ;
}
}
2021-11-23 09:46:17 -06:00
remoteCache . putIfAbsent ( key , value . getHotRodEntity ( ) ) ;
2021-08-17 06:30:52 -05:00
return value ;
}
2023-01-25 08:33:26 -06:00
private String getLockName ( String key ) {
return storedEntityDescriptor . getModelTypeClass ( ) . getName ( ) + " _ " + key ;
}
2021-08-17 06:30:52 -05:00
@Override
public V read ( String key ) {
Objects . requireNonNull ( key , " Key must be non-null " ) ;
2022-03-08 09:02:14 -06:00
K k = keyConverter . fromStringSafe ( key ) ;
2022-06-20 04:48:57 -05:00
2023-01-25 08:33:26 -06:00
if ( LockObjectsForModification . isEnabled ( session , storedEntityDescriptor . getModelTypeClass ( ) ) ) {
String lockName = getLockName ( key ) ;
HotRodLocksUtils . repeatPutIfAbsent ( locksCache , lockName , Duration . ofMillis ( lockTimeout ) , 50 , true ) ;
session . getTransactionManager ( ) . enlistAfterCompletion ( new AbstractKeycloakTransaction ( ) {
@Override
protected void commitImpl ( ) {
HotRodLocksUtils . removeWithInstanceIdentifier ( locksCache , lockName ) ;
}
@Override
protected void rollbackImpl ( ) {
HotRodLocksUtils . removeWithInstanceIdentifier ( locksCache , lockName ) ;
}
} ) ;
}
2022-06-21 04:41:42 -05:00
// Obtain value from Infinispan
2023-02-28 06:42:08 -06:00
MetadataValue < E > entityWithMetadata = remoteCache . getWithMetadata ( k ) ;
if ( entityWithMetadata = = null ) return null ;
// store entity version
LOG . tracef ( " Entity %s read in version %s " , key , entityWithMetadata . getVersion ( ) , getShortStackTrace ( ) ) ;
entityVersionCache . put ( k , entityWithMetadata . getVersion ( ) ) ;
2022-06-21 04:41:42 -05:00
// Create delegate that implements Map*Entity
2023-02-28 06:42:08 -06:00
return entityWithMetadata . getValue ( ) ! = null ? delegateProducer . apply ( entityWithMetadata . getValue ( ) ) : null ;
2021-08-17 06:30:52 -05:00
}
@Override
public V update ( V value ) {
2022-03-08 09:02:14 -06:00
K key = keyConverter . fromStringSafe ( value . getId ( ) ) ;
2022-06-21 04:41:42 -05:00
2022-08-29 09:12:28 -05:00
if ( isExpirableEntity ) {
Long lifespan = getLifespan ( value ) ;
if ( lifespan ! = null ) {
if ( lifespan > 0 ) {
2023-02-28 06:42:08 -06:00
if ( ! remoteCache . replaceWithVersion ( key , value . getHotRodEntity ( ) , entityVersionCache . get ( key ) , lifespan , TimeUnit . MILLISECONDS , - 1 , TimeUnit . MILLISECONDS ) ) {
throw new OptimisticLockException ( " Entity " + key + " with version " + entityVersionCache . get ( key ) + " already changed by a different transaction. " ) ;
}
2022-08-29 09:12:28 -05:00
} else {
2023-02-28 06:42:08 -06:00
if ( ! remoteCache . removeWithVersion ( key , entityVersionCache . get ( key ) ) ) {
throw new OptimisticLockException ( " Entity " + key + " with version " + entityVersionCache . get ( key ) + " already changed by a different transaction. " ) ;
}
2022-08-29 09:12:28 -05:00
LOG . warnf ( " Removing entity %s from storage due to negative/zero lifespan. " , key ) ;
}
2023-02-28 06:42:08 -06:00
return delegateProducer . apply ( value . getHotRodEntity ( ) ) ;
2022-08-29 09:12:28 -05:00
}
}
2023-02-28 06:42:08 -06:00
if ( ! remoteCache . replaceWithVersion ( key , value . getHotRodEntity ( ) , entityVersionCache . get ( key ) ) ) {
throw new OptimisticLockException ( " Entity " + key + " with version " + entityVersionCache . get ( key ) + " already changed by a different transaction. " ) ;
}
return delegateProducer . apply ( value . getHotRodEntity ( ) ) ;
2021-08-17 06:30:52 -05:00
}
@Override
public boolean delete ( String key ) {
2022-03-08 09:02:14 -06:00
K k = keyConverter . fromStringSafe ( key ) ;
2023-02-28 06:42:08 -06:00
Long entityVersion = entityVersionCache . get ( k ) ;
if ( entityVersion ! = null ) {
if ( ! remoteCache . removeWithVersion ( k , entityVersion ) ) {
throw new OptimisticLockException ( " Entity " + key + " with version " + entityVersion + " already changed by a different transaction. " ) ;
}
return true ;
}
2021-08-17 06:30:52 -05:00
return remoteCache . remove ( k ) ! = null ;
}
private static String toOrderString ( QueryParameters . OrderBy < ? > orderBy ) {
SearchableModelField < ? > field = orderBy . getModelField ( ) ;
String modelFieldName = IckleQueryMapModelCriteriaBuilder . getFieldName ( field ) ;
String orderString = orderBy . getOrder ( ) . equals ( QueryParameters . Order . ASCENDING ) ? " ASC " : " DESC " ;
return modelFieldName + " " + orderString ;
}
@Override
public Stream < V > read ( QueryParameters < M > queryParameters ) {
2023-01-25 08:33:26 -06:00
DefaultModelCriteria < M > dmc = queryParameters . getModelCriteriaBuilder ( ) ;
// Optimization if the criteria contains only one id
String id = ( String ) dmc . getSingleRestrictionArgument ( " id " ) ;
if ( id ! = null ) {
// We have a criteria that contains "id EQ 'some_key'". We can change this to reading only some_key using read method and then apply the rest of criteria.
MapModelCriteriaBuilder < K , V , M > mapMcb = dmc . flashToModelCriteriaBuilder ( new MapModelCriteriaBuilder < > ( keyConverter , fieldPredicates ) ) ;
V entity = read ( id ) ;
if ( entity = = null ) {
return Stream . empty ( ) ;
}
2023-02-28 06:42:08 -06:00
K k = keyConverter . fromString ( id ) ;
boolean fulfillsQueryCriteria = mapMcb . getKeyFilter ( ) . test ( k ) & & mapMcb . getEntityFilter ( ) . test ( entity ) ;
2023-01-25 08:33:26 -06:00
if ( ! fulfillsQueryCriteria ) {
// entity does not fulfill whole criteria, we can release lock now
2023-02-28 06:42:08 -06:00
if ( LockObjectsForModification . isEnabled ( session , storedEntityDescriptor . getModelTypeClass ( ) ) ) {
HotRodLocksUtils . removeWithInstanceIdentifier ( locksCache , getLockName ( id ) ) ;
entityVersionCache . remove ( k ) ;
}
2023-01-25 08:33:26 -06:00
return Stream . empty ( ) ;
}
return Stream . of ( entity ) ;
}
2023-02-28 06:42:08 -06:00
// workaround if the query contains us.clientId field, in which case don't read by id => read without optimistic locking.
// See https://issues.redhat.com/browse/ISPN-14537
if ( ! dmc . isEmpty ( ) & & dmc . partiallyEvaluate ( ( field , op , arg ) - > field = = UserSessionModel . SearchableFields . CLIENT_ID ) . toString ( ) . contains ( " __TRUE__ " ) ) {
Query < E > query = prepareQueryWithPrefixAndParameters ( null , queryParameters ) ;
CloseableIterator < E > iterator = paginateQuery ( query , queryParameters . getOffset ( ) ,
queryParameters . getLimit ( ) ) . iterator ( ) ;
return closing ( StreamSupport . stream ( Spliterators . spliteratorUnknownSize ( iterator , 0 ) , false ) )
. onClose ( iterator : : close )
. filter ( Objects : : nonNull ) // see https://github.com/keycloak/keycloak/issues/9271
. map ( this . delegateProducer ) ;
}
// Criteria does not contain only one id, we need to read ids without locking and then read entities one by one pessimistically or optimistically
2023-01-25 08:33:26 -06:00
Query < Object [ ] > query = prepareQueryWithPrefixAndParameters ( " SELECT id " , queryParameters ) ;
CloseableIterator < Object [ ] > iterator = paginateQuery ( query , queryParameters . getOffset ( ) ,
queryParameters . getLimit ( ) ) . iterator ( ) ;
return closing ( StreamSupport . stream ( Spliterators . spliteratorUnknownSize ( iterator , 0 ) , false ) )
. onClose ( iterator : : close )
// Extract ids from the result
. map ( a - > a [ 0 ] )
. map ( String . class : : cast )
2023-02-28 06:42:08 -06:00
// read by id => this will register the entity in an ISPN transaction
2023-01-25 08:33:26 -06:00
. map ( this : : read )
// Entity can be removed in the meanwhile, we need to check for null
. filter ( Objects : : nonNull ) ;
}
private < T > Query < T > prepareQueryWithPrefixAndParameters ( String prefix , QueryParameters < M > queryParameters ) {
2021-11-30 11:06:19 -06:00
IckleQueryMapModelCriteriaBuilder < E , M > iqmcb = queryParameters . getModelCriteriaBuilder ( )
2021-08-17 06:30:52 -05:00
. flashToModelCriteriaBuilder ( createCriteriaBuilder ( ) ) ;
2023-01-25 08:33:26 -06:00
String queryString = ( prefix ! = null ? prefix : " " ) + iqmcb . getIckleQuery ( ) ;
2021-08-17 06:30:52 -05:00
if ( ! queryParameters . getOrderBy ( ) . isEmpty ( ) ) {
queryString + = " ORDER BY " + queryParameters . getOrderBy ( ) . stream ( ) . map ( HotRodMapStorage : : toOrderString )
2023-01-25 08:33:26 -06:00
. collect ( Collectors . joining ( " , " ) ) ;
2021-08-17 06:30:52 -05:00
}
2023-01-25 08:33:26 -06:00
LOG . tracef ( " Preparing Ickle query: '%s'%s " , queryString , getShortStackTrace ( ) ) ;
2021-08-17 06:30:52 -05:00
QueryFactory queryFactory = Search . getQueryFactory ( remoteCache ) ;
2023-01-25 08:33:26 -06:00
Query < T > query = queryFactory . create ( queryString ) ;
2021-08-17 06:30:52 -05:00
query . setParameters ( iqmcb . getParameters ( ) ) ;
2023-01-25 08:33:26 -06:00
return query ;
2021-08-17 06:30:52 -05:00
}
@Override
public long getCount ( QueryParameters < M > queryParameters ) {
2021-11-30 11:06:19 -06:00
IckleQueryMapModelCriteriaBuilder < E , M > iqmcb = queryParameters . getModelCriteriaBuilder ( )
2021-08-17 06:30:52 -05:00
. flashToModelCriteriaBuilder ( createCriteriaBuilder ( ) ) ;
String queryString = iqmcb . getIckleQuery ( ) ;
LOG . tracef ( " Executing count Ickle query: %s " , queryString ) ;
QueryFactory queryFactory = Search . getQueryFactory ( remoteCache ) ;
2022-01-06 05:54:50 -06:00
Query < E > query = queryFactory . create ( queryString ) ;
2021-08-17 06:30:52 -05:00
query . setParameters ( iqmcb . getParameters ( ) ) ;
return query . execute ( ) . hitCount ( ) . orElse ( 0 ) ;
}
@Override
public long delete ( QueryParameters < M > queryParameters ) {
2022-12-16 11:55:31 -06:00
if ( queryParameters . getLimit ( ) ! = null | | queryParameters . getOffset ( ) ! = null ) {
throw new IllegalArgumentException ( " HotRod storage does not support pagination for delete query " ) ;
}
2021-08-17 06:30:52 -05:00
2023-01-25 08:33:26 -06:00
Query < Object [ ] > query = prepareQueryWithPrefixAndParameters ( " DELETE " , queryParameters ) ;
2022-11-08 08:57:06 -06:00
return query . executeStatement ( ) ;
2021-08-17 06:30:52 -05:00
}
2023-01-30 04:48:31 -06:00
@Override
public boolean exists ( String key ) {
Objects . requireNonNull ( key , " Key must be non-null " ) ;
K k = keyConverter . fromStringSafe ( key ) ;
return remoteCache . containsKey ( k ) ;
}
2021-11-30 11:06:19 -06:00
public IckleQueryMapModelCriteriaBuilder < E , M > createCriteriaBuilder ( ) {
return new IckleQueryMapModelCriteriaBuilder < > ( storedEntityDescriptor . getEntityTypeClass ( ) ) ;
2021-08-17 06:30:52 -05:00
}
@Override
public MapKeycloakTransaction < V , M > createTransaction ( KeycloakSession session ) {
2023-02-01 09:24:07 -06:00
// Here we return transaction that has no action because the returned transaction is enlisted to different
// phase than we need. Instead of tx returned by this method txWrapper is enlisted and executes all changes
// performed by the returned transaction.
return new NoActionHotRodTransactionWrapper < > ( ( ConcurrentHashMapKeycloakTransaction < K , V , M > ) txWrapper . getOrCreateTxForModel ( storedEntityDescriptor . getModelTypeClass ( ) , ( ) - > createTransactionInternal ( session ) ) ) ;
2021-08-17 06:30:52 -05:00
}
2022-07-11 01:54:39 -05:00
protected MapKeycloakTransaction < V , M > createTransactionInternal ( KeycloakSession session ) {
return new ConcurrentHashMapKeycloakTransaction < > ( this , keyConverter , cloner , fieldPredicates ) ;
}
2022-08-29 09:12:28 -05:00
// V must be an instance of ExpirableEntity
// returns null if expiration field is not set
// in certain cases can return 0 or negative number, which needs to be handled carefully when using as ISPN lifespan
private Long getLifespan ( V value ) {
Long expiration = ( ( ExpirableEntity ) value ) . getExpiration ( ) ;
return expiration ! = null ? expiration - Time . currentTimeMillis ( ) : null ;
}
2021-08-17 06:30:52 -05:00
}