Assign event listener to socket close event on open before attempting post-open logic

src/common/connection.ts

@@ -489,18 +489,6 @@ export class Connection extends EventEmitter {
       this._ws.on('error', error =>
         this.emit('error', 'websocket', error.message, error)
       )
-      // Finalize the connection and resolve all awaiting connect() requests
-      try {
-        this._retryConnectionBackoff.reset()
-        await this._subscribeToLedger()
-        this._startHeartbeatInterval()
-        this._connectionManager.resolveAllAwaiting()
-        this.emit('connected')
-      } catch (error) {
-        this._connectionManager.rejectAllAwaiting(error)
-        this.disconnect()
-        return
-      }
       // Handle a closed connection: reconnect if it was unexpected
       this._ws.once('close', code => {
         this._clearHeartbeatInterval()
@@ -524,6 +512,17 @@ export class Connection extends EventEmitter {
           }, retryTimeout)
         }
       })
+      // Finalize the connection and resolve all awaiting connect() requests
+      try {
+        this._retryConnectionBackoff.reset()
+        await this._subscribeToLedger()
+        this._startHeartbeatInterval()
+        this._connectionManager.resolveAllAwaiting()
+        this.emit('connected')
+      } catch (error) {
+        this._connectionManager.rejectAllAwaiting(error)
+        await this.disconnect().catch(() => {}) // Ignore this error, propagate the root cause.
+      }
     })
     return this._connectionManager.awaitConnection()
   }

test/connection-test.ts

@@ -590,6 +590,26 @@ describe('Connection', function() {
     }
   )
 
+  it('should clean up websocket connection if error after websocket is opened', async function() {
+    await this.api.disconnect();
+    // fail on connection
+    this.api.connection._subscribeToLedger = async () => {
+      throw new Error('error on _subscribeToLedger')
+    }
+    try {
+      await this.api.connect();
+    } catch (err) {
+      // _ws.close event listener should have cleaned up the socket when disconnect _ws.close is run on connection error
+      // do not fail on connection anymore
+      this.api.connection._subscribeToLedger = async () => {}
+      await new Promise((resolve, reject) => {
+        setTimeout(async () => {
+          this.api.connect().then(resolve).catch(reject); // should succeed and not throw websocket not cleaned up error
+        }, 1000)
+      })
+    }
+  })
+
   it('should try to reconnect on empty subscribe response on reconnect', function(done) {
     this.timeout(23000)
     this.api.on('error', error => {