Merge pull request #11961 from bstansberry/WFLY-11543

[WFLY-11543] Server-side services outside the transactions subsystem should not use ContextTransactionSynchronizationRegistry

datasources-agroal/src/main/java/org/wildfly/extension/datasources/agroal/DataSourceOperations.java

@@ -21,6 +21,8 @@
  */
 package org.wildfly.extension.datasources.agroal;
 
+import javax.transaction.TransactionSynchronizationRegistry;
+
 import io.agroal.api.AgroalDataSource;
 import io.agroal.api.configuration.supplier.AgroalConnectionFactoryConfigurationSupplier;
 import io.agroal.api.configuration.supplier.AgroalConnectionPoolConfigurationSupplier;
@@ -82,7 +84,10 @@ protected void performRuntime(OperationContext context, ModelNode operation, Mod
 
             CapabilityServiceBuilder<AgroalDataSource> serviceBuilder = context.getCapabilityServiceTarget().addCapability(AbstractDataSourceDefinition.DATA_SOURCE_CAPABILITY.fromBaseCapability(datasourceName), dataSourceService);
             serviceBuilder.addCapabilityRequirement(DriverDefinition.AGROAL_DRIVER_CAPABILITY.getDynamicName(driverName), Class.class, dataSourceService.getDriverInjector());
-
+            if (jta) {
+                // TODO add a Stage.MODEL requirement
+                serviceBuilder.addCapabilityRequirement("org.wildfly.transactions.transaction-synchronization-registry", TransactionSynchronizationRegistry.class, dataSourceService.getTransactionSynchronizationRegistryInjector());
+            }
             AbstractDataSourceOperations.setupElytronSecurity(context, factoryModel, dataSourceService, serviceBuilder);
 
             serviceBuilder.install();

datasources-agroal/src/main/java/org/wildfly/extension/datasources/agroal/DataSourceService.java

@@ -49,7 +49,6 @@
 import org.wildfly.security.credential.source.CredentialSource;
 import org.wildfly.security.password.interfaces.ClearPassword;
 import org.wildfly.transaction.client.ContextTransactionManager;
-import org.wildfly.transaction.client.ContextTransactionSynchronizationRegistry;
 
 import javax.security.auth.callback.Callback;
 import javax.security.auth.callback.NameCallback;
@@ -88,6 +87,7 @@ public class DataSourceService implements Service<AgroalDataSource>, Supplier<Ag
     private InjectedValue<Class> driverInjector = new InjectedValue<>();
     private InjectedValue<AuthenticationContext> authenticationContextInjector = new InjectedValue<>();
     private InjectedValue<ExceptionSupplier<CredentialSource, Exception>> credentialSourceSupplierInjector = new InjectedValue<>();
+    private InjectedValue<TransactionSynchronizationRegistry> transactionSynchronizationRegistryInjector = new InjectedValue<>();
 
     public DataSourceService(String dataSourceName, String jndiName, boolean jta, boolean connectable, boolean xa, AgroalDataSourceConfigurationSupplier dataSourceConfiguration) {
         this.dataSourceName = dataSourceName;
@@ -115,7 +115,7 @@ public void start(StartContext context) throws StartException {
 
         if (jta || xa) {
             TransactionManager transactionManager = ContextTransactionManager.getInstance();
-            TransactionSynchronizationRegistry transactionSynchronizationRegistry = ContextTransactionSynchronizationRegistry.getInstance();
+            TransactionSynchronizationRegistry transactionSynchronizationRegistry = transactionSynchronizationRegistryInjector.getValue();
 
             if (transactionManager == null || transactionSynchronizationRegistry == null) {
                 throw AgroalLogger.SERVICE_LOGGER.missingTransactionManager();
@@ -231,4 +231,8 @@ public InjectedValue<AuthenticationContext> getAuthenticationContextInjector() {
     public InjectedValue<ExceptionSupplier<CredentialSource, Exception>> getCredentialSourceSupplierInjector() {
         return credentialSourceSupplierInjector;
     }
+
+     InjectedValue<TransactionSynchronizationRegistry> getTransactionSynchronizationRegistryInjector() {
+        return transactionSynchronizationRegistryInjector;
+    }
 }

datasources-agroal/src/main/java/org/wildfly/extension/datasources/agroal/XADataSourceOperations.java

@@ -23,6 +23,8 @@
 
 import static org.jboss.as.controller.descriptions.ModelDescriptionConstants.OP_ADDR;
 
+import javax.transaction.TransactionSynchronizationRegistry;
+
 import io.agroal.api.AgroalDataSource;
 import io.agroal.api.configuration.supplier.AgroalConnectionFactoryConfigurationSupplier;
 import io.agroal.api.configuration.supplier.AgroalConnectionPoolConfigurationSupplier;
@@ -83,6 +85,8 @@ protected void performRuntime(OperationContext context, ModelNode operation, Mod
 
             CapabilityServiceBuilder<AgroalDataSource> serviceBuilder = context.getCapabilityServiceTarget().addCapability(AbstractDataSourceDefinition.DATA_SOURCE_CAPABILITY.fromBaseCapability(datasourceName), dataSourceService);
             serviceBuilder.addCapabilityRequirement(DriverDefinition.AGROAL_DRIVER_CAPABILITY.getDynamicName(driverName), Class.class, dataSourceService.getDriverInjector());
+            // TODO add a Stage.MODEL requirement
+            serviceBuilder.addCapabilityRequirement("org.wildfly.transactions.transaction-synchronization-registry", TransactionSynchronizationRegistry.class, dataSourceService.getTransactionSynchronizationRegistryInjector());
 
             AbstractDataSourceOperations.setupElytronSecurity(context, factoryModel, dataSourceService, serviceBuilder);
 

datasources-agroal/src/main/java/org/wildfly/extension/datasources/agroal/deployment/DataSourceDefinitionInjectionSource.java

@@ -21,17 +21,29 @@
  */
 package org.wildfly.extension.datasources.agroal.deployment;
 
+import static org.jboss.as.server.deployment.Attachments.MODULE;
+
+import java.sql.Driver;
+import java.time.Duration;
+import java.util.Map;
+import java.util.function.Supplier;
+
+import javax.sql.DataSource;
+import javax.transaction.TransactionSynchronizationRegistry;
+
 import io.agroal.api.configuration.AgroalConnectionFactoryConfiguration;
 import io.agroal.api.configuration.supplier.AgroalConnectionFactoryConfigurationSupplier;
 import io.agroal.api.configuration.supplier.AgroalConnectionPoolConfigurationSupplier;
 import io.agroal.api.configuration.supplier.AgroalDataSourceConfigurationSupplier;
 import io.agroal.api.security.NamePrincipal;
 import io.agroal.api.security.SimplePassword;
+import org.jboss.as.controller.capability.CapabilityServiceSupport;
 import org.jboss.as.ee.resource.definition.ResourceDefinitionInjectionSource;
 import org.jboss.as.naming.ManagedReferenceFactory;
 import org.jboss.as.naming.ServiceBasedNamingStore;
 import org.jboss.as.naming.deployment.ContextNames;
 import org.jboss.as.naming.service.BinderService;
+import org.jboss.as.server.deployment.Attachments;
 import org.jboss.as.server.deployment.DeploymentPhaseContext;
 import org.jboss.as.server.deployment.DeploymentUnitProcessingException;
 import org.jboss.msc.inject.Injector;
@@ -40,13 +52,6 @@
 import org.wildfly.extension.datasources.agroal.AgroalExtension;
 import org.wildfly.extension.datasources.agroal.logging.AgroalLogger;
 
-import javax.sql.DataSource;
-import java.sql.Driver;
-import java.time.Duration;
-import java.util.Map;
-
-import static org.jboss.as.server.deployment.Attachments.MODULE;
-
 /**
  * Injection source for a DataSource
  *
@@ -228,8 +233,16 @@ public void getResourceValue(ResolutionContext context, ServiceBuilder<?> servic
                     .addDependency(bindInfo.getParentContextServiceName(), ServiceBasedNamingStore.class, binderService.getNamingStoreInjector())
                     .install();
 
-        DataSourceDefinitionService dataSourceService = new DataSourceDefinitionService(bindInfo, transactional, dataSourceConfiguration);
-        phaseContext.getServiceTarget().addService(dataSourceServiceName).setInstance(dataSourceService).install();
+        ServiceBuilder svcBuilder = phaseContext.getServiceTarget().addService(dataSourceServiceName);
+        Supplier<TransactionSynchronizationRegistry> tsrSupplier = null;
+        if (transactional) {
+            CapabilityServiceSupport css = phaseContext.getDeploymentUnit().getAttachment(Attachments.CAPABILITY_SERVICE_SUPPORT);
+            ServiceName tsrName = css.getCapabilityServiceName("org.wildfly.transactions.transaction-synchronization-registry");
+            //noinspection unchecked
+            tsrSupplier = (Supplier<TransactionSynchronizationRegistry>) svcBuilder.requires(tsrName);
+        }
+        DataSourceDefinitionService dataSourceService = new DataSourceDefinitionService(bindInfo, transactional, dataSourceConfiguration, tsrSupplier);
+        svcBuilder.setInstance(dataSourceService).install();
 
         serviceBuilder.requires(bindInfo.getBinderServiceName());
         serviceBuilder.addDependency(dataSourceServiceName, ManagedReferenceFactory.class, injector);

datasources-agroal/src/main/java/org/wildfly/extension/datasources/agroal/deployment/DataSourceDefinitionService.java

@@ -36,11 +36,11 @@
 import org.wildfly.extension.datasources.agroal.logging.AgroalLogger;
 import org.wildfly.extension.datasources.agroal.logging.LoggingDataSourceListener;
 import org.wildfly.transaction.client.ContextTransactionManager;
-import org.wildfly.transaction.client.ContextTransactionSynchronizationRegistry;
 
 import javax.transaction.TransactionManager;
 import javax.transaction.TransactionSynchronizationRegistry;
 import java.sql.SQLException;
+import java.util.function.Supplier;
 
 /**
  * Defines an extension to provide DataSources based on the Agroal project
@@ -53,12 +53,15 @@ public class DataSourceDefinitionService implements Service<ManagedReferenceFact
     private final String jndiBinding;
     private final boolean transactional;
     private final AgroalDataSourceConfigurationSupplier dataSourceConfiguration;
+    private final Supplier<TransactionSynchronizationRegistry> tsrSupplier;
     private AgroalDataSource agroalDataSource;
 
-    public DataSourceDefinitionService(ContextNames.BindInfo bindInfo, boolean transactional, AgroalDataSourceConfigurationSupplier dataSourceConfiguration) {
+    public DataSourceDefinitionService(ContextNames.BindInfo bindInfo, boolean transactional, AgroalDataSourceConfigurationSupplier dataSourceConfiguration,
+                                       Supplier<TransactionSynchronizationRegistry> tsrSupplier) {
         this.dataSourceName = bindInfo.getParentContextServiceName().getSimpleName() + "." + bindInfo.getBindName().replace('/', '.');
         this.jndiBinding = bindInfo.getBindName();
         this.transactional = transactional;
+        this.tsrSupplier = tsrSupplier;
         this.dataSourceConfiguration = dataSourceConfiguration;
     }
 
@@ -76,7 +79,7 @@ public String getJndiViewInstanceValue() {
     public void start(StartContext context) throws StartException {
         if (transactional) {
             TransactionManager transactionManager = ContextTransactionManager.getInstance();
-            TransactionSynchronizationRegistry transactionSynchronizationRegistry = ContextTransactionSynchronizationRegistry.getInstance();
+            TransactionSynchronizationRegistry transactionSynchronizationRegistry = tsrSupplier != null ? tsrSupplier.get() : null;
 
             if (transactionManager == null || transactionSynchronizationRegistry == null) {
                 throw AgroalLogger.SERVICE_LOGGER.missingTransactionManager();

ejb3/src/main/java/org/jboss/as/ejb3/component/EJBComponent.java

@@ -33,6 +33,9 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Function;
 
 import javax.ejb.EJBHome;
 import javax.ejb.EJBLocalHome;
@@ -49,10 +52,6 @@
 import javax.transaction.TransactionSynchronizationRegistry;
 import javax.transaction.UserTransaction;
 
-import java.util.Set;
-import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.function.Function;
-
 import org.jboss.as.core.security.ServerSecurityManager;
 import org.jboss.as.ee.component.BasicComponent;
 import org.jboss.as.ee.component.ComponentView;
@@ -86,7 +85,6 @@
 import org.wildfly.security.authz.Roles;
 import org.wildfly.security.manager.WildFlySecurityManager;
 import org.wildfly.transaction.client.ContextTransactionManager;
-import org.wildfly.transaction.client.ContextTransactionSynchronizationRegistry;
 
 /**
  * @author <a href="mailto:cdewolf@redhat.com">Carlo de Wolf</a>
@@ -121,6 +119,7 @@ public abstract class EJBComponent extends BasicComponent implements ServerActiv
     private final InvocationMetrics invocationMetrics = new InvocationMetrics();
     private final EJBSuspendHandlerService ejbSuspendHandlerService;
     private final ShutDownInterceptorFactory shutDownInterceptorFactory;
+    private final TransactionSynchronizationRegistry transactionSynchronizationRegistry;
     private final UserTransaction userTransaction;
     private final ServerSecurityManager serverSecurityManager;
     private final ControlPoint controlPoint;
@@ -181,6 +180,7 @@ protected EJBComponent(final EJBComponentCreateService ejbComponentCreateService
         this.timeoutInterceptors = Collections.unmodifiableMap(ejbComponentCreateService.getTimeoutInterceptors());
         this.shutDownInterceptorFactory = ejbComponentCreateService.getShutDownInterceptorFactory();
         this.ejbSuspendHandlerService = ejbComponentCreateService.getEJBSuspendHandler();
+        this.transactionSynchronizationRegistry = ejbComponentCreateService.getTransactionSynchronizationRegistry();
         this.userTransaction = ejbComponentCreateService.getUserTransaction();
         this.serverSecurityManager = ejbComponentCreateService.getServerSecurityManager();
         this.controlPoint = ejbComponentCreateService.getControlPoint();
@@ -416,13 +416,8 @@ public TransactionManager getTransactionManager() {
         return ContextTransactionManager.getInstance();
     }
 
-    /**
-     * @deprecated Use {@link ContextTransactionSynchronizationRegistry#getInstance()} instead.
-     * @return the value of {@link ContextTransactionSynchronizationRegistry#getInstance()}
-     */
-    @Deprecated
     public TransactionSynchronizationRegistry getTransactionSynchronizationRegistry() {
-        return ContextTransactionSynchronizationRegistry.getInstance();
+        return transactionSynchronizationRegistry;
     }
 
     public int getTransactionTimeout(final MethodIntf methodIntf, final Method method) {

ejb3/src/main/java/org/jboss/as/ejb3/component/EJBComponentCreateService.java

@@ -22,11 +22,6 @@
 
 package org.jboss.as.ejb3.component;
 
-import javax.ejb.TimerService;
-import javax.ejb.TransactionAttributeType;
-import javax.ejb.TransactionManagementType;
-import javax.transaction.UserTransaction;
-
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
 import java.util.Collections;
@@ -39,6 +34,12 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Function;
 
+import javax.ejb.TimerService;
+import javax.ejb.TransactionAttributeType;
+import javax.ejb.TransactionManagementType;
+import javax.transaction.TransactionSynchronizationRegistry;
+import javax.transaction.UserTransaction;
+
 import org.jboss.as.core.security.ServerSecurityManager;
 import org.jboss.as.ee.component.BasicComponentCreateService;
 import org.jboss.as.ee.component.ComponentConfiguration;
@@ -97,6 +98,7 @@ public class EJBComponentCreateService extends BasicComponentCreateService {
     private final String distinctName;
     private final String policyContextID;
 
+    private final InjectedValue<TransactionSynchronizationRegistry> transactionSynchronizationRegistryValue = new InjectedValue<TransactionSynchronizationRegistry>();
     private final InjectedValue<ServerSecurityManager> serverSecurityManagerInjectedValue = new InjectedValue<>();
     private final InjectedValue<ControlPoint> controlPoint = new InjectedValue<>();
     private final InjectedValue<AtomicBoolean> exceptionLoggingEnabled = new InjectedValue<>();
@@ -327,6 +329,14 @@ UserTransaction getUserTransaction() {
         return LocalUserTransaction.getInstance();
     }
 
+    Injector<TransactionSynchronizationRegistry> getTransactionSynchronizationRegistryInjector() {
+        return transactionSynchronizationRegistryValue;
+    }
+
+    TransactionSynchronizationRegistry getTransactionSynchronizationRegistry() {
+        return transactionSynchronizationRegistryValue.getValue();
+    }
+
     public Injector<EJBSuspendHandlerService> getEJBSuspendHandlerInjector() {
         return this.ejbSuspendHandler;
     }

ejb3/src/main/java/org/jboss/as/ejb3/component/EJBComponentDescription.java

@@ -43,6 +43,7 @@
 import javax.ejb.TimerService;
 import javax.ejb.TransactionAttributeType;
 import javax.ejb.TransactionManagementType;
+import javax.transaction.TransactionSynchronizationRegistry;
 
 import org.jboss.as.controller.capability.CapabilityServiceSupport;
 import org.jboss.as.core.security.ServerSecurityManager;
@@ -568,7 +569,7 @@ public void configureDependency(final ServiceBuilder<?> serviceBuilder, final EJ
                         // add dependency on the local transaction provider
                         serviceBuilder.requires(support.getCapabilityServiceName("org.wildfly.transactions.global-default-local-provider"));
                         // add dependency on TransactionSynchronizationRegistry
-                        serviceBuilder.requires(support.getCapabilityServiceName("org.wildfly.transactions.transaction-synchronization-registry"));
+                        serviceBuilder.addDependency(support.getCapabilityServiceName("org.wildfly.transactions.transaction-synchronization-registry"), TransactionSynchronizationRegistry.class, ejbComponentCreateService.getTransactionSynchronizationRegistryInjector());
                     }
                 });