From 3afd0799be4bb9292d7a887bba1d73cc1fb9fbdd Mon Sep 17 00:00:00 2001
From: mincong <mincong.h@gmail.com>
Date: Sun, 17 Jul 2016 18:19:12 +0200
Subject: [PATCH] #25 add the afterChunk enhancement

---
 .../steps/afterChunk/AfterChunkBatchlet.java  | 76 +++++++++++++++++++
 .../META-INF/batch-jobs/mass-index.xml        | 11 ++-
 .../search/jsr352/MassIndexerIT.java          |  2 +-
 3 files changed, 87 insertions(+), 2 deletions(-)
 create mode 100644 core/src/main/java/org/hibernate/search/jsr352/internal/steps/afterChunk/AfterChunkBatchlet.java

diff --git a/core/src/main/java/org/hibernate/search/jsr352/internal/steps/afterChunk/AfterChunkBatchlet.java b/core/src/main/java/org/hibernate/search/jsr352/internal/steps/afterChunk/AfterChunkBatchlet.java
new file mode 100644
index 00000000000..e024cd6ee3d
--- /dev/null
+++ b/core/src/main/java/org/hibernate/search/jsr352/internal/steps/afterChunk/AfterChunkBatchlet.java
@@ -0,0 +1,76 @@
+/*
+ * Hibernate Search, full-text search for your domain model
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later
+ * See the lgpl.txt file in the root directory or <http://www.gnu.org/licenses/lgpl-2.1.html>.
+ */
+package org.hibernate.search.jsr352.internal.steps.afterChunk;
+
+import java.util.Set;
+
+import javax.batch.api.BatchProperty;
+import javax.batch.api.Batchlet;
+import javax.batch.runtime.context.JobContext;
+import javax.inject.Inject;
+import javax.inject.Named;
+import javax.naming.InitialContext;
+import javax.persistence.EntityManager;
+
+import org.hibernate.Session;
+import org.hibernate.search.backend.spi.BatchBackend;
+import org.hibernate.search.hcore.util.impl.ContextHelper;
+import org.hibernate.search.jsr352.internal.JobContextData;
+import org.jboss.logging.Logger;
+
+/**
+ * Enhancements after the chunk step "produceLuceneDoc" (lucene document
+ * production)
+ *
+ * @author Mincong Huang
+ */
+@Named
+public class AfterChunkBatchlet implements Batchlet {
+
+	private static final Logger logger = Logger.getLogger( AfterChunkBatchlet.class );
+	private final JobContext jobContext;
+
+	@Inject
+	@BatchProperty
+	private String persistenceUnitName;
+
+	@Inject
+	@BatchProperty
+	private boolean optimizeAtEnd;
+
+	@Inject
+	public AfterChunkBatchlet(JobContext jobContext) {
+		this.jobContext = jobContext;
+	}
+
+	@Override
+	public String process() throws Exception {
+
+		if ( this.optimizeAtEnd ) {
+
+			logger.info( "purging index for all entities ..." );
+			String path = "java:comp/env/" + persistenceUnitName;
+			EntityManager em = (EntityManager) InitialContext.doLookup( path );
+			Session session = em.unwrap( Session.class );
+			final BatchBackend backend = ContextHelper
+					.getSearchintegrator( session )
+					.makeBatchBackend( null );
+
+			logger.info( "optimizing all entities ..." );
+			JobContextData jobData = (JobContextData) jobContext.getTransientUserData();
+			Set<Class<?>> targetedClasses = jobData.getEntityClazzSet();
+			backend.optimize( targetedClasses );
+			backend.flush( targetedClasses );
+		}
+		return null;
+	}
+
+	@Override
+	public void stop() throws Exception {
+		// TODO Auto-generated method stub
+	}
+}
diff --git a/core/src/main/resources/META-INF/batch-jobs/mass-index.xml b/core/src/main/resources/META-INF/batch-jobs/mass-index.xml
index 8e93437a3b4..c9f69885f21 100644
--- a/core/src/main/resources/META-INF/batch-jobs/mass-index.xml
+++ b/core/src/main/resources/META-INF/batch-jobs/mass-index.xml
@@ -27,7 +27,7 @@
         </batchlet>
     </step>
 
-    <step id="produceLuceneDoc">
+    <step id="produceLuceneDoc" next="afterChunk">
         <chunk checkpoint-policy="custom">
             <reader ref="itemReader">
                 <properties>
@@ -69,4 +69,13 @@
             <reducer ref="partitionReducer" />
         </partition>
     </step>
+
+    <step id="afterChunk">
+        <batchlet ref="afterChunkBatchlet">
+            <properties>
+                <property name="optimizeAtEnd" value="#{jobParameters['optimizeAtEnd']}" />
+                <property name="persistenceUnitName" value="#{jobParameters['persistenceUnitName']}" />
+            </properties>
+        </batchlet>
+    </step>
 </job>
diff --git a/integrationtest/javaee-wildfly/src/test/java/org/hibernate/search/jsr352/MassIndexerIT.java b/integrationtest/javaee-wildfly/src/test/java/org/hibernate/search/jsr352/MassIndexerIT.java
index 66924fb02fa..2d9c99e51f3 100644
--- a/integrationtest/javaee-wildfly/src/test/java/org/hibernate/search/jsr352/MassIndexerIT.java
+++ b/integrationtest/javaee-wildfly/src/test/java/org/hibernate/search/jsr352/MassIndexerIT.java
@@ -49,7 +49,7 @@ public class MassIndexerIT {
 	private static final Logger logger = Logger.getLogger( MassIndexerIT.class );
 
 	private final boolean JOB_OPTIMIZE_AFTER_PURGE = true;
-	private final boolean JOB_OPTIMIZE_AT_END = false;
+	private final boolean JOB_OPTIMIZE_AT_END = true;
 	private final boolean JOB_PURGE_AT_START = true;
 	private final int JOB_FETCH_SIZE = 100 * 1000;
 	private final int JOB_MAX_RESULTS = 200 * 1000;