dl/lib/dataprovider/WebSocketDataProvider.js

DreamLab Libs

var WebSocket = require('ws');

var WebSocketServer = WebSocket.Server;
var ZooKeeper = require('zookeeper');

var AbstractDataProvider = require('./AbstractDataProvider.js').AbstractDataProvider;

var WebSocketDataProvider = function (params) {
    AbstractDataProvider.apply(this, arguments);

    this._wss = null;

    this._params = params;

    this._connections = [];
    this._readyCallback     = null;
    this._reloadedCallback  = null;
    this._ready = false;
    this._messageId = 0;
    this._lastUseConnIdx = -1;
    this._callbacks = {};

    /* lista zapiętych watcherów: {key: klucz, callback: callback} */
    this._watchers = [];
    this._requestQueue = [];
    this._ephemeralOwner = {};
    this._ephemeralKeys = {};
};

WebSocketDataProvider.prototype = Object.create(AbstractDataProvider.prototype);

WebSocketDataProvider.prototype._sendThrough = function (connIdx, method, params, callback) {
    var messageId = ++this._messageId;

    if (method !== 'get') {
        console.info('WebSocketDataProvider/_sendThrough ->', connIdx, messageId, method, params.key);
    }

    var data = JSON.stringify({
        id: messageId,
        method: method,
        params: params
    });

    var ws = this._connections[connIdx];
    var that = this;

    this._callbacks[connIdx][messageId] = callback;

    ws.send(data, { compress: !!this._params.compress, binary: !!this._params.binary }, function (err) {
        if (err) {
            console.error('WebSocketDataProvider/_sendThrough wsc=%s id=%d error: %s', connIdx, messageId, err);
            delete that._callbacks[connIdx][messageId];

            callback(err);
        }

        if (method !== 'get') {
            console.info('WebSocketDataProvider/_sendThrough done ->', connIdx, messageId, method, params.key);
        }
    });
};

WebSocketDataProvider.prototype._send = function (method, params, callback) {
    var connIdx = this._getConnectionIdx(params.key);
    this._sendThrough(connIdx, method, params, callback);
};

WebSocketDataProvider.prototype._getOpenConnections = function () {
    var openConnections = [];

    for (var i = 0, l = this._connections.length; i < l; i++) {
        if (this._connections[i] && this._connections[i].readyState == WebSocket.OPEN) {
            openConnections.push(i);
        }
    }

    return openConnections;
};

WebSocketDataProvider.prototype._getConnectionIdx = function (ephemeralKey) {
    if (ephemeralKey && this._ephemeralOwner.hasOwnProperty(ephemeralKey)) {
        return this._ephemeralOwner[ephemeralKey];
    }

    var openConnections = this._getOpenConnections();

    this._lastUseConnIdx = (this._lastUseConnIdx + 1) % openConnections.length;
    return openConnections[this._lastUseConnIdx];
};

