soa/service/zmq_endpoint.h

/* zmq_endpoint.h                                                  -*- C++ -*-
   Jeremy Barnes, 25 September 2012
   Copyright (c) 2012 Datacratic Inc.  All rights reserved.

   Endpoints for zeromq.
*/

#ifndef __service__zmq_endpoint_h__
#define __service__zmq_endpoint_h__

#include "named_endpoint.h"
#include "message_loop.h"
#include <set>
#include <type_traits>
#include "jml/utils/smart_ptr_utils.h"
#include "jml/utils/vector_utils.h"
#include <boost/make_shared.hpp>
#include "jml/arch/backtrace.h"
#include "jml/arch/timers.h"
#include "jml/arch/cmp_xchg.h"
#include "zmq_utils.h"

namespace Datacratic {


/*****************************************************************************/
/* ZMQ EVENT SOURCE                                                          */
/*****************************************************************************/

struct ZmqEventSource : public AsyncEventSource {

    typedef std::function<void (std::vector<std::string>)>
        AsyncMessageHandler;
    AsyncMessageHandler asyncMessageHandler;

    typedef std::function<std::vector<std::string> (std::vector<std::string>)>
        SyncMessageHandler;
    SyncMessageHandler syncMessageHandler;

    typedef std::mutex SocketLock;

    ZmqEventSource();

    ZmqEventSource(zmq::socket_t & socket, SocketLock * lock = nullptr);

    template<typename T>
    ZmqEventSource(zmq::socket_t & socket,
                   const T & handler,
                   SocketLock * lock = nullptr,
                   typename std::enable_if<!std::is_convertible<decltype(std::declval<T>()(std::declval<std::vector<std::string> >())),
                                                                std::vector<std::string> >::value, void>::type * = 0)
        : asyncMessageHandler(handler)
    {
        init(socket, lock);
    }

    template<typename T>
    ZmqEventSource(zmq::socket_t & socket,
                   const T & handler,
                   SocketLock * lock = nullptr,
                   typename std::enable_if<std::is_convertible<decltype(std::declval<T>()(std::declval<std::vector<std::string> >())),
                                                                std::vector<std::string> >::value, void>::type * = 0)
        : syncMessageHandler(handler)
    {
        init(socket, lock);
    }

    void init(zmq::socket_t & socket, SocketLock * lock = nullptr);

    virtual int selectFd() const;

    virtual bool poll() const;

    virtual bool processOne();

    virtual void handleMessage(const std::vector<std::string> & message);

    virtual std::vector<std::string>
    handleSyncMessage(const std::vector<std::string> & message);

    zmq::socket_t & socket() const
    {
        ExcAssert(socket_);
        return *socket_;
    }

    SocketLock * socketLock() const
    {
        return socketLock_;
    }

    zmq::socket_t * socket_;

    SocketLock * socketLock_;
};


/*****************************************************************************/
/* ZMQ BINARY EVENT SOURCE                                                   */
/*****************************************************************************/

struct ZmqBinaryEventSource : public AsyncEventSource {

    typedef std::function<void (std::vector<zmq::message_t> &&)>
        MessageHandler;
    MessageHandler messageHandler;

    ZmqBinaryEventSource()
        : socket_(0)
    {
        needsPoll = true;
    }

    ZmqBinaryEventSource(zmq::socket_t & socket,
                         MessageHandler messageHandler = MessageHandler())
        : messageHandler(std::move(messageHandler)),
          socket_(&socket)
    {
        needsPoll = true;
    }

    void init(zmq::socket_t & socket)
    {
        socket_ = &socket;
        needsPoll = true;
    }

    virtual int selectFd() const
    {
        int res = -1;
        size_t resSize = sizeof(int);
        socket().getsockopt(ZMQ_FD, &res, &resSize);
        if (res == -1)
            throw ML::Exception("no fd for zeromq socket");
        return res;
    }

    virtual bool poll() const
    {
        return getEvents(socket()).first;
    }

