path: root/contrib/native/client/src/clientlib/drillClientImpl.hpp

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */


#ifndef DRILL_CLIENT_IMPL_H
#define DRILL_CLIENT_IMPL_H

#include "drill/common.hpp"
// Define some BOOST defines
// WIN32_SHUTDOWN_ON_TIMEOUT is defined in "drill/common.hpp" for Windows 32 bit platform
#ifndef WIN32_SHUTDOWN_ON_TIMEOUT
#define BOOST_ASIO_ENABLE_CANCELIO
#endif //WIN32_SHUTDOWN_ON_TIMEOUT

#include <algorithm>
#include <queue>
#include <vector>

#include <boost/asio.hpp>
#if defined _WIN32  || defined _WIN64
//Windows header files redefine 'random'
#ifdef random
#undef random
#endif
#endif
#include <boost/asio/deadline_timer.hpp>
#include <boost/function.hpp>
#include <boost/thread.hpp>

#include "drill/drillClient.hpp"
#include "drill/preparedStatement.hpp"
#include "collectionsImpl.hpp"
#include "metadata.hpp"
#include "rpcMessage.hpp"
#include "utils.hpp"
#include "User.pb.h"
#include "UserBitShared.pb.h"
#include "saslAuthenticatorImpl.hpp"

namespace Drill {

class DrillClientImpl;

class DrillClientQueryHandle;

class DrillClientPrepareHandle;
class RecordBatch;

/*
 * Defines the interface used by DrillClient and implemented by DrillClientImpl and PooledDrillClientImpl
 * */
class DrillClientImplBase{
    public:
        DrillClientImplBase(){
        }

        virtual ~DrillClientImplBase(){
        }

        //Connect via Zookeeper or directly.
        //Makes an initial connection to a drillbit. successful connect adds the first drillbit to the pool.
        virtual connectionStatus_t connect(const char* connStr, DrillUserProperties* props)=0;

        // Test whether the client is active. Returns true if any one of the underlying connections is active
        virtual bool Active()=0;

        // Closes all open connections. 
        virtual void Close()=0;

        // Returns the last error encountered by any of the underlying executing queries or connections
        virtual DrillClientError* getError()=0;

        // Submits a query to a drillbit. 
        virtual DrillClientQueryResult* SubmitQuery(::exec::shared::QueryType t, const std::string& plan, pfnQueryResultsListener listener, void* listenerCtx)=0;
        virtual DrillClientPrepareHandle* PrepareQuery(const std::string& plan, pfnPreparedStatementListener listener, void* listenerCtx)=0;
        virtual DrillClientQueryResult* ExecuteQuery(const PreparedStatement& pstmt, pfnQueryResultsListener listener, void* listenerCtx)=0;

        //Waits as a connection has results pending
        virtual void waitForResults()=0;

        //Validates handshake at connect time.
        virtual connectionStatus_t validateHandshake(DrillUserProperties* props)=0;

        virtual void freeQueryResources(DrillClientQueryHandle* pQryHandle)=0;

        virtual meta::DrillMetadata* getMetadata() = 0;

        virtual void freeMetadata(meta::DrillMetadata* metadata) = 0;
};

/**
 * Base type for query handles
 */
class DrillClientQueryHandle{
    friend class DrillClientImpl;
    public:
    DrillClientQueryHandle(DrillClientImpl& client, int32_t coordId, const std::string& query, void* context):
        m_client(client),
        m_coordinationId(coordId),
        m_query(query),
		m_status(QRY_SUCCESS),
        m_bCancel(false),
        m_bHasError(false),
        m_pError(NULL),
		m_pApplicationContext(context){
    };

    virtual ~DrillClientQueryHandle(){
        clearAndDestroy();
    };

    virtual void cancel();
    bool isCancelled() const {return m_bCancel;};
    int32_t getCoordinationId() const { return m_coordinationId;}
    const std::string&  getQuery() const { return m_query;}

    bool hasError() const { return m_bHasError;}
    void resetError() { m_bHasError = false; }

    status_t getErrorStatus() const { return m_pError!=NULL?(status_t)m_pError->status:QRY_SUCCESS;}
    const DrillClientError* getError() const { return m_pError;}
    void setQueryStatus(status_t s){ m_status = s;}
    status_t getQueryStatus() const { return m_status;}
    inline DrillClientImpl& client() const { return m_client; };

