fix(ws): fail in-flight queries when the connection drops or times out

On the WebSocket transport a query waits for the socket.io ack of
'GET /v2/weblite/sql'. When the socket disconnected mid-query the ack
never fired and the error was only sent to the connect callback, which
had already run, so the query promise stayed pending forever.

- track the commands waiting for an ack and fail them with the
  disconnect/connection error when the socket goes away or close() is
  called
- report the connect callback only once, for the outcome of the
  connection attempt
- honour config.timeout on the WebSocket transport too

Fixes #283

Author: Gioee

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@sqlitecloud/drivers",
-  "version": "1.0.982",
+  "version": "1.0.1022",
   "description": "SQLiteCloud drivers for Typescript/Javascript in edge, web and node clients",
   "main": "./lib/index.js",
   "types": "./lib/index.d.ts",

src/drivers/connection-ws.ts

@@ -78,6 +78,9 @@ export class SQLiteCloudWebsocketConnection extends SQLiteCloudConnection {
   /** Socket.io used to communicated with SQLiteCloud server */
   private socket?: Socket
 
+  /** Commands sent to the gateway that are still waiting for their ack, failed if the socket goes away */
+  private pendingCommands?: Set<(error: Error) => void>
+
   /** True if connection is open */
   get connected(): boolean {
     return !!(this.socket && this.socket?.connected)
@@ -116,24 +119,34 @@ export class SQLiteCloudWebsocketConnection extends SQLiteCloudConnection {
         }
         this.socket = io(gatewayUrl, ioOpts)
 
+        // the connect callback reports the outcome of the connection attempt and is called only once,
+        // errors that happen later are reported to the commands that are waiting for a response
+        let connectReported = false
+        const reportConnect = (error: Error | null) => {
+          if (!connectReported) {
+            connectReported = true
+            callback?.call(this, error)
+          }
+        }
+        const fail = (error: Error) => {
+          this.closeWithError(error)
+          reportConnect(error)
+        }
+
         this.socket.on('connect', () => {
-          callback?.call(this, null)
+          reportConnect(null)
         })
 
         this.socket.on('disconnect', (reason, description) => {
-          this.close()
-          callback?.call(
-            this,
+          fail(
             (reason === 'parse error' && getAttachmentLimitError(description, websocketMaxAttachments)) ||
               new SQLiteCloudError('Disconnected', { errorCode: 'ERR_CONNECTION_ENDED', cause: reason })
           )
         })
 
         this.socket.on('connect_error', (error: any) => {
-          this.close()
           if (error?.message === 'parse error' || error?.cause === 'parse error') {
-            callback?.call(
-              this,
+            fail(
               getAttachmentLimitError(error?.description || error?.data || error?.cause, websocketMaxAttachments) ||
                 new SQLiteCloudError('Connection error', { errorCode: 'ERR_CONNECTION_ERROR', cause: error })
             )
@@ -148,12 +161,11 @@ export class SQLiteCloudWebsocketConnection extends SQLiteCloudConnection {
               message = error.context.responseText
             }
           }
-          callback?.call(this, new SQLiteCloudError(message, { errorCode: 'ERR_CONNECTION_ERROR' }))
+          fail(new SQLiteCloudError(message, { errorCode: 'ERR_CONNECTION_ERROR' }))
         })
 
         this.socket.on('error', (error: Error) => {
-          this.close()
-          callback?.call(this, new SQLiteCloudError('Connection error', { errorCode: 'ERR_CONNECTION_ERROR', cause: error }))
+          fail(new SQLiteCloudError('Connection error', { errorCode: 'ERR_CONNECTION_ERROR', cause: error }))
         })
       }
     } catch (error) {
@@ -175,6 +187,28 @@ export class SQLiteCloudWebsocketConnection extends SQLiteCloudConnection {
       commands = { query: commands }
     }
 
+    // the callback is called exactly once: with the gateway response, on timeout or when the connection drops
+    let settled = false
+    let timeout: ReturnType<typeof setTimeout> | undefined
+    const finish = (error: Error | null, results?: any) => {
+      if (settled) return
+      settled = true
+      clearTimeout(timeout)
+      this.pendingCommands?.delete(fail)
+      callback?.call(this, error, results)
+    }
+    const fail = (error: Error) => finish(error)
+
+    this.pendingCommands ||= new Set()
+    this.pendingCommands.add(fail)
+
+    const timeoutMs = this.config?.timeout ?? 0
+    if (timeoutMs > 0) {
+      timeout = setTimeout(() => {
+        finish(new SQLiteCloudError(`Request timed out after ${timeoutMs} ms`, { errorCode: 'ERR_CONNECTION_TIMEOUT' }))
+      }, timeoutMs)
+    }
+
     this.socket.emit(
       'GET /v2/weblite/sql',
       {
@@ -194,7 +228,7 @@ export class SQLiteCloudWebsocketConnection extends SQLiteCloudConnection {
           // for the `errors[]` shape; fall back to top-level for the legacy `error` shape.
           const errorFields = (gatewayError.meta && typeof gatewayError.meta === 'object') ? gatewayError.meta : gatewayError
           const error = new SQLiteCloudError(gatewayError.detail || gatewayError.message || 'Gateway error', { ...errorFields })
-          callback?.call(this, error)
+          finish(error)
         } else {
           const { metadata } = response
           const blobTransferFormat = getResponseBlobTransferFormat(response)
@@ -208,11 +242,11 @@ export class SQLiteCloudWebsocketConnection extends SQLiteCloudConnection {
               console.assert(Array.isArray(data), 'SQLiteCloudWebsocketConnection.transportCommands - data is not an array')
               // we can recreate a SQLiteCloudRowset from the response which we know to be an array of arrays
               const rowset = new SQLiteCloudRowset(metadata, data.flat())
-              callback?.call(this, null, rowset)
+              finish(null, rowset)
               return
             }
           }
-          callback?.call(this, null, data)
+          finish(null, data)
         }
       }
     )
@@ -222,6 +256,11 @@ export class SQLiteCloudWebsocketConnection extends SQLiteCloudConnection {
 
   /** Disconnect socket.io from server */
   public close(): this {
+    return this.closeWithError(new SQLiteCloudError('Connection closed', { errorCode: 'ERR_CONNECTION_CLOSED' }))
+  }
+
+  /** Disconnect socket.io from server and fail the commands still waiting for a response with the given error */
+  private closeWithError(error: Error): this {
     console.assert(this.socket !== null, 'SQLiteCloudWebsocketConnection.close - connection already closed')
     if (this.socket) {
       this.socket?.removeAllListeners()
@@ -230,6 +269,10 @@ export class SQLiteCloudWebsocketConnection extends SQLiteCloudConnection {
     }
 
     this.operations.clear()
+
+    const pendingCommands = this.pendingCommands
+    this.pendingCommands = undefined
+    pendingCommands?.forEach(fail => fail(error))
     return this
   }
 }

src/drivers/types.ts

@@ -68,7 +68,7 @@ export interface SQLiteCloudConfig {
   /** Connect using plain TCP port, without TLS encryption, NOT RECOMMENDED, TEST ONLY */
   insecure?: boolean
 
-  /** Optional query timeout passed directly to TLS socket */
+  /** Optional query timeout in milliseconds, applies to both TLS and WebSocket connections */
   timeout?: number
   /** Name of database to open */
   database?: string

test/connection-ws-unit.test.ts

@@ -259,4 +259,125 @@ describe('websocket transport helpers', () => {
     })
     expect((callback.mock.calls[0]?.[0] as Error).message).toContain('websocketMaxAttachments')
   })
+
+  describe('in-flight commands', () => {
+    function connectMockedSocket(config: Record<string, any> = {}) {
+      const handlers: Record<string, (...args: any[]) => void> = {}
+      const socket = {
+        connected: true,
+        on: jest.fn((event: string, handler: (...args: any[]) => void) => {
+          handlers[event] = handler
+        }),
+        emit: jest.fn(),
+        removeAllListeners: jest.fn(),
+        close: jest.fn()
+      }
+      const mockedIo = io as jest.MockedFunction<typeof io>
+      mockedIo.mockReturnValue(socket as any)
+
+      const connectCallback = jest.fn()
+      const connection = Object.create(SQLiteCloudWebsocketConnection.prototype) as any
+      connection.operations = { clear: jest.fn() }
+      connection.connectTransport(
+        { connectionstring: 'sqlitecloud://host.sqlite.cloud/database?apikey=secret', host: 'host.sqlite.cloud', apikey: 'secret', ...config },
+        connectCallback
+      )
+      handlers.connect?.()
+
+      return { connection, socket, handlers, connectCallback }
+    }
+
+    afterEach(() => {
+      jest.useRealTimers()
+    })
+
+    it('should fail the command in flight when the connection drops', () => {
+      const { connection, socket, handlers, connectCallback } = connectMockedSocket()
+      const commandCallback = jest.fn()
+
+      connection.transportCommands('SELECT 1', commandCallback)
+      expect(socket.emit).toHaveBeenCalledTimes(1)
+      expect(commandCallback).not.toHaveBeenCalled()
+
+      handlers.disconnect?.('transport error')
+
+      expect(commandCallback).toHaveBeenCalledTimes(1)
+      expect(commandCallback.mock.calls[0]?.[0]).toMatchObject({ message: 'Disconnected', errorCode: 'ERR_CONNECTION_ENDED', cause: 'transport error' })
+      expect(socket.close).toHaveBeenCalled()
+      expect(connection.connected).toBe(false)
+
+      // the connect callback already reported a successful connection and must not be called again
+      expect(connectCallback).toHaveBeenCalledTimes(1)
+      expect(connectCallback).toHaveBeenCalledWith(null)
+    })
+
+    it('should report a disconnection before connect to the connect callback', () => {
+      const handlers: Record<string, (...args: any[]) => void> = {}
+      const socket = {
+        connected: false,
+        on: jest.fn((event: string, handler: (...args: any[]) => void) => {
+          handlers[event] = handler
+        }),
+        removeAllListeners: jest.fn(),
+        close: jest.fn()
+      }
+      ;(io as jest.MockedFunction<typeof io>).mockReturnValue(socket as any)
+
+      const connectCallback = jest.fn()
+      const connection = Object.create(SQLiteCloudWebsocketConnection.prototype) as any
+      connection.operations = { clear: jest.fn() }
+      connection.connectTransport({ connectionstring: 'sqlitecloud://host.sqlite.cloud/database', host: 'host.sqlite.cloud' }, connectCallback)
+
+      handlers.disconnect?.('transport close')
+
+      expect(connectCallback).toHaveBeenCalledTimes(1)
+      expect(connectCallback.mock.calls[0]?.[0]).toMatchObject({ errorCode: 'ERR_CONNECTION_ENDED' })
+    })
+
+    it('should fail the command in flight when the connection is closed', () => {
+      const { connection } = connectMockedSocket()
+      const commandCallback = jest.fn()
+
+      connection.transportCommands('SELECT 1', commandCallback)
+      connection.close()
+
+      expect(commandCallback).toHaveBeenCalledTimes(1)
+      expect(commandCallback.mock.calls[0]?.[0]).toMatchObject({ errorCode: 'ERR_CONNECTION_CLOSED' })
+    })
+
+    it('should fail the command in flight after the configured timeout and ignore a late response', () => {
+      jest.useFakeTimers()
+      const { connection, socket } = connectMockedSocket({ timeout: 1000 })
+      const commandCallback = jest.fn()
+
+      connection.transportCommands('SELECT 1', commandCallback)
+
+      jest.advanceTimersByTime(999)
+      expect(commandCallback).not.toHaveBeenCalled()
+
+      jest.advanceTimersByTime(1)
+      expect(commandCallback).toHaveBeenCalledTimes(1)
+      expect(commandCallback.mock.calls[0]?.[0]).toMatchObject({ errorCode: 'ERR_CONNECTION_TIMEOUT' })
+
+      const ack = socket.emit.mock.calls[0]?.[2] as (response: any) => void
+      ack({ data: 1 })
+      expect(commandCallback).toHaveBeenCalledTimes(1)
+    })
+
+    it('should not fail a completed command when the connection drops later', () => {
+      jest.useFakeTimers()
+      const { connection, socket, handlers } = connectMockedSocket({ timeout: 1000 })
+      const commandCallback = jest.fn()
+
+      connection.transportCommands('SELECT 1', commandCallback)
+      const ack = socket.emit.mock.calls[0]?.[2] as (response: any) => void
+      ack({ data: 1 })
+
+      handlers.disconnect?.('transport close')
+      jest.advanceTimersByTime(1000)
+
+      expect(commandCallback).toHaveBeenCalledTimes(1)
+      expect(commandCallback).toHaveBeenCalledWith(null, 1)
+    })
+  })
 })