Custom Thread Groups - Add support for IteratingController implement (undera#714)

* Add support for IteratingController implement

* Increase version number

Author: 3dgiordano

plugins/casutg/pom.xml

@@ -12,7 +12,7 @@
 
     <groupId>kg.apc</groupId>
     <artifactId>jmeter-plugins-casutg</artifactId>
-    <version>3.0</version>
+    <version>3.1</version>
 
     <name>Custom Thread Groups</name>
     <description>Custom Thread Groups</description>

plugins/casutg/src/main/java/com/blazemeter/jmeter/control/VirtualUserController.java

@@ -5,40 +5,47 @@
 import com.blazemeter.jmeter.threads.arrivals.ArrivalsThreadGroup;
 import com.blazemeter.jmeter.threads.concurrency.ConcurrencyThreadGroup;
 import org.apache.jmeter.control.GenericController;
+import org.apache.jmeter.control.IteratingController;
 import org.apache.jmeter.control.NextIsNullException;
+import org.apache.jmeter.engine.event.LoopIterationEvent;
 import org.apache.jmeter.samplers.Sampler;
 import org.apache.jmeter.threads.JMeterContextService;
 import org.apache.jmeter.threads.JMeterThread;
 import org.slf4j.LoggerFactory;
 import org.slf4j.Logger;
 
-public class VirtualUserController extends GenericController {
+public class VirtualUserController extends GenericController implements IteratingController {
     private static final Logger log = LoggerFactory.getLogger(VirtualUserController.class);
     private boolean hasArrived = false;
     protected AbstractDynamicThreadGroup owner;
-
+    private boolean breakLoop;
     private long iterationNo = 0;
 
     @Override
     public Sampler next() {
-        if (owner.isLimitReached()) {
-            setDone(true);
-        } else if (!hasArrived) {
-            if (owner.isLimitReached()) {
-                throw new IllegalStateException("Should not have more iterations");
-            }
-            hasArrived = true;
-            iterationNo++;
-            if (owner instanceof ArrivalsThreadGroup) {
-                getOwnerAsArrivals().arrivalFact(JMeterContextService.getContext().getThread(), iterationNo);
-                if (!owner.isRunning()) {
-                    setDone(true);
-                    return null;
+        updateIterationIndex(owner.getName(), (int) iterationNo);
+        try {
+            if (breakLoop || owner.isLimitReached()) {
+                setDone(true);
+            } else if (!hasArrived) {
+                if (owner.isLimitReached()) {
+                    throw new IllegalStateException("Should not have more iterations");
+                }
+                hasArrived = true;
+                incrementLoopCount();
+                updateIterationIndex(owner.getName(), (int) iterationNo);
+                if (owner instanceof ArrivalsThreadGroup) {
+                    getOwnerAsArrivals().arrivalFact(JMeterContextService.getContext().getThread(), iterationNo);
+                    if (!owner.isRunning()) {
+                        setDone(true);
+                        return null;
+                    }
                 }
             }
+            return super.next();
+        } finally {
+            updateIterationIndex(owner.getName(), (int) iterationNo);
         }
-
-        return super.next();
     }
 
     private boolean moveToPool(JMeterThread thread) {
@@ -57,6 +64,12 @@ protected void reInitialize() {
         hasArrived = false;
     }
 
+    @Override
+    protected void setDone(boolean done) {
+        resetBreakLoop();
+        super.setDone(done);
+    }
+
     @Override
     protected Sampler nextIsNull() throws NextIsNullException {
         JMeterThread thread = JMeterContextService.getContext().getThread();
@@ -87,10 +100,23 @@ protected Sampler nextIsNull() throws NextIsNullException {
         }
     }
 
+    protected void incrementLoopCount() {
+        iterationNo++;
+    }
+
+    protected void resetLoopCount() {
+        iterationNo = 0;
+    }
+
     public void setOwner(AbstractDynamicThreadGroup owner) {
         this.owner = owner;
     }
 
+    @Override
+    protected int getIterCount() {
+        return (int) (iterationNo + 1);
+    }
+
     public void startNextLoop() {
         JMeterThread thread = JMeterContextService.getContext().getThread();
         if (owner instanceof ArrivalsThreadGroup) {
@@ -104,7 +130,37 @@ public void startNextLoop() {
         }
     }
 
+    @Override
+    public void breakLoop() {
+        breakLoop = true;
+        setFirst(true);
+        resetCurrent();
+        resetLoopCount();
+        recoverRunningVersion();
+    }
+
+    private void resetBreakLoop() {
+        if(breakLoop) {
+            breakLoop = false;
+        }
+    }
+
     private ArrivalsThreadGroup getOwnerAsArrivals() {
         return (ArrivalsThreadGroup) owner;
     }
+
+    @Override
+    public void removed() {
+        super.removed();
+    }
+
+    @Override
+    public void iterationStart(LoopIterationEvent loopIterationEvent) {
+        if (log.isDebugEnabled()) {
+            log.debug("iterationStart called on {} with source {} and iteration {}", owner.getName(),
+                loopIterationEvent.getSource(), loopIterationEvent.getIteration());
+        }
+        reInitialize();
+        resetLoopCount();
+    }
 }