Merge pull request #1068 from amfa-team/feature/websocket-timeouts

Support WebSocket timeouts limits & fix connectedAt context attribute

README.md

@@ -128,6 +128,8 @@ All CLI options are optional:
 --useChildProcesses         Run handlers in a child process
 --useWorkerThreads          Uses worker threads for handlers. Requires node.js v11.7.0 or higher
 --websocketPort             WebSocket port to listen on. Default: 3001
+--webSocketHardTimeout      Set WebSocket hard timeout in seconds to reproduce AWS limits (https://docs.aws.amazon.com/apigateway/latest/developerguide/limits.html#apigateway-execution-service-websocket-limits-table). Default: 7200 (2 hours)
+--webSocketIdleTimeout      Set WebSocket idle timeout in seconds to reproduce AWS limits (https://docs.aws.amazon.com/apigateway/latest/developerguide/limits.html#apigateway-execution-service-websocket-limits-table). Default: 600 (10 minutes)
 --useDocker                 Run handlers in a docker container.
 --layersDir                 The directory layers should be stored in. Default: ${codeDir}/.serverless-offline/layers'
 --dockerReadOnly            Marks if the docker code layer should be read only. Default: true

src/config/commandOptions.js

@@ -78,6 +78,14 @@ export default {
   websocketPort: {
     usage: 'Websocket port to listen on. Default: 3001',
   },
+  webSocketHardTimeout: {
+    usage:
+      'Set WebSocket hard timeout in seconds to reproduce AWS limits (https://docs.aws.amazon.com/apigateway/latest/developerguide/limits.html#apigateway-execution-service-websocket-limits-table). Default: 7200 (2 hours)',
+  },
+  webSocketIdleTimeout: {
+    usage:
+      'Set WebSocket idle timeout in seconds to reproduce AWS limits (https://docs.aws.amazon.com/apigateway/latest/developerguide/limits.html#apigateway-execution-service-websocket-limits-table). Default: 600 (10 minutes)',
+  },
   useDocker: {
     usage: 'Uses docker for node/python/ruby/provided',
   },

src/config/defaultOptions.js

@@ -22,6 +22,8 @@ export default {
   useChildProcesses: false,
   useWorkerThreads: false,
   websocketPort: 3001,
+  webSocketHardTimeout: 7200,
+  webSocketIdleTimeout: 600,
   useDocker: false,
   layersDir: null,
   dockerReadOnly: true,

src/events/websocket/WebSocketClients.js

@@ -20,6 +20,8 @@ export default class WebSocketClients {
   #options = null
   #webSocketRoutes = new Map()
   #websocketsApiRouteSelectionExpression = null
+  #idleTimeouts = new WeakMap()
+  #hardTimeouts = new WeakMap()
 
   constructor(serverless, options, lambda) {
     this.#lambda = lambda
@@ -32,6 +34,8 @@ export default class WebSocketClients {
   _addWebSocketClient(client, connectionId) {
     this.#clients.set(client, connectionId)
     this.#clients.set(connectionId, client)
+    this._onWebSocketUsed(connectionId)
+    this._addHardTimeout(client, connectionId)
   }
 
   _removeWebSocketClient(client) {
@@ -47,6 +51,36 @@ export default class WebSocketClients {
     return this.#clients.get(connectionId)
   }
 
+  _addHardTimeout(client, connectionId) {
+    const timeoutId = setTimeout(() => {
+      debugLog(`timeout:hard:${connectionId}`)
+      client.close(1001, 'Going away')
+    }, this.#options.webSocketHardTimeout * 1000)
+    this.#hardTimeouts.set(client, timeoutId)
+  }
+
+  _clearHardTimeout(client) {
+    const timeoutId = this.#hardTimeouts.get(client)
+    clearTimeout(timeoutId)
+  }
+
+  _onWebSocketUsed(connectionId) {
+    const client = this._getWebSocketClient(connectionId)
+    this._clearIdleTimeout(client)
+    debugLog(`timeout:idle:${connectionId}:reset`)
+
+    const timeoutId = setTimeout(() => {
+      debugLog(`timeout:idle:${connectionId}:trigger`)
+      client.close(1001, 'Going away')
+    }, this.#options.webSocketIdleTimeout * 1000)
+    this.#idleTimeouts.set(client, timeoutId)
+  }
+
+  _clearIdleTimeout(client) {
+    const timeoutId = this.#idleTimeouts.get(client)
+    clearTimeout(timeoutId)
+  }
+
   async _processEvent(websocketClient, connectionId, route, event) {
     let functionKey = this.#webSocketRoutes.get(route)
 
@@ -136,6 +170,9 @@ export default class WebSocketClients {
         connectionId,
       ).create()
 
+      this._clearHardTimeout(webSocketClient)
+      this._clearIdleTimeout(webSocketClient)
+
       this._processEvent(
         webSocketClient,
         connectionId,
@@ -152,6 +189,7 @@ export default class WebSocketClients {
       debugLog(`route:${route} on connection=${connectionId}`)
 
       const event = new WebSocketEvent(connectionId, route, message).create()
+      this._onWebSocketUsed(connectionId)
 
       this._processEvent(webSocketClient, connectionId, route, event)
     })
@@ -179,6 +217,7 @@ export default class WebSocketClients {
     const client = this._getWebSocketClient(connectionId)
 
     if (client) {
+      this._onWebSocketUsed(connectionId)
       client.send(payload)
       return true
     }

src/events/websocket/lambda-events/WebSocketRequestContext.js

@@ -2,23 +2,36 @@ import { createUniqueId, formatToClfTime } from '../../../utils/index.js'
 
 const { now } = Date
 
+const connectedAt = new Map()
+
 export default class WebSocketRequestContext {
   #connectionId = null
   #eventType = null
   #route = null
+  #connectedAt = null
 
   constructor(eventType, route, connectionId) {
     this.#connectionId = connectionId
     this.#eventType = eventType
     this.#route = route
+
+    if (eventType === 'CONNECT') {
+      connectedAt.set(connectionId, now())
+    }
+
+    this.#connectedAt = connectedAt.get(connectionId)
+
+    if (eventType === 'DISCONNECT') {
+      connectedAt.delete(connectionId)
+    }
   }
 
   create() {
     const timeEpoch = now()
 
     const requestContext = {
       apiId: 'private',
-      connectedAt: now(), // TODO this is probably not correct, and should be the initial connection time?
+      connectedAt: this.#connectedAt,
       connectionId: this.#connectionId,
       domainName: 'localhost',
       eventType: this.#eventType,