Handle timeouts more gracefully by allowing the application to shutdown

couscous.yml

@@ -114,6 +114,10 @@ menu:
                     text: Databases
                     url: /docs/environment/database.html
                     title: Using a database from AWS Lambda
+                timeouts:
+                    text: Timeouts
+                    url: /docs/environment/timeouts.html
+                    title: Configure and handle timeouts
                 custom-domains:
                     text: Custom domains
                     url: /docs/environment/custom-domains.html

docs/environment/database.md

@@ -54,6 +54,8 @@ When possible, an alternative to NAT Gateways is to split the work done by a lam
 
 Finally, another free alternative to NAT Gateway is to access AWS services by creating "*private VPC endpoints*": this is possible for S3, API Gateway, [and more](https://docs.aws.amazon.com/en_pv/vpc/latest/userguide/vpc-endpoints-access.html).
 
+Read more in the section about [timeouts](/docs/environment/timeouts.md).
+
 ## Creating a database
 
 On the [RDS console](https://console.aws.amazon.com/rds/home):

docs/environment/timeouts.md

@@ -0,0 +1,65 @@
+---
+title: Timeouts
+current_menu: timeouts
+introduction: Configure and handle timeouts.
+---
+
+When a Lambda function times out, it is like the power to the computer is suddenly
+just turned off. This does not give the application a chance to shutdown properly.
+This often leaves you without any logs and the problem could be hard to fix.
+
+Bref will throw an `LambdaTimeout` exception just before the Lambda actually times
+out. This will allow your application to actually shutdown.
+
+This feature is enabled automatically for the `php-xx` layer and the `console` layer.
+The `php-xx-fpm` layer needs to opt-in by adding the following to `index.php`.
+
+```php
+if (isset($_SERVER['LAMBDA_TASK_ROOT'])) {
+    \Bref\Timeout\Timeout::enable();
+}
+```
+
+## Configuration
+
+You may configure this behavior with the `BREF_TIMEOUT` environment variable. To
+always trigger an exception after 10 seconds, set `BREF_TIMEOUT=10`. To disable
+Bref throwing an exception use value `BREF_TIMEOUT=-1`. To automatically set the
+timeout just a hair shorter than the Lambda timeout, use `BREF_TIMEOUT=0`.
+
+## Catching the exception
+
+If you are using a framework, then the framework is probably catching all exceptions
+and displays an error page for the users. You may of course catch the exception
+yourself:
+
+```php
+<?php
+
+require dirname(__DIR__) . '/vendor/autoload.php';
+
+use Bref\Context\Context;
+use Bref\Timeout\LambdaTimeout;
+
+class Handler implements \Bref\Event\Handler
+{
+    public function handle($event, Context $context)
+    {
+        try {
+            $this->generateResponse();
+        } catch (LambdaTimeout $e) {
+            echo 'Oops, sorry. We spent too much time on this.';
+        } catch (\Throwable $e) {
+            echo 'Some unexpected error happened.';
+        }
+    }
+
+    private function generateResponse()
+    {
+        $pi = // ...
+        echo 'Pi is '.$pi;
+    }
+}
+
+return new Handler();
+```

runtime/layers/fpm/bootstrap

@@ -25,7 +25,8 @@ if (getenv('BREF_DOWNLOAD_VENDOR')) {
     require $appRoot . '/vendor/autoload.php';
 }
 
-$lambdaRuntime = LambdaRuntime::fromEnvironmentVariable();
+// Get a LambdaRuntime and disable timeout exceptions.
+$lambdaRuntime = LambdaRuntime::fromEnvironmentVariable(-1);
 
 $handlerFile = $appRoot . '/' . getenv('_HANDLER');
 if (! is_file($handlerFile)) {

src/Runtime/LambdaRuntime.php

@@ -5,6 +5,7 @@
 use Bref\Context\Context;
 use Bref\Context\ContextBuilder;
 use Bref\Event\Handler;
+use Bref\Timeout\Timeout;
 use Exception;
 use Psr\Http\Server\RequestHandlerInterface;
 
@@ -42,19 +43,33 @@ final class LambdaRuntime
     /** @var Invoker */
     private $invoker;
 
-    public static function fromEnvironmentVariable(): self
+    /** @var int seconds */
+    private $timeout;
+
+    public static function fromEnvironmentVariable(?int $timeout = null): self
     {
-        return new self((string) getenv('AWS_LAMBDA_RUNTIME_API'));
+        return new self((string) getenv('AWS_LAMBDA_RUNTIME_API'), $timeout ?? (int) getenv('BREF_TIMEOUT'));
     }
 
-    public function __construct(string $apiUrl)
+    /**
+     * @param int $timeout number of seconds before a TimeoutException is thrown.
+     *                     Value -1 means "disabled". Value 0 means "auto", this will
+     *                     set the timeout just a bit shorter than the Lambda timeout.
+     */
+    public function __construct(string $apiUrl, int $timeout = 0)
     {
         if ($apiUrl === '') {
             die('At the moment lambdas can only be executed in an Lambda environment');
         }
 
         $this->apiUrl = $apiUrl;
         $this->invoker = new Invoker;
+        $this->timeout = $timeout;
+
+        if ($timeout >= 0 && ! Timeout::init()) {
+            // If we fail to initialize
+            $this->timeout = -1;
+        }
     }
 
     public function __destruct()
@@ -96,6 +111,13 @@ public function processNextEvent($handler): void
         [$event, $context] = $this->waitNextInvocation();
         \assert($context instanceof Context);
 
+        if ($this->timeout > 0) {
+            Timeout::timeoutAfter($this->timeout);
+        } elseif ($this->timeout === 0 && 0 < $context->getRemainingTimeInMillis()) {
+            // Throw exception one second before Lambda pulls the plug.
+            Timeout::timeoutAfter(max(1, (int) floor($context->getRemainingTimeInMillis() / 1000) - 1));
+        }
+
         $this->ping();
 
         try {
@@ -104,6 +126,8 @@ public function processNextEvent($handler): void
             $this->sendResponse($context->getAwsRequestId(), $result);
         } catch (\Throwable $e) {
             $this->signalFailure($context->getAwsRequestId(), $e);
+        } finally {
+            Timeout::reset();
         }
     }
 

src/Timeout/LambdaTimeout.php

@@ -0,0 +1,12 @@
+<?php declare(strict_types=1);
+
+namespace Bref\Timeout;
+
+/**
+ * The application took too long to produce a response. This exception is thrown
+ * to give the application a chance to flush logs and shut it self down before
+ * the power to AWS Lambda is disconnected.
+ */
+class LambdaTimeout extends \RuntimeException
+{
+}

src/Timeout/Timeout.php

@@ -0,0 +1,92 @@
+<?php declare(strict_types=1);
+
+namespace Bref\Timeout;
+
+/**
+ * Helper class to trigger an exception just before the Lamba times out. This
+ * will give the application a chance to shut down.
+ */
+final class Timeout
+{
+    /** @var bool */
+    private static $initialized = false;
+
+    /**
+     * Read environment variables and setup timeout exception.
+     */
+    public static function enable(): void
+    {
+        if (isset($_SERVER['BREF_TIMEOUT'])) {
+            $timeout = (int) $_SERVER['BREF_TIMEOUT'];
+            if ($timeout === -1) {
+                return;
+            }
+
+            if ($timeout > 0) {
+                self::timeoutAfter($timeout);
+
+                return;
+            }
+
+            // else if 0, continue
+        }
+
+        if (isset($_SERVER['LAMBDA_INVOCATION_CONTEXT'])) {
+            $context = json_decode($_SERVER['LAMBDA_INVOCATION_CONTEXT'], true, 512, JSON_THROW_ON_ERROR);
+            $deadlineMs = $context['deadlineMs'];
+            $remainingTime = $deadlineMs - intval(microtime(true) * 1000);
+
+            self::timeoutAfter((int) floor($remainingTime / 1000));
+
+            return;
+        }
+
+        throw new \LogicException('Could not find value for bref timeout. Are we running on Lambda?');
+    }
+
+    /**
+     * Setup custom handler for SIGTERM. One need to call Timeout::timoutAfter()
+     * to make an exception to be thrown.
+     *
+     * @return bool true if successful.
+     */
+    public static function init(): bool
+    {
+        if (self::$initialized) {
+            return true;
+        }
+
+        if (! function_exists('pcntl_async_signals')) {
+            trigger_error('Could not enable timeout exceptions because pcntl extension is not enabled.');
+            return false;
+        }
+
+        pcntl_async_signals(true);
+        pcntl_signal(SIGALRM, function (): void {
+            throw new LambdaTimeout('Maximum AWS Lambda execution time reached');
+        });
+
+        self::$initialized = true;
+
+        return true;
+    }
+
+    /**
+     * Set a timer to throw an exception.
+     */
+    public static function timeoutAfter(int $seconds): void
+    {
+        self::init();
+        pcntl_alarm($seconds);
+    }
+
+    /**
+     * Reset timeout.
+     */
+    public static function reset(): void
+    {
+        if (self::$initialized) {
+            pcntl_alarm(0);
+        }
+    }
+}

tests/Runtime/LambdaRuntimeTest.php

@@ -35,7 +35,7 @@ protected function setUp(): void
     {
         ob_start();
         Server::start();
-        $this->runtime = new LambdaRuntime('localhost:8126');
+        $this->runtime = new LambdaRuntime('localhost:8126', -1);
     }
 
     protected function tearDown(): void
@@ -44,6 +44,29 @@ protected function tearDown(): void
         ob_end_clean();
     }
 
+    public function testFromEnvironmentVariable()
+    {
+        $getTimeout = function ($runtime) {
+            $reflectionProp = (new \ReflectionObject($runtime))->getProperty('timeout');
+            $reflectionProp->setAccessible(true);
+
+            return $reflectionProp->getValue($runtime);
+        };
+
+        putenv('AWS_LAMBDA_RUNTIME_API=foo');
+        putenv('BREF_TIMEOUT'); // unset
+        $this->assertEquals(0, $getTimeout(LambdaRuntime::fromEnvironmentVariable()));
+        $this->assertEquals(-1, $getTimeout(LambdaRuntime::fromEnvironmentVariable(-1)));
+        $this->assertEquals(0, $getTimeout(LambdaRuntime::fromEnvironmentVariable(0)));
+        $this->assertEquals(10, $getTimeout(LambdaRuntime::fromEnvironmentVariable(10)));
+
+        putenv('BREF_TIMEOUT=5');
+        $this->assertEquals(5, $getTimeout(LambdaRuntime::fromEnvironmentVariable()));
+        $this->assertEquals(-1, $getTimeout(LambdaRuntime::fromEnvironmentVariable(-1)));
+        $this->assertEquals(0, $getTimeout(LambdaRuntime::fromEnvironmentVariable(0)));
+        $this->assertEquals(10, $getTimeout(LambdaRuntime::fromEnvironmentVariable(10)));
+    }
+
     public function test basic behavior()
     {
         $this->givenAnEvent(['Hello' => 'world!']);