    inline void* getApplicationContext() const { return m_pApplicationContext; }

    protected:

    virtual void signalError(DrillClientError* pErr);
    virtual void clearAndDestroy();

    private:
    DrillClientImpl& m_client;

    int32_t m_coordinationId;
    std::string m_query;
    status_t m_status;
    bool m_bCancel;
    bool m_bHasError;

    const DrillClientError* m_pError;

    void* m_pApplicationContext;
};

template<typename Listener, typename ListenerValue>
class DrillClientBaseHandle: public DrillClientQueryHandle {
    friend class DrillClientImpl;
    public:
    DrillClientBaseHandle(DrillClientImpl& client, int32_t coordId, const std::string& query, Listener listener, void* context):
    	DrillClientQueryHandle(client, coordId, query, context),
		m_pApplicationListener(listener){
    };

    virtual ~DrillClientBaseHandle(){
        clearAndDestroy();
    };

    inline Listener getApplicationListener() const { return m_pApplicationListener; }


    protected:
    virtual status_t notifyListener(ListenerValue v, DrillClientError* pErr);

    virtual void signalError(DrillClientError* pErr);
    void setHasError(bool hasError) { m_bHasError = hasError; }

    private:
    Listener m_pApplicationListener;
};

class DrillClientQueryResult: public DrillClientBaseHandle<pfnQueryResultsListener, RecordBatch*>{
    friend class DrillClientImpl;
    public:
    DrillClientQueryResult(DrillClientImpl& client, int32_t coordId, const std::string& query, pfnQueryResultsListener listener, void* listenerCtx):
    	DrillClientBaseHandle<pfnQueryResultsListener, RecordBatch*>(client, coordId, query, listener, listenerCtx),
        m_numBatches(0),
        m_columnDefs(new std::vector<Drill::FieldMetadata*>),
        m_bIsQueryPending(true),
        m_bIsLastChunk(false),
        m_bHasSchemaChanged(false),
        m_bHasData(false),
        m_queryState(exec::shared::QueryResult_QueryState_STARTING),
        m_pQueryId(NULL),
        m_pSchemaListener(NULL) {
    };

    ~DrillClientQueryResult(){
        this->clearAndDestroy();
    };

    // get data asynchronously
    void registerSchemaChangeListener(pfnSchemaListener l){
        m_pSchemaListener=l;
    }

    // Synchronous call to get data. Caller assumes ownership of the record batch
    // returned and it is assumed to have been consumed.
    RecordBatch*  getNext();
    // Synchronous call to get a look at the next Record Batch. This
    // call does not move the current pointer forward. Repeated calls
    // to peekNext return the same value until getNext is called.
    RecordBatch*  peekNext();
    // Blocks until data is available.
    void waitForData();

    // placeholder to return an empty col def vector when calls are made out of order.
    static FieldDefPtr s_emptyColDefs;

    FieldDefPtr getColumnDefs() {
        boost::lock_guard<boost::mutex> bufferLock(this->m_schemaMutex);
        return this->m_columnDefs;
    }

    bool hasSchemaChanged() const {return this->m_bHasSchemaChanged;};

    void setQueryId(exec::shared::QueryId* q){this->m_pQueryId=q;}
    exec::shared::QueryId& getQueryId() const { return *(this->m_pQueryId); }

    void setQueryState(exec::shared::QueryResult_QueryState s){ m_queryState = s;}
    exec::shared::QueryResult_QueryState getQueryState() const { return m_queryState;}
    void setIsQueryPending(bool isPending){
        boost::lock_guard<boost::mutex> cvLock(this->m_cvMutex);
        m_bIsQueryPending=isPending;
    }
    protected:
    virtual status_t notifyListener(RecordBatch* batch, DrillClientError* pErr);
    virtual void signalError(DrillClientError* pErr);
    virtual void clearAndDestroy();

    private:
    status_t setupColumnDefs(exec::shared::QueryData* pQueryData);
    status_t defaultQueryResultsListener(void* ctx, RecordBatch* b, DrillClientError* err);
    // Construct a DrillClientError object, set the appropriate state and signal any listeners, condition variables.
    // Also used when a query is cancelled or when a query completed response is received.
    // Error object is now owned by the DrillClientQueryResult object.
    void signalComplete();

    size_t m_numBatches; // number of record batches received so far

