feat: add onReconnect hook to wsReconnect

Author: simone-sanfratello

README.md

@@ -229,7 +229,13 @@ The default implementation forwards the `cookie` header.
 
 ## `wsReconnect`
 
-The `wsReconnect` option contains the configuration for the WebSocket reconnection feature; is an object with the following properties:
+**Experimental.** (default: `disabled`)
+
+Reconnection feature detects and closes broken connections and reconnects automatically, see [how to detect and close broken connections](https://github.com/websockets/ws#how-to-detect-and-close-broken-connections).
+The connection is considered broken if the target does not respond to the ping messages or no data is received from the target.
+
+The `wsReconnect` option contains the configuration for the WebSocket reconnection feature.
+To enable the feature, set the `wsReconnect` option to an object with the following properties:
 
 - `pingInterval`: The interval between ping messages in ms (default: `30_000`).
 - `maxReconnectionRetries`: The maximum number of reconnection retries (`1` to `Infinity`, default: `Infinity`).
@@ -238,9 +244,7 @@ The `wsReconnect` option contains the configuration for the WebSocket reconnecti
 - `connectionTimeout`: The timeout for establishing the connection in ms (default: `5_000`).
 - `reconnectOnClose`: Whether to reconnect on close, as long as the connection from the related client to the proxy is active (default: `false`).
 - `logs`: Whether to log the reconnection process (default: `false`).
-
-Reconnection feature detects and closes broken connections and reconnects automatically, see [how to detect and close broken connections](https://github.com/websockets/ws#how-to-detect-and-close-broken-connections).
-The connection is considered broken if the target does not respond to the ping messages or no data is received from the target.
+- `onReconnect`: A hook function that is called when the connection is reconnected `async onReconnect(oldSocket, newSocket)` (default: `undefined`).
 
 ## Benchmarks
 

index.js

@@ -112,7 +112,7 @@ function proxyWebSockets (source, target) {
   /* c8 ignore stop */
 }
 
-async function reconnect (logger, source, wsReconnectOptions, targetParams) {
+async function reconnect (logger, source, wsReconnectOptions, oldTarget, targetParams) {
   const { url, subprotocols, optionsWs } = targetParams
 
   let attempts = 0
@@ -138,21 +138,24 @@ async function reconnect (logger, source, wsReconnectOptions, targetParams) {
   }
 
   wsReconnectOptions.logs && logger.info({ target: targetParams.url, attempts }, 'proxy ws reconnected')
+  wsReconnectOptions.onReconnect(oldTarget, target)
   proxyWebSocketsWithReconnection(logger, source, target, wsReconnectOptions, targetParams)
 }
 
-function proxyWebSocketsWithReconnection (logger, source, target, options, targetParams, fromReconnection = false) {
+function proxyWebSocketsWithReconnection (logger, source, target, options, targetParams) {
   function close (code, reason) {
     target.pingTimer && clearTimeout(source.pingTimer)
     target.pingTimer = undefined
 
     // reconnect target as long as the source connection is active
     if (source.isAlive && (target.broken || options.reconnectOnClose)) {
+      // clean up the target and related source listeners
       target.isAlive = false
       target.removeAllListeners()
       // need to specify the listeners to remove
       removeSourceListeners(source)
-      reconnect(logger, source, options, targetParams)
+
+      reconnect(logger, source, options, target, targetParams)
       return
     }
 

src/options.js

@@ -7,6 +7,9 @@ const DEFAULT_RECONNECT_DECAY = 1.5
 const DEFAULT_CONNECTION_TIMEOUT = 5_000
 const DEFAULT_RECONNECT_ON_CLOSE = false
 const DEFAULT_LOGS = false
+const DEFAULT_ON_RECONNECT = noop
+
+function noop () {}
 
 function validateOptions (options) {
   if (!options.upstream && !options.websocket && !((options.upstream === '' || options.wsUpstream === '') && options.replyOptions && typeof options.replyOptions.getUpstream === 'function')) {
@@ -50,6 +53,11 @@ function validateOptions (options) {
       throw new Error('wsReconnect.logs must be a boolean')
     }
     wsReconnect.logs = wsReconnect.logs ?? DEFAULT_LOGS
+
+    if (wsReconnect.onReconnect !== undefined && typeof wsReconnect.onReconnect !== 'function') {
+      throw new Error('wsReconnect.onReconnect must be a function')
+    }
+    wsReconnect.onReconnect = wsReconnect.onReconnect ?? DEFAULT_ON_RECONNECT
   }
 
   return options
@@ -63,5 +71,6 @@ module.exports = {
   DEFAULT_RECONNECT_DECAY,
   DEFAULT_CONNECTION_TIMEOUT,
   DEFAULT_RECONNECT_ON_CLOSE,
-  DEFAULT_LOGS
+  DEFAULT_LOGS,
+  DEFAULT_ON_RECONNECT
 }

test/options.js

@@ -3,7 +3,7 @@
 const { test } = require('node:test')
 const assert = require('node:assert')
 const { validateOptions } = require('../src/options')
-const { DEFAULT_PING_INTERVAL, DEFAULT_MAX_RECONNECTION_RETRIES, DEFAULT_RECONNECT_INTERVAL, DEFAULT_RECONNECT_DECAY, DEFAULT_CONNECTION_TIMEOUT, DEFAULT_RECONNECT_ON_CLOSE, DEFAULT_LOGS } = require('../src/options')
+const { DEFAULT_PING_INTERVAL, DEFAULT_MAX_RECONNECTION_RETRIES, DEFAULT_RECONNECT_INTERVAL, DEFAULT_RECONNECT_DECAY, DEFAULT_CONNECTION_TIMEOUT, DEFAULT_RECONNECT_ON_CLOSE, DEFAULT_LOGS, DEFAULT_ON_RECONNECT } = require('../src/options')
 
 test('validateOptions', (t) => {
   const requiredOptions = {
@@ -41,11 +41,26 @@ test('validateOptions', (t) => {
   assert.throws(() => validateOptions({ ...requiredOptions, wsReconnect: { logs: '1' } }), /wsReconnect.logs must be a boolean/)
   assert.doesNotThrow(() => validateOptions({ ...requiredOptions, wsReconnect: { logs: true } }))
 
+  assert.throws(() => validateOptions({ ...requiredOptions, wsReconnect: { onReconnect: '1' } }), /wsReconnect.onReconnect must be a function/)
+  assert.doesNotThrow(() => validateOptions({ ...requiredOptions, wsReconnect: { onReconnect: () => { } } }))
+
   // set all values
-  assert.doesNotThrow(() => validateOptions({ ...requiredOptions, wsReconnect: { pingInterval: 1, maxReconnectionRetries: 1, reconnectInterval: 100, reconnectDecay: 1, connectionTimeout: 1, reconnectOnClose: true, logs: true } }))
+  assert.doesNotThrow(() => validateOptions({
+    ...requiredOptions,
+    wsReconnect: {
+      pingInterval: 1,
+      maxReconnectionRetries: 1,
+      reconnectInterval: 100,
+      reconnectDecay: 1,
+      connectionTimeout: 1,
+      reconnectOnClose: true,
+      logs: true,
+      onReconnect: () => { }
+    }
+  }))
 
   // get default values
-  assert.deepEqual(validateOptions({ ...requiredOptions, wsReconnect: { } }), {
+  assert.deepEqual(validateOptions({ ...requiredOptions, wsReconnect: {} }), {
     ...requiredOptions,
     wsReconnect: {
       pingInterval: DEFAULT_PING_INTERVAL,
@@ -54,7 +69,8 @@ test('validateOptions', (t) => {
       reconnectDecay: DEFAULT_RECONNECT_DECAY,
       connectionTimeout: DEFAULT_CONNECTION_TIMEOUT,
       reconnectOnClose: DEFAULT_RECONNECT_ON_CLOSE,
-      logs: DEFAULT_LOGS
+      logs: DEFAULT_LOGS,
+      onReconnect: DEFAULT_ON_RECONNECT
     }
   })
 })

test/ws-reconnect.js

@@ -73,7 +73,8 @@ async function createServices ({ t, wsReconnectOptions, wsTargetOptions, wsServe
     },
     proxy,
     client,
-    loggerSpy
+    loggerSpy,
+    logger
   }
 }
 
@@ -144,7 +145,7 @@ test('should not reconnect after max retries', async (t) => {
   await waitForLogMessage(loggerSpy, 'proxy ws failed to reconnect! No more retries')
 })
 
-test('should reconnect on regular target connection close', async (t) => {
+test('should not reconnect when the target connection is closed and reconnectOnClose is off', async (t) => {
   const wsReconnectOptions = { pingInterval: 200, reconnectInterval: 100, maxReconnectionRetries: 1, reconnectOnClose: false, logs: true }
 
   const { target, loggerSpy } = await createServices({ t, wsReconnectOptions })
@@ -154,7 +155,7 @@ test('should reconnect on regular target connection close', async (t) => {
       socket.pong()
     })
 
-    await wait(1_000)
+    await wait(500)
     socket.close()
   })
 
@@ -197,3 +198,50 @@ test('should reconnect retrying after a few failures', async (t) => {
   await createTargetServer(t, { autoPong: true }, targetPort)
   await waitForLogMessage(loggerSpy, 'proxy ws reconnected')
 })
+
+test('should reconnect when the target connection is closed gracefully and reconnectOnClose is on', async (t) => {
+  const wsReconnectOptions = { pingInterval: 200, reconnectInterval: 100, maxReconnectionRetries: 1, reconnectOnClose: true, logs: true }
+
+  const { target, loggerSpy } = await createServices({ t, wsReconnectOptions })
+
+  target.ws.on('connection', async (socket) => {
+    socket.on('ping', async () => {
+      socket.pong()
+    })
+
+    await wait(500)
+    socket.close()
+  })
+
+  await waitForLogMessage(loggerSpy, 'proxy ws target close event')
+  await waitForLogMessage(loggerSpy, 'proxy ws reconnected')
+})
+
+test('should call onReconnect hook function when the connection is reconnected', async (t) => {
+  const onReconnect = (oldSocket, newSocket) => {
+    logger.info('onReconnect called')
+  }
+  const wsReconnectOptions = {
+    pingInterval: 100,
+    reconnectInterval: 100,
+    maxReconnectionRetries: 1,
+    reconnectOnClose: true,
+    logs: true,
+    onReconnect
+  }
+
+  const { target, loggerSpy, logger } = await createServices({ t, wsReconnectOptions })
+
+  target.ws.on('connection', async (socket) => {
+    socket.on('ping', async () => {
+      socket.pong()
+    })
+
+    await wait(500)
+    socket.close()
+  })
+
+  await waitForLogMessage(loggerSpy, 'proxy ws target close event')
+  await waitForLogMessage(loggerSpy, 'proxy ws reconnected')
+  await waitForLogMessage(loggerSpy, 'onReconnect called')
+})