[WFLY-3424] Correctly implement the ArtifactFactory to work better wi…

…th CDI.

[WFLY-8942] Setup the correct context for @requestScope beans.

batch/extension-jberet/pom.xml

@@ -43,6 +43,15 @@
             <groupId>org.jboss.spec.javax.batch</groupId>
             <artifactId>jboss-batch-api_1.0_spec</artifactId>
         </dependency>
+
+        <dependency>
+            <groupId>org.wildfly</groupId>
+            <artifactId>wildfly-weld-spi</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.jboss.weld</groupId>
+            <artifactId>weld-core-impl</artifactId>
+        </dependency>
         <dependency>
             <groupId>org.wildfly.core</groupId>
             <artifactId>wildfly-controller</artifactId>

batch/extension-jberet/src/main/java/org/wildfly/extension/batch/jberet/BatchServiceNames.java

@@ -28,6 +28,17 @@ public static ServiceName batchEnvironmentServiceName(final DeploymentUnit deplo
         return deploymentUnit.getServiceName().append("batch").append("environment");
     }
 
+    /**
+     * Creates a service name for the {@linkplain org.jberet.spi.ArtifactFactory artifact factory} service.
+     *
+     * @param deploymentUnit the deployment unit to create the service name for
+     *
+     * @return the service name
+     */
+    public static ServiceName batchArtifactFactoryServiceName(final DeploymentUnit deploymentUnit) {
+        return deploymentUnit.getServiceName().append("batch").append("artifact").append("factory");
+    }
+
     /**
      * Creates the service name used for the bean manager on the deployment.
      *

batch/extension-jberet/src/main/java/org/wildfly/extension/batch/jberet/deployment/ArtifactFactoryService.java

@@ -0,0 +1,172 @@
+/*
+ * JBoss, Home of Professional Open Source.
+ * Copyright 2015, Red Hat, Inc., and individual contributors
+ * as indicated by the @author tags. See the copyright.txt file in the
+ * distribution for a full listing of individual contributors.
+ *
+ * This is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU Lesser General Public License as
+ * published by the Free Software Foundation; either version 2.1 of
+ * the License, or (at your option) any later version.
+ *
+ * This software is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with this software; if not, write to the Free
+ * Software Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA
+ * 02110-1301 USA, or see the FSF site: http://www.fsf.org.
+ */
+
+package org.wildfly.extension.batch.jberet.deployment;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+import javax.enterprise.context.Dependent;
+import javax.enterprise.context.RequestScoped;
+import javax.enterprise.context.spi.CreationalContext;
+import javax.enterprise.inject.spi.Bean;
+import javax.enterprise.inject.spi.BeanManager;
+
+import org.jberet.creation.AbstractArtifactFactory;
+import org.jberet.spi.ArtifactFactory;
+import org.jboss.msc.service.Service;
+import org.jboss.msc.service.StartContext;
+import org.jboss.msc.service.StartException;
+import org.jboss.msc.service.StopContext;
+import org.jboss.msc.value.InjectedValue;
+import org.jboss.weld.bean.builtin.BeanManagerProxy;
+import org.jboss.weld.context.RequestContext;
+import org.jboss.weld.context.unbound.UnboundLiteral;
+import org.jboss.weld.manager.BeanManagerImpl;
+import org.wildfly.extension.batch.jberet._private.BatchLogger;
+
+/**
+ * ArtifactFactory for Java EE runtime environment.
+ *
+ * @author <a href="mailto:jperkins@redhat.com">James R. Perkins</a>
+ */
+public class ArtifactFactoryService extends AbstractArtifactFactory implements Service<ArtifactFactory>, WildFlyArtifactFactory {
+    private final InjectedValue<BeanManager> beanManagerInjector = new InjectedValue<>();
+
+    private final Map<Object, Holder> contexts = Collections.synchronizedMap(new HashMap<>());
+    private volatile BeanManager beanManager;
+
+    @Override
+    public void destroy(final Object instance) {
+        final Holder holder = contexts.remove(instance);
+        if (holder == null) {
+            // This bean was not created via CDI, we need to invoke the JBeret cleanup
+            super.destroy(instance);
+        } else {
+            // Only release the context for @Dependent beans, Weld should take care of the other scopes
+            if (holder.isDependent()) {
+                // We're not destroying the bean here because weld should handle that for use when we release the
+                // CreationalContext
+                holder.context.release();
+            }
+        }
+    }
+
+    @Override
+    public Class<?> getArtifactClass(final String ref, final ClassLoader classLoader) {
+        final Bean<?> bean = getBean(ref, getBeanManager());
+        return bean == null ? null : bean.getBeanClass();
+    }
+
+    @Override
+    public Object create(final String ref, Class<?> cls, final ClassLoader classLoader) throws Exception {
+        final BeanManager beanManager = getBeanManager();
+        final Bean<?> bean = getBean(ref, beanManager);
+        if (bean == null) {
+            return null;
+        }
+        final CreationalContext<?> context = beanManager.createCreationalContext(bean);
+        final Object result = beanManager.getReference(bean, bean.getBeanClass(), context);
+        contexts.put(result, new Holder(bean, context));
+        return result;
+    }
+
+    @Override
+    public void start(final StartContext context) throws StartException {
+        beanManager = beanManagerInjector.getOptionalValue();
+    }
+
+    @Override
+    public void stop(final StopContext context) {
+        beanManager = null;
+        synchronized (contexts) {
+            for (Holder holder : contexts.values()) {
+                // Only release the context for @Dependent beans, Weld should take care of the other scopes
+                if (holder.isDependent()) {
+                    holder.context.release();
+                }
+            }
+            contexts.clear();
+        }
+    }
+
+    @Override
+    public ArtifactFactory getValue() throws IllegalStateException, IllegalArgumentException {
+        return this;
+    }
+
+    @Override
+    public ContextHandle createContextHandle() {
+        final BeanManagerImpl beanManager = getBeanManager();
+        return () -> {
+            if (beanManager == null || beanManager.isContextActive(RequestScoped.class)) {
+                return () -> {
+                };
+            }
+            final RequestContext requestContext = beanManager.instance().select(RequestContext.class, UnboundLiteral.INSTANCE).get();
+            requestContext.activate();
+            return () -> {
+                requestContext.invalidate();
+                requestContext.deactivate();
+            };
+        };
+    }
+
+    public InjectedValue<BeanManager> getBeanManagerInjector() {
+        return beanManagerInjector;
+    }
+
+    private BeanManagerImpl getBeanManager() {
+        final BeanManager beanManager = this.beanManager;
+        return beanManager == null ? null : BeanManagerProxy.unwrap(beanManager);
+    }
+
+    private static Bean<?> getBean(final String ref, final BeanManager beanManager) {
+        if (beanManager == null) {
+            return null;
+        }
+        BatchLogger.LOGGER.tracef("Looking up bean reference for '%s'", ref);
+        final Set<Bean<?>> beans = beanManager.getBeans(ref);
+        final Bean<?> bean = beanManager.resolve(beans);
+        if (bean != null) {
+            BatchLogger.LOGGER.tracef("Found bean '%s' for reference '%s'", bean, ref);
+        } else {
+            BatchLogger.LOGGER.tracef("No bean found for reference '%s;'", ref);
+        }
+        return bean;
+    }
+
+    private static class Holder {
+        final Bean<?> bean;
+        final CreationalContext<?> context;
+
+        private Holder(final Bean<?> bean, final CreationalContext<?> context) {
+            this.bean = bean;
+            this.context = context;
+        }
+
+        boolean isDependent() {
+            return Dependent.class.equals(bean.getScope());
+        }
+    }
+}