    // Vector of Buffers holding data returned by the server
    // Each data buffer is decoded into a RecordBatch
    std::vector<ByteBuf_t> m_dataBuffers;
    std::queue<RecordBatch*> m_recordBatches;
    FieldDefPtr m_columnDefs;

    // Mutex to protect schema definitions
    boost::mutex m_schemaMutex;
    // Mutex for Cond variable for read write to batch vector
    boost::mutex m_cvMutex;
    // Condition variable to signal arrival of more data. Condition variable is signaled
    // if the recordBatches queue is not empty
    boost::condition_variable m_cv;

    // state
    // if m_bIsQueryPending is true, we continue to wait for results
    bool m_bIsQueryPending;
    bool m_bIsLastChunk;
    bool m_bHasSchemaChanged;
    bool m_bHasData;

    // state in the last query result received from the server.
    exec::shared::QueryResult_QueryState m_queryState;

    exec::shared::QueryId* m_pQueryId;

    // Schema change listener
    pfnSchemaListener m_pSchemaListener;
};

class DrillClientPrepareHandle: public DrillClientBaseHandle<pfnPreparedStatementListener, PreparedStatement*>, public PreparedStatement {
    public:
    DrillClientPrepareHandle(DrillClientImpl& client, int32_t coordId, const std::string& query, pfnPreparedStatementListener listener, void* listenerCtx):
    	DrillClientBaseHandle<pfnPreparedStatementListener, PreparedStatement*>(client, coordId, query, listener, listenerCtx),
		PreparedStatement(),
        m_columnDefs(new std::vector<Drill::FieldMetadata*>) {
    };

    // PreparedStatement overrides
	virtual std::size_t getNumFields() const { return m_columnDefs->size(); }
	virtual const Drill::FieldMetadata& getFieldMetadata(std::size_t index) const { return *m_columnDefs->at(index);}

    protected:
    virtual void clearAndDestroy();

    private:
    friend class DrillClientImpl;
    status_t setupPreparedStatement(const exec::user::PreparedStatement& pstmt);

    FieldDefPtr m_columnDefs;
    ::exec::user::PreparedStatementHandle m_preparedStatementHandle;
};

template<typename Listener, typename MetaType, typename MetaImpl, typename MetadataResult>
class DrillClientMetadataResult: public DrillClientBaseHandle<Listener, const DrillCollection<MetaType>*> {
public:
    DrillClientMetadataResult(DrillClientImpl& client, int32_t coordId, const std::string& query, Listener listener, void* listenerCtx):
    	DrillClientBaseHandle<Listener, const DrillCollection<MetaType>*>(client, coordId, query, listener, listenerCtx) {}

    void attachMetadataResult(MetadataResult* result) { this->m_pMetadata.reset(result); }

private:
    friend class DrillClientImpl;

    // Meta informations returned to the user, linked to the handle
    DrillVector<MetaType, MetaImpl> m_meta;

