feat: timeout optimizations

middyjs · Dec 10, 2023 · f8a3523 · f8a3523
1 parent 8d0a139
commit f8a3523
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 27 deletions.
diff --git a/packages/core/__benchmarks__/index.js b/packages/core/__benchmarks__/index.js
@@ -1,5 +1,6 @@
 import { Bench } from 'tinybench'
 import middy from '../index.js'
+// import middyNext from '../index.next.js'
 
 const bench = new Bench({ time: 1_000 })
 
@@ -45,7 +46,7 @@ middlewaresAsync.fill(middlewareAsync())
 const warmAsyncMiddlewareHandler = middy()
   .use(middlewaresAsync)
   .handler(baseHandler)
-const warmTimeoutHandler = middy({ timeoutEarlyInMillis: 0 }).handler(
+const warmDisableTimeoutHandler = middy({ timeoutEarlyInMillis: 0 }).handler(
   baseHandler
 )
 
@@ -60,40 +61,41 @@ const warmTimeoutHandler = middy({ timeoutEarlyInMillis: 0 }).handler(
 //   baseHandler
 // )
 
+const event = {}
 await bench
-  .add('Cold Invocation', async (event = {}) => {
+  .add('Cold Invocation', async () => {
     const coldHandler = middy().handler(baseHandler)
     await coldHandler(event, context)
   })
-  .add('Cold Invocation with middleware', async (event = {}) => {
+  .add('Cold Invocation with middleware', async () => {
     const middlewares = new Array(25)
     middlewares.fill(middleware())
     const coldHandler = middy().use(middlewares).handler(baseHandler)
     await coldHandler(event, context)
   })
-  .add('Warm Invocation', async (event = {}) => {
+  .add('Warm Invocation', async () => {
     await warmHandler(event, context)
   })
   // .add('Warm Invocation * next', async (event = {}) => {
   //   await warmNextHandler(event, context)
   // })
-  .add('Warm Async Invocation', async (event = {}) => {
+  .add('Warm Async Invocation', async () => {
     await warmAsyncHandler(event, context)
   })
-  .add('Warm Invocation without Timeout', async (event = {}) => {
-    await warmTimeoutHandler(event, context)
+  .add('Warm Invocation with disabled Timeout', async () => {
+    await warmDisableTimeoutHandler(event, context)
   })
-  // .add('Warm Invocation with Timeout * next', async (event = {}) => {
+  // .add('Warm Invocation with disabled Timeout * next', async (event = {}) => {
   //   await warmNextTimeoutHandler(event, context)
   // })
   // TODO StreamifyResponse
-  .add('Warm Invocation with middleware', async (event = {}) => {
+  .add('Warm Invocation with middleware', async () => {
     await warmMiddlewareHandler(event, context)
   })
   // .add('Warm Invocation with middleware * next', async (event = {}) => {
   //   await warmNextMiddlewareHandler(event, context)
   // })
-  .add('Warm Invocation with async middleware', async (event = {}) => {
+  .add('Warm Invocation with async middleware', async () => {
     await warmAsyncMiddlewareHandler(event, context)
   })
 

diff --git a/packages/core/index.js b/packages/core/index.js
@@ -1,7 +1,7 @@
 /* global awslambda */
 import { Readable } from 'node:stream'
 import { pipeline } from 'node:stream/promises'
-import { setTimeout } from 'node:timers/promises'
+import { setTimeout } from 'node:timers'
 
 const defaultLambdaHandler = () => {}
 const defaultPlugin = {
@@ -137,7 +137,7 @@ const runRequest = async (
   onErrorMiddlewares,
   plugin
 ) => {
-  let timeoutAbort
+  let timeoutID
   // context.getRemainingTimeInMillis checked for when AWS context missing (tests, containers)
   const timeoutEarly =
     plugin.timeoutEarly && request.context.getRemainingTimeInMillis
@@ -149,39 +149,45 @@ const runRequest = async (
     if (typeof request.response === 'undefined') {
       plugin.beforeHandler?.()
 
-      // Note: signal.abort is slow ~6000ns
+      // Can't manually abort and timeout with same AbortSignal
+      // https://developer.mozilla.org/en-US/docs/Web/API/AbortSignal/timeout_static
       const handlerAbort = new AbortController()
       const promises = [
         lambdaHandler(request.event, request.context, {
           signal: handlerAbort.signal
         })
       ]
 
-      // Can't manually abort and timeout with same AbortSignal
-      // https://developer.mozilla.org/en-US/docs/Web/API/AbortSignal/timeout_static
+      // clearTimeout pattern is 10x faster than using AbortController
+      // Note: signal.abort is slow ~6000ns
       if (timeoutEarly) {
-        timeoutAbort = new AbortController()
-        promises.push(
-          setTimeout(
-            request.context.getRemainingTimeInMillis() -
-              plugin.timeoutEarlyInMillis,
-            undefined,
-            { signal: timeoutAbort.signal }
-          ).then(() => {
+        let timeoutResolve
+        const timeoutPromise = new Promise((resolve, reject) => {
+          timeoutResolve = () => {
             handlerAbort.abort()
-            return plugin.timeoutEarlyResponse()
-          })
+            try {
+              resolve(plugin.timeoutEarlyResponse())
+            } catch (e) {
+              reject(e)
+            }
+          }
+        })
+        timeoutID = setTimeout(
+          timeoutResolve,
+          request.context.getRemainingTimeInMillis() -
+            plugin.timeoutEarlyInMillis
         )
+        promises.push(timeoutPromise)
       }
       request.response = await Promise.race(promises)
-      timeoutAbort?.abort() // lambdaHandler may not support .then()
+      clearTimeout(timeoutID)
 
       plugin.afterHandler?.()
       await runMiddlewares(request, afterMiddlewares, plugin)
     }
   } catch (e) {
     // timeout should be aborted when errors happen in handler
-    timeoutAbort?.abort()
+    clearTimeout(timeoutID)
 
     // Reset response changes made by after stack before error thrown
     request.response = undefined