    virtual bool processOne()
    {
        ExcAssert(socket_);

        std::vector<zmq::message_t> messages;

        int64_t more = 1;
        size_t more_size = sizeof (more);

        while (more) {
            zmq::message_t message;
            bool got = socket_->recv(&message, messages.empty() ? ZMQ_NOBLOCK: 0);
            if (!got) return false;  // no first part available
            messages.emplace_back(std::move(message));
            socket_->getsockopt(ZMQ_RCVMORE, &more, &more_size);
        }

        handleMessage(std::move(messages));

        return poll();
    }

    virtual void handleMessage(std::vector<zmq::message_t> && message)
    {
        if (messageHandler)
            messageHandler(std::move(message));
        else throw ML::Exception("need to override handleMessage");
    }

    zmq::socket_t & socket() const
    {
        ExcAssert(socket_);
        return *socket_;
    }

    zmq::socket_t * socket_;

};


/*****************************************************************************/
/* ZMQ BINARY TYPED EVENT SOURCE                                             */
/*****************************************************************************/

template<typename Arg>
struct ZmqBinaryTypedEventSource: public AsyncEventSource {

    typedef std::function<void (Arg)> MessageHandler;
    MessageHandler messageHandler;

    ZmqBinaryTypedEventSource()
        : socket_(0)
    {
        needsPoll = true;
    }

    ZmqBinaryTypedEventSource(zmq::socket_t & socket,
                              MessageHandler messageHandler = MessageHandler())
        : messageHandler(std::move(messageHandler)),
          socket_(&socket)
    {
        needsPoll = true;
    }

    void init(zmq::socket_t & socket)
    {
        socket_ = &socket;
        needsPoll = true;
    }

    virtual int selectFd() const
    {
        int res = -1;
        size_t resSize = sizeof(int);
        socket().getsockopt(ZMQ_FD, &res, &resSize);
        if (res == -1)
            throw ML::Exception("no fd for zeromq socket");
        return res;
    }

    virtual bool poll() const
    {
        return getEvents(socket()).first;
    }

#if 0
    template<typename Arg, int Index>
    const Arg &
    getArg(const std::vector<zmq::message_t> & messages,
           const ML::InPosition<Arg, Index> * arg)
    {
        auto & m = messages.at(Index);
        ExcAssertEqual(m.size(), sizeof(Arg));
        return * reinterpret_cast<const Arg *>(msg.data());
    }
#endif

    virtual bool processOne()
    {
        zmq::message_t message;
        ExcAssert(socket_);
        bool got = socket_->recv(&message, ZMQ_NOBLOCK);
        if (!got) return false;

        handleMessage(* reinterpret_cast<const Arg *>(message.data()));

        return poll();
    }

    virtual void handleMessage(const Arg & arg)
    {
        if (messageHandler)
            messageHandler(arg);
        else throw ML::Exception("handleMessage not done");
    }

    zmq::socket_t & socket() const
    {
        ExcAssert(socket_);
        return *socket_;
    }

    zmq::socket_t * socket_;
};


/*****************************************************************************/
/* ZMQ TYPED EVENT SOURCE                                                    */
/*****************************************************************************/

template<typename T>
struct ZmqTypedEventSource: public ZmqEventSource {
    typedef std::function<void (T &&, std::string address)>
      OnMessage;

    OnMessage onMessage;
    bool routable;
    std::string messageTopic;

    ZmqTypedEventSource()
    {
    }

    ZmqTypedEventSource(zmq::socket_t & socket,
                        bool routable,
                        const std::string & messageTopic)
    {
        init(routable, messageTopic);
    }

    void init(zmq::socket_t & socket,
              bool routable,
              const std::string messageTopic)
    {
        this->routable = routable;
        this->messageTopic = messageTopic;
    }