    // to keep a reference on the underlying metadata object, and
    // make sure it's clean when this handle is destroyed
    boost::shared_ptr<MetadataResult> m_pMetadata;

};

class DrillClientCatalogResult: public DrillClientMetadataResult<Metadata::pfnCatalogMetadataListener, meta::CatalogMetadata, meta::DrillCatalogMetadata, exec::user::GetCatalogsResp> {
    friend class DrillClientImpl;
public:
    DrillClientCatalogResult(DrillClientImpl& client, int32_t coordId, Metadata::pfnCatalogMetadataListener listener, void* listenerCtx):
    	DrillClientMetadataResult<Metadata::pfnCatalogMetadataListener, meta::CatalogMetadata, meta::DrillCatalogMetadata, exec::user::GetCatalogsResp>(client, coordId, "getCatalog", listener, listenerCtx) {}
};

class DrillClientSchemaResult: public DrillClientMetadataResult<Metadata::pfnSchemaMetadataListener, meta::SchemaMetadata, meta::DrillSchemaMetadata, exec::user::GetSchemasResp> {
    friend class DrillClientImpl;
public:
    DrillClientSchemaResult(DrillClientImpl& client, int32_t coordId, Metadata::pfnSchemaMetadataListener listener, void* listenerCtx):
    	DrillClientMetadataResult<Metadata::pfnSchemaMetadataListener, meta::SchemaMetadata, meta::DrillSchemaMetadata, exec::user::GetSchemasResp>(client, coordId, "getSchemas", listener, listenerCtx) {}
};

class DrillClientTableResult: public DrillClientMetadataResult<Metadata::pfnTableMetadataListener, meta::TableMetadata, meta::DrillTableMetadata, exec::user::GetTablesResp> {
    friend class DrillClientImpl;
public:
    DrillClientTableResult(DrillClientImpl& client, int32_t coordId, Metadata::pfnTableMetadataListener listener, void* listenerCtx):
    	DrillClientMetadataResult<Metadata::pfnTableMetadataListener, meta::TableMetadata, meta::DrillTableMetadata, exec::user::GetTablesResp>(client, coordId, "getTables", listener, listenerCtx) {}
};

class DrillClientColumnResult: public DrillClientMetadataResult<Metadata::pfnColumnMetadataListener, meta::ColumnMetadata, meta::DrillColumnMetadata, exec::user::GetColumnsResp> {
    friend class DrillClientImpl;
    public:
    DrillClientColumnResult(DrillClientImpl& client, int32_t coordId, Metadata::pfnColumnMetadataListener listener, void* listenerCtx):
    	DrillClientMetadataResult<Metadata::pfnColumnMetadataListener, meta::ColumnMetadata, meta::DrillColumnMetadata, exec::user::GetColumnsResp>(client, coordId, "getColumns", listener, listenerCtx) {}
};


class DrillClientImpl : public DrillClientImplBase{
    public:
        DrillClientImpl():
            m_coordinationId(1),
            m_handshakeVersion(0),
            m_handshakeStatus(exec::user::SUCCESS),
            m_bIsConnected(false),
            m_saslAuthenticator(NULL),
            m_saslDone(false),
            m_pendingRequests(0),
            m_pError(NULL),
            m_pListenerThread(NULL),
            m_pWork(NULL),
            m_socket(m_io_service),
            m_deadlineTimer(m_io_service),
            m_heartbeatTimer(m_io_service),
            m_rbuf(NULL),
            m_wbuf(MAX_SOCK_RD_BUFSIZE),
			m_bIsDirectConnection(false)
    {
        m_coordinationId=rand()%1729+1;
    };

        ~DrillClientImpl(){
            //TODO: Cleanup.
            //Free any record batches or buffers remaining
            //Cancel any pending requests
            //Clear and destroy DrillClientQueryResults vector?
            if(this->m_pWork!=NULL){
                delete this->m_pWork;
                this->m_pWork = NULL;
            }
            if(this->m_saslAuthenticator!=NULL){
                delete this->m_saslAuthenticator;
                this->m_saslAuthenticator = NULL;
            }

            m_heartbeatTimer.cancel();
            m_deadlineTimer.cancel();
            m_io_service.stop();
            boost::system::error_code ignorederr;
            m_socket.shutdown(boost::asio::ip::tcp::socket::shutdown_both, ignorederr);
            m_socket.close();
            if(m_rbuf!=NULL){
                Utils::freeBuffer(m_rbuf, MAX_SOCK_RD_BUFSIZE); m_rbuf=NULL;
            }
            if(m_pError!=NULL){
                delete m_pError; m_pError=NULL;
            }
            //Terminate and free the heartbeat thread
            //if(this->m_pHeartbeatThread!=NULL){
            //    this->m_pHeartbeatThread->interrupt();
            //    this->m_pHeartbeatThread->join();
            //    delete this->m_pHeartbeatThread;
            //    this->m_pHeartbeatThread = NULL;
            //}
            //Terminate and free the listener thread
            if(this->m_pListenerThread!=NULL){
                this->m_pListenerThread->interrupt();
                this->m_pListenerThread->join();
                delete this->m_pListenerThread;
                this->m_pListenerThread = NULL;
            }
        };

        //Connect via Zookeeper or directly
        connectionStatus_t connect(const char* connStr, DrillUserProperties* props);
        // test whether the client is active
        bool Active();
        void Close() ;
        DrillClientError* getError(){ return m_pError;}
        DrillClientQueryResult* SubmitQuery(::exec::shared::QueryType t, const std::string& plan, pfnQueryResultsListener listener, void* listenerCtx);
        DrillClientPrepareHandle* PrepareQuery(const std::string& plan, pfnPreparedStatementListener listener, void* listenerCtx);
        DrillClientQueryResult* ExecuteQuery(const PreparedStatement& pstmt, pfnQueryResultsListener listener, void* listenerCtx);

