Make channels and connections Promise aware for close operations

Author: ali-ince

src/internal/browser/browser-channel.js

@@ -21,6 +21,13 @@ import HeapBuffer from './browser-buf'
 import { newError } from '../../error'
 import { ENCRYPTION_OFF, ENCRYPTION_ON } from '../util'
 
+// Just to be sure that these values are with us even after WebSocket is injected
+// for tests.
+const WS_CONNECTING = 0
+const WS_OPEN = 1
+const WS_CLOSING = 2
+const WS_CLOSED = 3
+
 /**
  * Create a new WebSocketChannel to be used in web browsers.
  * @access private
@@ -131,13 +138,19 @@ export default class WebSocketChannel {
 
   /**
    * Close the connection
-   * @param {function} cb - Function to call on close.
+   * @returns {Promise} A promise that will be resolved after channel is closed
    */
-  close (cb = () => null) {
-    this._open = false
-    this._clearConnectionTimeout()
-    this._ws.close()
-    this._ws.onclose = cb
+  close () {
+    return new Promise((resolve, reject) => {
+      if (this._ws && this._ws.readyState != WS_CLOSED) {
+        this._open = false
+        this._clearConnectionTimeout()
+        this._ws.onclose = () => resolve()
+        this._ws.close()
+      } else {
+        resolve()
+      }
+    })
   }
 
   /**
@@ -151,7 +164,7 @@ export default class WebSocketChannel {
       const webSocket = this._ws
 
       return setTimeout(() => {
-        if (webSocket.readyState !== WebSocket.OPEN) {
+        if (webSocket.readyState !== WS_OPEN) {
           this._connectionTimeoutFired = true
           webSocket.close()
         }

src/internal/connection-channel.js

@@ -438,9 +438,9 @@ export default class ChannelConnection extends Connection {
 
   /**
    * Call close on the channel.
-   * @param {function} cb - Function to call on close.
+   * @returns {Promise<void>} - A promise that will be resolved when the underlying channel is closed.
    */
-  close (cb = () => null) {
+  close () {
     if (this._log.isDebugEnabled()) {
       this._log.debug(`${this} closing`)
     }
@@ -451,11 +451,10 @@ export default class ChannelConnection extends Connection {
       this._protocol.prepareToClose()
     }
 
-    this._ch.close(() => {
+    return this._ch.close().then(() => {
       if (this._log.isDebugEnabled()) {
         this._log.debug(`${this} closed`)
       }
-      cb()
     })
   }
 

src/internal/connection-delegate.js

@@ -83,8 +83,8 @@ export default class DelegateConnection extends Connection {
     return this._delegate.resetAndFlush()
   }
 
-  close (cb = () => null) {
-    return this._delegate.close(cb)
+  close () {
+    return this._delegate.close()
   }
 
   _release () {

src/internal/connection.js

@@ -106,9 +106,10 @@ export default class Connection {
 
   /**
    * Call close on the channel.
-   * @param {function} cb - Function to call on close.
+   * @returns {Promise<void>} - A promise that will be resolved when the connection is closed.
+   *
    */
-  close (cb = () => null) {
+  close () {
     throw new Error('not implemented')
   }
 

src/internal/node/node-channel.js

@@ -334,16 +334,18 @@ export default class NodeChannel {
 
   /**
    * Close the connection
-   * @param {function} cb - Function to call on close.
+   * @returns {Promise} A promise that will be resolved after channel is closed
    */
-  close (cb = () => null) {
-    this._open = false
-    if (this._conn) {
-      this._conn.end()
-      this._conn.removeListener('end', this._handleConnectionTerminated)
-      this._conn.on('end', cb)
-    } else {
-      cb()
-    }
+  close () {
+    return new Promise((resolve, reject) => {
+      if (this._open) {
+        this._open = false
+        this._conn.removeListener('end', this._handleConnectionTerminated)
+        this._conn.on('end', () => resolve())
+        this._conn.end()
+      } else {
+        resolve()
+      }
+    })
   }
 }

test/browser/karma-firefox.conf.js

@@ -36,7 +36,7 @@ module.exports = function (config) {
     reporters: ['spec'],
     port: 9876, // karma web server port
     colors: true,
-    logLevel: config.LOG_DEBUG,
+    logLevel: config.LOG_ERROR,
     browsers: ['FirefoxHeadless'],
     autoWatch: false,
     singleRun: true,

test/internal/browser/browser-channel.test.js

@@ -24,6 +24,11 @@ import { setTimeoutMock } from '../timers-util'
 import { ENCRYPTION_OFF, ENCRYPTION_ON } from '../../../src/internal/util'
 import ServerAddress from '../../../src/internal/server-address'
 
+const WS_CONNECTING = 0
+const WS_OPEN = 1
+const WS_CLOSING = 2
+const WS_CLOSED = 3
+
 /* eslint-disable no-global-assign */
 describe('#unit WebSocketChannel', () => {
   let OriginalWebSocket
@@ -158,6 +163,44 @@ describe('#unit WebSocketChannel', () => {
     testWarningInMixedEnvironment(ENCRYPTION_OFF, 'https')
   })
 
+  it('should resolve close if websocket is already closed', () => {
+    WebSocket = () => {
+      return {
+        readyState: WS_CLOSED
+      }
+    }
+
+    const address = ServerAddress.fromUrl('bolt://localhost:8989')
+    const channelConfig = new ChannelConfig(address, {}, SERVICE_UNAVAILABLE)
+
+    const channel = new WebSocketChannel(channelConfig)
+
+    return expectAsync(channel.close()).toBeResolved()
+  })
+
+  it('should resolve close when websocket is closed', () => {
+    WebSocket = () => {
+      const ws = {
+        readyState: WS_OPEN,
+        onclose: () => {}
+      }
+
+      ws.close = () => {
+        ws.readyState = WS_CLOSED
+        ws.onclose()
+      }
+
+      return ws
+    }
+
+    const address = ServerAddress.fromUrl('bolt://localhost:8989')
+    const channelConfig = new ChannelConfig(address, {}, SERVICE_UNAVAILABLE)
+
+    const channel = new WebSocketChannel(channelConfig)
+
+    return expectAsync(channel.close()).toBeResolved()
+  })
+
   function testFallbackToLiteralIPv6 (boltAddress, expectedWsAddress) {
     // replace real WebSocket with a function that throws when IPv6 address is used
     WebSocket = url => {

test/internal/node/node-channel.test.js

@@ -0,0 +1,64 @@
+/**
+ * Copyright (c) 2002-2019 "Neo4j,"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+import NodeChannel from '../../../src/internal/node/node-channel'
+import ChannelConfig from '../../../src/internal/channel-config'
+import { SERVICE_UNAVAILABLE } from '../../../src/error'
+import ServerAddress from '../../../src/internal/server-address'
+import { connect } from 'net'
+
+describe('#unit NodeChannel', () => {
+  it('should resolve close if websocket is already closed', () => {
+    const address = ServerAddress.fromUrl('bolt://localhost:9999')
+    const channelConfig = new ChannelConfig(address, {}, SERVICE_UNAVAILABLE)
+    const channel = new NodeChannel(channelConfig)
+
+    // Modify the connection to be closed
+    channel._open = false
+
+    return expectAsync(channel.close()).toBeResolved()
+  })
+
+  it('should resolve close when websocket is connected', () => {
+    const channel = createMockedChannel(true)
+
+    return expectAsync(channel.close()).toBeResolved()
+  })
+})
+
+function createMockedChannel (connected) {
+  let endCallback = null
+  const address = ServerAddress.fromUrl('bolt://localhost:9999')
+  const channelConfig = new ChannelConfig(address, {}, SERVICE_UNAVAILABLE)
+  const channel = new NodeChannel(channelConfig)
+  const socket = {
+    end: () => {
+      channel._open = false
+      endCallback()
+    },
+    removeListener: () => {},
+    on: (key, cb) => {
+      if (key === 'end') {
+        endCallback = cb
+      }
+    }
+  }
+  channel._conn = socket
+  channel._open = connected
+  return channel
+}