    virtual void handleMessage(const std::vector<std::string> & message)
    {
        int expectedSize = routable + 2;
        if (message.size() != expectedSize)
            throw ML::Exception("unexpected message size in ZmqTypedMessageSink");

        int i = routable;
        if (message[i + 1] != messageTopic)
            throw ML::Exception("unexpected messake kind in ZmqTypedMessageSink");

        std::istringstream stream(message[i + 2]);
        ML::DB::Store_Reader store(stream);
        T result;
        store >> result;

        handleTypedMessage(std::move(result), routable ? message[0] : "");
    }

    virtual void handleTypedMessage(T && message, const std::string & address)
    {
        if (onMessage)
            onMessage(message, address);
        else
            throw ML::Exception("need to override handleTypedMessage or assign "
                                "to onMessage");
    }
};


/*****************************************************************************/
/* ZMQ SOCKET MONITOR                                                        */
/*****************************************************************************/

struct ZmqSocketMonitor : public ZmqBinaryTypedEventSource<zmq_event_t> {

    ZmqSocketMonitor(zmq::context_t & context);

    ~ZmqSocketMonitor()
    {
        shutdown();
    }

    void shutdown();

    void init(zmq::socket_t & socketToMonitor, int events = ZMQ_EVENT_ALL);

    typedef std::function<void (std::string, int, zmq_event_t)> EventHandler;

    // Success handlers
    EventHandler connectHandler, bindHandler, acceptHandler;

    // Socket event handlers
    EventHandler closeHandler, disconnectHandler;

    // Failure handlers
    EventHandler connectFailureHandler, acceptFailureHandler,
        bindFailureHandler, closeFailureHandler;

    // Retry handlers
    EventHandler connectRetryHandler;

    // Catch all handler, for when other handlers aren't registered
    EventHandler defaultHandler;

    virtual int handleEvent(const zmq_event_t & event);

private:
    typedef std::mutex Lock;
    mutable Lock lock;

    std::string connectedUri;

    std::unique_ptr<zmq::socket_t> monitorEndpoint;

    zmq::socket_t * monitoredSocket;
};


/*****************************************************************************/
/* ZEROMQ NAMED ENDPOINT                                                     */
/*****************************************************************************/

struct ZmqNamedEndpoint : public NamedEndpoint, public MessageLoop {

    ZmqNamedEndpoint(std::shared_ptr<zmq::context_t> context);

    ~ZmqNamedEndpoint()
    {
        shutdown();
    }

    void init(std::shared_ptr<ConfigurationService> config,
              int socketType,
              const std::string & endpointName);

    void shutdown()
    {
        MessageLoop::shutdown();

        if (socket_) {
            unbindAll();
            socket_.reset();
        }

        //ML::sleep(0.1);
        monitor.shutdown();
    }

    std::string bindTcp(PortRange const & portRange = PortRange(), std::string host = "");

    void bind(const std::string & address)
    {
        if (!socket_)
            throw ML::Exception("need to call ZmqNamedEndpoint::init() before "
                                "bind");

        std::unique_lock<Lock> guard(lock);
        socket_->bind(address);
        boundAddresses[address];
    }

    void unbindAll()
    {
        std::unique_lock<Lock> guard(lock);
        ExcAssert(socket_);
        for (auto addr: boundAddresses)
            socket_->tryUnbind(addr.first);
    }

    template<typename... Args>
    void sendMessage(Args&&... args)
    {
        using namespace std;
        std::unique_lock<Lock> guard(lock);
        ExcAssert(socket_);
        Datacratic::sendMessage(*socket_, std::forward<Args>(args)...);
    }

    void sendMessage(const std::vector<std::string> & message)
    {
        using namespace std;
        std::unique_lock<Lock> guard(lock);
        ExcAssert(socket_);
        Datacratic::sendAll(*socket_, message);
    }

    void sendMessage(std::vector<zmq::message_t> && message)
    {
        using namespace std;
        std::unique_lock<Lock> guard(lock);
        ExcAssert(socket_);
        for (unsigned i = 0;  i < message.size();  ++i) {
            socket_->send(message[i],
                          i == message.size() - 1
                          ? 0 : ZMQ_SNDMORE);
        }
    }

