Add graceful max execution timeout.

d-ph · d-ph · commit 2037c2a70b9c · 2017-01-22T18:01:50.000Z
Allows for running consumer up to certain time, which guarantees, that
resources are cleaned up periodically and unobtrusively.
diff --git a/CHANGELOG b/CHANGELOG
@@ -1,3 +1,6 @@
+- 2017-01-22
+    * Add `graceful_max_execution_timeout`
+
 - 2015-02-07
     * Added possibility to set serialize/unserialize function for rpc servers/clients
 
diff --git a/DependencyInjection/Configuration.php b/DependencyInjection/Configuration.php
@@ -138,6 +138,8 @@ protected function addConsumers(ArrayNodeDefinition $node)
                             ->scalarNode('callback')->isRequired()->end()
                             ->scalarNode('idle_timeout')->end()
                             ->scalarNode('idle_timeout_exit_code')->end()
+                            ->scalarNode('graceful_max_execution_timeout')->end()
+                            ->scalarNode('graceful_max_execution_timeout_exit_code')->end()
                             ->scalarNode('auto_setup_fabric')->defaultTrue()->end()
                             ->arrayNode('qos_options')
                                 ->canBeUnset()
diff --git a/DependencyInjection/OldSoundRabbitMqExtension.php b/DependencyInjection/OldSoundRabbitMqExtension.php
@@ -200,6 +200,18 @@ protected function loadConsumers()
             if (isset($consumer['idle_timeout_exit_code'])) {
                 $definition->addMethodCall('setIdleTimeoutExitCode', array($consumer['idle_timeout_exit_code']));
             }
+            if (isset($consumer['graceful_max_execution_timeout'])) {
+                $definition->addMethodCall(
+                    'setGracefulMaxExecutionDateTimeFromSecondsInTheFuture',
+                    array($consumer['graceful_max_execution_timeout'])
+                );
+            }
+            if (isset($consumer['graceful_max_execution_timeout_exit_code'])) {
+                $definition->addMethodCall(
+                    'setGracefulMaxExecutionTimeoutExitCode',
+                    array($consumer['graceful_max_execution_timeout_exit_code'])
+                );
+            }
             if (!$consumer['auto_setup_fabric']) {
                 $definition->addMethodCall('disableAutoSetupFabric');
             }
