kitejs/src/consul.ts

the rpc framework Kite for Node.js

import { Address } from './protocol/protocol';
import * as Debug from 'debug';
import * as NConsul from 'consul';
import { Lookup } from './consul/lookup';
import { Map } from './types/lang';
import { assign } from './util';

const debug = Debug('consul');

export class Consul {
    useCheckEnv: boolean = true;
    options: Map<any>;

    public get(serAddress: Address, service: string, callback, env: string = 'prod', idc ?: string, action?: string, requestOptions ?: Map<any>) {
        if (!service) {
            return callback(new Error('unknown service name.'), null);
        }

        if (typeof env !== 'string') {
            this.useCheckEnv = false;
        }

        debug(`start_request host: ${serAddress.getHost} port: ${serAddress.getPort()}`);

        this.initOptions(serAddress, service, env, idc);
        
        if (action === 'catelog') {
            this.requestWithCatelog(callback);
        } else {
            this.requestWithLookup(requestOptions, callback);
        }
    }

    private initOptions(serAddress: Address, service: string, env: string = 'prod', idc ?: string) {
        this.options = {
            address: serAddress,
            service: service,
            env: env,
            idc: idc
        };
    }

    private requestWithCatelog(callback) {
        const consul  = NConsul({
            host: this.options.address.getHost(),
            port: this.options.address.getPort()
        });
        
        var __this = this;

        consul.catalog.service.nodes(
            {
                service: __this.options.service
            }, function (err, nodes) {
                if (err) return callback(err);

                var hosts: Array<Address> = [];
                nodes.forEach(node => {
                    if (__this.useCheckEnv && !__this.checkEnv(node.ServiceTags, __this.options.env)) {
                        return;
                    }

                    hosts.push(new Address(
                        node.Address || node.ServiceAddress,
                        node.ServicePort
                    ));
                });

                callback(null, hosts);
            }
        );     
    }

    private requestWithLookup(requestOptions: Map<any>, callback) {
        let debug = this.options.env !== 'prod';
        let lookup = new Lookup({debug: debug});
        let __this = this;

        lookup.request(assign({
            service: this.options.service,
            host: this.options.address.host,
            port: this.options.address.port
        }, requestOptions || {}), function (err, nodes) {
            if (err) {
                callback(err);
                return;
            }
            
            var hosts: Array<Address> = [];
            nodes.forEach(node => {
                if (__this.useCheckEnv && node.Tags && node.Tags.env) {
                    if (node.Tags.env !== __this.options.env) {
                        return;
                    }
                }

                hosts.push(new Address(
                    node.Host,
                    node.Port
                ));
            });

            callback(null, hosts);
        });
    }

    /**
     * tags maybe is null, @TODO
     */
    private checkEnv(tags: Array<string>, env: string) {
        if (Array.isArray(tags)) {
            for (let i = 0; i < tags.length; i++) {
                let tag = tags[i];
                if (tag.indexOf('env:') === 0 && tag === ('env:' + env)) {
                    return true;
                }
            }

            return false;
        }
        
        return true;
    }
}

export var consul = new Consul();