WebSocketDataProvider.prototype._addConnection = function (ws) {
    var that = this;

    // find first empty slot - null
    var idx = this._connections.indexOf(null);

    if (idx > -1) {
        this._connections[idx] = ws;
    } else {
        idx = this._connections.length;
        this._connections.push(ws);
    }

    this._callbacks[idx] = {};
    this._ephemeralKeys[idx] = [];

    console.info('DataProvider[WebSocket]: Entered connected state', idx);

    ws.on('close', function () {
        console.info('DataProvider[WebSocket]: Entered closed state', idx);

        that._removeConnection(this);

        // when one connection is closed it may hold some ephemeral nodes that must be registered again
        if (that.isConnected()) {
            that._reloadedCallback();
            that._assignWatchers();
        }
    });

    ws.on('message', function (message) {
        try {
            message = JSON.parse(message);
        } catch (ex) {
            console.error('DataProvider[WebSocket]: Failed to parse message');
            return;
        }

        if (message.method && message.method === 'notify') {
            var params = message.params;

            console.info('DataProvider[WebSocket]: On watch notify: %s wsc: %d', params.key, idx);

            for (var i = 0, l = that._watchers.length; i < l; i++) {
                var watcher = that._watchers[i];

                if (watcher.key === params.key && watcher.connIdx === idx) {
                    console.info('DataProvider[WebSocket]: On watch notify: %s, launching...', watcher.key);

                    process.nextTick((function (_watcher, _rc, _error, _data, _stat) {
                        return function () {
                            if (!params.rc) {
                                _watcher.callback(null, _data, _stat);
                            } else {
                                _watcher.callback(_rc, _error);
                            }
                        };
                    })(watcher, params.rc, params.error, params.data, params.stat));
                }
            }

            return;
        } else if (message.method && message.method === 'reloaded') {
            console.info('DataProvider[WebSocket]: Reloaded wsc: %d', idx);

            return that._reloadedCallback();

        } else if (!message.id) {
            console.warn('DataProvider[WebSocket]: Message has no id');
            return;
        }

        //console.log('DataProvider[WebSocket] <-', message.id);

        if (!that._callbacks[idx][message.id]) {
            console.warn('DataProvider[WebSocket]: No callback for message id:', message.id);
            return;
        }

        process.nextTick(function () {
            var cb = that._callbacks[idx][message.id];
            delete that._callbacks[idx][message.id];

            cb(message.rc, message.error, message.data, message.stat);
        });
    });
};

WebSocketDataProvider.prototype._removeConnection = function (ws) {
    var index = this._connections.indexOf(ws);

    console.info('DataProvider[WebSocket]: removing connection', index);

    // remove ephemeral owner for all keys owned by dropped connection
    for (var i = 0, l = this._ephemeralKeys[index].length; i < l; i++) {
        delete this._ephemeralOwner[this._ephemeralKeys[index][i]];
    }

    // remove list of keys owned by dropped connection
    delete this._ephemeralKeys[index];

    // invalidate all waiting callbacks
    var cbKeys = Object.keys(this._callbacks[index]);
    var cb;

    console.info('WebSocketDataProvider/_removeConnection invalidating %d callbacks wsc=%d', cbKeys.length, index);

    for (var j = 0, jl = cbKeys.length; j < jl; j++) {
        var msgId = cbKeys[j];

        cb = this._callbacks[index][msgId];
        delete this._callbacks[index][msgId];

        cb(ZooKeeper.ZCONNECTIONLOSS);
    }

    delete this._callbacks[index];

    // mark watches as orphaned
    for (var k = 0, lk = this._watchers.length; k < lk; k++) {
        var watcher = this._watchers[k];

        if (watcher.connIdx === index) {
            console.info('WebSocketDataProvider/_removeConnection remove watcher connection wsc=%d key=%s',
                         index, watcher.key);

            watcher.connIdx = null;
        }
    }

    this._connections[index] = null;
};

WebSocketDataProvider.prototype.init = function (readyCallback, reloadedCallback) {
    var that = this;

    this._readyCallback = function () {
        console.log('DataProvider[WebSocket]: ready callback');
        readyCallback(null, true);
    };

    this._reloadedCallback = function () {
        console.log('DataProvider[WebSocket]: reloaded callback');
        reloadedCallback(null, true);
    };

    console.info('DataProvider[WebSocket]: init');

    this._wss = new WebSocketServer(this._params);

    this._wss.on('connection', function (ws) {
        that._addConnection(ws);

        that._assignWatchers();
        that._reloadedCallback();
        that._flushQueue();

        /* pierwsze uruchomienie */
        if (!that._ready) {
            that._readyCallback();
        }

        that._ready = true;
    });
};

