Merge pull request #2952 from mposolda/master

KEYCLOAK-3149 DB update triggered before DBLock was acquired

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);
+        factory.register(new ProviderEventListener() {
-        AuthorizationProvider authorization = providerFactory.create(factory.create());
+
-        authorization.getStoreFactory().getPolicyStore().findByType(getId()).forEach(this::update);
+            @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

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;
 
 /**
@@ -157,6 +159,12 @@ public class DefaultJpaConnectionProviderFactory implements JpaConnectionProvide
 	                            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);
 	                        } else if (databaseSchema.equals("validate")) {

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

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);
+    }
 }

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")) {

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");
@@ -124,6 +126,11 @@ public class MongoDBLockProvider implements DBLockProvider {
         }
     }
 
+    @Override
+    public boolean hasLock() {
+        return factory.hasLock();
+    }
+
     @Override
     public boolean supportsForcedUnlock() {
         return true;

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);
+    }
+
 }

server-spi/src/main/java/org/keycloak/models/dblock/DBLockManager.java

@@ -15,24 +15,19 @@
  * limitations under the License.
  */
 
-package org.keycloak.services.managers;
+package org.keycloak.models.dblock;
 
+import org.jboss.logging.Logger;
 import org.keycloak.models.KeycloakSession;
-import org.keycloak.models.KeycloakSessionFactory;
-import org.keycloak.models.KeycloakSessionTask;
 import org.keycloak.models.RealmProvider;
 import org.keycloak.models.RealmProviderFactory;
-import org.keycloak.models.dblock.DBLockProvider;
-import org.keycloak.models.dblock.DBLockProviderFactory;
-import org.keycloak.models.utils.KeycloakModelUtils;
-import org.keycloak.services.ServicesLogger;
 
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
  */
 public class DBLockManager {
 
-    protected static final ServicesLogger logger = ServicesLogger.ROOT_LOGGER;
+    protected static final Logger logger = Logger.getLogger(DBLockManager.class);
 
     private final KeycloakSession session;
 
@@ -45,7 +40,7 @@ public class DBLockManager {
         if (Boolean.getBoolean("keycloak.dblock.forceUnlock")) {
             DBLockProvider lock = getDBLock();
             if (lock.supportsForcedUnlock()) {
-                logger.forcedReleaseDBLock();
+                logger.warn("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.");
                 lock.releaseLock();
             } else {
                 throw new IllegalStateException("Forced unlock requested, but provider " + lock + " doesn't support it");

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

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();
 }

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;

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;