    zmq::socket_t & getSocketUnsafe() const
    {
        ExcAssert(socket_);
        return *socket_;
    }

    typedef std::function<void (std::vector<zmq::message_t> &&)>
        RawMessageHandler;
    RawMessageHandler rawMessageHandler;

    typedef std::function<void (std::vector<std::string> &&)> MessageHandler;
    MessageHandler messageHandler;

    virtual void handleRawMessage(std::vector<zmq::message_t> && message)
    {
        if (rawMessageHandler)
            rawMessageHandler(std::move(message));
        else {
            std::vector<std::string> msg2;
            for (unsigned i = 0;  i < message.size();  ++i) {
                msg2.push_back(message[i].toString());
            }
            handleMessage(std::move(msg2));
        }
    }

    virtual void handleMessage(std::vector<std::string> && message)
    {
        if (messageHandler)
            messageHandler(std::move(message));
        else throw ML::Exception("need to override handleRawMessage or "
                                 "handleMessage");
    }

    typedef std::function<void (std::string bindAddress)>
        ConnectionEventHandler;

    ConnectionEventHandler acceptEventHandler;

    ConnectionEventHandler disconnectEventHandler;

    ConnectionEventHandler closeEventHandler;

    virtual void handleAcceptEvent(std::string boundAddress)
    {
        if (acceptEventHandler)
            acceptEventHandler(boundAddress);
    }

    virtual void handleDisconnectEvent(std::string boundAddress)
    {
        if (disconnectEventHandler)
            disconnectEventHandler(boundAddress);
    }

    virtual void handleCloseEvent(std::string boundAddress)
    {
        if (closeEventHandler)
            closeEventHandler(boundAddress);
    }

    size_t numBoundAddresses() const
    {
        std::unique_lock<Lock> guard(lock);
        return boundAddresses.size();
    }

    size_t numActiveConnections(std::string addr = "") const
    {
        std::unique_lock<Lock> guard(lock);
        if (addr == "") {
            size_t result = 0;
            for (auto & addr: boundAddresses)
                result += addr.second.connectedFds.size();
            return result;
        }
        else {
            auto it = boundAddresses.find(addr);
            if (it == boundAddresses.end())
                return 0;
            return it->second.connectedFds.size();
        }
    }

private:
    typedef std::mutex Lock;
    mutable Lock lock;

    std::shared_ptr<zmq::context_t> context_;
    std::shared_ptr<zmq::socket_t> socket_;

    ZmqSocketMonitor monitor;

    struct AddressInfo {
        AddressInfo()
            : listeningFd(-1)
        {
        }

        int listeningFd;

        std::set<int> connectedFds;
    };

    std::map<std::string, AddressInfo> boundAddresses;

    int socketType;
};


/*****************************************************************************/
/* ZMQ NAMED CLIENT BUS                                                      */
/*****************************************************************************/

struct ZmqNamedClientBus: public ZmqNamedEndpoint {

    ZmqNamedClientBus(std::shared_ptr<zmq::context_t> context,
                      double deadClientDelay = 5.0)
        : ZmqNamedEndpoint(context), deadClientDelay(deadClientDelay)
    {
    }

    void init(std::shared_ptr<ConfigurationService> config,
              const std::string & endpointName)
    {
        ZmqNamedEndpoint::init(config, ZMQ_XREP, endpointName);
        addPeriodic("ZmqNamedClientBus::checkClient", 1.0,
                    [=] (uint64_t v) { this->onCheckClient(v); });
    }

    virtual ~ZmqNamedClientBus()
    {
        shutdown();
    }

    void shutdown()
    {
        MessageLoop::shutdown();
        ZmqNamedEndpoint::shutdown();
    }

    double deadClientDelay;

    std::function<void (std::string)> onConnection;

    std::function<void (std::string)> onDisconnection;



    template<typename... Args>
    void sendMessage(const std::string & address,
                     const std::string & topic,
                     Args&&... args)
    {
        ZmqNamedEndpoint::sendMessage(address, topic,
                                      std::forward<Args>(args)...);
    }

