Handle timeouts more gracefully by throwing exception near the hard end

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::timoutAfter($this->timeout);
+        } elseif ($this->timeout === 0 && 0 < $context->getRemainingTimeInMillis()) {
+            // Throw exception one second before Lambda pulls the plug.
+            Timeout::timoutAfter(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::timoutAfter($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::timoutAfter((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 timoutAfter(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=10');
+        $this->assertEquals(10, $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!']);

tests/Timeout/TimeoutTest.php

@@ -0,0 +1,99 @@
+<?php declare(strict_types=1);
+
+namespace Timeout;
+
+use Bref\Timeout\LambdaTimeout;
+use Bref\Timeout\Timeout;
+use PHPUnit\Framework\TestCase;
+
+class TimeoutTest extends TestCase
+{
+    public static function setUpBeforeClass(): void
+    {
+        if (! function_exists('pcntl_async_signals')) {
+            self::markTestSkipped('PCNTL extension is not enabled.');
+        }
+    }
+
+    protected function setUp(): void
+    {
+        parent::setUp();
+        unset($_SERVER['LAMBDA_INVOCATION_CONTEXT']);
+        unset($_SERVER['BREF_TIMEOUT']);
+    }
+
+    protected function tearDown(): void
+    {
+        Timeout::reset();
+        parent::tearDown();
+    }
+
+    public function testEnableWithoutContext()
+    {
+        $this->expectException(\LogicException::class);
+        Timeout::enable();
+    }
+
+    public function testEnableWithBrefTimeoutInactive()
+    {
+        $_SERVER['BREF_TIMEOUT'] = -1;
+        $_SERVER['LAMBDA_INVOCATION_CONTEXT'] = json_encode(['deadlineMs' => (time() + 30) * 1000]);
+
+        Timeout::enable();
+        $timeout = pcntl_alarm(0);
+        $this->assertSame(0, $timeout, 'Timeout should not be active when BREF_TIMEOUT=-1');
+    }
+
+    public function testEnableWithBrefTimeout()
+    {
+        $_SERVER['BREF_TIMEOUT'] = 10;
+
+        Timeout::enable();
+        $timeout = pcntl_alarm(0);
+        $this->assertSame(10, $timeout, 'BREF_TIMEOUT=10 should have effect');
+    }
+
+    public function testEnableWithBrefTimeoutAndContext()
+    {
+        $_SERVER['BREF_TIMEOUT'] = 10;
+        $_SERVER['LAMBDA_INVOCATION_CONTEXT'] = json_encode(['deadlineMs' => (time() + 30) * 1000]);
+
+        Timeout::enable();
+        $timeout = pcntl_alarm(0);
+        $this->assertSame(10, $timeout, 'BREF_TIMEOUT=10 should have effect over context');
+    }
+
+    public function testEnableWithBrefTimeoutZeroAndContext()
+    {
+        $_SERVER['BREF_TIMEOUT'] = 0;
+        $_SERVER['LAMBDA_INVOCATION_CONTEXT'] = json_encode(['deadlineMs' => (time() + 30) * 1000]);
+
+        Timeout::enable();
+        $timeout = pcntl_alarm(0);
+        $this->assertEqualsWithDelta(30, $timeout, 1, 'BREF_TIMEOUT=0 should fallback to context');
+    }
+
+    public function testEnableWithContext()
+    {
+        $_SERVER['LAMBDA_INVOCATION_CONTEXT'] = json_encode(['deadlineMs' => (time() + 30) * 1000]);
+
+        Timeout::enable();
+        $timeout = pcntl_alarm(0);
+        $this->assertEqualsWithDelta(30, $timeout, 1);
+    }
+
+    public function testTimeoutAfter()
+    {
+        $start = microtime(true);
+        Timeout::timoutAfter(2);
+        try {
+            sleep(4);
+            $this->fail('We expect a LambdaTimeout before we reach this line');
+        } catch (LambdaTimeout $e) {
+            $time = 1000 * (microtime(true) - $start);
+            $this->assertEqualsWithDelta(2000, $time, 200, 'We must wait about 2 seconds');
+        } catch (\Throwable $e) {
+            $this->fail('It must throw a LambdaTimeout.');
+        }
+    }
+}