feat: close

Author: shaoshuai0102

lib/client.js

@@ -19,6 +19,8 @@ const publish = symbols.publish;
 const invoke = symbols.invoke;
 const subInfo = symbols.subInfo;
 const pubInfo = symbols.pubInfo;
+const closeHandler = symbols.closeHandler;
+const close = symbols.close;
 
 class ClusterClient extends Base {
   /**
@@ -38,6 +40,7 @@ class ClusterClient extends Base {
     this[subInfo] = new Map();
     this[pubInfo] = new Map();
 
+    this[closeHandler] = this[closeHandler].bind(this);
     this[init]();
   }
 
@@ -84,12 +87,7 @@ class ClusterClient extends Base {
       utils.delegateEvents(this[innerClient], this);
 
       // re init when connection is close
-      this[innerClient].on('close', () => {
-        this[logger].warn('[ClusterClient#%s] %s closed, and try to init it again', name, this[innerClient].isLeader ? 'leader' : 'follower');
-        this[isReady] = false;
-        this.ready(false);
-        this[init]();
-      });
+      this[innerClient].on('close', this[closeHandler]);
 
       // wait leader/follower ready
       yield this[innerClient].ready();
@@ -165,6 +163,29 @@ class ClusterClient extends Base {
     debug('[ClusterClient#%s] invoke method: %s, args: %j', this._options.name, method, args);
     return this[innerClient].invoke(method, args, callback);
   }
+
+  [closeHandler]() {
+    this[logger].warn('[ClusterClient#%s] %s closed, and try to init it again', this._options.name, this[innerClient].isLeader ? 'leader' : 'follower');
+    this[isReady] = false;
+    this.ready(false);
+    this[init]();
+  }
+
+  [close]() {
+    return co(function* () {
+      // close after ready, in case of innerClient is initializing
+      yield this.ready();
+
+      const client = this[innerClient];
+      // prevent re-initializing
+      client.removeListener('close', this[closeHandler]);
+
+      if (client.close) {
+        yield utils.callFn(client.close.bind(client));
+      }
+      this.removeAllListeners();
+    }.bind(this));
+  }
 }
 
 module.exports = ClusterClient;

lib/connection.js

@@ -81,10 +81,19 @@ class Connection extends Base {
       this.emit('error', err);
     }
 
+    this._clearInvokes();
     this.emit('close');
     this.removeAllListeners();
   }
 
+  _clearInvokes() {
+    for (const id of this._invokes.keys()) {
+      const req = this._invokes.get(id);
+      clearTimeout(req.timer);
+    }
+    this._invokes.clear();
+  }
+
   _onReadable() {
     try {
       let remaining = false;

lib/index.js

@@ -1,6 +1,7 @@
 'use strict';
 
 const is = require('is-type-of');
+const symbols = require('./symbol');
 const logger = require('./default_logger');
 const transcode = require('./default_transcode');
 const ClusterClient = require('./client');
@@ -203,3 +204,13 @@ class ClientWrapper {
 module.exports = function createWrapper(clientClass, options) {
   return new ClientWrapper(clientClass, options);
 };
+
+/**
+ * Close a ClusterClient
+ *
+ * @param {Object} client - ClusterClient instance to be closed
+ * @return {Promise} returns a promise which will be resolved after fully closed
+ */
+module.exports.close = function(client) {
+  return client[symbols.close]();
+};

lib/leader.js

@@ -6,6 +6,7 @@ const is = require('is-type-of');
 const Base = require('sdk-base');
 const utils = require('./utils');
 const random = require('utility').random;
+const ClusterServer = require('./server');
 const Connection = require('./connection');
 const Request = require('./protocol/request');
 
@@ -43,18 +44,12 @@ class Leader extends Base {
       this._realClient.ready(() => this.ready(true));
     }
 
+    this._handleClose = this._handleClose.bind(this);
+    this._handleConnection = this._handleConnection.bind(this);
+
     // subscribe its own channel
-    this._server.on(`${this._options.name}_connection`, socket => this._handleConnection(socket));
-    this._server.once('close', () => {
-      this.logger.info('[Loader:%s] leader server is closed', this._options.name);
-      // close the real client
-      if (this._realClient && is.function(this._realClient.close)) {
-        this._realClient.close();
-      }
-      clearInterval(this._heartbeatTimer);
-      this._heartbeatTimer = null;
-      this.emit('close');
-    });
+    this._server.on(`${this._options.name}_connection`, this._handleConnection);
+    this._server.once('close', this._handleClose);
 
     this._heartbeatTimer = setInterval(() => {
       const now = Date.now();
@@ -80,6 +75,7 @@ class Leader extends Base {
   subscribe(reg, listener) {
     const transcode = this._transcode;
     const conn = Object.create(Base.prototype, {
+      isMock: { value: true },
       key: { value: `mock_conn_${Date.now()}` },
       lastActiveTime: {
         get() {
@@ -302,6 +298,55 @@ class Leader extends Base {
   _errorHandler(err) {
     setImmediate(() => this.emit('error', err));
   }
+
+  * _handleClose() {
+    this.logger.info('[Loader:%s] leader server is closed', this._options.name);
+    // close the real client
+    if (this._realClient && is.function(this._realClient.close)) {
+      // support common function, generatorFunction, and function returning a promise
+      yield utils.callFn(this._realClient.close.bind(this._realClient));
+    }
+    clearInterval(this._heartbeatTimer);
+    this._heartbeatTimer = null;
+    this.emit('close');
+  }
+
+  close() {
+    return co(function* () {
+      // 1. stop listening to server channel
+      this._server.removeListener(`${this._options.name}_connection`, this._handleConnection);
+
+      // 2. close all mock connections
+      for (const [ key, conn ] of this._connections.entries()) {
+        if (conn.isMock) this._connections.delete(key);
+      }
+
+      // 3. wait all followers close
+      yield new Promise((resolve, reject) => {
+        if (this._connections.size === 0) return resolve();
+
+        for (const conn of this._connections.values()) {
+          conn.once('close', () => {
+            if (this._connections.size === 0) return resolve();
+          });
+        }
+
+        setTimeout(() => {
+          reject(new Error(`[Leader#${this._options.name}] close failed: follower connections are still not closed after 30s`));
+        }, 30000);
+      });
+
+      // 4. close server
+      //    CANNOT close server directly by server.close(), other cluster clients may be using it
+      this._server.removeListener('close', this._handleClose);
+      yield ClusterServer.close(this._options.name, this._server);
+
+      // 5. close real client
+      yield this._handleClose();
+
+      this.removeAllListeners();
+    }.bind(this));
+  }
 }
 
 module.exports = Leader;

lib/server.js

@@ -94,15 +94,19 @@ class ClusterServer extends Base {
   }
 
   close() {
-    if (this.isClosed) {
-      return;
-    }
+    return new Promise((resolve, reject) => {
+      if (this.isClosed) return resolve();
 
-    this._server.close();
+      this._server.close(err => {
+        if (err) return reject(err);
+        resolve();
+      });
 
-    for (const socket of this._sockets.values()) {
-      socket.destroy();
-    }
+      // sockets must be closed manually, otherwise server.close callback will never be called
+      for (const socket of this._sockets.values()) {
+        socket.destroy();
+      }
+    });
   }
 
   _handleSocket(socket) {
@@ -188,6 +192,27 @@ class ClusterServer extends Base {
     }
   }
 
+  static* close(name, server) {
+    const port = server._port;
+
+    // remove from typeSet, so other client can occupy
+    typeSet.delete(`${name}@${port}`);
+
+    let listening = false;
+    for (const key of typeSet.values()) {
+      if (key.endsWith(`@${port}`)) {
+        listening = true;
+        break;
+      }
+    }
+
+    // close server if no one is listening on this port any more
+    if (!listening) {
+      const server = serverMap.get(port);
+      yield server && server.close();
+    }
+  }
+
   /**
    * Wait for Leader Startup
    *
@@ -203,7 +228,7 @@ class ClusterServer extends Base {
 
       // if timeout, throw error
       if (Date.now() - start > timeout) {
-        throw new Error(`[ClusterClient] leader dose not be active in ${timeout}ms on port:${port}`);
+        throw new Error(`[ClusterClient] leader does not be active in ${timeout}ms on port:${port}`);
       }
       if (!connect) {
         yield sleep(3000);

lib/symbol.js

@@ -9,3 +9,5 @@ exports.publish = Symbol.for('ClusterClient#publish');
 exports.invoke = Symbol.for('ClusterClient#invoke');
 exports.subInfo = Symbol.for('ClusterClient#subInfo');
 exports.pubInfo = Symbol.for('ClusterClient#pubInfo');
+exports.closeHandler = Symbol.for('ClusterClient#closeHandler');
+exports.close = Symbol.for('ClusterClient#close');

lib/utils.js

@@ -1,5 +1,6 @@
 'use strict';
 
+const is = require('is-type-of');
 const stringify = require('json-stringify-safe');
 
 const MAX_INT_HIGH = Math.pow(2, 21);
@@ -69,3 +70,24 @@ exports.handleLong = function(val) {
   }
   return val.toNumber();
 };
+
+/**
+ * call a function, support common function, generator function, or a function returning promise
+ *
+ * @param {Function} fn - common function, generator function, or a function returning promise
+ * @param {Array} args - args as fn() paramaters
+ * @return {*} data returned by fn
+ */
+exports.callFn = function* (fn, args) {
+  args = args || [];
+  if (!is.function(fn)) return;
+  if (is.generatorFunction(fn)) {
+    return yield fn(...args);
+  }
+  const r = fn(...args);
+  if (is.promise(r)) {
+    return yield r;
+  }
+  return r;
+};
+

package.json

@@ -50,7 +50,6 @@
     "mm": "^2.0.0",
     "npminstall": "^2.13.3",
     "pedding": "^1.1.0",
-    "power-assert": "^1.4.2",
     "supertest": "^2.0.1"
   },
   "engines": {