From 126d3d92aa2d1dacc0a492beec323029fa3aadfe Mon Sep 17 00:00:00 2001
From: tvs <tvansloten@xebia.com>
Date: Tue, 5 Feb 2019 22:00:22 +0100
Subject: [PATCH] Enable function invocation timeout using Promise.race

---
 bin/cmd.js   |  9 +++++----
 lib/serve.js | 27 ++++++++++++++++++++++-----
 2 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/bin/cmd.js b/bin/cmd.js
index 6cc8889e..91f80c28 100755
--- a/bin/cmd.js
+++ b/bin/cmd.js
@@ -18,7 +18,8 @@ program.version(pkg.version);
 program
   .option("-c --config <webpack-config>", "additional webpack configuration")
   .option("-p --port <port>", "port to serve from (default: 9000)")
-  .option("-s --static", "serve pre-built lambda files")
+  .option("-t --timeout <timeout>", "function invocation timeout in seconds (default: 10)")
+  .option("-s --static", "serve pre-built lambda files");
 
 program
   .command("serve <dir>")
@@ -26,8 +27,8 @@ program
   .action(function(cmd, options) {
     console.log("netlify-lambda: Starting server");
     var static = Boolean(program.static);
-    var server = serve.listen(program.port || 9000, static);
-    if (static) return // early terminate, don't build
+    var server = serve.listen(program.port || 9000, static, Number(program.timeout) || 10);
+    if (static) return; // early terminate, don't build
     build.watch(cmd, program.config, function(err, stats) {
       if (err) {
         console.error(err);
@@ -75,4 +76,4 @@ var NO_COMMAND_SPECIFIED = program.args.length === 0;
 if (NO_COMMAND_SPECIFIED) {
   // user did not supply args, so show --help
   program.help();
-}
\ No newline at end of file
+}
diff --git a/lib/serve.js b/lib/serve.js
index 9d702ca2..bafc543d 100644
--- a/lib/serve.js
+++ b/lib/serve.js
@@ -14,6 +14,18 @@ function handleErr(err, response) {
   return;
 }
 
+function handleInvocationTimeout(response, timeout) {
+  response.statusCode = 500;
+  response.write(`Function invocation took longer than ${timeout} seconds.`);
+  response.end();
+  console.log(
+    `Your lambda function took longer than ${timeout} seconds to finish.
+If you need a longer execution time, you can increase the timeout using the -t or --timeout flag.
+Please note that default function invocation is 10 seconds, check our documentation for more information (https://www.netlify.com/docs/functions/#custom-deployment-options).
+`
+  );
+}
+
 function createCallback(response) {
   return function callback(err, lambdaResponse) {
     if (err) {
@@ -50,7 +62,7 @@ function promiseCallback(promise, callback) {
   if (typeof promise.then !== 'function') return;
   if (typeof callback !== 'function') return;
 
-  promise.then(
+  return promise.then(
     function(data) {
       callback(null, data);
     },
@@ -80,7 +92,7 @@ function buildClientContext(headers) {
   }
 }
 
-function createHandler(dir, static) {
+function createHandler(dir, static, timeout) {
   return function(request, response) {
     // handle proxies without path re-writes (http-servr)
     var cleanPath = request.path.replace(/^\/.netlify\/functions/, '');
@@ -123,11 +135,16 @@ function createHandler(dir, static) {
       { clientContext: buildClientContext(request.headers) || {} },
       callback
     );
-    promiseCallback(promise, callback);
+    Promise.race([
+      promiseCallback(promise, callback),
+      setTimeout(function() {
+        handleInvocationTimeout(response, timeout)
+      }, timeout * 1000)
+    ])
   };
 }
 
-exports.listen = function(port, static) {
+exports.listen = function(port, static, timeout) {
   var config = conf.load();
   var app = express();
   var dir = config.build.functions || config.build.Functions;
@@ -142,7 +159,7 @@ exports.listen = function(port, static) {
   app.get('/favicon.ico', function(req, res) {
     res.status(204).end();
   });
-  app.all('*', createHandler(dir, static));
+  app.all('*', createHandler(dir, static, timeout));
 
   app.listen(port, function(err) {
     if (err) {