        void waitForResults();
        connectionStatus_t validateHandshake(DrillUserProperties* props);
        void freeQueryResources(DrillClientQueryHandle* pQryHandle){
            delete pQryHandle;
        };
        
        const exec::user::RpcEndpointInfos& getServerInfos() const { return m_serverInfos; }

        meta::DrillMetadata* getMetadata();

        void freeMetadata(meta::DrillMetadata* metadata);

    private:
        friend class meta::DrillMetadata;
        friend class DrillClientQueryHandle;
        friend class DrillClientQueryResult;
        friend class PooledDrillClientImpl;

        struct compareQueryId{
            bool operator()(const exec::shared::QueryId* q1, const exec::shared::QueryId* q2) const {
                return q1->part1()<q2->part1() || (q1->part1()==q2->part1() && q1->part2() < q2->part2());
            }
        };

        // Direct connection to a drillbit
        // host can be name or ip address, port can be port number or name of service in /etc/services
        connectionStatus_t connect(const char* host, const char* port);
        void startHeartbeatTimer();// start or restart the heartbeat timer
        connectionStatus_t sendHeartbeat(); // send a heartbeat to the server
        void handleHeartbeatTimeout(const boost::system::error_code & err); // send a heartbeat. If send fails, broadcast error, close connection and bail out.

        int32_t getNextCoordinationId(){ return ++m_coordinationId; };
        // send synchronous messages
        //connectionStatus_t recvSync(rpc::InBoundRpcMessage& msg);
        connectionStatus_t sendSync(rpc::OutBoundRpcMessage& msg);
        // handshake
        connectionStatus_t recvHandshake();
        void handleHandshake(ByteBuf_t b, const boost::system::error_code& err, std::size_t bytes_transferred );
        void handleHShakeReadTimeout(const boost::system::error_code & err);
        // starts the listener thread that receives responses/messages from the server
        void startMessageListener(); 
        // Query results
        void getNextResult();
        status_t readMsg(
                ByteBuf_t _buf,
                AllocatedBufferPtr* allocatedBuffer,
                rpc::InBoundRpcMessage& msg);
        status_t processQueryResult(AllocatedBufferPtr allocatedBuffer, const rpc::InBoundRpcMessage& msg);
        status_t processQueryData(AllocatedBufferPtr allocatedBuffer, const rpc::InBoundRpcMessage& msg);
        status_t processCancelledQueryResult( exec::shared::QueryId& qid, exec::shared::QueryResult* qr);
        status_t processQueryId(AllocatedBufferPtr allocatedBuffer, const rpc::InBoundRpcMessage& msg );
        status_t processPreparedStatement(AllocatedBufferPtr allocatedBuffer, const rpc::InBoundRpcMessage& msg );
        status_t processCatalogsResult(AllocatedBufferPtr allocatedBuffer, const rpc::InBoundRpcMessage& msg );
        status_t processSchemasResult(AllocatedBufferPtr allocatedBuffer, const rpc::InBoundRpcMessage& msg );
        status_t processTablesResult(AllocatedBufferPtr allocatedBuffer, const rpc::InBoundRpcMessage& msg );
        status_t processColumnsResult(AllocatedBufferPtr allocatedBuffer, const rpc::InBoundRpcMessage& msg );
        DrillClientQueryResult* findQueryResult(const exec::shared::QueryId& qid);
        status_t processQueryStatusResult( exec::shared::QueryResult* qr,
                DrillClientQueryResult* pDrillClientQueryResult);
        void handleReadTimeout(const boost::system::error_code & err);
        void handleRead(ByteBuf_t _buf, const boost::system::error_code & err, size_t bytes_transferred) ;
        status_t validateDataMessage(const rpc::InBoundRpcMessage& msg, const exec::shared::QueryData& qd, std::string& valError);
        status_t validateResultMessage(const rpc::InBoundRpcMessage& msg, const exec::shared::QueryResult& qr, std::string& valError);
        connectionStatus_t handleConnError(connectionStatus_t status, const std::string& msg);
        status_t handleQryCancellation(status_t status, DrillClientQueryResult* pQueryResult);
        status_t handleQryError(status_t status, const std::string& msg, DrillClientQueryHandle* pQueryHandle);
        status_t handleQryError(status_t status, const exec::shared::DrillPBError& e, DrillClientQueryHandle* pQueryHandle);
        // handle query state indicating query is COMPLETED or CANCELED
        // (i.e., COMPLETED or CANCELED)
        status_t handleTerminatedQryState(status_t status,
                const std::string& msg,
                DrillClientQueryResult* pQueryResult);
        void broadcastError(DrillClientError* pErr);
        void removeQueryHandle(DrillClientQueryHandle* pQueryHandle);
        void removeQueryResult(DrillClientQueryResult* pQueryResult);
        void sendAck(const rpc::InBoundRpcMessage& msg, bool isOk);
        void sendCancel(const exec::shared::QueryId* pQueryId);