diff --git a/README.md b/README.md
@@ -393,6 +393,29 @@ consumers:
         idle_timeout_exit_code: 0
 ```
 
+#### Graceful max execution timeout ####
+
+If you'd like your consumer to be running up to certain time and then gracefully exit, then set the `graceful_max_execution_timeout` in seconds.
+"Gracefully exit" means, that the consumer will exit either after the currently running task or immediatelly, when waiting for new tasks.
+The `graceful_max_execution_timeout_exit_code` specifies what exit code should be returned by the consumer when the graceful max execution timeout occurs. Without specifying it, the consumer will exit with status `0`.
+
+This feature is great in conjuction with supervisord, which together can allow for periodical memory leaks cleanup, connection with database/rabbitmq renewal and more.
+
+```yaml
+consumers:
+    upload_picture:
+        connection:             default
+        exchange_options:       {name: 'upload-picture', type: direct}
+        queue_options:          {name: 'upload-picture'}
+        callback:               upload_picture_service
+        
+        # 30 minutes
+        graceful_max_execution_timeout: 1800
+         
+        # default is 0
+        graceful_max_execution_timeout_exit_code: 10 
+```
+
 #### Fair dispatching ####
 
 > You might have noticed that the dispatching still doesn't work exactly as we want. For example in a situation with two workers, when all odd messages are heavy and even messages are light, one worker will be constantly busy and the other one will do hardly any work. Well, RabbitMQ doesn't know anything about that and will still dispatch messages evenly.
diff --git a/RabbitMq/Consumer.php b/RabbitMq/Consumer.php
@@ -13,11 +13,25 @@
 
 class Consumer extends BaseConsumer
 {
+    const TIMEOUT_TYPE_IDLE = 'idle';
+    const TIMEOUT_TYPE_GRACEFUL_MAX_EXECUTION = 'graceful-max-execution';
+
     /**
      * @var int $memoryLimit
      */
     protected $memoryLimit = null;
 
+    /**
+     * @var \DateTime|null DateTime after which the consumer will gracefully exit. "Gracefully" means, that
+     *      any currently running consumption will not be interrupted.
+     */
+    protected $gracefulMaxExecutionDateTime;
+
+    /**
+     * @var int Exit code used, when consumer is closed by the Graceful Max Execution Timeout feature.
+     */
+    protected $gracefulMaxExecutionTimeoutExitCode = 0;
+
     /**
      * Set the memory limit
      *
@@ -42,6 +56,7 @@ public function getMemoryLimit()
      * Consume the message
      *
      * @param int $msgAmount
+     * @return int
      */
     public function consume($msgAmount)
     {
@@ -52,10 +67,27 @@ public function consume($msgAmount)
         while (count($this->getChannel()->callbacks)) {
             $this->dispatchEvent(OnConsumeEvent::NAME, new OnConsumeEvent($this));
             $this->maybeStopConsumer();
+
+            /*
+             * Be careful not to trigger ::wait() with 0 or less seconds, when
+             * graceful max execution timeout is being used.
+             */
+            $waitTimeout = $this->chooseWaitTimeout();
+            if (
+                $waitTimeout['timeoutType'] === self::TIMEOUT_TYPE_GRACEFUL_MAX_EXECUTION
+                && $waitTimeout['seconds'] < 1
+            ) {
+                return $this->getGracefulMaxExecutionTimeoutExitCode();
+            }
+
             if (!$this->forceStop) {
                 try {
-                    $this->getChannel()->wait(null, false, $this->getIdleTimeout());
+                    $this->getChannel()->wait(null, false, $waitTimeout['seconds']);
                 } catch (AMQPTimeoutException $e) {
+                    if (self::TIMEOUT_TYPE_GRACEFUL_MAX_EXECUTION === $waitTimeout['timeoutType']) {
+                        return $this->getGracefulMaxExecutionTimeoutExitCode();
+                    }
+
                     $idleEvent = new OnIdleEvent($this);
                     $this->dispatchEvent(OnIdleEvent::NAME, $idleEvent);
 
@@ -69,6 +101,8 @@ public function consume($msgAmount)
                 }
             }
         }
+
+        return 0;
     }
 
     /**
@@ -176,4 +210,92 @@ protected function isRamAlmostOverloaded()
 
         return $memoryManager->isRamAlmostOverloaded($this->getMemoryLimit(), '5M');
     }
+
+    /**
+     * @param \DateTime|null $dateTime
+     */
+    public function setGracefulMaxExecutionDateTime($dateTime = null)
+    {
+        $this->gracefulMaxExecutionDateTime = $dateTime;
+    }
+
+    /**
+     * @param int $secondsInTheFuture
+     */
+    public function setGracefulMaxExecutionDateTimeFromSecondsInTheFuture($secondsInTheFuture)
+    {
+        $this->setGracefulMaxExecutionDateTime(new \DateTime("+{$secondsInTheFuture} seconds"));
+    }
+
+    /**
+     * @param int $exitCode
+     */
+    public function setGracefulMaxExecutionTimeoutExitCode($exitCode)
+    {
+        $this->gracefulMaxExecutionTimeoutExitCode = $exitCode;
+    }
+
+    /**
+     * @return \DateTime|null
+     */
+    public function getGracefulMaxExecutionDateTime()
+    {
+        return $this->gracefulMaxExecutionDateTime;
+    }
+
+    /**
+     * @return int
+     */
+    public function getGracefulMaxExecutionTimeoutExitCode()
+    {
+        return $this->gracefulMaxExecutionTimeoutExitCode;
+    }
+
+    /**
+     * Choose the timeout to use for the $this->getChannel()->wait() method.
+     *
+     * @return array Of structure
+     *  {
+     *      timeoutType: string; // one of self::TIMEOUT_TYPE_*
+     *      seconds: int;
+     *  }
+     */
+    private function chooseWaitTimeout()
+    {
+        if ($this->gracefulMaxExecutionDateTime) {
+            $allowedExecutionDateInterval = $this->gracefulMaxExecutionDateTime->diff(new \DateTime());
+            $allowedExecutionSeconds =  $allowedExecutionDateInterval->days * 86400
+                + $allowedExecutionDateInterval->h * 3600
+                + $allowedExecutionDateInterval->i * 60
+                + $allowedExecutionDateInterval->s;
+
+            if (!$allowedExecutionDateInterval->invert) {
+                $allowedExecutionSeconds *= -1;
+            }
+
+            /*
+             * Respect the idle timeout if it's set and if it's less than
+             * the remaining allowed execution.
+             */
+            if (
+                $this->getIdleTimeout()
+                && $this->getIdleTimeout() < $allowedExecutionSeconds
+            ) {
+                return array(
+                    'timeoutType' => self::TIMEOUT_TYPE_IDLE,
+                    'seconds' => $this->getIdleTimeout(),
+                );
+            }
+
+            return array(
+                'timeoutType' => self::TIMEOUT_TYPE_GRACEFUL_MAX_EXECUTION,
+                'seconds' => $allowedExecutionSeconds,
+            );
+        }
+
+        return array(
+            'timeoutType' => self::TIMEOUT_TYPE_IDLE,
+            'seconds' => $this->getIdleTimeout(),
+        );
+    }
 }
diff --git a/Tests/RabbitMq/ConsumerTest.php b/Tests/RabbitMq/ConsumerTest.php
@@ -255,4 +255,66 @@ public function testShouldAllowContinueConsumptionAfterIdleTimeout()
         $this->setExpectedException('PhpAmqpLib\Exception\AMQPTimeoutException');
         $consumer->consume(10);
     }
+
+    public function testGracefulMaxExecutionTimeoutExitCode()
+    {
+        // set up amqp connection
+        $amqpConnection = $this->prepareAMQPConnection();
+        // set up amqp channel
+        $amqpChannel = $this->prepareAMQPChannel();
+        $amqpChannel->expects($this->atLeastOnce())
+            ->method('getChannelId')
+            ->with()
+            ->willReturn(true);
+        $amqpChannel->expects($this->once())
+            ->method('basic_consume')
+            ->withAnyParameters()
+            ->willReturn(true);
+
+        // set up consumer
+        $consumer = $this->getConsumer($amqpConnection, $amqpChannel);
+        // disable autosetup fabric so we do not mock more objects
+        $consumer->disableAutoSetupFabric();
+        $consumer->setChannel($amqpChannel);
+
+        $consumer->setGracefulMaxExecutionDateTimeFromSecondsInTheFuture(60);
+        $consumer->setGracefulMaxExecutionTimeoutExitCode(10);
+        $amqpChannel->callbacks = array('graceful_max_execution_timeout_test');
+
+        $amqpChannel->expects($this->exactly(1))
+            ->method('wait')
+            ->willThrowException(new AMQPTimeoutException());
+
+        $this->assertTrue(10 == $consumer->consume(1));
+    }
+
+    public function testGracefulMaxExecutionWontWaitIfPastTheTimeout()
+    {
+        // set up amqp connection
+        $amqpConnection = $this->prepareAMQPConnection();
+        // set up amqp channel
+        $amqpChannel = $this->prepareAMQPChannel();
+        $amqpChannel->expects($this->atLeastOnce())
+            ->method('getChannelId')
+            ->with()
+            ->willReturn(true);
+        $amqpChannel->expects($this->once())
+            ->method('basic_consume')
+            ->withAnyParameters()
+            ->willReturn(true);
+
+        // set up consumer
+        $consumer = $this->getConsumer($amqpConnection, $amqpChannel);
+        // disable autosetup fabric so we do not mock more objects
+        $consumer->disableAutoSetupFabric();
+        $consumer->setChannel($amqpChannel);
+
+        $consumer->setGracefulMaxExecutionDateTimeFromSecondsInTheFuture(0);
+        $amqpChannel->callbacks = array('graceful_max_execution_timeout_test');
+
+        $amqpChannel->expects($this->never())
+            ->method('wait');
+
+        $consumer->consume(1);
+    }
 }
