keycloak-memoizeit

diff --git a/authz/policy/drools/src/main/java/org/keycloak/authorization/policy/provider/drools/DroolsPolicyProviderFactory.java b/authz/policy/drools/src/main/java/org/keycloak/authorization/policy/provider/drools/DroolsPolicyProviderFactory.java
index 3df4210..0e1e208 100644
--- a/authz/policy/drools/src/main/java/org/keycloak/authorization/policy/provider/drools/DroolsPolicyProviderFactory.java
+++ b/authz/policy/drools/src/main/java/org/keycloak/authorization/policy/provider/drools/DroolsPolicyProviderFactory.java
@@ -9,6 +9,9 @@ import org.keycloak.authorization.policy.provider.PolicyProviderAdminService;
 import org.keycloak.authorization.policy.provider.PolicyProviderFactory;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.KeycloakSessionFactory;
+import org.keycloak.models.utils.PostMigrationEvent;
+import org.keycloak.provider.ProviderEvent;
+import org.keycloak.provider.ProviderEventListener;
 import org.keycloak.provider.ProviderFactory;
 import org.kie.api.KieServices;
 import org.kie.api.KieServices.Factory;
@@ -61,9 +64,19 @@ public class DroolsPolicyProviderFactory implements PolicyProviderFactory {
 
     @Override
     public void postInit(KeycloakSessionFactory factory) {
-        ProviderFactory<AuthorizationProvider> providerFactory = factory.getProviderFactory(AuthorizationProvider.class);
-        AuthorizationProvider authorization = providerFactory.create(factory.create());
-        authorization.getStoreFactory().getPolicyStore().findByType(getId()).forEach(this::update);
+        factory.register(new ProviderEventListener() {
+
+            @Override
+            public void onEvent(ProviderEvent event) {
+                // Ensure the initialization is done after DB upgrade is finished
+                if (event instanceof PostMigrationEvent) {
+                    ProviderFactory<AuthorizationProvider> providerFactory = factory.getProviderFactory(AuthorizationProvider.class);
+                    AuthorizationProvider authorization = providerFactory.create(factory.create());
+                    authorization.getStoreFactory().getPolicyStore().findByType(getId()).forEach(DroolsPolicyProviderFactory.this::update);
+                }
+            }
+
+        });
     }
 
     @Override

diff --git a/model/jpa/src/main/java/org/keycloak/connections/jpa/DefaultJpaConnectionProviderFactory.java b/model/jpa/src/main/java/org/keycloak/connections/jpa/DefaultJpaConnectionProviderFactory.java
index b5621ef..784438c 100755
--- a/model/jpa/src/main/java/org/keycloak/connections/jpa/DefaultJpaConnectionProviderFactory.java
+++ b/model/jpa/src/main/java/org/keycloak/connections/jpa/DefaultJpaConnectionProviderFactory.java
@@ -37,7 +37,9 @@ import org.keycloak.connections.jpa.updater.JpaUpdaterProvider;
 import org.keycloak.connections.jpa.util.JpaUtils;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.KeycloakSessionFactory;
+import org.keycloak.models.dblock.DBLockProvider;
 import org.keycloak.provider.ServerInfoAwareProviderFactory;
+import org.keycloak.models.dblock.DBLockManager;
 import org.keycloak.timer.TimerProvider;
 
 /**
@@ -156,6 +158,12 @@ public class DefaultJpaConnectionProviderFactory implements JpaConnectionProvide
 	                        if (updater == null) {
 	                            throw new RuntimeException("Can't update database: JPA updater provider not found");
 	                        }
+
+                            // Check if having DBLock before trying to initialize hibernate
+                            DBLockProvider dbLock = new DBLockManager(session).getDBLock();
+                            if (!dbLock.hasLock()) {
+                                throw new IllegalStateException("Trying to update database, but don't have a DB lock acquired");
+                            }
 	
 	                        if (databaseSchema.equals("update")) {
                                 updater.update(connection, schema);

diff --git a/model/jpa/src/main/java/org/keycloak/connections/jpa/updater/liquibase/lock/LiquibaseDBLockProvider.java b/model/jpa/src/main/java/org/keycloak/connections/jpa/updater/liquibase/lock/LiquibaseDBLockProvider.java
index f44641e..5874084 100644
--- a/model/jpa/src/main/java/org/keycloak/connections/jpa/updater/liquibase/lock/LiquibaseDBLockProvider.java
+++ b/model/jpa/src/main/java/org/keycloak/connections/jpa/updater/liquibase/lock/LiquibaseDBLockProvider.java
@@ -91,6 +91,7 @@ public class LiquibaseDBLockProvider implements DBLockProvider {
         while (maxAttempts > 0) {
             try {
                 lockService.waitForLock();
+                factory.setHasLock(true);
                 this.maxAttempts = DEFAULT_MAX_ATTEMPTS;
                 return;
             } catch (LockRetryException le) {
@@ -111,6 +112,12 @@ public class LiquibaseDBLockProvider implements DBLockProvider {
     public void releaseLock() {
         lockService.releaseLock();
         lockService.reset();
+        factory.setHasLock(false);
+    }
+
+    @Override
+    public boolean hasLock() {
+        return factory.hasLock();
     }
 
     @Override

diff --git a/model/jpa/src/main/java/org/keycloak/connections/jpa/updater/liquibase/lock/LiquibaseDBLockProviderFactory.java b/model/jpa/src/main/java/org/keycloak/connections/jpa/updater/liquibase/lock/LiquibaseDBLockProviderFactory.java
index 3026f7d..ef1fae1 100644
--- a/model/jpa/src/main/java/org/keycloak/connections/jpa/updater/liquibase/lock/LiquibaseDBLockProviderFactory.java
+++ b/model/jpa/src/main/java/org/keycloak/connections/jpa/updater/liquibase/lock/LiquibaseDBLockProviderFactory.java
@@ -17,6 +17,8 @@
 
 package org.keycloak.connections.jpa.updater.liquibase.lock;
 
+import java.util.concurrent.atomic.AtomicBoolean;
+
 import org.jboss.logging.Logger;
 import org.keycloak.Config;
 import org.keycloak.common.util.Time;
@@ -33,6 +35,9 @@ public class LiquibaseDBLockProviderFactory implements DBLockProviderFactory {
 
     private long lockWaitTimeoutMillis;
 
+    // True if this node has a lock acquired
+    private AtomicBoolean hasLock = new AtomicBoolean(false);
+
     protected long getLockWaitTimeoutMillis() {
         return lockWaitTimeoutMillis;
     }
@@ -68,4 +73,12 @@ public class LiquibaseDBLockProviderFactory implements DBLockProviderFactory {
     public String getId() {
         return "jpa";
     }
+
+    public boolean hasLock() {
+        return hasLock.get();
+    }
+
+    public void setHasLock(boolean hasLock) {
+        this.hasLock.set(hasLock);
+    }
 }

diff --git a/model/mongo/src/main/java/org/keycloak/connections/mongo/DefaultMongoConnectionFactoryProvider.java b/model/mongo/src/main/java/org/keycloak/connections/mongo/DefaultMongoConnectionFactoryProvider.java
index 491ad71..7fc3439 100755
--- a/model/mongo/src/main/java/org/keycloak/connections/mongo/DefaultMongoConnectionFactoryProvider.java
+++ b/model/mongo/src/main/java/org/keycloak/connections/mongo/DefaultMongoConnectionFactoryProvider.java
@@ -22,6 +22,7 @@ import java.net.UnknownHostException;
 import java.util.Collections;
 import java.util.LinkedHashMap;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import javax.net.ssl.SSLSocketFactory;
 
@@ -33,6 +34,8 @@ import org.keycloak.connections.mongo.impl.context.TransactionMongoStoreInvocati
 import org.keycloak.connections.mongo.updater.MongoUpdaterProvider;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.KeycloakSessionFactory;
+import org.keycloak.models.dblock.DBLockManager;
+import org.keycloak.models.dblock.DBLockProvider;
 import org.keycloak.provider.ServerInfoAwareProviderFactory;
 
 import com.mongodb.DB;
@@ -170,6 +173,11 @@ public class DefaultMongoConnectionFactoryProvider implements MongoConnectionPro
                 throw new RuntimeException("Can't update database: Mongo updater provider not found");
             }
 
+            DBLockProvider dbLock = new DBLockManager(session).getDBLock();
+            if (!dbLock.hasLock()) {
+                throw new IllegalStateException("Trying to update database, but don't have a DB lock acquired");
+            }
+
             if (databaseSchema.equals("update")) {
                 mongoUpdater.update(session, db);
             } else if (databaseSchema.equals("validate")) {

diff --git a/model/mongo/src/main/java/org/keycloak/connections/mongo/lock/MongoDBLockProvider.java b/model/mongo/src/main/java/org/keycloak/connections/mongo/lock/MongoDBLockProvider.java
index 1b36889..a11f729 100644
--- a/model/mongo/src/main/java/org/keycloak/connections/mongo/lock/MongoDBLockProvider.java
+++ b/model/mongo/src/main/java/org/keycloak/connections/mongo/lock/MongoDBLockProvider.java
@@ -91,6 +91,7 @@ public class MongoDBLockProvider implements DBLockProvider {
             WriteResult wr = db.getCollection(DB_LOCK_COLLECTION).update(query, update, true, false);
             if (wr.getN() == 1) {
                 logger.debugf("Successfully acquired DB lock");
+                factory.setHasLock(true);
                 return true;
             } else {
                 return false;
@@ -115,6 +116,7 @@ public class MongoDBLockProvider implements DBLockProvider {
         try {
             WriteResult wr = db.getCollection(DB_LOCK_COLLECTION).update(query, update, true, false);
             if (wr.getN() > 0) {
+                factory.setHasLock(false);
                 logger.debugf("Successfully released DB lock");
             } else {
                 logger.warnf("Attempt to release DB lock, but nothing was released");
@@ -125,6 +127,11 @@ public class MongoDBLockProvider implements DBLockProvider {
     }
 
     @Override
+    public boolean hasLock() {
+        return factory.hasLock();
+    }
+
+    @Override
     public boolean supportsForcedUnlock() {
         return true;
     }

diff --git a/model/mongo/src/main/java/org/keycloak/connections/mongo/lock/MongoDBLockProviderFactory.java b/model/mongo/src/main/java/org/keycloak/connections/mongo/lock/MongoDBLockProviderFactory.java
index 7bd6e02..64f65f6 100644
--- a/model/mongo/src/main/java/org/keycloak/connections/mongo/lock/MongoDBLockProviderFactory.java
+++ b/model/mongo/src/main/java/org/keycloak/connections/mongo/lock/MongoDBLockProviderFactory.java
@@ -17,6 +17,8 @@
 
 package org.keycloak.connections.mongo.lock;
 
+import java.util.concurrent.atomic.AtomicBoolean;
+
 import com.mongodb.DB;
 import org.jboss.logging.Logger;
 import org.keycloak.Config;
@@ -37,6 +39,9 @@ public class MongoDBLockProviderFactory implements DBLockProviderFactory {
     private long lockRecheckTimeMillis;
     private long lockWaitTimeoutMillis;
 
+    // True if this node has a lock acquired
+    private AtomicBoolean hasLock = new AtomicBoolean(false);
+
     protected long getLockRecheckTimeMillis() {
         return lockRecheckTimeMillis;
     }
@@ -81,4 +86,13 @@ public class MongoDBLockProviderFactory implements DBLockProviderFactory {
     public String getId() {
         return "mongo";
     }
+
+    public boolean hasLock() {
+        return hasLock.get();
+    }
+
+    public void setHasLock(boolean hasLock) {
+        this.hasLock.set(hasLock);
+    }
+
 }

diff --git a/server-spi/src/main/java/org/keycloak/models/dblock/DBLockProvider.java b/server-spi/src/main/java/org/keycloak/models/dblock/DBLockProvider.java
index d9dc131..d7d6053 100644
--- a/server-spi/src/main/java/org/keycloak/models/dblock/DBLockProvider.java
+++ b/server-spi/src/main/java/org/keycloak/models/dblock/DBLockProvider.java
@@ -39,6 +39,13 @@ public interface DBLockProvider extends Provider {
      */
     void releaseLock();
 
+    /**
+     * Check if I have lock
+     *
+     * @return
+     */
+    boolean hasLock();
+
 
     /**
      * @return true if provider supports forced unlock at startup

diff --git a/services/src/main/java/org/keycloak/services/resources/KeycloakApplication.java b/services/src/main/java/org/keycloak/services/resources/KeycloakApplication.java
index d6d44d4..818d571 100644
--- a/services/src/main/java/org/keycloak/services/resources/KeycloakApplication.java
+++ b/services/src/main/java/org/keycloak/services/resources/KeycloakApplication.java
@@ -26,7 +26,7 @@ import org.keycloak.exportimport.ExportImportManager;
 import org.keycloak.migration.MigrationModelManager;
 import org.keycloak.models.*;
 import org.keycloak.models.dblock.DBLockProvider;
-import org.keycloak.services.managers.DBLockManager;
+import org.keycloak.models.dblock.DBLockManager;
 import org.keycloak.models.utils.PostMigrationEvent;
 import org.keycloak.models.utils.RepresentationToModel;
 import org.keycloak.representations.idm.RealmRepresentation;

diff --git a/services/src/main/java/org/keycloak/services/ServicesLogger.java b/services/src/main/java/org/keycloak/services/ServicesLogger.java
index c50191f..fb71a2e 100644
--- a/services/src/main/java/org/keycloak/services/ServicesLogger.java
+++ b/services/src/main/java/org/keycloak/services/ServicesLogger.java
@@ -402,8 +402,4 @@ public interface ServicesLogger extends BasicLogger {
     @LogMessage(level = ERROR)
     @Message(id=90, value="Failed to close ProviderSession")
     void failedToCloseProviderSession(@Cause Throwable t);
-
-    @LogMessage(level = WARN)
-    @Message(id=91, value="Forced release of DB lock at startup requested by System property. Make sure to not use this in production environment! And especially when more cluster nodes are started concurrently.")
-    void forcedReleaseDBLock();
 }

diff --git a/testsuite/integration/src/test/java/org/keycloak/testsuite/model/DBLockTest.java b/testsuite/integration/src/test/java/org/keycloak/testsuite/model/DBLockTest.java
index 14fde53..2537332 100644
--- a/testsuite/integration/src/test/java/org/keycloak/testsuite/model/DBLockTest.java
+++ b/testsuite/integration/src/test/java/org/keycloak/testsuite/model/DBLockTest.java
@@ -24,12 +24,11 @@ import java.util.concurrent.atomic.AtomicInteger;
 import org.jboss.logging.Logger;
 import org.junit.Assert;
 import org.junit.Before;
-import org.junit.Ignore;
 import org.junit.Test;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.KeycloakSessionFactory;
 import org.keycloak.models.KeycloakSessionTask;
-import org.keycloak.services.managers.DBLockManager;
+import org.keycloak.models.dblock.DBLockManager;
 import org.keycloak.models.dblock.DBLockProvider;
 import org.keycloak.models.dblock.DBLockProviderFactory;
 import org.keycloak.models.utils.KeycloakModelUtils;