    virtual void handleMessage(std::vector<std::string> && message)
    {
        using namespace std;
        //cerr << "ZmqNamedClientBus got message " << message << endl;

        const std::string & agent = message.at(0);
        const std::string & topic = message.at(1);

        if (topic == "HEARTBEAT") {
            // Not the first message from the client
            auto it = clientInfo.find(agent);
            if (it == clientInfo.end()) {
                // Disconnection then reconnection
                if (onConnection)
                    onConnection(agent);
                it = clientInfo.insert(make_pair(agent, ClientInfo())).first;
            }
            it->second.lastHeartbeat = Date::now();
            sendMessage(agent, "HEARTBEAT");
        }
        else if (topic == "HELLO") {
            // First message from client
            auto it = clientInfo.find(agent);
            if (it == clientInfo.end()) {
                // New connection
                if (onConnection)
                    onConnection(agent);
                it = clientInfo.insert(make_pair(agent, ClientInfo())).first;
            }
            else {
                // Client must have disappeared then reappared without us
                // noticing.
                // Do this disconnection, then the reconnection
                if (onDisconnection)
                    onDisconnection(agent);
                if (onConnection)
                    onConnection(agent);
            }
            it->second.lastHeartbeat = Date::now();
            sendMessage(agent, "HEARTBEAT");
        }
        else {
            handleClientMessage(message);
        }

#if 0
        cerr << "poll() returned " << poll() << endl;

        std::vector<std::string> msg
            = recvAllNonBlocking(socket());

        while (!msg.empty()) {
            cerr << "*** GOT FURTHER MESSAGE " << msg << endl;
            msg = recvAllNonBlocking(socket());
        }

        cerr << "poll() returned " << poll() << endl;
#endif
    }

    typedef std::function<void (std::vector<std::string>)>
    ClientMessageHandler;
    ClientMessageHandler clientMessageHandler;

    virtual void handleClientMessage(const std::vector<std::string> & message)
    {
        if (clientMessageHandler)
            clientMessageHandler(message);
        else {
            throw ML::Exception("need to assign to onClientMessage "
                                "or override handleClientMessage for message "
                                + message.at(1));
        }
#if 0
        using namespace std;
        cerr << "ZmqNamedClientBus handleClientMessage " << message << endl;
        throw ML::Exception("handleClientMessage");
#endif
    }

private:
    void onCheckClient(uint64_t numEvents)
    {
        Date now = Date::now();
        Date expiry = now.plusSeconds(-deadClientDelay);

        std::vector<std::string> deadClients;

        for (auto & c: clientInfo)
            if (c.second.lastHeartbeat < expiry)
                deadClients.push_back(c.first);

        for (auto d: deadClients) {
            if (onDisconnection)
                onDisconnection(d);
            clientInfo.erase(d);
        }
    }

    struct ClientInfo {
        ClientInfo()
            : lastHeartbeat(Date::now())
        {
        }

        Date lastHeartbeat;
    };

    std::map<std::string, ClientInfo> clientInfo;
};


enum ConnectionStyle {
    CS_ASYNCHRONOUS,  
    CS_SYNCHRONOUS,   
    CS_MUST_SUCCEED   
};




/*****************************************************************************/
/* ZEROMQ NAMED PROXY                                                        */
/*****************************************************************************/

// THIS SHOULD BE REPLACED BY ZmqNamedSocket

struct ZmqNamedProxy: public MessageLoop {

    ZmqNamedProxy();

    ZmqNamedProxy(std::shared_ptr<zmq::context_t> context);

    ~ZmqNamedProxy()
    {
        shutdown();
    }

    void shutdown()
    {
        MessageLoop::shutdown();
        if(socket_) {
            std::lock_guard<ZmqEventSource::SocketLock> guard(socketLock_);
            socket_.reset();
        }
    }

    bool isConnected() const { return connectionState == CONNECTED; }