        template<typename Handle>
        Handle* sendMsg(boost::function<Handle*(int32_t)> handleFactory, ::exec::user::RpcType type, const ::google::protobuf::Message& msg);

        // metadata requests
        DrillClientCatalogResult* getCatalogs(const std::string& catalogPattern, Metadata::pfnCatalogMetadataListener listener, void* listenerCtx);
        DrillClientSchemaResult* getSchemas(const std::string& catalogPattern, const std::string& schemaPattern, Metadata::pfnSchemaMetadataListener listener, void* listenerCtx);
        DrillClientTableResult* getTables(const std::string& catalogPattern, const std::string& schemaPattern, const std::string& tablePattern, const std::vector<std::string>* tableTypes, Metadata::pfnTableMetadataListener listener, void* listenerCtx);
        DrillClientColumnResult* getColumns(const std::string& catalogPattern, const std::string& schemaPattern, const std::string& tablePattern, const std::string& columnPattern, Metadata::pfnColumnMetadataListener listener, void* listenerCtx);

        // SASL exchange
        connectionStatus_t handleAuthentication(const DrillUserProperties *userProperties);
        void initiateAuthentication();
        void sendSaslResponse(const exec::shared::SaslMessage& response);
        void processSaslChallenge(AllocatedBufferPtr allocatedBuffer, const rpc::InBoundRpcMessage& msg);
        void finishAuthentication();

        void shutdownSocket();

        int32_t m_coordinationId;
        int32_t m_handshakeVersion;
        exec::user::HandshakeStatus m_handshakeStatus;
        std::string m_handshakeErrorId;
        std::string m_handshakeErrorMsg;
        exec::user::RpcEndpointInfos m_serverInfos;
        bool m_bIsConnected;

        std::vector<std::string> m_serverAuthMechanisms;
        SaslAuthenticatorImpl* m_saslAuthenticator;
        int m_saslResultCode;
        bool m_saslDone;
        boost::mutex m_saslMutex; // mutex to protect m_saslDone
        boost::condition_variable m_saslCv; // to signal completion of SASL exchange

        std::string m_connectStr; 

        // 
        // number of outstanding read requests.
        // handleRead will keep asking for more results as long as this number is not zero.
        size_t m_pendingRequests;
        //mutex to protect m_pendingRequests
        boost::mutex m_prMutex;

        // Error Object. NULL if no error. Set if the error is valid for ALL running queries.
        // All the query result objects will
        // also have the error object set.
        // If the error is query specific, only the query results object will have the error set.
        DrillClientError* m_pError;

        //Started after the connection is established and sends heartbeat messages after {heartbeat frequency} seconds
        //The thread is killed on disconnect.
        //boost::thread * m_pHeartbeatThread;

        // for boost asio
        boost::thread * m_pListenerThread;
        boost::asio::io_service m_io_service;
        // the work object prevent io_service running out of work
        boost::asio::io_service::work * m_pWork;
        boost::asio::ip::tcp::socket m_socket;
        boost::asio::deadline_timer m_deadlineTimer; // to timeout async queries that never return
        boost::asio::deadline_timer m_heartbeatTimer; // to send heartbeat messages

        std::string m_connectedHost; // The hostname and port the socket is connected to.

        //for synchronous messages, like validate handshake
        ByteBuf_t m_rbuf; // buffer for receiving synchronous messages
        DataBuf m_wbuf; // buffer for sending synchronous message

        // Mutex to protect drill client operations
        boost::mutex m_dcMutex;

        // Map of coordination id to Query handles.
        std::map<int, DrillClientQueryHandle*> m_queryHandles;