WebSocketDataProvider.prototype.get = function (key, callback) {
    if (!this.isConnected()) {
        this._deferRequest(arguments.callee.bind(this, key, callback), callback);
        return;
    }

    var that = this;

    this._send('get', { key: key }, function (rc, error, data, stat) {
        switch (rc) {
            case ZooKeeper.ZOK:
                callback(null, data, stat);
                break;
            case ZooKeeper.ZCONNECTIONLOSS:
            case ZooKeeper.ZOPERATIONTIMEOUT:
                console.info('DataProvider[WebSocket]:', rc, ' - defering request', 'GET', key);
                that._deferRequest(that.get.bind(that, key, callback), callback);
                break;
            default:
                console.warn('DataProvider[WebSocket]:', 'Unhandled error', 'GET', rc, error, key);
                callback(rc, error);
        }
    });
};

WebSocketDataProvider.prototype.set = function (key, value, options, callback) {
    if (!this.isConnected()) {
        this._deferRequest(arguments.callee.bind(this, key, value, options, callback), callback);
        return;
    }

    options = options || {};

    var that = this;
    var flags = options.flags || 0;
    var ownerIdx = this._getConnectionIdx(key);
    var ephemeral = flags & 1; // ZOO_EPHEMERAL

    this._sendThrough(ownerIdx, 'set', {
        key: key,
        value: value,
        options: options
    }, function (rc, error, stat) {
        switch (rc) {
            case ZooKeeper.ZOK:
                if (ephemeral && stat && stat.createdInThisSession && !that._ephemeralOwner[key]) {
                    that._ephemeralOwner[key] = ownerIdx;
                    that._ephemeralKeys[ownerIdx].push(key);
                }

                callback(null, stat);
                break;
            case ZooKeeper.ZCONNECTIONLOSS: /* stracilismy polaczenie do ZK */
            case ZooKeeper.ZOPERATIONTIMEOUT:
                console.warn('DataProvider[WebSocket]:', rc, ' - defering request', 'SET', key);
                that._deferRequest(that.set.bind(that, key, value, options, callback), callback);
                break;
            default:
                console.warn('DataProvider[WebSocket]:', 'Unhandled error', 'SET', rc, error, key);
                callback(rc, error);
        }
    });
};

WebSocketDataProvider.prototype.sync = function (key, callback) {
    if (!this.isConnected()) {
        this._deferRequest(arguments.callee.bind(this, key, callback), callback);
        return;
    }

    var that = this;

    this._send('sync', { key: key }, function (rc, error) {
        switch (rc) {
            case ZooKeeper.ZOK:
                callback(null, null);
                break;
            case ZooKeeper.ZCONNECTIONLOSS:
            case ZooKeeper.ZOPERATIONTIMEOUT:
                console.warn('DataProvider[WebSocket]:', rc, ' - defering request', 'SYNC', key);
                that._deferRequest(that.sync.bind(that, key, callback), callback);
                break;
            default:
                console.warn('DataProvider[WebSocket]:', 'Unhandled error', 'SYNC', rc, error, key);
                callback(rc, error);
        }
    });
};

WebSocketDataProvider.prototype.remove = function (key, options, callback) {
    if (!this.isConnected()) {
        this._deferRequest(arguments.callee.bind(this, key, options, callback), callback);
        return;
    }

    var that = this;

    this._send('remove', { key: key, options: options }, function (rc, error) {
        switch (rc) {
            case ZooKeeper.ZOK:
                callback(null, null);
                break;
            case ZooKeeper.ZCONNECTIONLOSS:
            case ZooKeeper.ZOPERATIONTIMEOUT:
                console.warn('DataProvider[WebSocket]:', rc, ' - defering request', 'REMOVE', key);
                that._deferRequest(that.remove.bind(that, key, options, callback), callback);
                break;
            default:
                console.warn('DataProvider[WebSocket]:', 'Unhandled error', 'REMOVE', rc, error, key);
                callback(rc, error);
        }
    });
};

WebSocketDataProvider.prototype.exists = function (key, callback) {
    if (!this.isConnected()) {
        this._deferRequest(arguments.callee.bind(this, key, callback), callback);
        return;
    }

    var that = this;

    this._send('exists', { key: key }, function (rc, error, stat) {
        switch (rc) {
            case ZooKeeper.ZOK:
                callback(null, stat);
                break;
            case ZooKeeper.ZCONNECTIONLOSS:
            case ZooKeeper.ZOPERATIONTIMEOUT:
                console.warn('DataProvider[WebSocket]:', rc, ' - defering request', 'EXISTS', key);
                that._deferRequest(that.exists.bind(that, key, callback), callback);
                break;
            default:
                console.warn('DataProvider[WebSocket]:', 'Unhandled error', 'EXISTS', rc, error, key);
                callback(rc, error);
        }
    });
};

