[WFLY-11543] Remove use of ContextTransactionSynchronizationRegistry …

…in agroal

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;
@@ -31,6 +33,7 @@
 import org.jboss.as.controller.OperationContext;
 import org.jboss.as.controller.OperationFailedException;
 import org.jboss.as.controller.OperationStepHandler;
+import org.jboss.as.txn.service.TxnServices;
 import org.jboss.dmr.ModelNode;
 import org.jboss.msc.service.ServiceName;
 
@@ -82,7 +85,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 replace with capability injection
+                serviceBuilder.addDependency(TxnServices.JBOSS_TXN_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;
@@ -34,6 +36,7 @@
 import org.jboss.as.controller.OperationFailedException;
 import org.jboss.as.controller.OperationStepHandler;
 import org.jboss.as.controller.PathAddress;
+import org.jboss.as.txn.service.TxnServices;
 import org.jboss.dmr.ModelNode;
 import org.jboss.msc.service.ServiceName;
 
@@ -83,6 +86,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 replace with capability injection
+            serviceBuilder.addDependency(TxnServices.JBOSS_TXN_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,6 +21,16 @@
  */
 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;
@@ -34,19 +44,13 @@
 import org.jboss.as.naming.service.BinderService;
 import org.jboss.as.server.deployment.DeploymentPhaseContext;
 import org.jboss.as.server.deployment.DeploymentUnitProcessingException;
+import org.jboss.as.txn.service.TxnServices;
 import org.jboss.msc.inject.Injector;
 import org.jboss.msc.service.ServiceBuilder;
 import org.jboss.msc.service.ServiceName;
 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 +232,15 @@ 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) {
+            // TODO replace with capability injection
+            //noinspection unchecked
+            tsrSupplier = (Supplier<TransactionSynchronizationRegistry>) svcBuilder.requires(TxnServices.JBOSS_TXN_SYNCHRONIZATION_REGISTRY);
+        }
+        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();