    typedef std::function<void (std::string)> ConnectionHandler;

    ConnectionHandler connectHandler;

    virtual void onConnect(const std::string & source)
    {
        if (connectHandler)
            connectHandler(source);
    }

    ConnectionHandler disconnectHandler;

    virtual void onDisconnect(const std::string & source)
    {
        if (disconnectHandler)
            disconnectHandler(source);
    }

    void init(std::shared_ptr<ConfigurationService> config,
              int socketType,
              const std::string & identity = "");

    bool connect(const std::string & endpointName,
                 ConnectionStyle style = CS_ASYNCHRONOUS);

    bool connectToServiceClass(const std::string & serviceClass,
                               const std::string & endpointName,
                               ConnectionStyle style = CS_ASYNCHRONOUS);

    bool onConfigChange(ConfigurationService::ChangeType change,
                        const std::string & key,
                        const Json::Value & newValue);

    zmq::socket_t & socket() const
    {
        ExcAssert(socket_);
        return *socket_;
    }

    ZmqEventSource::SocketLock * socketLock() const
    {
        return &socketLock_;
    }

    template<typename... Args>
    void sendMessage(Args&&... args)
    {
        std::lock_guard<ZmqEventSource::SocketLock> guard(socketLock_);

        ExcCheckNotEqual(connectionState, NOT_CONNECTED,
                "sending on an unconnected socket: " + endpointName);

        if (connectionState == CONNECTION_PENDING) {
            std::cerr << ("dropping message for " + endpointName + "\n");
            return;
        }

        Datacratic::sendMessage(socket(), std::forward<Args>(args)...);
    }

    void disconnect()
    {
        if (connectionState == NOT_CONNECTED) return;

        {
            std::lock_guard<ZmqEventSource::SocketLock> guard(socketLock_);

            if (connectionState == CONNECTED)
                socket_->disconnect(connectedUri);

            connectionState = NOT_CONNECTED;
        }

        onDisconnect(connectedUri);
    }


protected:
    ConfigurationService::Watch serviceWatch, endpointWatch;
    std::shared_ptr<ConfigurationService> config;
    std::shared_ptr<zmq::context_t> context_;
    std::shared_ptr<zmq::socket_t> socket_;

    mutable ZmqEventSource::SocketLock socketLock_;

    enum ConnectionType {
        NO_CONNECTION,        
        CONNECT_DIRECT,       
        CONNECT_TO_CLASS,     
    } connectionType;

    enum ConnectionState {
        NOT_CONNECTED,      // connect() was not called
        CONNECTION_PENDING, // connect() was called but service is not available
        CONNECTED           // connect() was called and the socket was connected
    } connectionState;

    void onServiceNodeChange(const std::string & path,
                             ConfigurationService::ChangeType change);
    void onEndpointNodeChange(const std::string & path,
                              ConfigurationService::ChangeType change);

    std::string serviceClass;      
    std::string endpointName;      
    std::string connectedService;  
    std::string connectedUri;      
};


/*****************************************************************************/
/* ZEROMQ NAMED CLIENT BUS PROXY                                             */
/*****************************************************************************/

struct ZmqNamedClientBusProxy : public ZmqNamedProxy {

    ZmqNamedClientBusProxy()
        : timeout(2.0)
    {
    }

    ZmqNamedClientBusProxy(std::shared_ptr<zmq::context_t> context)
        : ZmqNamedProxy(context), timeout(2.0)
    {
    }

    ~ZmqNamedClientBusProxy()
    {
        shutdown();
    }

