Add b | batches option for batch command consume (#41)

* Added -b to add a specific number of batches to be consumed set in command

Author: andreea-anamaria

Command/BatchConsumerCommand.php

@@ -16,6 +16,9 @@ final class BatchConsumerCommand extends BaseRabbitMqCommand
      */
     protected $consumer;
 
+    /** @var int */
+    protected $amount;
+
     public function stopConsumer()
     {
         if ($this->consumer instanceof BatchConsumer) {
@@ -36,6 +39,7 @@ protected function configure()
         $this
             ->setName('rabbitmq:batch:consumer')
             ->addArgument('name', InputArgument::REQUIRED, 'Consumer Name')
+            ->addOption('batches', 'b', InputOption::VALUE_OPTIONAL, 'Number of batches to consume', 0)
             ->addOption('route', 'r', InputOption::VALUE_OPTIONAL, 'Routing Key', '')
             ->addOption('memory-limit', 'l', InputOption::VALUE_OPTIONAL, 'Allowed memory for this process', null)
             ->addOption('debug', 'd', InputOption::VALUE_NONE, 'Enable Debugging')
@@ -52,7 +56,7 @@ protected function configure()
      *
      * @return  integer                         0 if everything went fine, or an error code
      *
-     * @throws  \InvalidArgumentException       When the number of messages to consume is less than 0
+     * @throws  \InvalidArgumentException       When the number of batches to consume is less than 0
      * @throws  \BadFunctionCallException       When the pcntl is not installed and option -s is true
      */
     protected function execute(InputInterface $input, OutputInterface $output)
@@ -74,9 +78,15 @@ protected function execute(InputInterface $input, OutputInterface $output)
             define('AMQP_DEBUG', (bool) $input->getOption('debug'));
         }
 
+        $this->amount = (int) $input->getOption('batches');
+
+        if (0 > $this->amount) {
+            throw new \InvalidArgumentException("The -b option should be null or greater than 0");
+        }
+
         $this->initConsumer($input);
 
-        return $this->consumer->consume();
+        return $this->consumer->consume($this->amount);
     }
 
     /**

README.md

@@ -913,6 +913,8 @@ How to run the following batch consumer:
 ```
 
 Important: BatchConsumers will not have the -m|messages option available
+Important: BatchConsumers can also have the -b|batches option available if you want to only consume a specific number of batches and then stop the consumer.
+! Give the number of the batches only if you want the consumer to stop after those batch messages were consumed.! 
 
 ### STDIN Producer ###
 

RabbitMq/BatchConsumer.php

@@ -59,6 +59,11 @@ class BatchConsumer extends BaseAmqp implements DequeuerInterface
      */
     protected $batchCounter = 0;
 
+    /**
+     * @var int
+     */
+    protected $batchAmount = 0;
+
     /**
      * @var int
      */
@@ -70,6 +75,9 @@ class BatchConsumer extends BaseAmqp implements DequeuerInterface
      */
     protected $gracefulMaxExecutionDateTime;
 
+    /** @var int */
+    private $target;
+
     /**
      * @param \DateTime|null $dateTime
      */
@@ -98,8 +106,10 @@ public function setCallback($callback)
         return $this;
     }
 
-    public function consume()
+    public function consume(int $batchAmount = 0)
     {
+        $this->target = $batchAmount;
+
         $this->setupConsumer();
 
         while (count($this->getChannel()->callbacks)) {
@@ -117,6 +127,7 @@ public function consume()
             } catch (AMQPTimeoutException $e) {
                 if (!$this->isEmptyBatch()) {
                     $this->batchConsume();
+                    $this->maybeStopConsumer();
                 } elseif ($this->keepAlive === true) {
                     continue;
                 } elseif (null !== $this->getIdleTimeoutExitCode()) {
@@ -174,6 +185,7 @@ private function batchConsume()
             throw $e;
         }
 
+        $this->batchAmount++;
         $this->resetBatch();
     }
 
@@ -211,6 +223,7 @@ private function handleProcessFlag($deliveryTag, $processFlag)
             // Remove message from queue only if callback return not false
             $this->getMessageChannel($deliveryTag)->basic_ack($deliveryTag);
         }
+
     }
 
     /**
@@ -359,7 +372,7 @@ protected function maybeStopConsumer()
             pcntl_signal_dispatch();
         }
 
-        if ($this->forceStop) {
+        if ($this->forceStop || ($this->batchAmount == $this->target && $this->target > 0)) {
             $this->stopConsuming();
         }