From c17827e030fec150bff22035cfccf8b38315c820 Mon Sep 17 00:00:00 2001
From: Nicholas Smith <nickewansmith@gmail.com>
Date: Tue, 28 Jan 2020 09:08:47 -0500
Subject: [PATCH] Assign event listener to socket close event on open before
 attempting post-open logic (#1186)

* Assign event listener to socket close event on open before attempting to execute post-connection logic, prottects possible unhandled rejection in disconnect
* Remove try/catch for linting failure
* Up timeout for CI failure
* Feedback emit error for disconnection failure on connect failure
* Feedback ignore error on disconnect, propagate root cause
* Feedback add extra test check for expected error on connect
* Feedback remove setTimeout for await reconnect
---
 src/common/connection.ts | 23 +++++++++++------------
 test/connection-test.ts  | 18 ++++++++++++++++++
 2 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/src/common/connection.ts b/src/common/connection.ts
index b667384d..5fe1f19c 100644
--- a/src/common/connection.ts
+++ b/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()
   }
diff --git a/test/connection-test.ts b/test/connection-test.ts
index 22c23ad2..304e43cb 100644
--- a/test/connection-test.ts
+++ b/test/connection-test.ts
@@ -590,6 +590,24 @@ 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();
+      throw new Error('expected connect() to reject, but it resolved')
+    } catch (err) {
+      assert(err.message === 'error on _subscribeToLedger');
+      // _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 this.api.connection.reconnect();
+    }
+  })
+
   it('should try to reconnect on empty subscribe response on reconnect', function(done) {
     this.timeout(23000)
     this.api.on('error', error => {