Modify driver.close() and related internals to return Promise

Author: ali-ince

README.md

@@ -48,7 +48,7 @@ var neo4j = require('neo4j-driver')
 Driver instance should be closed when Node.js application exits:
 
 ```javascript
-driver.close()
+driver.close() // returns a Promise
 ```
 
 otherwise application shutdown might hang or it might exit with a non-zero exit code.
@@ -87,7 +87,7 @@ WebSockets when the page is unloaded. However, driver instance should be explici
 is not the same as the lifetime of the web page:
 
 ```javascript
-driver.close()
+driver.close() // returns a Promise
 ```
 
 ## Usage examples
@@ -104,7 +104,7 @@ var driver = neo4j.driver(
 
 // Close the driver when application exits.
 // This closes all used network connections.
-driver.close()
+await driver.close()
 ```
 
 ### Acquiring a Session
@@ -189,7 +189,7 @@ session
       console.log(record.get('name'))
     },
     onCompleted: () => {
-      session.close()
+      session.close() // returns a Promise
     },
     onError: error => {
       console.log(error)

src/driver.js

@@ -174,8 +174,9 @@ class Driver {
   close () {
     this._log.info(`Driver ${this._id} closing`)
     if (this._connectionProvider) {
-      this._connectionProvider.close()
+      return this._connectionProvider.close()
     }
+    return Promise.resolve()
   }
 
   /**

src/internal/browser/browser-channel.js

@@ -38,7 +38,11 @@ export default class WebSocketChannel {
    * @param {ChannelConfig} config - configuration for this channel.
    * @param {function(): string} protocolSupplier - function that detects protocol of the web page. Should only be used in tests.
    */
-  constructor (config, protocolSupplier = detectWebPageProtocol) {
+  constructor (
+    config,
+    protocolSupplier = detectWebPageProtocol,
+    socketFactory = url => new WebSocket(url)
+  ) {
     this._open = true
     this._pending = []
     this._error = null
@@ -51,14 +55,14 @@ export default class WebSocketChannel {
       return
     }
 
-    this._ws = createWebSocket(scheme, config.address)
+    this._ws = createWebSocket(scheme, config.address, socketFactory)
     this._ws.binaryType = 'arraybuffer'
 
     let self = this
     // All connection errors are not sent to the error handler
     // we must also check for dirty close calls
     this._ws.onclose = function (e) {
-      if (!e.wasClean) {
+      if (e && !e.wasClean) {
         self._handleConnectionError()
       }
     }
@@ -142,7 +146,7 @@ export default class WebSocketChannel {
    */
   close () {
     return new Promise((resolve, reject) => {
-      if (this._ws && this._ws.readyState != WS_CLOSED) {
+      if (this._ws && this._ws.readyState !== WS_CLOSED) {
         this._open = false
         this._clearConnectionTimeout()
         this._ws.onclose = () => resolve()
@@ -187,11 +191,11 @@ export default class WebSocketChannel {
   }
 }
 
-function createWebSocket (scheme, address) {
+function createWebSocket (scheme, address, socketFactory) {
   const url = scheme + '://' + address.asHostPort()
 
   try {
-    return new WebSocket(url)
+    return socketFactory(url)
   } catch (error) {
     if (isIPv6AddressIssueOnWindows(error, address)) {
       // WebSocket in IE and Edge browsers on Windows do not support regular IPv6 address syntax because they contain ':'.
@@ -208,7 +212,7 @@ function createWebSocket (scheme, address) {
       // That is why here we "catch" SyntaxError and rewrite IPv6 address if needed.
 
       const windowsFriendlyUrl = asWindowsFriendlyIPv6Address(scheme, address)
-      return new WebSocket(windowsFriendlyUrl)
+      return socketFactory(windowsFriendlyUrl)
     } else {
       throw error
     }

src/internal/connection-channel.js

@@ -440,7 +440,7 @@ export default class ChannelConnection extends Connection {
    * Call close on the channel.
    * @returns {Promise<void>} - A promise that will be resolved when the underlying channel is closed.
    */
-  close () {
+  async close () {
     if (this._log.isDebugEnabled()) {
       this._log.debug(`${this} closing`)
     }
@@ -451,11 +451,11 @@ export default class ChannelConnection extends Connection {
       this._protocol.prepareToClose()
     }
 
-    return this._ch.close().then(() => {
-      if (this._log.isDebugEnabled()) {
-        this._log.debug(`${this} closed`)
-      }
-    })
+    await this._ch.close()
+
+    if (this._log.isDebugEnabled()) {
+      this._log.debug(`${this} closed`)
+    }
   }
 
   toString () {

src/internal/connection-delegate.js

@@ -92,6 +92,6 @@ export default class DelegateConnection extends Connection {
       this._delegate._errorHandler = this._originalErrorHandler
     }
 
-    this._delegate._release()
+    return this._delegate._release()
   }
 }

src/internal/connection-provider-pooled.js

@@ -98,22 +98,16 @@ export default class PooledConnectionProvider extends ConnectionProvider {
    */
   _destroyConnection (conn) {
     delete this._openConnections[conn.id]
-    conn.close()
+    return conn.close()
   }
 
-  close () {
-    try {
-      // purge all idle connections in the connection pool
-      this._connectionPool.purgeAll()
-    } finally {
-      // then close all connections driver has ever created
-      // it is needed to close connections that are active right now and are acquired from the pool
-      for (let connectionId in this._openConnections) {
-        if (this._openConnections.hasOwnProperty(connectionId)) {
-          this._openConnections[connectionId].close()
-        }
-      }
-    }
+  async close () {
+    // purge all idle connections in the connection pool
+    await this._connectionPool.close()
+
+    // then close all connections driver has ever created
+    // it is needed to close connections that are active right now and are acquired from the pool
+    await Promise.all(Object.values(this._openConnections).map(c => c.close()))
   }
 
   static _installIdleObserverOnConnection (conn, observer) {

src/internal/connection-provider-routing.js

@@ -155,7 +155,9 @@ export default class RoutingConnectionProvider extends PooledConnectionProvider
       )
     }
 
-    this._connectionPool.purge(address)
+    // We're firing and forgetting this operation explicitly and listening for any
+    // errors to avoid unhandled promise rejection
+    this._connectionPool.purge(address).catch(() => {})
   }
 
   forgetWriter (address, database) {
@@ -223,10 +225,9 @@ export default class RoutingConnectionProvider extends PooledConnectionProvider
           currentRoutingTable
         )
       })
-      .then(newRoutingTable => {
+      .then(newRoutingTable =>
         this._applyRoutingTableIfPossible(currentRoutingTable, newRoutingTable)
-        return newRoutingTable
-      })
+      )
   }
 
   _fetchRoutingTableFromKnownRoutersFallbackToSeedRouter (
@@ -249,10 +250,9 @@ export default class RoutingConnectionProvider extends PooledConnectionProvider
           currentRoutingTable
         )
       })
-      .then(newRoutingTable => {
+      .then(newRoutingTable =>
         this._applyRoutingTableIfPossible(currentRoutingTable, newRoutingTable)
-        return newRoutingTable
-      })
+      )
   }
 
   _fetchRoutingTableUsingKnownRouters (knownRouters, currentRoutingTable) {
@@ -381,7 +381,7 @@ export default class RoutingConnectionProvider extends PooledConnectionProvider
       })
   }
 
-  _applyRoutingTableIfPossible (currentRoutingTable, newRoutingTable) {
+  async _applyRoutingTableIfPossible (currentRoutingTable, newRoutingTable) {
     if (!newRoutingTable) {
       // none of routing servers returned valid routing table, throw exception
       throw newError(
@@ -396,12 +396,14 @@ export default class RoutingConnectionProvider extends PooledConnectionProvider
       this._useSeedRouter = true
     }
 
-    this._updateRoutingTable(newRoutingTable)
+    await this._updateRoutingTable(newRoutingTable)
+
+    return newRoutingTable
   }
 
-  _updateRoutingTable (newRoutingTable) {
+  async _updateRoutingTable (newRoutingTable) {
     // close old connections to servers not present in the new routing table
-    this._connectionPool.keepAll(newRoutingTable.allServers())
+    await this._connectionPool.keepAll(newRoutingTable.allServers())
 
     // make this driver instance aware of the new table
     this._routingTables[newRoutingTable.database] = newRoutingTable

src/internal/node/node-channel.js

@@ -278,6 +278,7 @@ export default class NodeChannel {
   }
 
   _handleConnectionTerminated () {
+    this._open = false
     this._error = newError(
       'Connection was closed by server',
       this._connectionErrorCode
@@ -338,13 +339,23 @@ export default class NodeChannel {
    */
   close () {
     return new Promise((resolve, reject) => {
+      const cleanup = () => {
+        if (!this._conn.destroyed) {
+          this._conn.destroy()
+        }
+
+        resolve()
+      }
+
       if (this._open) {
         this._open = false
         this._conn.removeListener('end', this._handleConnectionTerminated)
-        this._conn.on('end', () => resolve())
+        this._conn.on('end', () => cleanup())
+        this._conn.on('close', () => cleanup())
         this._conn.end()
+        this._conn.destroy()
       } else {
-        resolve()
+        cleanup()
       }
     })
   }