batch/extension-jberet/src/main/java/org/wildfly/extension/batch/jberet/deployment/BatchEnvironmentProcessor.java

@@ -26,6 +26,7 @@
 import javax.transaction.TransactionManager;
 
 import org.jberet.repository.JobRepository;
+import org.jberet.spi.ArtifactFactory;
 import org.jberet.spi.ContextClassLoaderJobOperatorContextSelector;
 import org.jberet.spi.JobExecutor;
 import org.jberet.spi.JobOperatorContext;
@@ -43,6 +44,7 @@
 import org.jboss.as.txn.service.TxnServices;
 import org.jboss.modules.Module;
 import org.jboss.msc.service.ServiceBuilder;
+import org.jboss.msc.service.ServiceName;
 import org.jboss.msc.service.ServiceTarget;
 import org.jboss.msc.value.ImmediateValue;
 import org.wildfly.extension.batch.jberet.BatchConfiguration;
@@ -127,11 +129,18 @@ public void deploy(final DeploymentPhaseContext phaseContext) throws DeploymentU
             serviceBuilder.addDependency(support.getCapabilityServiceName(Capabilities.BATCH_CONFIGURATION_CAPABILITY.getName()), BatchConfiguration.class, service.getBatchConfigurationInjector());
             serviceBuilder.addDependency(TxnServices.JBOSS_TXN_TRANSACTION_MANAGER, TransactionManager.class, service.getTransactionManagerInjector());
 
