MapGlobalLockProvider.java
/*
* Copyright 2023 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.lock;
import org.keycloak.common.util.Retry;
import org.keycloak.common.util.Time;
import org.keycloak.models.KeycloakSession;
import org.keycloak.models.KeycloakSessionTaskWithResult;
import org.keycloak.models.locking.GlobalLockProvider;
import org.keycloak.models.locking.LockAcquiringTimeoutException;
import org.keycloak.models.map.common.DeepCloner;
import org.keycloak.models.map.storage.MapStorage;
import org.keycloak.models.map.storage.ModelCriteriaBuilder;
import org.keycloak.models.map.storage.QueryParameters;
import org.keycloak.models.map.storage.criteria.DefaultModelCriteria;
import org.keycloak.models.utils.KeycloakModelUtils;
import java.net.InetAddress;
import java.net.UnknownHostException;
import java.time.Duration;
import java.time.Instant;
import java.util.Optional;
import java.util.function.Supplier;
import static org.keycloak.models.map.storage.criteria.DefaultModelCriteria.criteria;
/**
* Implementing a {@link GlobalLockProvider} based on a map storage.
* This requires the map store to support the entity type {@link MapLockEntity}. One of the stores which supports
* this is the JPA Map Store. The store needs to support the uniqueness of entries in the lock area, see
* {@link #lock(String)} for details.
*
* @author Alexander Schwartz
*/
public class MapGlobalLockProvider implements GlobalLockProvider {
private final KeycloakSession session;
private final long defaultTimeoutMilliseconds;
private MapStorage<MapLockEntity, MapLockEntity> store;
/**
* The lockStoreSupplier allows the store to be initialized lazily and only when needed: As this provider is initialized
* for both the outer and the inner transactions, and the store is needed only for the inner transactions.
*/
private final Supplier<MapStorage<MapLockEntity, MapLockEntity>> lockStoreSupplier;
public MapGlobalLockProvider(KeycloakSession session, long defaultTimeoutMilliseconds, Supplier<MapStorage<MapLockEntity, MapLockEntity>> lockStoreSupplier) {
this.defaultTimeoutMilliseconds = defaultTimeoutMilliseconds;
this.session = session;
this.lockStoreSupplier = lockStoreSupplier;
}
@Override
public <V> V withLock(String lockName, Duration timeToWaitForLock, KeycloakSessionTaskWithResult<V> task) throws LockAcquiringTimeoutException {
MapLockEntity[] lockEntity = {null};
try {
if (timeToWaitForLock == null) {
// Set default timeout if null provided
timeToWaitForLock = Duration.ofMillis(defaultTimeoutMilliseconds);
}
String[] keycloakInstanceIdentifier = {null};
Instant[] timeWhenAcquired = {null};
try {
Retry.executeWithBackoff(i -> lockEntity[0] = KeycloakModelUtils.runJobInTransactionWithResult(this.session.getKeycloakSessionFactory(),
innerSession -> {
MapGlobalLockProvider provider = (MapGlobalLockProvider) innerSession.getProvider(GlobalLockProvider.class);
// even if the call to provider.lock() succeeds, due to concurrency one can only be sure after a commit that all DB constraints have been met
return provider.lock(lockName);
}), (iteration, t) -> {
if (t instanceof LockAcquiringTimeoutException) {
LockAcquiringTimeoutException ex = (LockAcquiringTimeoutException) t;
keycloakInstanceIdentifier[0] = ex.getKeycloakInstanceIdentifier();
timeWhenAcquired[0] = ex.getTimeWhenAcquired();
}
}, timeToWaitForLock, 500);
} catch (RuntimeException ex) {
if (!(ex instanceof LockAcquiringTimeoutException)) {
throw new LockAcquiringTimeoutException(lockName, keycloakInstanceIdentifier[0], timeWhenAcquired[0], ex);
}
throw ex;
}
return KeycloakModelUtils.runJobInTransactionWithResult(this.session.getKeycloakSessionFactory(), task);
} finally {
if (lockEntity[0] != null) {
KeycloakModelUtils.runJobInTransaction(this.session.getKeycloakSessionFactory(), innerSession -> {
MapGlobalLockProvider provider = (MapGlobalLockProvider) innerSession.getProvider(GlobalLockProvider.class);
provider.unlock(lockEntity[0]);
});
}
}
}
@Override
public void forceReleaseAllLocks() {
KeycloakModelUtils.runJobInTransaction(this.session.getKeycloakSessionFactory(), innerSession -> {
MapGlobalLockProvider provider = (MapGlobalLockProvider) innerSession.getProvider(GlobalLockProvider.class);
provider.releaseAllLocks();
});
}
@Override
public void close() {
}
private void prepareTx() {
if (store == null) {
this.store = lockStoreSupplier.get();
}
}
/**
* Create a {@link MapLockEntity} for the provided <code>lockName</code>.
* The underlying store must ensure that a lock with the given name can be created only once in the store.
* This constraint needs to be checked either at the time of creation, or at the latest when the transaction
* is committed. If such a constraint violation is detected at the time of the transaction commit, it should
* throw an exception and the transaction should roll back.
* <p/>
* The JPA Map Store implements this with a unique index, which is checked by the database both at the time of
* insertion and at the time the transaction is committed.
*/
private MapLockEntity lock(String lockName) {
prepareTx();
DefaultModelCriteria<MapLockEntity> mcb = criteria();
mcb = mcb.compare(MapLockEntity.SearchableFields.NAME, ModelCriteriaBuilder.Operator.EQ, lockName);
Optional<MapLockEntity> entry = store.read(QueryParameters.withCriteria(mcb)).findFirst();
if (entry.isEmpty()) {
MapLockEntity entity = DeepCloner.DUMB_CLONER.newInstance(MapLockEntity.class);
entity.setName(lockName);
entity.setKeycloakInstanceIdentifier(getKeycloakInstanceIdentifier());
entity.setTimeAcquired(Time.currentTimeMillis());
return store.create(entity);
} else {
throw new LockAcquiringTimeoutException(lockName, entry.get().getKeycloakInstanceIdentifier(), Instant.ofEpochMilli(entry.get().getTimeAcquired()));
}
}
/**
* Unlock the previously created lock.
* Will fail if the lock doesn't exist, or has a different owner.
*/
private void unlock(MapLockEntity lockEntity) {
prepareTx();
MapLockEntity readLockEntity = store.read(lockEntity.getId());
if (readLockEntity == null) {
throw new RuntimeException("didn't find lock - someone else unlocked it?");
} else if (!lockEntity.isLockUnchanged(readLockEntity)) {
// this case is there for stores which might re-use IDs or derive it from the name of the entity (like the file store map store does in some cases).
throw new RuntimeException(String.format("Lock owned by different instance: Lock [%s] acquired by keycloak instance [%s] at the time [%s]",
readLockEntity.getName(), readLockEntity.getKeycloakInstanceIdentifier(), readLockEntity.getTimeAcquired()));
} else {
store.delete(readLockEntity.getId());
}
}
private void releaseAllLocks() {
prepareTx();
DefaultModelCriteria<MapLockEntity> mcb = criteria();
store.delete(QueryParameters.withCriteria(mcb));
}
private static String getKeycloakInstanceIdentifier() {
long pid = ProcessHandle.current().pid();
String hostname;
try {
hostname = InetAddress.getLocalHost().getHostName();
} catch (UnknownHostException e) {
hostname = "unknown-host";
}
String threadName = Thread.currentThread().getName();
return threadName + "#" + pid + "@" + hostname;
}
}