        // Map of query id to query result for currently executing queries
        std::map<exec::shared::QueryId*, DrillClientQueryResult*, compareQueryId> m_queryResults;

        // Condition variable to signal completion of all queries. 
        boost::condition_variable m_cv;

        bool m_bIsDirectConnection;
};

inline bool DrillClientImpl::Active() {
    return this->m_bIsConnected;
}


/* *
 *  Provides the same public interface as a DrillClientImpl but holds a pool of DrillClientImpls.
 *  Every submitQuery uses a different DrillClientImpl to distribute the load.
 *  DrillClient can use this class instead of DrillClientImpl to get better load balancing.
 * */
class PooledDrillClientImpl : public DrillClientImplBase{
    public:
        PooledDrillClientImpl():
        	m_lastConnection(-1),
			m_queriesExecuted(0),
			m_maxConcurrentConnections(DEFAULT_MAX_CONCURRENT_CONNECTIONS),
			m_bIsDirectConnection(false),
			m_pError(NULL),
			m_pUserProperties() {
            char* maxConn=std::getenv(MAX_CONCURRENT_CONNECTIONS_ENV);
            if(maxConn!=NULL){
                m_maxConcurrentConnections=atoi(maxConn);
            }
        }

        ~PooledDrillClientImpl(){
            for(std::vector<DrillClientImpl*>::iterator it = m_clientConnections.begin(); it != m_clientConnections.end(); ++it){
                delete *it;
            }
            m_clientConnections.clear();
            if(m_pError!=NULL){ delete m_pError; m_pError=NULL;}
        }

        //Connect via Zookeeper or directly.
        //Makes an initial connection to a drillbit. successful connect adds the first drillbit to the pool.
        connectionStatus_t connect(const char* connStr, DrillUserProperties* props);

        // Test whether the client is active. Returns true if any one of the underlying connections is active
        bool Active();

        // Closes all open connections. 
        void Close() ;

        // Returns the last error encountered by any of the underlying executing queries or connections
        DrillClientError* getError();

        // Submits a query to a drillbit. If more than one query is to be sent, we may choose a
        // a different drillbit in the pool. No more than m_maxConcurrentConnections will be allowed.
        // Connections once added to the pool will be removed only when the DrillClient is closed.
        DrillClientQueryResult* SubmitQuery(::exec::shared::QueryType t, const std::string& plan, pfnQueryResultsListener listener, void* listenerCtx);

        DrillClientPrepareHandle* PrepareQuery(const std::string& plan, pfnPreparedStatementListener listener, void* listenerCtx);
        DrillClientQueryResult* ExecuteQuery(const PreparedStatement& pstmt, pfnQueryResultsListener listener, void* listenerCtx);

        //Waits as long as any one drillbit connection has results pending
        void waitForResults();

        //Validates handshake only against the first drillbit connected to.
        connectionStatus_t validateHandshake(DrillUserProperties* props);

        void freeQueryResources(DrillClientQueryHandle* pQueryHandle);

        int getDrillbitCount() const { return m_drillbits.size();};
        
        meta::DrillMetadata* getMetadata();

        void freeMetadata(meta::DrillMetadata* metadata);

    private:
        
        std::string m_connectStr; 
        std::string m_lastQuery;
        
        // A list of all the current client connections. We choose a new one for every query. 
        // When picking a drillClientImpl to use, we see how many queries each drillClientImpl
        // is currently executing. If none,  
        std::vector<DrillClientImpl*> m_clientConnections; 
		boost::mutex m_poolMutex; // protect access to the vector
        // Use this to decide which drillbit to select next from the list of drillbits.
        size_t m_lastConnection;
		boost::mutex m_cMutex;

        // Number of queries executed so far. Can be used to select a new Drillbit from the pool.
        size_t m_queriesExecuted;

        size_t m_maxConcurrentConnections;

        bool m_bIsDirectConnection;

        DrillClientError* m_pError;

        connectionStatus_t handleConnError(connectionStatus_t status, std::string msg);
        // get a connection from the pool or create a new one. Return NULL if none is found
        DrillClientImpl* getOneConnection();

        std::vector<std::string> m_drillbits;

        boost::shared_ptr<DrillUserProperties> m_pUserProperties;//Keep a copy of user properties
};

} // namespace Drill

#endif