    void init(std::shared_ptr<ConfigurationService> config,
              const std::string & identity = "")
    {
        ZmqNamedProxy::init(config, ZMQ_XREQ, identity);

        auto doMessage = [=] (const std::vector<std::string> & message)
            {
                const std::string & topic = message.at(0);
                if (topic == "HEARTBEAT")
                    this->lastHeartbeat = Date::now();
                else handleMessage(message);
            };

        addSource("ZmqNamedClientBusProxy::doMessage",
                  std::make_shared<ZmqEventSource>(socket(), doMessage, socketLock()));
 
        auto doHeartbeat = [=] (int64_t skipped)
            {
                if (connectionState != CONNECTED) return;

                sendMessage("HEARTBEAT");

                auto now = Date::now();
                auto end = now.plusSeconds(-timeout);
                //if(lastHeartbeat < end) {
                    //std::cerr << "no heartbeat for " << timeout << "s... should be disconnecting from " << connectedUri << std::endl;
                    //disconnect();
                //}
            };

        addPeriodic("ZmqNamedClientBusProxy::doHeartbeat", 1.0, doHeartbeat);
   }

    void shutdown()
    {
        MessageLoop::shutdown();
        ZmqNamedProxy::shutdown();
    }

    virtual void onConnect(const std::string & where)
    {
        lastHeartbeat = Date::now();

        sendMessage("HELLO");

        if (connectHandler)
            connectHandler(where);
    }

    virtual void onDisconnect(const std::string & where)
    {
        if (disconnectHandler)
            disconnectHandler(where);
    }

    ZmqEventSource::AsyncMessageHandler messageHandler;

    virtual void handleMessage(const std::vector<std::string> & message)
    {
        if (messageHandler)
            messageHandler(message);
        else
            throw ML::Exception("need to override on messageHandler or handleMessage");
    }

    Date lastHeartbeat;
    double timeout;
};


/*****************************************************************************/
/* ZEROMQ MULTIPLE NAMED CLIENT BUS PROXY                                    */
/*****************************************************************************/

struct ZmqMultipleNamedClientBusProxy: public MessageLoop {

    ZmqMultipleNamedClientBusProxy()
        : zmqContext(new zmq::context_t(1))
    {
        connected = false;
    }

    ZmqMultipleNamedClientBusProxy(std::shared_ptr<zmq::context_t> context)
        : zmqContext(context)
    {
        connected = false;
    }

    ~ZmqMultipleNamedClientBusProxy()
    {
        shutdown();
    }

    void init(std::shared_ptr<ConfigurationService> config,
              const std::string & identity = "")
    {
        this->config = config;
        this->identity = identity;
    }

    void shutdown()
    {
        MessageLoop::shutdown();
        for (auto & c: connections)
            if (c.second)
                c.second->shutdown();
    }

    template<typename... Args>
    void sendMessage(const std::string & recipient,
                     const std::string & topic,
                     Args&&... args) const
    {
        std::unique_lock<Lock> guard(connectionsLock);
        auto it = connections.find(recipient);
        if (it == connections.end()) {
            throw ML::Exception("attempt to deliver " + topic + " message to unknown client "
                                + recipient);
        }
        it->second->sendMessage(topic, std::forward<Args>(args)...);
    }

    void connectAllServiceProviders(const std::string & serviceClass,
                                    const std::string & endpointName)
    {
        if (connected)
            throw ML::Exception("alread connected to service providers");

        this->serviceClass = serviceClass;
        this->endpointName = endpointName;

        serviceProvidersWatch.init([=] (const std::string & path,
                                        ConfigurationService::ChangeType change)
                                   {
                                       onServiceProvidersChanged("serviceClass/" + serviceClass);
                                   });

        onServiceProvidersChanged("serviceClass/" + serviceClass);
        // std::cerr << "++++after call to onServiceProvidersChanged " << std::endl;
        connected = true;
    }


    void connectSingleServiceProvider(const std::string & service);

    void connectUri(const std::string & zmqUri);

    size_t connectionCount() const
    {
        std::unique_lock<Lock> guard(connectionsLock);
        return connections.size();
    }

    typedef std::function<void (std::string)> ConnectionHandler;

    ConnectionHandler connectHandler;

    virtual void onConnect(const std::string & source)
    {
        if (connectHandler)
            connectHandler(source);
    }

    ConnectionHandler disconnectHandler;