WebSocketDataProvider.prototype.watch = function (key, callback) {
    console.info('DataProvider[WebSocket]: creating watch on %s', key);

    // register data for future reconnects / notifications
    var watcher = {
        key: key,
        callback: callback,
        connIdx: undefined // <- not yet set
    };

    this._watchers.push(watcher);
    this._watch(key, watcher);
};

WebSocketDataProvider.prototype._watch = function (key, watcher) {
    if (!this.isConnected()) {
        this._deferRequest(arguments.callee.bind(this, key, watcher));
        return;
    }

    var connIdx = this._getConnectionIdx(key);
    var that = this;

    // stick key to one connection to prevent multiple through different connections
    if (!this._ephemeralOwner.hasOwnProperty(key)) {
        this._ephemeralOwner[key] = connIdx;
        this._ephemeralKeys[connIdx].push(key);
    }

    watcher.connIdx = connIdx;

    console.info('DataProvider[WebSocket]: watching %s on connection %d', key, connIdx);

    this._sendThrough(connIdx, 'watch', { key: key }, function (rc, error, data, stat) {
        switch (rc) {
            case ZooKeeper.ZOK:
                console.info('DataProvider[WebSocket]: watch %s created on %d', key, connIdx);
                watcher.callback(null, data, stat);
                break;
            case ZooKeeper.ZCONNECTIONLOSS:
            case ZooKeeper.ZOPERATIONTIMEOUT:
                console.info('DataProvider[WebSocket]:', rc, ' - defering request', 'WATCH', key);
                that._deferRequest(that._watch.bind(that, key, watcher));
                break;
            default:
                console.warn('DataProvider[WebSocket]:', 'Unhandled error', 'WATCH', rc, key);
                watcher.callback(rc, error);
        }
    });
};

WebSocketDataProvider.prototype._assignWatchers = function () {
    console.info('DataProvider[WebSocket]: assigning watches after reconnect');
    var tmpWatchers = this._watchers;
    this._watchers = [];

    for (var i = 0, len = tmpWatchers.length; i < len; i++) {
        if (tmpWatchers[i].connIdx === null) {
            console.info('DataProvider[WebSocket]: reassign watch', tmpWatchers[i].key);
            this.watch(tmpWatchers[i].key, tmpWatchers[i].callback);
        } else {
            console.info('DataProvider[WebSocket]: restore watch wsc=%s key=%s',
                         tmpWatchers[i].connIdx, tmpWatchers[i].key);
            this._watchers.push(tmpWatchers[i]);
        }
    }
};

WebSocketDataProvider.prototype.isConnected = function () {
    return this._getOpenConnections().length > 0;
};

WebSocketDataProvider.prototype._deferRequest = function (fnc, timeoutCb) {
    this._requestQueue.push(fnc);

    if (timeoutCb) {
        var that = this;

        setTimeout(function () {
            var idx = that._requestQueue.indexOf(fnc);

            if (idx === -1) {
                return;
            }

            that._requestQueue.splice(idx, 1);

            return timeoutCb(new Error('Timed out'));
        }, WebSocketDataProvider.DEFER_TIMEOUT);
    }
};

WebSocketDataProvider.prototype._flushQueue = function () {
    console.info('DataProvider[WebSocket]: Flushing queue', this._requestQueue.length);

    var requestQueue = this._requestQueue.slice();
    this._requestQueue = [];

    while (requestQueue.length) {
        requestQueue.shift()();
    }
};

WebSocketDataProvider.DEFER_TIMEOUT = 60000;

exports.WebSocketDataProvider = WebSocketDataProvider;