+            final ServiceName artifactFactoryServiceName = BatchServiceNames.batchArtifactFactoryServiceName(deploymentUnit);
+            final ArtifactFactoryService artifactFactoryService = new ArtifactFactoryService();
+            final ServiceBuilder<ArtifactFactory> artifactFactoryServiceBuilder = serviceTarget.addService(artifactFactoryServiceName, artifactFactoryService);
+
             // Register the bean manager if this is a CDI deployment
             if (WeldDeploymentMarker.isPartOfWeldDeployment(deploymentUnit)) {
                 BatchLogger.LOGGER.tracef("Adding BeanManager service dependency for deployment %s", deploymentUnit.getName());
-                serviceBuilder.addDependency(BatchServiceNames.beanManagerServiceName(deploymentUnit), BeanManager.class, service.getBeanManagerInjector());
+                artifactFactoryServiceBuilder.addDependency(BatchServiceNames.beanManagerServiceName(deploymentUnit), BeanManager.class,
+                        artifactFactoryService.getBeanManagerInjector());
             }
+            artifactFactoryServiceBuilder.install();
+            serviceBuilder.addDependency(artifactFactoryServiceName, WildFlyArtifactFactory.class, service.getArtifactFactoryInjector());
 
             // No deployment defined repository, use the default
             if (jobRepositoryName != null) {

batch/extension-jberet/src/main/java/org/wildfly/extension/batch/jberet/deployment/BatchEnvironmentService.java

@@ -23,7 +23,6 @@
 package org.wildfly.extension.batch.jberet.deployment;
 
 import java.util.Properties;
-import javax.enterprise.inject.spi.BeanManager;
 import javax.transaction.TransactionManager;
 
 import org.jberet.repository.JobRepository;
@@ -52,7 +51,7 @@ public class BatchEnvironmentService implements Service<SecurityAwareBatchEnviro
 
     private static final Properties PROPS = new Properties();
 
-    private final InjectedValue<BeanManager> beanManagerInjector = new InjectedValue<>();
+    private final InjectedValue<WildFlyArtifactFactory> artifactFactoryInjector = new InjectedValue<>();
     private final InjectedValue<JobExecutor> jobExecutorInjector = new InjectedValue<>();
     private final InjectedValue<TransactionManager> transactionManagerInjector = new InjectedValue<>();
     private final InjectedValue<RequestController> requestControllerInjector = new InjectedValue<>();
@@ -86,7 +85,7 @@ public synchronized void start(final StartContext context) throws StartException
             jobRepository = batchConfiguration.getDefaultJobRepository();
         }
 
-        this.batchEnvironment = new WildFlyBatchEnvironment(beanManagerInjector.getOptionalValue(),
+        this.batchEnvironment = new WildFlyBatchEnvironment(artifactFactoryInjector.getValue(),
                 jobExecutor, transactionManagerInjector.getValue(),
                 jobRepository, jobXmlResolver);
 
@@ -113,8 +112,8 @@ public synchronized SecurityAwareBatchEnvironment getValue() throws IllegalState
         return batchEnvironment;
     }
 
-    public InjectedValue<BeanManager> getBeanManagerInjector() {
-        return beanManagerInjector;
+    public InjectedValue<WildFlyArtifactFactory> getArtifactFactoryInjector() {
+        return artifactFactoryInjector;
     }
 
     public InjectedValue<JobExecutor> getJobExecutorInjector() {
@@ -139,19 +138,19 @@ public InjectedValue<BatchConfiguration> getBatchConfigurationInjector() {
 
     private class WildFlyBatchEnvironment implements BatchEnvironment, SecurityAwareBatchEnvironment {
 
-        private final ArtifactFactory artifactFactory;
+        private final WildFlyArtifactFactory artifactFactory;
         private final JobExecutor jobExecutor;
         private final TransactionManager transactionManager;
         private final JobRepository jobRepository;
         private final JobXmlResolver jobXmlResolver;
 
-        WildFlyBatchEnvironment(final BeanManager beanManager,
+        WildFlyBatchEnvironment(final WildFlyArtifactFactory artifactFactory,
                                 final JobExecutor jobExecutor,
                                 final TransactionManager transactionManager,
                                 final JobRepository jobRepository,
                                 final JobXmlResolver jobXmlResolver) {
             this.jobXmlResolver = jobXmlResolver;
-            artifactFactory = new WildFlyArtifactFactory(beanManager);
+            this.artifactFactory = artifactFactory;
             this.jobExecutor = jobExecutor;
             this.transactionManager = transactionManager;
             this.jobRepository = jobRepository;
@@ -229,7 +228,8 @@ private ContextHandle createContextHandle() {
             // If the TCCL is null, use the deployments ModuleClassLoader
             final ClassLoaderContextHandle classLoaderContextHandle = (tccl == null ? new ClassLoaderContextHandle(classLoader) : new ClassLoaderContextHandle(tccl));
             // Class loader handle must be first so the TCCL is set before the other handles execute
-            return new ContextHandle.ChainedContextHandle(classLoaderContextHandle, new NamespaceContextHandle(), new SecurityContextHandle());
+            return new ContextHandle.ChainedContextHandle(classLoaderContextHandle, new NamespaceContextHandle(),
+                    new SecurityContextHandle(), artifactFactory.createContextHandle());
         }
     }
 }