    virtual void onDisconnect(const std::string & source)
    {
        if (disconnectHandler)
            disconnectHandler(source);
    }

    typedef std::function<void (std::string, std::vector<std::string>)>
    MessageHandler;

    MessageHandler messageHandler;

    virtual void handleMessage(const std::string & source,
                               const std::vector<std::string> & message)
    {
        if (messageHandler)
            messageHandler(source, message);
        else
            throw ML::Exception("need to override on messageHandler or handleMessage");
    }

private:
    bool connected;

    std::shared_ptr<zmq::context_t> zmqContext;

    std::shared_ptr<ConfigurationService> config;

    std::string serviceClass;

    std::string endpointName;

    std::string identity;

    typedef ML::Spinlock Lock;

    mutable Lock connectionsLock;

    std::map<std::string, std::unique_ptr<ZmqNamedClientBusProxy> > connections;

    ConfigurationService::Watch serviceProvidersWatch;

    void onServiceProvidersChanged(const std::string & path)
    {
        using namespace std;
        //cerr << "onServiceProvidersChanged(" << path << ")" << endl;

        // The list of service providers has changed

        vector<string> children
            = config->getChildren(path, serviceProvidersWatch);
        for (auto c: children) {
            Json::Value value = config->getJson(path + "/" + c);
            std::string name = value["serviceName"].asString();
            std::string path = value["servicePath"].asString();

            watchServiceProvider(name, path);
        }
    }

    struct OnConnectCallback
    {
        OnConnectCallback(const ConnectionHandler& fn, std::string name) :
            fn(fn), name(name), state(DEFER)
        {}

        void release()
        {
            State old = state;

            ExcAssertNotEqual(old, CALL);

            // If the callback wasn't triggered while we were holding the lock
            // then trigger it the next time we see it.
            if (old == DEFER && ML::cmp_xchg(state, old, CALL)) return;

            ExcAssertEqual(old, DEFERRED);
            fn(name);
        }

        void operator() (std::string)
        {
            State old = state;
            ExcAssertNotEqual(old, DEFERRED);

            // If we're still in the locked section then trigger the callback
            // when release is called.
            if (old == DEFER && ML::cmp_xchg(state, old, DEFERRED)) return;

            // We're out of the locked section so just trigger the callback.
            ExcAssertEqual(old, CALL);
            fn(name);
        }

    private:

        ConnectionHandler fn;
        std::string name;

        enum State {
            DEFER,    // We're holding the lock so defer an incoming callback.
            DEFERRED, // We were called while holding the lock.
            CALL      // We were not called while holding the lock.
        } state;
    };

    void watchServiceProvider(const std::string & name, const std::string & path)
    {
        // Protects the connections map... I think.
        std::unique_lock<Lock> guard(connectionsLock);

        auto & c = connections[name];

        // already connected
        if (c) return;

        try {
            std::unique_ptr<ZmqNamedClientBusProxy> newClient
                (new ZmqNamedClientBusProxy(zmqContext));
            newClient->init(config, identity);

            // The connect call below could trigger this callback while we're
            // holding the connectionsLock which is a big no-no. This fancy
            // wrapper ensures that it's only called after we call its release
            // function.
            newClient->connectHandler = OnConnectCallback(connectHandler, name);

            newClient->disconnectHandler = [=] (std::string s)
                {
                    // TODO: chain in so that we know it's not around any more
                    this->onDisconnect(s);
                };

            newClient->connect(path + "/" + endpointName);
            newClient->messageHandler = [=] (const std::vector<std::string> & msg)
                {
                    this->handleMessage(name, msg);
                };
            //newClient->debug(true);

            c.reset(newClient.release());

            // Add it to our message loop so that it can process messages
            addSource("ZmqMultipleNamedClientBusProxy child " + name, *c);

            guard.unlock();
            c->connectHandler.target<OnConnectCallback>()->release();

        } catch (...) {
            connections.erase(name);
            throw;
        }
    }

};



} // namespace Datacratic

#endif /* __service__zmq_endpoint_h__ */