From 5213f9d83cbf258e3a6ac2ab9c576db0c799c4ab Mon Sep 17 00:00:00 2001 From: Alexander Schwartz Date: Mon, 2 Nov 2026 13:54:36 +0100 Subject: [PATCH 1/4] Support KEYCLOAK_JDBC_PING2 in custom JGroups stacks by injecting the connection factory Closes keycloak/keycloak#54539 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Alexander Schwartz --- .../impl/embedded/JGroupsConfigurator.java | 111 +++++++++--- .../jgroups/protocol/JdbcPing2Test.java | 3 +- .../JGroupsConfiguratorCustomStackTest.java | 171 ++++++++++++++++++ 3 files changed, 259 insertions(+), 26 deletions(-) create mode 100644 model/infinispan/src/test/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfiguratorCustomStackTest.java diff --git a/model/infinispan/src/main/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfigurator.java b/model/infinispan/src/main/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfigurator.java index 8529cc0e884f..112fbef6ba90 100644 --- a/model/infinispan/src/main/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfigurator.java +++ b/model/infinispan/src/main/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfigurator.java @@ -25,6 +25,7 @@ import java.sql.ResultSet; import java.util.ArrayList; import java.util.Arrays; +import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Objects; @@ -240,8 +241,10 @@ private static void configureDiscovery(ConfigurationBuilderHolder holder, Keyclo var stackXmlAttribute = transportStackOf(holder); if (stackXmlAttribute.isModified() && !isJdbcPingStack(stackXmlAttribute.get())) { logger.debugf("Custom stack configured (%s). JDBC_PING discovery disabled.", stackXmlAttribute.get()); - if (Profile.isFeatureEnabled(Profile.Feature.STATELESS)) { - if (Objects.equals("test", stackXmlAttribute.get())) { + var tableName = JpaUtils.getTableNameForNativeQuery("JGROUPS_PING", session.getProvider(JpaConnectionProvider.class).getEntityManager()); + boolean jdbcPing = injectConnectionFactoryIntoCustomStack(holder, session.getKeycloakSessionFactory(), tableName, stackXmlAttribute.get()); + if (!jdbcPing && Profile.isFeatureEnabled(Profile.Feature.STATELESS)) { + if (isInJvmTestStack(holder, stackXmlAttribute.get())) { // TODO: Remove this case once the model tests have been migrated logger.error("The stateless feature must hast JDBC_PING discovery enabled"); } else { @@ -274,6 +277,48 @@ private static void configureDiscovery(ConfigurationBuilderHolder holder, Keyclo JGroupsConfigurator.logger.info("JGroups JDBC_PING discovery enabled."); } + /** + * If a custom stack uses {@link KEYCLOAK_JDBC_PING2}, wraps it in a configurator that injects the + * {@link JpaConnectionProviderFactory}, as the protocol is unable to connect to the database without it. + * The stack is registered under a new name as Infinispan rejects duplicate stack names. + * The JGroups address is not pre-registered in the database, as the stack is fully controlled by the user. + * + * @return {@code true} if the stack uses {@link KEYCLOAK_JDBC_PING2} + */ + static boolean injectConnectionFactoryIntoCustomStack(ConfigurationBuilderHolder holder, KeycloakSessionFactory sessionFactory, String tableName, String stackName) { + var configurator = transportOf(holder).build().transport().jgroups().configurator(stackName); + var combined = configurator.getProtocolStack(); + if (combined.stream().noneMatch(p -> KEYCLOAK_JDBC_PING2.class.getName().equals(p.getProtocolName()))) { + return false; + } + // Fail fast instead of starting with a protocol that has no database connection, or with a silently altered stack + if (!(configurator instanceof EmbeddedJGroupsChannelConfigurator embedded)) { + throw new IllegalStateException("%s is only supported in a stack defined in the XML configuration, not in stack '%s' defined by a file" + .formatted(KEYCLOAK_JDBC_PING2.class.getSimpleName(), stackName)); + } + var sites = embedded.getUncombinedRemoteSites(); + if ((sites != null && !sites.getRemoteSites().isEmpty()) + || combined.stream().anyMatch(p -> p.getProtocolName().contains("RELAY2"))) { + throw new IllegalStateException("Cross-site replication (RELAY2 or remote sites) is not supported in stack '%s' using %s" + .formatted(stackName, KEYCLOAK_JDBC_PING2.class.getSimpleName())); + } + var stack = new ArrayList(combined.size()); + for (var protocol : combined) { + var properties = new HashMap<>(protocol.getProperties()); + if (KEYCLOAK_JDBC_PING2.class.getName().equals(protocol.getProtocolName())) { + // values set in the XML take precedence + jdbcPingProperties(tableName).forEach(properties::putIfAbsent); + } + stack.add(new ProtocolConfiguration(protocol.getProtocolName(), properties)); + } + var connectionFactory = (JpaConnectionProviderFactory) sessionFactory.getProviderFactory(JpaConnectionProvider.class); + var wrappedName = stackName + "-keycloak"; + holder.addJGroupsStack(new JpaFactoryAwareJGroupsChannelConfigurator(sessionFactory, wrappedName, stack, embedded.getUncombinedRemoteSites(), null, connectionFactory, null), null); + transportOf(holder).stack(wrappedName); + logger.infof("Custom stack '%s' uses %s. Injected the database connection factory.", stackName, KEYCLOAK_JDBC_PING2.class.getSimpleName()); + return true; + } + /** * Generate the next sequence of the address, and place it into the JGROUPS_PING table so other nodes can see it. * If we are the first = smallest entry, the other nodes will wait for us to become a coordinator @@ -337,29 +382,32 @@ private static Address insertSequenceInTable(JpaConnectionProvider cp, String cl return address; } + private static Map jdbcPingProperties(String tableName) { + return Map.ofEntries( + // Leave initialize_sql blank as table is already created by Keycloak + Map.entry("initialize_sql", ""), + // Explicitly specify clear and select_all SQL to ensure "cluster_name" column is used, as the default + // "cluster" cannot be used with Oracle DB as it's a reserved word. + Map.entry("clear_sql", String.format("DELETE from %s WHERE cluster_name=?", tableName)), + Map.entry("delete_single_sql", String.format("DELETE from %s WHERE address=?", tableName)), + Map.entry("insert_single_sql", String.format("INSERT INTO %s (address, name, cluster_name, ip, coord, last_update, coordinated_by) values (?, ?, ?, ?, ?, ?, ?)", tableName)), + Map.entry("select_all_pingdata_sql", String.format("SELECT address, name, ip, coord, coordinated_by, last_update FROM %s WHERE cluster_name=?", tableName)), + Map.entry("select_other_clusters_count_sql", String.format("SELECT COUNT(*) FROM %s WHERE cluster_name != ? AND last_update >= ?", tableName)), + Map.entry("allow_multiple_clusters", String.valueOf(Profile.isFeatureEnabled(Profile.Feature.STATELESS))), + // This guarantees cleanup of stale data + Map.entry("remove_all_data_on_view_change", "true"), + // This guarantees that merging happens even after the info writer completed + Map.entry("write_data_on_find", "true"), + Map.entry("register_shutdown_hook", "false") + ); + } + private static List getProtocolConfigurations(String tableName, boolean tracingEnabled, boolean fdSockEnabled) { var list = new ArrayList(4); - list.add(new ProtocolConfiguration(KEYCLOAK_JDBC_PING2.class.getName(), - Map.ofEntries( - // Leave initialize_sql blank as table is already created by Keycloak - Map.entry("initialize_sql", ""), - // Explicitly specify clear and select_all SQL to ensure "cluster_name" column is used, as the default - // "cluster" cannot be used with Oracle DB as it's a reserved word. - Map.entry("clear_sql", String.format("DELETE from %s WHERE cluster_name=?", tableName)), - Map.entry("delete_single_sql", String.format("DELETE from %s WHERE address=?", tableName)), - Map.entry("insert_single_sql", String.format("INSERT INTO %s (address, name, cluster_name, ip, coord, last_update, coordinated_by) values (?, ?, ?, ?, ?, ?, ?)", tableName)), - Map.entry("select_all_pingdata_sql", String.format("SELECT address, name, ip, coord, coordinated_by, last_update FROM %s WHERE cluster_name=?", tableName)), - Map.entry("select_other_clusters_count_sql", String.format("SELECT COUNT(*) FROM %s WHERE cluster_name != ? AND last_update >= ?", tableName)), - Map.entry("allow_multiple_clusters", String.valueOf(Profile.isFeatureEnabled(Profile.Feature.STATELESS))), - // This guarantees cleanup of stale data - Map.entry("remove_all_data_on_view_change", "true"), - // This guarantees that merging happens even after the info writer completed - Map.entry("write_data_on_find", "true"), - Map.entry("register_shutdown_hook", "false"), - Map.entry("stack.combine", "REPLACE"), - Map.entry("stack.position", "MPING") - )) - ); + var pingProperties = new HashMap<>(jdbcPingProperties(tableName)); + pingProperties.put("stack.combine", "REPLACE"); + pingProperties.put("stack.position", "MPING"); + list.add(new ProtocolConfiguration(KEYCLOAK_JDBC_PING2.class.getName(), pingProperties)); list.add(new ProtocolConfiguration(TCP.class.getSimpleName(), // Using per-destination bundler is important to not block the communication to other nodes @@ -413,6 +461,13 @@ private static void validateTlsAvailable(ConfigurationBuilderHolder holder) { } } + // The model tests use a stack named "test" that has no network transport, as all nodes are in the same JVM. + static boolean isInJvmTestStack(ConfigurationBuilderHolder holder, String stackName) { + return "test".equals(stackName) + && transportOf(holder).build().transport().jgroups().configurator(stackName).getProtocolStack().stream() + .anyMatch(p -> p.getProtocolName().contains("SHARED_LOOPBACK")); + } + private static boolean isJdbcPingStack(String stackName) { return "jdbc-ping".equals(stackName); } @@ -431,7 +486,11 @@ private static class JpaFactoryAwareJGroupsChannelConfigurator extends EmbeddedJ private final Address address; public JpaFactoryAwareJGroupsChannelConfigurator(KeycloakSessionFactory factory, String name, List stack, JpaConnectionProviderFactory jpaConnectionProviderFactory, Address address) { - super(name, stack, null, "tcp"); + this(factory, name, stack, null, "tcp", jpaConnectionProviderFactory, address); + } + + public JpaFactoryAwareJGroupsChannelConfigurator(KeycloakSessionFactory factory, String name, List stack, RemoteSites remoteSites, String parent, JpaConnectionProviderFactory jpaConnectionProviderFactory, Address address) { + super(name, stack, remoteSites, parent); this.factory = factory; this.jpaConnectionProviderFactory = Objects.requireNonNull(jpaConnectionProviderFactory); this.address = address; @@ -439,7 +498,9 @@ public JpaFactoryAwareJGroupsChannelConfigurator(KeycloakSessionFactory factory, @Override protected JChannel amendChannel(JChannel channel) { - channel.addAddressGenerator(() -> address); + if (address != null) { + channel.addAddressGenerator(() -> address); + } return super.amendChannel(channel); } diff --git a/model/infinispan/src/test/java/org/keycloak/jgroups/protocol/JdbcPing2Test.java b/model/infinispan/src/test/java/org/keycloak/jgroups/protocol/JdbcPing2Test.java index b0dd4443412c..3af15b669a84 100644 --- a/model/infinispan/src/test/java/org/keycloak/jgroups/protocol/JdbcPing2Test.java +++ b/model/infinispan/src/test/java/org/keycloak/jgroups/protocol/JdbcPing2Test.java @@ -64,7 +64,8 @@ public class JdbcPing2Test { public static final String PROTOCOL_STACK = "jdbc-h2.xml"; static { - ClassConfigurator.addProtocol((short) 1026, KEYCLOAK_JDBC_PING2_FOR_TESTING.class); + // must not clash with the IDs registered in JGroupsConfigurator (1025, 1026), as tests share a JVM + ClassConfigurator.addProtocol((short) 1099, KEYCLOAK_JDBC_PING2_FOR_TESTING.class); } /** diff --git a/model/infinispan/src/test/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfiguratorCustomStackTest.java b/model/infinispan/src/test/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfiguratorCustomStackTest.java new file mode 100644 index 000000000000..0456df952662 --- /dev/null +++ b/model/infinispan/src/test/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfiguratorCustomStackTest.java @@ -0,0 +1,171 @@ +/* + * Copyright 2026 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.spi.infinispan.impl.embedded; + +import java.lang.reflect.Proxy; +import java.sql.SQLException; +import java.util.concurrent.atomic.AtomicBoolean; + +import org.keycloak.common.Profile; +import org.keycloak.connections.jpa.JpaConnectionProviderFactory; +import org.keycloak.jgroups.protocol.KEYCLOAK_JDBC_PING2; +import org.keycloak.models.KeycloakSessionFactory; + +import org.infinispan.configuration.global.TransportConfiguration; +import org.infinispan.configuration.parsing.ConfigurationBuilderHolder; +import org.infinispan.configuration.parsing.ParserRegistry; +import org.infinispan.remoting.transport.jgroups.EmbeddedJGroupsChannelConfigurator; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +/** + * Tests that a custom stack extending TCP and removing FD_SOCK2 gets the database connection factory injected into {@link KEYCLOAK_JDBC_PING2}. + */ +public class JGroupsConfiguratorCustomStackTest { + + @BeforeClass + public static void initProfile() { + Profile.defaults(); + } + + @AfterClass + public static void resetProfile() { + Profile.reset(); + } + + private static final String PING_PROTOCOL = ""; + + private static ConfigurationBuilderHolder parse(String discovery) { + return new ParserRegistry().parse(""" + + + + + + %s + + + + + + + """.formatted(discovery)); + } + + private static KeycloakSessionFactory sessionFactory(AtomicBoolean called) { + var jpaFactory = Proxy.newProxyInstance( + JpaConnectionProviderFactory.class.getClassLoader(), + new Class[]{JpaConnectionProviderFactory.class}, + (proxy, method, args) -> { + if ("getConnection".equals(method.getName())) { + called.set(true); + } + return null; + }); + return (KeycloakSessionFactory) Proxy.newProxyInstance( + KeycloakSessionFactory.class.getClassLoader(), + new Class[]{KeycloakSessionFactory.class}, + (proxy, method, args) -> "getProviderFactory".equals(method.getName()) ? jpaFactory : null); + } + + // Exposes the protected connection lookup + private static class TestPing extends KEYCLOAK_JDBC_PING2 { + void connect() { + try { + getConnection(); + } catch (SQLException e) { + // the stub factory returns no connection; only the factory invocation matters here + } + } + } + + @Test + public void testFactoryInjectedIntoCustomStack() throws Exception { + var holder = parse(PING_PROTOCOL); + var called = new AtomicBoolean(); + + assertTrue(JGroupsConfigurator.injectConnectionFactoryIntoCustomStack(holder, sessionFactory(called), "JGROUPS_PING", "custom")); + + var transport = holder.getGlobalConfigurationBuilder().transport(); + assertEquals("custom-keycloak", transport.attributes().attribute(TransportConfiguration.STACK).get()); + + var configurator = (EmbeddedJGroupsChannelConfigurator) transport.build().transport().jgroups().configurator("custom-keycloak"); + var protocols = configurator.getProtocolStack(); + assertTrue(protocols.stream().anyMatch(p -> p.getProtocolName().equals(KEYCLOAK_JDBC_PING2.class.getName()))); + assertTrue(protocols.stream().noneMatch(p -> p.getProtocolName().contains("FD_SOCK2"))); + assertTrue(protocols.stream().noneMatch(p -> p.getProtocolName().contains("MPING"))); + var verifySuspect = protocols.stream().filter(p -> p.getProtocolName().equals("VERIFY_SUSPECT2")).findFirst().orElseThrow(); + assertEquals("5000", verifySuspect.getProperties().get("timeout")); + + var ping = protocols.stream().filter(p -> p.getProtocolName().equals(KEYCLOAK_JDBC_PING2.class.getName())).findFirst().orElseThrow(); + assertEquals("DELETE from JGROUPS_PING WHERE cluster_name=?", ping.getProperties().get("clear_sql")); + assertEquals("", ping.getProperties().get("initialize_sql")); + // explicit XML values win over the defaults + assertEquals("false", ping.getProperties().get("write_data_on_find")); + var ping2 = new TestPing(); + configurator.afterCreation(ping2); + ping2.connect(); + assertTrue("The connection factory should be used by the protocol", called.get()); + } + + @Test + public void testInJvmTestStackExemptionIsNarrow() { + var inJvm = new ParserRegistry().parse(""" + + + + + + + + + + + """); + assertTrue(JGroupsConfigurator.isInJvmTestStack(inJvm, "test")); + + var networked = parse(""); + networked.getGlobalConfigurationBuilder().transport().stack("custom"); + assertFalse(JGroupsConfigurator.isInJvmTestStack(networked, "custom")); + } + + @Test + public void testStackWithoutKeycloakJdbcPingUntouched() { + var holder = parse(""); + + assertFalse(JGroupsConfigurator.injectConnectionFactoryIntoCustomStack(holder, sessionFactory(new AtomicBoolean()), "JGROUPS_PING", "custom")); + + var transport = holder.getGlobalConfigurationBuilder().transport(); + assertEquals("custom", transport.attributes().attribute(TransportConfiguration.STACK).get()); + assertFalse(transport.build().transport().jgroups().stacks().stream().anyMatch(s -> s.name().equals("custom-keycloak"))); + } + + @Test + public void testCrossSiteRejected() { + var holder = parse(PING_PROTOCOL + ""); + + assertThrows(IllegalStateException.class, + () -> JGroupsConfigurator.injectConnectionFactoryIntoCustomStack(holder, sessionFactory(new AtomicBoolean()), "JGROUPS_PING", "custom")); + } +} From fb660f816a87274b5207a57f94da727d7c97bd40 Mon Sep 17 00:00:00 2001 From: Alexander Schwartz Date: Tue, 3 Nov 2026 19:56:22 +0100 Subject: [PATCH 2/4] Update SHARED_LOOPBACK protocol check in JGroupsConfigurator Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Alexander Schwartz --- .../spi/infinispan/impl/embedded/JGroupsConfigurator.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/model/infinispan/src/main/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfigurator.java b/model/infinispan/src/main/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfigurator.java index 112fbef6ba90..d1187b69ebd6 100644 --- a/model/infinispan/src/main/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfigurator.java +++ b/model/infinispan/src/main/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfigurator.java @@ -465,7 +465,8 @@ private static void validateTlsAvailable(ConfigurationBuilderHolder holder) { static boolean isInJvmTestStack(ConfigurationBuilderHolder holder, String stackName) { return "test".equals(stackName) && transportOf(holder).build().transport().jgroups().configurator(stackName).getProtocolStack().stream() - .anyMatch(p -> p.getProtocolName().contains("SHARED_LOOPBACK")); + .anyMatch(p -> "SHARED_LOOPBACK".equals(p.getProtocolName()) + || "org.jgroups.protocols.SHARED_LOOPBACK".equals(p.getProtocolName())); } private static boolean isJdbcPingStack(String stackName) { From 6c7034e79a80972f8dcc6013a86b4e1620345fce Mon Sep 17 00:00:00 2001 From: Alexander Schwartz Date: Tue, 3 Nov 2026 19:56:30 +0100 Subject: [PATCH 3/4] Choose an unused name for the wrapped custom stack Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Alexander Schwartz --- .../impl/embedded/JGroupsConfigurator.java | 5 +++++ .../JGroupsConfiguratorCustomStackTest.java | 21 +++++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/model/infinispan/src/main/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfigurator.java b/model/infinispan/src/main/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfigurator.java index d1187b69ebd6..0c2672c525ef 100644 --- a/model/infinispan/src/main/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfigurator.java +++ b/model/infinispan/src/main/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfigurator.java @@ -312,7 +312,12 @@ static boolean injectConnectionFactoryIntoCustomStack(ConfigurationBuilderHolder stack.add(new ProtocolConfiguration(protocol.getProtocolName(), properties)); } var connectionFactory = (JpaConnectionProviderFactory) sessionFactory.getProviderFactory(JpaConnectionProvider.class); + var existingNames = transportOf(holder).build().transport().jgroups().stacks().stream() + .map(s -> s.name()).toList(); var wrappedName = stackName + "-keycloak"; + for (int i = 2; existingNames.contains(wrappedName); i++) { + wrappedName = stackName + "-keycloak-" + i; + } holder.addJGroupsStack(new JpaFactoryAwareJGroupsChannelConfigurator(sessionFactory, wrappedName, stack, embedded.getUncombinedRemoteSites(), null, connectionFactory, null), null); transportOf(holder).stack(wrappedName); logger.infof("Custom stack '%s' uses %s. Injected the database connection factory.", stackName, KEYCLOAK_JDBC_PING2.class.getSimpleName()); diff --git a/model/infinispan/src/test/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfiguratorCustomStackTest.java b/model/infinispan/src/test/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfiguratorCustomStackTest.java index 0456df952662..83eb1bb22e9e 100644 --- a/model/infinispan/src/test/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfiguratorCustomStackTest.java +++ b/model/infinispan/src/test/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfiguratorCustomStackTest.java @@ -129,6 +129,27 @@ public void testFactoryInjectedIntoCustomStack() throws Exception { assertTrue("The connection factory should be used by the protocol", called.get()); } + @Test + public void testWrapperNameAvoidsExistingStack() { + var holder = new ParserRegistry().parse(""" + + + + %s + + + + + + + + """.formatted(PING_PROTOCOL)); + + assertTrue(JGroupsConfigurator.injectConnectionFactoryIntoCustomStack(holder, sessionFactory(new AtomicBoolean()), "JGROUPS_PING", "custom")); + + assertEquals("custom-keycloak-2", holder.getGlobalConfigurationBuilder().transport().attributes().attribute(TransportConfiguration.STACK).get()); + } + @Test public void testInJvmTestStackExemptionIsNarrow() { var inJvm = new ParserRegistry().parse(""" From be574566325e59c476f2574b2fdd1fcd44689bb7 Mon Sep 17 00:00:00 2001 From: Alexander Schwartz Date: Wed, 4 Nov 2026 20:01:24 +0100 Subject: [PATCH 4/4] Pre-register the JGroups address for custom stacks using KEYCLOAK_JDBC_PING2 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Alexander Schwartz --- .../impl/embedded/JGroupsConfigurator.java | 22 +++++++++++++------ .../JGroupsConfiguratorCustomStackTest.java | 8 +++---- 2 files changed, 19 insertions(+), 11 deletions(-) diff --git a/model/infinispan/src/main/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfigurator.java b/model/infinispan/src/main/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfigurator.java index 0c2672c525ef..55f6750fc05d 100644 --- a/model/infinispan/src/main/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfigurator.java +++ b/model/infinispan/src/main/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfigurator.java @@ -30,6 +30,7 @@ import java.util.Map; import java.util.Objects; import java.util.function.Consumer; +import java.util.function.Supplier; import javax.net.ssl.KeyManager; import javax.net.ssl.SSLContext; import javax.net.ssl.SSLParameters; @@ -242,7 +243,8 @@ private static void configureDiscovery(ConfigurationBuilderHolder holder, Keyclo if (stackXmlAttribute.isModified() && !isJdbcPingStack(stackXmlAttribute.get())) { logger.debugf("Custom stack configured (%s). JDBC_PING discovery disabled.", stackXmlAttribute.get()); var tableName = JpaUtils.getTableNameForNativeQuery("JGROUPS_PING", session.getProvider(JpaConnectionProvider.class).getEntityManager()); - boolean jdbcPing = injectConnectionFactoryIntoCustomStack(holder, session.getKeycloakSessionFactory(), tableName, stackXmlAttribute.get()); + boolean jdbcPing = injectConnectionFactoryIntoCustomStack(holder, session.getKeycloakSessionFactory(), tableName, stackXmlAttribute.get(), + () -> registerJGroupsAddress(session.getKeycloakSessionFactory(), transportOf(holder).attributes().attribute(CLUSTER_NAME).get())); if (!jdbcPing && Profile.isFeatureEnabled(Profile.Feature.STATELESS)) { if (isInJvmTestStack(holder, stackXmlAttribute.get())) { // TODO: Remove this case once the model tests have been migrated @@ -268,24 +270,30 @@ private static void configureDiscovery(ConfigurationBuilderHolder holder, Keyclo String clusterName = transportOf(holder).attributes().attribute(CLUSTER_NAME).get(); - Address address = Retry.call(ignored -> KeycloakModelUtils.runJobInTransactionWithResult(session.getKeycloakSessionFactory(), - s -> prepareJGroupsAddress(s, clusterName)), - 50, 10); + Address address = registerJGroupsAddress(session.getKeycloakSessionFactory(), clusterName); holder.addJGroupsStack(new JpaFactoryAwareJGroupsChannelConfigurator(session.getKeycloakSessionFactory(), stackName, stack, connectionFactory, address), null); transportOf(holder).stack(stackName); JGroupsConfigurator.logger.info("JGroups JDBC_PING discovery enabled."); } + private static Address registerJGroupsAddress(KeycloakSessionFactory sessionFactory, String clusterName) { + return Retry.call(ignored -> KeycloakModelUtils.runJobInTransactionWithResult(sessionFactory, + s -> prepareJGroupsAddress(s, clusterName)), + 50, 10); + } + /** * If a custom stack uses {@link KEYCLOAK_JDBC_PING2}, wraps it in a configurator that injects the * {@link JpaConnectionProviderFactory}, as the protocol is unable to connect to the database without it. * The stack is registered under a new name as Infinispan rejects duplicate stack names. - * The JGroups address is not pre-registered in the database, as the stack is fully controlled by the user. + * The JGroups address is pre-registered in the database like for the default stack, as this prevents multiple + * coordinators when nodes start concurrently. * + * @param addressSupplier creates and registers the address of this node, and is only called when the stack uses {@link KEYCLOAK_JDBC_PING2} * @return {@code true} if the stack uses {@link KEYCLOAK_JDBC_PING2} */ - static boolean injectConnectionFactoryIntoCustomStack(ConfigurationBuilderHolder holder, KeycloakSessionFactory sessionFactory, String tableName, String stackName) { + static boolean injectConnectionFactoryIntoCustomStack(ConfigurationBuilderHolder holder, KeycloakSessionFactory sessionFactory, String tableName, String stackName, Supplier
addressSupplier) { var configurator = transportOf(holder).build().transport().jgroups().configurator(stackName); var combined = configurator.getProtocolStack(); if (combined.stream().noneMatch(p -> KEYCLOAK_JDBC_PING2.class.getName().equals(p.getProtocolName()))) { @@ -318,7 +326,7 @@ static boolean injectConnectionFactoryIntoCustomStack(ConfigurationBuilderHolder for (int i = 2; existingNames.contains(wrappedName); i++) { wrappedName = stackName + "-keycloak-" + i; } - holder.addJGroupsStack(new JpaFactoryAwareJGroupsChannelConfigurator(sessionFactory, wrappedName, stack, embedded.getUncombinedRemoteSites(), null, connectionFactory, null), null); + holder.addJGroupsStack(new JpaFactoryAwareJGroupsChannelConfigurator(sessionFactory, wrappedName, stack, embedded.getUncombinedRemoteSites(), null, connectionFactory, addressSupplier.get()), null); transportOf(holder).stack(wrappedName); logger.infof("Custom stack '%s' uses %s. Injected the database connection factory.", stackName, KEYCLOAK_JDBC_PING2.class.getSimpleName()); return true; diff --git a/model/infinispan/src/test/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfiguratorCustomStackTest.java b/model/infinispan/src/test/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfiguratorCustomStackTest.java index 83eb1bb22e9e..186ff8fd5dfa 100644 --- a/model/infinispan/src/test/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfiguratorCustomStackTest.java +++ b/model/infinispan/src/test/java/org/keycloak/spi/infinispan/impl/embedded/JGroupsConfiguratorCustomStackTest.java @@ -105,7 +105,7 @@ public void testFactoryInjectedIntoCustomStack() throws Exception { var holder = parse(PING_PROTOCOL); var called = new AtomicBoolean(); - assertTrue(JGroupsConfigurator.injectConnectionFactoryIntoCustomStack(holder, sessionFactory(called), "JGROUPS_PING", "custom")); + assertTrue(JGroupsConfigurator.injectConnectionFactoryIntoCustomStack(holder, sessionFactory(called), "JGROUPS_PING", "custom", () -> null)); var transport = holder.getGlobalConfigurationBuilder().transport(); assertEquals("custom-keycloak", transport.attributes().attribute(TransportConfiguration.STACK).get()); @@ -145,7 +145,7 @@ public void testWrapperNameAvoidsExistingStack() { """.formatted(PING_PROTOCOL)); - assertTrue(JGroupsConfigurator.injectConnectionFactoryIntoCustomStack(holder, sessionFactory(new AtomicBoolean()), "JGROUPS_PING", "custom")); + assertTrue(JGroupsConfigurator.injectConnectionFactoryIntoCustomStack(holder, sessionFactory(new AtomicBoolean()), "JGROUPS_PING", "custom", () -> null)); assertEquals("custom-keycloak-2", holder.getGlobalConfigurationBuilder().transport().attributes().attribute(TransportConfiguration.STACK).get()); } @@ -175,7 +175,7 @@ public void testInJvmTestStackExemptionIsNarrow() { public void testStackWithoutKeycloakJdbcPingUntouched() { var holder = parse(""); - assertFalse(JGroupsConfigurator.injectConnectionFactoryIntoCustomStack(holder, sessionFactory(new AtomicBoolean()), "JGROUPS_PING", "custom")); + assertFalse(JGroupsConfigurator.injectConnectionFactoryIntoCustomStack(holder, sessionFactory(new AtomicBoolean()), "JGROUPS_PING", "custom", () -> null)); var transport = holder.getGlobalConfigurationBuilder().transport(); assertEquals("custom", transport.attributes().attribute(TransportConfiguration.STACK).get()); @@ -187,6 +187,6 @@ public void testCrossSiteRejected() { var holder = parse(PING_PROTOCOL + ""); assertThrows(IllegalStateException.class, - () -> JGroupsConfigurator.injectConnectionFactoryIntoCustomStack(holder, sessionFactory(new AtomicBoolean()), "JGROUPS_PING", "custom")); + () -> JGroupsConfigurator.injectConnectionFactoryIntoCustomStack(holder, sessionFactory(new AtomicBoolean()), "JGROUPS_PING", "custom", () -> null)); } }