text: /*
   * 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.
   */
  
  
  #include "drill/common.hpp"
  #include <queue>
  #include <string.h>
  #include <boost/asio.hpp>
  #include <boost/bind.hpp>
  #include <boost/date_time/posix_time/posix_time.hpp>
  #include <boost/date_time/posix_time/posix_time_duration.hpp>
  #include <boost/lexical_cast.hpp>
  #include <boost/thread.hpp>
  #ifdef _WIN32
  #include <zookeeper.h>
  #else
  #include <zookeeper/zookeeper.h>
  #endif
  #include <boost/assign.hpp>
  
  #include "drill/drillClient.hpp"
  #include "drill/recordBatch.hpp"
  #include "drillClientImpl.hpp"
  #include "errmsgs.hpp"
  #include "logger.hpp"
  #include "rpcEncoder.hpp"
  #include "rpcDecoder.hpp"
  #include "rpcMessage.hpp"
  #include "utils.hpp"
  
  #include "GeneralRPC.pb.h"
  #include "UserBitShared.pb.h"
  
  namespace Drill{
  
  static std::map<exec::shared::QueryResult_QueryState, status_t> QUERYSTATE_TO_STATUS_MAP = boost::assign::map_list_of
      (exec::shared::QueryResult_QueryState_STARTING, QRY_PENDING)
      (exec::shared::QueryResult_QueryState_RUNNING, QRY_RUNNING)
      (exec::shared::QueryResult_QueryState_COMPLETED, QRY_COMPLETED)
      (exec::shared::QueryResult_QueryState_CANCELED, QRY_CANCELED)
      (exec::shared::QueryResult_QueryState_FAILED, QRY_FAILED)
      ;
  
  RpcEncoder DrillClientImpl::s_encoder;
  RpcDecoder DrillClientImpl::s_decoder;
  
  std::string debugPrintQid(const exec::shared::QueryId& qid){
      return std::string("[")+boost::lexical_cast<std::string>(qid.part1()) +std::string(":") + boost::lexical_cast<std::string>(qid.part2())+std::string("] ");
  }
  
  void setSocketTimeout(boost::asio::ip::tcp::socket& socket, int32_t timeout){
  #if defined _WIN32
      int32_t timeoutMsecs=timeout*1000;
      setsockopt(socket.native(), SOL_SOCKET, SO_RCVTIMEO, (const char*)&timeoutMsecs, sizeof(timeoutMsecs));
      setsockopt(socket.native(), SOL_SOCKET, SO_SNDTIMEO, (const char*)&timeoutMsecs, sizeof(timeoutMsecs));
  #else
      struct timeval tv;
      tv.tv_sec  = timeout;
      tv.tv_usec = 0;
      int e=0;
      e=setsockopt(socket.native(), SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
      e=setsockopt(socket.native(), SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv));
  #endif
  }
  
  
  void DrillClientImpl::parseConnectStr(const char* connectStr,
          std::string& pathToDrill,
          std::string& protocol,
          std::string& hostPortStr){
      char u[MAX_CONNECT_STR+1];
      strncpy(u,connectStr, MAX_CONNECT_STR); u[MAX_CONNECT_STR]=0;
      char* z=strtok(u, "=");
      char* c=strtok(NULL, "/");
      char* p=strtok(NULL, "");
  
      if(p!=NULL) pathToDrill=std::string("/")+p;
      protocol=z; hostPortStr=c;
      return;
  }
  
  connectionStatus_t DrillClientImpl::connect(const char* connStr){
      std::string pathToDrill, protocol, hostPortStr;
      std::string host;
      std::string port;
      if(!this->m_bIsConnected){
          parseConnectStr(connStr, pathToDrill, protocol, hostPortStr);
          if(!strcmp(protocol.c_str(), "zk")){
              ZookeeperImpl zook;
              if(zook.connectToZookeeper(hostPortStr.c_str(), pathToDrill.c_str())!=0){
                  return handleConnError(CONN_ZOOKEEPER_ERROR, getMessage(ERR_CONN_ZOOKEEPER, zook.getError().c_str()));
              }
              zook.debugPrint();
              exec::DrillbitEndpoint e=zook.getEndPoint();
              host=boost::lexical_cast<std::string>(e.address());
              port=boost::lexical_cast<std::string>(e.user_port());
              zook.close();
          }else if(!strcmp(protocol.c_str(), "local")){
              char tempStr[MAX_CONNECT_STR+1];
              strncpy(tempStr, hostPortStr.c_str(), MAX_CONNECT_STR); tempStr[MAX_CONNECT_STR]=0;
              host=strtok(tempStr, ":");
              port=strtok(NULL, "");
          }else{
              return handleConnError(CONN_INVALID_INPUT, getMessage(ERR_CONN_UNKPROTO, protocol.c_str()));
          }
          return this->connect(host.c_str(), port.c_str());
      }
      return CONN_SUCCESS;
  }
  
  connectionStatus_t DrillClientImpl::connect(const char* host, const char* port){
      using boost::asio::ip::tcp;
      tcp::endpoint endpoint;
      try{
          tcp::resolver resolver(m_io_service);
          tcp::resolver::query query(tcp::v4(), host, port);
          tcp::resolver::iterator iter = resolver.resolve(query);
          tcp::resolver::iterator end;
          while (iter != end){
              endpoint = *iter++;
              DRILL_LOG(LOG_TRACE) << endpoint << std::endl;
          }
          boost::system::error_code ec;
          m_socket.connect(endpoint, ec);
          if(ec){
              return handleConnError(CONN_FAILURE, getMessage(ERR_CONN_FAILURE, host, port, ec.message().c_str()));
          }
  
      }catch(std::exception e){
          // Handle case when the hostname cannot be resolved. "resolve" is hard-coded in boost asio resolver.resolve
          if (!strcmp(e.what(), "resolve")) {
              return handleConnError(CONN_HOSTNAME_RESOLUTION_ERROR, getMessage(ERR_CONN_EXCEPT, e.what()));
          }
          return handleConnError(CONN_FAILURE, getMessage(ERR_CONN_EXCEPT, e.what()));
      }
  
      // set socket keep alive
      boost::asio::socket_base::keep_alive keepAlive(true);
      m_socket.set_option(keepAlive);
  	// set no_delay
      boost::asio::ip::tcp::no_delay noDelay(true);
      m_socket.set_option(noDelay);
  
      //
      // We put some OS dependent code here for timing out a socket. Mostly, this appears to
      // do nothing. Should we leave it in there?
      //
      setSocketTimeout(m_socket, DrillClientConfig::getSocketTimeout());
  
      return CONN_SUCCESS;
  }
  
  void DrillClientImpl::startHeartbeatTimer(){
      DRILL_LOG(LOG_TRACE) << "Started new heartbeat timer with "
          << DrillClientConfig::getHeartbeatFrequency() << " seconds." << std::endl;
      m_heartbeatTimer.expires_from_now(boost::posix_time::seconds(DrillClientConfig::getHeartbeatFrequency()));
      m_heartbeatTimer.async_wait(boost::bind(
                  &DrillClientImpl::handleHeartbeatTimeout,
                  this,
                  boost::asio::placeholders::error
                  ));
  	    startMessageListener(); // start this thread early so we don't have the timer blocked
  }
  
  connectionStatus_t DrillClientImpl::sendHeartbeat(){
  	connectionStatus_t status=CONN_SUCCESS;
      exec::rpc::Ack ack;
      ack.set_ok(true);
      OutBoundRpcMessage heartbeatMsg(exec::rpc::PING, exec::user::ACK/*can be anything */, 0, &ack);
  	boost::lock_guard<boost::mutex> prLock(this->m_prMutex);
  	boost::lock_guard<boost::mutex> lock(m_dcMutex);
      DRILL_LOG(LOG_TRACE) << "Heartbeat sent." << std::endl;
      status=sendSync(heartbeatMsg);
      status=status==CONN_SUCCESS?status:CONN_DEAD;
      //If the server sends responses to a heartbeat, we need to increment the pending requests counter.
      if(m_pendingRequests++==0){
          getNextResult(); // async wait for results
      }
      return status;
  }
  
  void DrillClientImpl::resetHeartbeatTimer(){
      m_heartbeatTimer.cancel();
      DRILL_LOG(LOG_TRACE) << "Reset Heartbeat timer." << std::endl;
      startHeartbeatTimer();
  }
  
  
  
  void DrillClientImpl::handleHeartbeatTimeout(const boost::system::error_code & err){
      DRILL_LOG(LOG_TRACE) << "DrillClientImpl:: Heartbeat timer expired." << std::endl;
      if(err != boost::asio::error::operation_aborted){
          // Check whether the deadline has passed.
          DRILL_LOG(LOG_TRACE) << "DrillClientImpl::Heartbeat Timer -  Expires at: " 
              << to_simple_string(m_heartbeatTimer.expires_at())
              << " and time now is: "
              << to_simple_string(boost::asio::deadline_timer::traits_type::now())
              << std::endl;
              ;
          if (m_heartbeatTimer.expires_at() <= boost::asio::deadline_timer::traits_type::now()){
              // The deadline has passed.
              m_heartbeatTimer.expires_at(boost::posix_time::pos_infin);
              if(sendHeartbeat()==CONN_SUCCESS){
                  startHeartbeatTimer();
              }else{
                  // Close connection.
                  DRILL_LOG(LOG_TRACE) << "DrillClientImpl:: No heartbeat. Closing connection.";
                  shutdownSocket();
              }
          }
      }
      return;
  }
  
  
  void DrillClientImpl::Close() {
      shutdownSocket();
  }
  
  
  connectionStatus_t DrillClientImpl::sendSync(OutBoundRpcMessage& msg){
      DrillClientImpl::s_encoder.Encode(m_wbuf, msg);
      boost::system::error_code ec;
      size_t s=m_socket.write_some(boost::asio::buffer(m_wbuf), ec);
      if(!ec && s!=0){
          return CONN_SUCCESS;
      }else{
          return handleConnError(CONN_FAILURE, getMessage(ERR_CONN_WFAIL, ec.message().c_str()));
      }
  }
  
  connectionStatus_t DrillClientImpl::recvHandshake(){
      if(m_rbuf==NULL){
          m_rbuf = Utils::allocateBuffer(MAX_SOCK_RD_BUFSIZE);
      }
  
      m_io_service.reset();
      if (DrillClientConfig::getHandshakeTimeout() > 0){
          m_deadlineTimer.expires_from_now(boost::posix_time::seconds(DrillClientConfig::getHandshakeTimeout()));
          m_deadlineTimer.async_wait(boost::bind(
                      &DrillClientImpl::handleHShakeReadTimeout,
                      this,
                      boost::asio::placeholders::error
                      ));
          DRILL_LOG(LOG_TRACE) << "Started new handshake wait timer with "
                  << DrillClientConfig::getHandshakeTimeout() << " seconds." << std::endl;
      }
  
      async_read(
              this->m_socket,
              boost::asio::buffer(m_rbuf, LEN_PREFIX_BUFLEN),
              boost::bind(
                  &DrillClientImpl::handleHandshake,
                  this,
                  m_rbuf,
                  boost::asio::placeholders::error,
                  boost::asio::placeholders::bytes_transferred)
              );
      DRILL_LOG(LOG_DEBUG) << "DrillClientImpl::recvHandshake: async read waiting for server handshake response.\n";
      m_io_service.run();
      if(m_rbuf!=NULL){
          Utils::freeBuffer(m_rbuf, MAX_SOCK_RD_BUFSIZE); m_rbuf=NULL;
      }
  #ifdef WIN32_SHUTDOWN_ON_TIMEOUT
      if (m_pError != NULL) {
          return static_cast<connectionStatus_t>(m_pError->status);
      }
  #endif // WIN32_SHUTDOWN_ON_TIMEOUT
      startHeartbeatTimer();
  
      return CONN_SUCCESS;
  }
  
  void DrillClientImpl::handleHandshake(ByteBuf_t _buf,
          const boost::system::error_code& err,
          size_t bytes_transferred) {
      boost::system::error_code error=err;
      // cancel the timer
      m_deadlineTimer.cancel();
      DRILL_LOG(LOG_TRACE) << "Deadline timer cancelled." << std::endl;
      if(!error){
          InBoundRpcMessage msg;
          uint32_t length = 0;
          int bytes_read = DrillClientImpl::s_decoder.LengthDecode(m_rbuf, &length);
          if(length>0){
              size_t leftover = LEN_PREFIX_BUFLEN - bytes_read;
              ByteBuf_t b=m_rbuf + LEN_PREFIX_BUFLEN;
              size_t bytesToRead=length - leftover;
              while(1){
                  size_t dataBytesRead=m_socket.read_some(
                          boost::asio::buffer(b, bytesToRead),
                          error);
                  if(err) break;
                  DRILL_LOG(LOG_TRACE) << "Handshake Message: actual bytes read = " << dataBytesRead << std::endl;
                  if(dataBytesRead==bytesToRead) break;
                  bytesToRead-=dataBytesRead;
                  b+=dataBytesRead;
              }
              DrillClientImpl::s_decoder.Decode(m_rbuf+bytes_read, length, msg);
          }else{
              DRILL_LOG(LOG_TRACE) << "DrillClientImpl::handleHandshake: ERR_CONN_RDFAIL. No handshake.\n";
              handleConnError(CONN_FAILURE, getMessage(ERR_CONN_RDFAIL, "No handshake"));
              return;
          }
          exec::user::BitToUserHandshake b2u;
          b2u.ParseFromArray(msg.m_pbody.data(), msg.m_pbody.size());
          this->m_handshakeVersion=b2u.rpc_version();
          this->m_handshakeStatus=b2u.status();
          this->m_handshakeErrorId=b2u.errorid();
          this->m_handshakeErrorMsg=b2u.errormessage();
  
      }else{
          // boost error
          if(error==boost::asio::error::eof){ // Server broke off the connection
              handleConnError(CONN_HANDSHAKE_FAILED,
                  getMessage(ERR_CONN_NOHSHAKE, DRILL_RPC_VERSION));
          }else{
              handleConnError(CONN_FAILURE, getMessage(ERR_CONN_RDFAIL, error.message().c_str()));
          }
          return;
      }
      return;
  }
  
  void DrillClientImpl::handleHShakeReadTimeout(const boost::system::error_code & err){
      // if err == boost::asio::error::operation_aborted) then the caller cancelled the timer.
      if(err != boost::asio::error::operation_aborted){
          // Check whether the deadline has passed.
          if (m_deadlineTimer.expires_at() <= boost::asio::deadline_timer::traits_type::now()){
              // The deadline has passed.
              m_deadlineTimer.expires_at(boost::posix_time::pos_infin);
              DRILL_LOG(LOG_TRACE) << "DrillClientImpl::HandleHShakeReadTimeout: Deadline timer expired; ERR_CONN_HSHAKETIMOUT.\n";
              handleConnError(CONN_HANDSHAKE_TIMEOUT, getMessage(ERR_CONN_HSHAKETIMOUT));
              m_io_service.stop();
              boost::system::error_code ignorederr;
              m_socket.shutdown(boost::asio::ip::tcp::socket::shutdown_both, ignorederr);
          }
      }
      return;
  }
  
  connectionStatus_t DrillClientImpl::validateHandshake(DrillUserProperties* properties){
  
      DRILL_LOG(LOG_TRACE) << "validateHandShake\n";
  
      exec::user::UserToBitHandshake u2b;
      u2b.set_channel(exec::shared::USER);
      u2b.set_rpc_version(DRILL_RPC_VERSION);
      u2b.set_support_listening(true);
      u2b.set_support_timeout(true);
  
      if(properties != NULL && properties->size()>0){
          std::string username;
          std::string err;
          if(!properties->validate(err)){
              DRILL_LOG(LOG_INFO) << "Invalid user input:" << err << std::endl;
          }
          exec::user::UserProperties* userProperties = u2b.mutable_properties();
  
          std::map<char,int>::iterator it;
          for(size_t i=0; i<properties->size(); i++){
              std::map<std::string,uint32_t>::const_iterator it=DrillUserProperties::USER_PROPERTIES.find(properties->keyAt(i));
              if(it==DrillUserProperties::USER_PROPERTIES.end()){
                  DRILL_LOG(LOG_WARNING) << "Connection property ("<< properties->keyAt(i) 
                      << ") is unknown and is being skipped" << std::endl;
                  continue;
              }
              if(IS_BITSET((*it).second,USERPROP_FLAGS_SERVERPROP)){
                  exec::user::Property* connProp = userProperties->add_properties();
                  connProp->set_key(properties->keyAt(i));
                  connProp->set_value(properties->valueAt(i));
                  //Username(but not the password) also needs to be set in UserCredentials
                  if(IS_BITSET((*it).second,USERPROP_FLAGS_USERNAME)){
                      exec::shared::UserCredentials* creds = u2b.mutable_credentials();
                      username=properties->valueAt(i);
                      creds->set_user_name(username);
                      //u2b.set_credentials(&creds);
                  }
                  if(IS_BITSET((*it).second,USERPROP_FLAGS_PASSWORD)){
                      DRILL_LOG(LOG_INFO) <<  properties->keyAt(i) << ": ********** " << std::endl;
                  }else{
                      DRILL_LOG(LOG_INFO) << properties->keyAt(i) << ":" << properties->valueAt(i) << std::endl;
                  }
              }// Server properties
          }
      }
  
      {
          boost::lock_guard<boost::mutex> lock(this->m_dcMutex);
          uint64_t coordId = this->getNextCoordinationId();
  
          OutBoundRpcMessage out_msg(exec::rpc::REQUEST, exec::user::HANDSHAKE, coordId, &u2b);
          sendSync(out_msg);
          DRILL_LOG(LOG_TRACE) << "Sent handshake request message. Coordination id: " << coordId << "\n";
      }
  
      connectionStatus_t ret = recvHandshake();
      if(ret!=CONN_SUCCESS){
          return ret;
      }
      if(this->m_handshakeStatus != exec::user::SUCCESS){
          switch(this->m_handshakeStatus){
              case exec::user::RPC_VERSION_MISMATCH:
                  DRILL_LOG(LOG_TRACE) << "Invalid rpc version.  Expected "
                      << DRILL_RPC_VERSION << ", actual "<< m_handshakeVersion << "." << std::endl;
                  return handleConnError(CONN_BAD_RPC_VER,
                          getMessage(ERR_CONN_BAD_RPC_VER, DRILL_RPC_VERSION,
                              m_handshakeVersion,
                              this->m_handshakeErrorId.c_str(),
                              this->m_handshakeErrorMsg.c_str()));
              case exec::user::AUTH_FAILED:
                  DRILL_LOG(LOG_TRACE) << "Authentication failed." << std::endl;
                  return handleConnError(CONN_AUTH_FAILED,
                          getMessage(ERR_CONN_AUTHFAIL,
                              this->m_handshakeErrorId.c_str(),
                              this->m_handshakeErrorMsg.c_str()));
              case exec::user::UNKNOWN_FAILURE:
                  DRILL_LOG(LOG_TRACE) << "Unknown error during handshake." << std::endl;
                  return handleConnError(CONN_HANDSHAKE_FAILED,
                          getMessage(ERR_CONN_UNKNOWN_ERR,
                              this->m_handshakeErrorId.c_str(),
                              this->m_handshakeErrorMsg.c_str()));
              default:
                  break;
          }
      }
      // reset io_service after handshake is validated before running queries
      m_io_service.reset();
      return CONN_SUCCESS;
  }
  
  
  FieldDefPtr DrillClientQueryResult::s_emptyColDefs( new (std::vector<Drill::FieldMetadata*>));
  
  void DrillClientImpl::startMessageListener() {
      if(this->m_pListenerThread==NULL){
          // Stopping the io_service from running out-of-work
          if(m_io_service.stopped()){
              DRILL_LOG(LOG_DEBUG) << "DrillClientImpl::startMessageListener: io_service is stopped. Restarting." <<std::endl;
              m_io_service.reset();
          }
          this->m_pWork = new boost::asio::io_service::work(m_io_service);
          this->m_pListenerThread = new boost::thread(boost::bind(&boost::asio::io_service::run,
                      &this->m_io_service));
          DRILL_LOG(LOG_DEBUG) << "DrillClientImpl::startMessageListener: Starting listener thread: "
              << this->m_pListenerThread << std::endl;
      }
  }
  
  DrillClientQueryResult* DrillClientImpl::SubmitQuery(::exec::shared::QueryType t,
          const std::string& plan,
          pfnQueryResultsListener l,
          void* lCtx){
      exec::user::RunQuery query;
      query.set_results_mode(exec::user::STREAM_FULL);
      query.set_type(t);
      query.set_plan(plan);
  
      uint64_t coordId;
      DrillClientQueryResult* pQuery=NULL;
      {
          boost::lock_guard<boost::mutex> prLock(this->m_prMutex);
          boost::lock_guard<boost::mutex> dcLock(this->m_dcMutex);
          coordId = this->getNextCoordinationId();
          OutBoundRpcMessage out_msg(exec::rpc::REQUEST, exec::user::RUN_QUERY, coordId, &query);
          sendSync(out_msg);
  
          pQuery = new DrillClientQueryResult(this, coordId);
          pQuery->registerListener(l, lCtx);
          bool sendRequest=false;
          this->m_queryIds[coordId]=pQuery;
  
          DRILL_LOG(LOG_DEBUG)  << "Sent query request. Coordination id = " << coordId << std::endl;
  
          if(m_pendingRequests++==0){
              sendRequest=true;
          }else{
              DRILL_LOG(LOG_DEBUG) << "Queueing query request to server" << std::endl;
              DRILL_LOG(LOG_DEBUG) << "Number of pending requests = " << m_pendingRequests << std::endl;
          }
          if(sendRequest){
              DRILL_LOG(LOG_DEBUG) << "Sending query request. Number of pending requests = "
                  << m_pendingRequests << std::endl;
              getNextResult(); // async wait for results
          }
      }
  
      //run this in a new thread
      startMessageListener();
  
      return pQuery;
  }
  
  void DrillClientImpl::getNextResult(){
  
      // This call is always made from within a function where the mutex has already been acquired
      //boost::lock_guard<boost::mutex> lock(this->m_dcMutex);
  
      {
          boost::unique_lock<boost::mutex> memLock(AllocatedBuffer::s_memCVMutex);
          DRILL_LOG(LOG_TRACE) << "Read blocked waiting for memory." << std::endl;
          while(AllocatedBuffer::s_isBufferLimitReached){
              AllocatedBuffer::s_memCV.wait(memLock);
          }
      }
      
      //use free, not delete to free
      ByteBuf_t readBuf = Utils::allocateBuffer(LEN_PREFIX_BUFLEN);
      if (DrillClientConfig::getQueryTimeout() > 0){
          DRILL_LOG(LOG_TRACE) << "Started new query wait timer with "
                  << DrillClientConfig::getQueryTimeout() << " seconds." << std::endl;
          m_deadlineTimer.expires_from_now(boost::posix_time::seconds(DrillClientConfig::getQueryTimeout()));
          m_deadlineTimer.async_wait(boost::bind(
              &DrillClientImpl::handleReadTimeout,
              this,
              boost::asio::placeholders::error
              ));
      }
  
      resetHeartbeatTimer();
  
      async_read(
              this->m_socket,
              boost::asio::buffer(readBuf, LEN_PREFIX_BUFLEN),
              boost::bind(
                  &DrillClientImpl::handleRead,
                  this,
                  readBuf,
                  boost::asio::placeholders::error,
                  boost::asio::placeholders::bytes_transferred)
              );
      DRILL_LOG(LOG_DEBUG) << "DrillClientImpl::getNextResult: async_read from the server\n";
  }
  
  void DrillClientImpl::waitForResults(){
      // The listener thread never exists because it may be sending/receiving a heartbeat. Before the heartbeat was introduced
      // we could check if the listener thread has exited to tell if the queries are done. We can no longer do so now. We check
      // a condition variable instead
      {
          boost::unique_lock<boost::mutex> cvLock(this->m_dcMutex);
          //if no more data, return NULL;
          while(this->m_pendingRequests>0) {
              this->m_cv.wait(cvLock);
          }
      }
  }
  
  status_t DrillClientImpl::readMsg(ByteBuf_t _buf,
          AllocatedBufferPtr* allocatedBuffer,
          InBoundRpcMessage& msg,
          boost::system::error_code& error){
  
      DRILL_LOG(LOG_TRACE) << "DrillClientImpl::readMsg: Read message from buffer "
          <<  reinterpret_cast<int*>(_buf) << std::endl;
      size_t leftover=0;
      uint32_t rmsgLen;
      AllocatedBufferPtr currentBuffer;
      *allocatedBuffer=NULL;
      {
          // We need to protect the readLength and read buffer, and the pending requests counter,
          // but we don't have to keep the lock while we decode the rest of the buffer.
          boost::lock_guard<boost::mutex> lock(this->m_dcMutex);
          int bytes_read = DrillClientImpl::s_decoder.LengthDecode(_buf, &rmsgLen);
          DRILL_LOG(LOG_TRACE) << "len bytes = " << bytes_read << std::endl;
          DRILL_LOG(LOG_TRACE) << "rmsgLen = " << rmsgLen << std::endl;
  
          if(rmsgLen>0){
              leftover = LEN_PREFIX_BUFLEN - bytes_read;
              // Allocate a buffer
              currentBuffer=new AllocatedBuffer(rmsgLen);
              DRILL_LOG(LOG_TRACE) << "DrillClientImpl::readMsg: Allocated and locked buffer: [ "
                  << currentBuffer << ", size = " << rmsgLen << " ]\n";
              if(currentBuffer==NULL){
                  Utils::freeBuffer(_buf, LEN_PREFIX_BUFLEN);
                  return handleQryError(QRY_CLIENT_OUTOFMEM, getMessage(ERR_QRY_OUTOFMEM), NULL);
              }
              *allocatedBuffer=currentBuffer;
              if(leftover){
                  memcpy(currentBuffer->m_pBuffer, _buf + bytes_read, leftover);
              }
              DRILL_LOG(LOG_TRACE) << "reading data (rmsgLen - leftover) : "
                  << (rmsgLen - leftover) << std::endl;
              ByteBuf_t b=currentBuffer->m_pBuffer + leftover;
              size_t bytesToRead=rmsgLen - leftover;
                
              while(1){
                  size_t dataBytesRead=this->m_socket.read_some(
                          boost::asio::buffer(b, bytesToRead),
                          error);
                  if(error) break;
                  DRILL_LOG(LOG_TRACE) << "Data Message: actual bytes read = " << dataBytesRead << std::endl;
                  if(dataBytesRead==bytesToRead) break;
                  bytesToRead-=dataBytesRead;
                  b+=dataBytesRead;
              }
              
              if(!error){
                  // read data successfully
                  DrillClientImpl::s_decoder.Decode(currentBuffer->m_pBuffer, rmsgLen, msg);
                  DRILL_LOG(LOG_TRACE) << "Done decoding chunk. Coordination id: " <<msg.m_coord_id<< std::endl;
              }else{
                  Utils::freeBuffer(_buf, LEN_PREFIX_BUFLEN);
                  return handleQryError(QRY_COMM_ERROR,
                          getMessage(ERR_QRY_COMMERR, error.message().c_str()), NULL);
              }
          }else{
              // got a message with an invalid read length.
              Utils::freeBuffer(_buf, LEN_PREFIX_BUFLEN);
              return handleQryError(QRY_INTERNAL_ERROR, getMessage(ERR_QRY_INVREADLEN), NULL);
          }
      }
      DRILL_LOG(LOG_TRACE) << "DrillClientImpl::readMsg: Free buffer "
          <<  reinterpret_cast<int*>(_buf) << std::endl;
      Utils::freeBuffer(_buf, LEN_PREFIX_BUFLEN);
      return QRY_SUCCESS;
  }
  
  status_t DrillClientImpl::processQueryResult(AllocatedBufferPtr  allocatedBuffer, InBoundRpcMessage& msg ){
      DrillClientQueryResult* pDrillClientQueryResult=NULL;
      status_t ret=QRY_SUCCESS;
      exec::shared::QueryId qid;
      sendAck(msg, true);
      {
          boost::lock_guard<boost::mutex> lock(this->m_dcMutex);
          exec::shared::QueryResult qr;
  
          DRILL_LOG(LOG_DEBUG) << "Processing Query Result " << std::endl;
          qr.ParseFromArray(msg.m_pbody.data(), msg.m_pbody.size());
          DRILL_LOG(LOG_TRACE) << qr.DebugString() << std::endl;
          
          qid.CopyFrom(qr.query_id());
          
          if (qr.has_query_state() &&
                  qr.query_state() != exec::shared::QueryResult_QueryState_RUNNING &&
                  qr.query_state() != exec::shared::QueryResult_QueryState_STARTING) {
              pDrillClientQueryResult=findQueryResult(qid);
              //Queries that have been cancelled or whose resources are freed before completion 
              //do not have a DrillClientQueryResult object. We need not handle the terminal message 
              //in that case since all it does is to free resources (and they have already been freed)
              if(pDrillClientQueryResult!=NULL){
                  //Validate the RPC message
                  std::string valErr;
                  if( (ret=validateResultMessage(msg, qr, valErr)) != QRY_SUCCESS){
                      delete allocatedBuffer;
                      DRILL_LOG(LOG_TRACE) << "DrillClientImpl::processQueryResult: ERR_QRY_INVRPC." << std::endl;
                      return handleQryError(ret, getMessage(ERR_QRY_INVRPC, valErr.c_str()), pDrillClientQueryResult);
                  }
                  ret=processQueryStatusResult(&qr, pDrillClientQueryResult);
              }else{
                  // We've received the final message for a query that has been cancelled
                  // or for which the resources have been freed. We no longer need to listen
                  // for more incoming messages for such a query.
                  DRILL_LOG(LOG_TRACE) << "DrillClientImpl::processQueryResult:" << debugPrintQid(qid)<< " completed."<< std::endl;
                  m_pendingRequests--;
                  DRILL_LOG(LOG_TRACE) << "DrillClientImpl::processQueryResult: pending requests is " << m_pendingRequests<< std::endl;
                  ret=QRY_CANCELED;
              }
              delete allocatedBuffer;
              //return ret;
          }else{
              // Normal query results come back with query_state not set.
              // Actually this is not strictly true. The query state is set to
              // 0(i.e. PENDING), but protobuf thinks this means the value is not set.
              DRILL_LOG(LOG_TRACE) << "DrillClientImpl::processQueryResult: Query State was not set.\n";
          }
      }
      DRILL_LOG(LOG_TRACE) << "DrillClientImpl::processQueryResult: " << m_pendingRequests << " requests pending." << std::endl;
      if(m_pendingRequests==0){
          // signal any waiting client that it can exit because there are no more any query results to arrive.
          // We keep the heartbeat going though.
          m_cv.notify_one();
      }
      return ret;
  }
  
  status_t DrillClientImpl::processQueryData(AllocatedBufferPtr  allocatedBuffer, InBoundRpcMessage& msg ){
      DrillClientQueryResult* pDrillClientQueryResult=NULL;
      status_t ret=QRY_SUCCESS;
      exec::shared::QueryId qid;
      // Be a good client and send ack as early as possible.
      // Drillbit pushed the query result to the client, the client should send ack
      // whenever it receives the message
      sendAck(msg, true);
      RecordBatch* pRecordBatch=NULL;
      {
          boost::lock_guard<boost::mutex> lock(this->m_dcMutex);
          exec::shared::QueryData* qr = new exec::shared::QueryData; //Record Batch will own this object and free it up.
  
          DRILL_LOG(LOG_DEBUG) << "Processing Query Data " << std::endl;
          qr->ParseFromArray(msg.m_pbody.data(), msg.m_pbody.size());
          DRILL_LOG(LOG_TRACE) << qr->DebugString() << std::endl;
  
          qid.CopyFrom(qr->query_id());
          if(qid.part1()==0){
              DRILL_LOG(LOG_TRACE) << "DrillClientImpl::processQueryData: QID=0. Ignore and return QRY_SUCCESS." << std::endl;
              delete allocatedBuffer;
              return QRY_SUCCESS;
          }
  
          pDrillClientQueryResult=findQueryResult(qid);
          if(pDrillClientQueryResult==NULL){
              DRILL_LOG(LOG_TRACE) << "Cleaning up resources allocated for canceled query (" 
                                   << debugPrintQid(qid) << ")." << std::endl;
              delete qr;
              delete allocatedBuffer;
              return ret;
          }
          
          //Validate the RPC message
          std::string valErr;
          if( (ret=validateDataMessage(msg, *qr, valErr)) != QRY_SUCCESS){
              delete allocatedBuffer;
              delete qr;
              DRILL_LOG(LOG_TRACE) << "DrillClientImpl::processQueryData: ERR_QRY_INVRPC.\n";
              pDrillClientQueryResult->setQueryStatus(ret);
              return handleQryError(ret, getMessage(ERR_QRY_INVRPC, valErr.c_str()), pDrillClientQueryResult);
          }
  
          //Build Record Batch here
          DRILL_LOG(LOG_DEBUG) << "Building record batch for Query Id - " << debugPrintQid(qr->query_id()) << std::endl;
  
          pRecordBatch= new RecordBatch(qr, allocatedBuffer,  msg.m_dbody);
          pDrillClientQueryResult->m_numBatches++;
  
          DRILL_LOG(LOG_TRACE) << "Allocated new Record batch." << (void*)pRecordBatch << std::endl;
          pRecordBatch->build();
          DRILL_LOG(LOG_DEBUG) << debugPrintQid(qr->query_id())<<"recordBatch.numRecords "
              << pRecordBatch->getNumRecords()  << std::endl;
          DRILL_LOG(LOG_DEBUG) << debugPrintQid(qr->query_id())<<"recordBatch.numFields "
              << pRecordBatch->getNumFields()  << std::endl;
  
          ret=pDrillClientQueryResult->setupColumnDefs(qr);
          if(ret==QRY_SUCCESS_WITH_INFO){
              pRecordBatch->schemaChanged(true);
          }
  
          pDrillClientQueryResult->setIsQueryPending(true);
          pfnQueryResultsListener pResultsListener=pDrillClientQueryResult->m_pResultsListener;
          if(pDrillClientQueryResult->m_bIsLastChunk){
              DRILL_LOG(LOG_DEBUG) << debugPrintQid(*pDrillClientQueryResult->m_pQueryId)
                  <<  "Received last batch. " << std::endl;
              ret=QRY_NO_MORE_DATA;
          }
          pDrillClientQueryResult->setQueryStatus(ret);
          if(pResultsListener!=NULL){
              ret = pResultsListener(pDrillClientQueryResult, pRecordBatch, NULL);
          }else{
              //Use a default callback that is called when a record batch is received
              ret = pDrillClientQueryResult->defaultQueryResultsListener(pDrillClientQueryResult,
                      pRecordBatch, NULL);
          }
      } // release lock
      if(ret==QRY_FAILURE){
          sendCancel(&qid);
          // Do not decrement pending requests here. We have sent a cancel and we may still receive results that are
          // pushed on the wire before the cancel is processed.
          pDrillClientQueryResult->setIsQueryPending(false);
          DRILL_LOG(LOG_DEBUG) << "Client app cancelled query." << std::endl;
          pDrillClientQueryResult->setQueryStatus(ret);
          clearMapEntries(pDrillClientQueryResult);
          return ret;
      }
      return ret;
  }
  
  status_t DrillClientImpl::processQueryId(AllocatedBufferPtr allocatedBuffer, InBoundRpcMessage& msg ){
      DrillClientQueryResult* pDrillClientQueryResult=NULL;
      DRILL_LOG(LOG_DEBUG) << "Processing Query Handle with coordination id:" << msg.m_coord_id << std::endl;
      status_t ret=QRY_SUCCESS;
  
      boost::lock_guard<boost::mutex> lock(m_dcMutex);
      std::map<int,DrillClientQueryResult*>::iterator it;
      for(it=this->m_queryIds.begin();it!=this->m_queryIds.end();it++){
          std::string qidString = it->second->m_pQueryId!=NULL?debugPrintQid(*it->second->m_pQueryId):std::string("NULL");
          DRILL_LOG(LOG_DEBUG) << "DrillClientImpl::processQueryId: m_queryIds: coordinationId: " << it->first
          << " QueryId: "<< qidString << std::endl;
      }
      if(msg.m_coord_id==0){
          DRILL_LOG(LOG_TRACE) << "DrillClientImpl::processQueryId: m_coord_id=0. Ignore and return QRY_SUCCESS." << std::endl;
          return QRY_SUCCESS;
      }
      it=this->m_queryIds.find(msg.m_coord_id);
      if(it!=this->m_queryIds.end()){
          pDrillClientQueryResult=(*it).second;
          exec::shared::QueryId *qid = new exec::shared::QueryId;
          DRILL_LOG(LOG_TRACE)  << "Received Query Handle " << msg.m_pbody.size() << std::endl;
          qid->ParseFromArray(msg.m_pbody.data(), msg.m_pbody.size());
          DRILL_LOG(LOG_DEBUG) << "Query Id - " << debugPrintQid(*qid) << std::endl;
          m_queryResults[qid]=pDrillClientQueryResult;
          //save queryId allocated here so we can free it later
          pDrillClientQueryResult->setQueryId(qid);
      }else{
          delete allocatedBuffer;
          return handleQryError(QRY_INTERNAL_ERROR, getMessage(ERR_QRY_INVQUERYID), NULL);
      }
      delete allocatedBuffer;
      return ret;
  }
  
  DrillClientQueryResult* DrillClientImpl::findQueryResult(exec::shared::QueryId& qid){
      DrillClientQueryResult* pDrillClientQueryResult=NULL;
      DRILL_LOG(LOG_DEBUG) << "Searching for Query Id - " << debugPrintQid(qid) << std::endl;
      std::map<exec::shared::QueryId*, DrillClientQueryResult*, compareQueryId>::iterator it;
      DRILL_LOG(LOG_TRACE) << "DrillClientImpl::processQueryResult: m_queryResults size: " << m_queryResults.size() << std::endl;
      if(m_queryResults.size() != 0){
          for(it=m_queryResults.begin(); it!=m_queryResults.end(); it++){
              DRILL_LOG(LOG_TRACE) << "DrillClientImpl::findQueryResult: m_QueryResult ids: [" << it->first->part1() << ":"
                  << it->first->part2() << "]\n";
          }
      }
      it=this->m_queryResults.find(&qid);
      if(it!=this->m_queryResults.end()){
          pDrillClientQueryResult=(*it).second;
          DRILL_LOG(LOG_DEBUG) << "Drill Client Query Result Query Id - " <<
              debugPrintQid(*pDrillClientQueryResult->m_pQueryId) << std::endl;
      }
      return pDrillClientQueryResult;
  }
  
  status_t DrillClientImpl::processQueryStatusResult(exec::shared::QueryResult* qr,
          DrillClientQueryResult* pDrillClientQueryResult){
      status_t ret = QUERYSTATE_TO_STATUS_MAP[qr->query_state()];
      if(pDrillClientQueryResult!=NULL){
          pDrillClientQueryResult->setQueryStatus(ret);
          pDrillClientQueryResult->setQueryState(qr->query_state());
      }
      switch(qr->query_state()) {
          case exec::shared::QueryResult_QueryState_FAILED:
              {
                  // get the error message from protobuf and handle errors
                  ret=handleQryError(ret, qr->error(0), pDrillClientQueryResult);
              }
              break;
              // m_pendingRequests should be decremented when the query is
              // completed
          case exec::shared::QueryResult_QueryState_CANCELED:
              {
                  ret=handleTerminatedQryState(ret,
                          getMessage(ERR_QRY_CANCELED),
                          pDrillClientQueryResult);
                  m_pendingRequests--;
              }
              break;
          case exec::shared::QueryResult_QueryState_COMPLETED:
              {
                  //Not clean to call the handleTerminatedQryState method
                  //because it signals an error to the listener.
                  //The ODBC driver expects this though and the sync API
                  //handles this (luckily).
                  ret=handleTerminatedQryState(ret,
                          getMessage(ERR_QRY_COMPLETED),
                          pDrillClientQueryResult);
                  m_pendingRequests--;
              }
              break;
          default:
              {
                  DRILL_LOG(LOG_TRACE) << "DrillClientImpl::processQueryStatusResult: Unknown Query State.\n";
                  ret=handleQryError(QRY_INTERNAL_ERROR,
                          getMessage(ERR_QRY_UNKQRYSTATE),
                          pDrillClientQueryResult);
              }
              break;
      }
      return ret;
  }
  
  void DrillClientImpl::handleReadTimeout(const boost::system::error_code & err){
      // if err == boost::asio::error::operation_aborted) then the caller cancelled the timer.
      if(err != boost::asio::error::operation_aborted){
  
          // Check whether the deadline has passed.
          if (m_deadlineTimer.expires_at() <= boost::asio::deadline_timer::traits_type::now()){
              // The deadline has passed.
              DRILL_LOG(LOG_TRACE) << "DrillClientImpl::handleReadTimeout: Deadline timer expired; ERR_QRY_TIMOUT. \n";
              handleQryError(QRY_TIMEOUT, getMessage(ERR_QRY_TIMOUT), NULL);
              // There is no longer an active deadline. The expiry is set to positive
              // infinity so that the timer never expires until a new deadline is set.
              // Note that at this time, the caller is not in a (async) wait for the timer.
              m_deadlineTimer.expires_at(boost::posix_time::pos_infin);
              // Cancel all pending async IOs.
              // The cancel call _MAY_ not work on all platforms. To be a little more reliable we need
              // to have the BOOST_ASIO_ENABLE_CANCELIO macro (as well as the BOOST_ASIO_DISABLE_IOCP macro?)
              // defined. To be really sure, we need to close the socket. Closing the socket is a bit
              // drastic and we will defer that till a later release.
  #ifdef WIN32_SHUTDOWN_ON_TIMEOUT
              boost::system::error_code ignorederr;
              m_socket.shutdown(boost::asio::ip::tcp::socket::shutdown_both, ignorederr);
  #else // NOT WIN32_SHUTDOWN_ON_TIMEOUT
              m_socket.cancel();
  #endif // WIN32_SHUTDOWN_ON_TIMEOUT
          }
      }
      return;
  }
  
  void DrillClientImpl::handleRead(ByteBuf_t _buf,
          const boost::system::error_code& err,
          size_t bytes_transferred) {
      boost::system::error_code error=err;
      DRILL_LOG(LOG_TRACE) << "DrillClientImpl::handleRead: Handle Read from buffer "
          <<  reinterpret_cast<int*>(_buf) << std::endl;
      if(DrillClientConfig::getQueryTimeout() > 0){
          // Cancel the timeout if handleRead is called
          DRILL_LOG(LOG_TRACE) << "DrillClientImpl::handleRead: Cancel deadline timer.\n";
          m_deadlineTimer.cancel();
      }
      if(!error){
          InBoundRpcMessage msg;
          boost::lock_guard<boost::mutex> lock(this->m_prMutex);
  
          DRILL_LOG(LOG_TRACE) << "Getting new message" << std::endl;
          AllocatedBufferPtr allocatedBuffer=NULL;
  
          if(readMsg(_buf, &allocatedBuffer, msg, error)!=QRY_SUCCESS){
              if(m_pendingRequests!=0){
                  boost::lock_guard<boost::mutex> lock(this->m_dcMutex);
                  getNextResult();
              }
              return;
          }
  
          if(!error && msg.m_mode==exec::rpc::PONG){ //heartbeat response. Throw it away
              m_pendingRequests--;
              delete allocatedBuffer;
              DRILL_LOG(LOG_TRACE) << "Received heartbeat from server. " <<  std::endl;
              if(m_pendingRequests!=0){
                  boost::lock_guard<boost::mutex> lock(this->m_dcMutex);
                  getNextResult();
              }else{
  				boost::unique_lock<boost::mutex> cvLock(this->m_dcMutex);
                  DRILL_LOG(LOG_TRACE) << "No more results expected from server. " <<  std::endl;
  				m_cv.notify_one();
              }
              return;
          }else if(!error && msg.m_rpc_type==exec::user::QUERY_RESULT){
              status_t s = processQueryResult(allocatedBuffer, msg);
              if(s !=QRY_SUCCESS && s!= QRY_NO_MORE_DATA){
                  if(m_pendingRequests!=0){
                      boost::lock_guard<boost::mutex> lock(this->m_dcMutex);
                      getNextResult();
                  }
                  return;
              }
          }else if(!error && msg.m_rpc_type==exec::user::QUERY_DATA){
              if(processQueryData(allocatedBuffer, msg)!=QRY_SUCCESS){
                  if(m_pendingRequests!=0){
                      boost::lock_guard<boost::mutex> lock(this->m_dcMutex);
                      getNextResult();
                  }
                  return;
              }
          }else if(!error && msg.m_rpc_type==exec::user::QUERY_HANDLE){
              if(processQueryId(allocatedBuffer, msg)!=QRY_SUCCESS){
                  if(m_pendingRequests!=0){
                      boost::lock_guard<boost::mutex> lock(this->m_dcMutex);
                      getNextResult();
                  }
                  return;
              }
          }else if(!error && msg.m_rpc_type==exec::user::ACK){
              // Cancel requests will result in an ACK sent back.
              // Consume silently
              delete allocatedBuffer;
              if(m_pendingRequests!=0){
                  boost::lock_guard<boost::mutex> lock(this->m_dcMutex);
                  getNextResult();
              }
              return;
          }else{
              boost::lock_guard<boost::mutex> lock(this->m_dcMutex);
              if(error){
                  // We have a socket read error, but we do not know which query this is for.
                  // Signal ALL pending queries that they should stop waiting.
                  delete allocatedBuffer;
                  DRILL_LOG(LOG_TRACE) << "read error: " << error << std::endl;
                  handleQryError(QRY_COMM_ERROR, getMessage(ERR_QRY_COMMERR, error.message().c_str()), NULL);
                  return;
              }else{
                  // If not QUERY_RESULT, then we think something serious has gone wrong?
                  // In one case when the client hung, we observed that the server was sending a handshake request to the client
                  // We should properly handle these handshake requests/responses
                  if(msg.has_rpc_type() && msg.m_rpc_type==exec::user::HANDSHAKE){
                      if(msg.has_mode() && msg.m_mode==exec::rpc::REQUEST){
                          DRILL_LOG(LOG_TRACE) << "DrillClientImpl::handleRead: Handshake request from server. Send response.\n";
                          exec::user::UserToBitHandshake u2b;
                          u2b.set_channel(exec::shared::USER);
                          u2b.set_rpc_version(DRILL_RPC_VERSION);
                          u2b.set_support_listening(true);
                          OutBoundRpcMessage out_msg(exec::rpc::RESPONSE, exec::user::HANDSHAKE, msg.m_coord_id, &u2b);
                          sendSync(out_msg);
                          DRILL_LOG(LOG_TRACE) << "DrillClientImpl::handleRead: Handshake response sent.\n";
                      }else{
                          DRILL_LOG(LOG_TRACE) << "DrillClientImpl::handleRead: Handshake response from server. Ignore.\n";
                      }
                  }else{
                      DRILL_LOG(LOG_TRACE) << "DrillClientImpl::handleRead: ERR_QRY_INVRPCTYPE. "
                          << "QueryResult returned " << msg.m_rpc_type << std::endl;
                      handleQryError(QRY_INTERNAL_ERROR, getMessage(ERR_QRY_INVRPCTYPE, msg.m_rpc_type), NULL);
                  }
                  delete allocatedBuffer;
                  return;
              }
          }
          {
              boost::lock_guard<boost::mutex> lock(this->m_dcMutex);
              getNextResult();
          }
      }else{
          // boost error
          Utils::freeBuffer(_buf, LEN_PREFIX_BUFLEN);
          boost::lock_guard<boost::mutex> lock(this->m_dcMutex);
          DRILL_LOG(LOG_TRACE) << "DrillClientImpl::handleRead: ERR_QRY_COMMERR. "
              "Boost Communication Error: " << error.message() << std::endl;
          handleQryError(QRY_COMM_ERROR, getMessage(ERR_QRY_COMMERR, error.message().c_str()), NULL);
          return;
      }
      return;
  }
  
  status_t DrillClientImpl::validateDataMessage(InBoundRpcMessage& msg, exec::shared::QueryData& qd, std::string& valErr){
      if(msg.m_mode == exec::rpc::RESPONSE_FAILURE){
          valErr=getMessage(ERR_QRY_RESPFAIL);
          return QRY_FAILURE;
      }
      if(qd.def().carries_two_byte_selection_vector() == true){
          valErr=getMessage(ERR_QRY_SELVEC2);
          return QRY_FAILURE;
      }
      return QRY_SUCCESS;
  }
  
  status_t DrillClientImpl::validateResultMessage(InBoundRpcMessage& msg, exec::shared::QueryResult& qr, std::string& valErr){
      if(msg.m_mode == exec::rpc::RESPONSE_FAILURE){
          valErr=getMessage(ERR_QRY_RESPFAIL);
          return QRY_FAILURE;
      }
      if(qr.query_state()==exec::shared::QueryResult_QueryState_CANCELED){
          valErr=getMessage(ERR_QRY_CANCELED);
          return QRY_FAILURE;
      }
      return QRY_SUCCESS;
  }
  
  connectionStatus_t DrillClientImpl::handleConnError(connectionStatus_t status, std::string msg){
      DrillClientError* pErr = new DrillClientError(status, DrillClientError::CONN_ERROR_START+status, msg);
      m_pendingRequests=0;
      if(!m_queryIds.empty()){
          // set query error only if queries are running
          broadcastError(pErr);
      }else{
          if(m_pError!=NULL){ delete m_pError; m_pError=NULL;}
          m_pError=pErr;
      }
      return status;
  }
  
  status_t DrillClientImpl::handleQryError(status_t status, std::string msg, DrillClientQueryResult* pQueryResult){
      DrillClientError* pErr = new DrillClientError(status, DrillClientError::QRY_ERROR_START+status, msg);
      // set query error only if queries are running
      if(pQueryResult!=NULL){
          m_pendingRequests--;
          pQueryResult->signalError(pErr);
      }else{
          m_pendingRequests=0;
          broadcastError(pErr);
      }
      return status;
  }
  
  status_t DrillClientImpl::handleQryError(status_t status,
          const exec::shared::DrillPBError& e,
          DrillClientQueryResult* pQueryResult){
      assert(pQueryResult!=NULL);
      DrillClientError* pErr =  DrillClientError::getErrorObject(e);
      pQueryResult->signalError(pErr);
      m_pendingRequests--;
      return status;
  }
  
  void DrillClientImpl::broadcastError(DrillClientError* pErr){
      if(pErr!=NULL){
          std::map<int, DrillClientQueryResult*>::iterator iter;
          if(!m_queryIds.empty()){
              for(iter = m_queryIds.begin(); iter != m_queryIds.end(); iter++) {
                  DrillClientError* err=new DrillClientError(pErr->status, pErr->errnum, pErr->msg);
                  iter->second->signalError(err);
              }
          }
          delete pErr;
      }
      // We have an error at the connection level. Cancel the heartbeat. 
      // And close the connection 
      m_heartbeatTimer.cancel();
      m_pendingRequests=0;
      m_cv.notify_one();
      shutdownSocket();
      return;
  }
  
  // The implementation is similar to handleQryError
  status_t DrillClientImpl::handleTerminatedQryState(
          status_t status,
          std::string msg,
          DrillClientQueryResult* pQueryResult){
      assert(pQueryResult!=NULL);
      if(status==QRY_COMPLETED){
          pQueryResult->signalComplete();
      }else{
          // set query error only if queries did not complete successfully
          DrillClientError* pErr = new DrillClientError(status, DrillClientError::QRY_ERROR_START+status, msg);
          pQueryResult->signalError(pErr);
      }
      return status;
  }
  
  
  void DrillClientImpl::clearMapEntries(DrillClientQueryResult* pQueryResult){
      std::map<int, DrillClientQueryResult*>::iterator iter;
      boost::lock_guard<boost::mutex> lock(m_dcMutex);
      if(!m_queryIds.empty()){
          for(iter=m_queryIds.begin(); iter!=m_queryIds.end(); iter++) {
              if(pQueryResult==(DrillClientQueryResult*)iter->second){
                  m_queryIds.erase(iter->first);
                  break;
              }
          }
      }
      if(!m_queryResults.empty()){
          std::map<exec::shared::QueryId*, DrillClientQueryResult*, compareQueryId>::iterator it;
          for(it=m_queryResults.begin(); it!=m_queryResults.end(); it++) {
              if(pQueryResult==(DrillClientQueryResult*)it->second){
                  m_queryResults.erase(it->first);
                  break;
              }
          }
      }
  }
  
  void DrillClientImpl::sendAck(InBoundRpcMessage& msg, bool isOk){
      exec::rpc::Ack ack;
      ack.set_ok(isOk);
      OutBoundRpcMessage ack_msg(exec::rpc::RESPONSE, exec::user::ACK, msg.m_coord_id, &ack);
      boost::lock_guard<boost::mutex> lock(m_dcMutex);
      sendSync(ack_msg);
      DRILL_LOG(LOG_TRACE) << "ACK sent" << std::endl;
  }
  
  void DrillClientImpl::sendCancel(exec::shared::QueryId* pQueryId){
      boost::lock_guard<boost::mutex> lock(m_dcMutex);
      uint64_t coordId = this->getNextCoordinationId();
      OutBoundRpcMessage cancel_msg(exec::rpc::REQUEST, exec::user::CANCEL_QUERY, coordId, pQueryId);
      sendSync(cancel_msg);
      DRILL_LOG(LOG_TRACE) << "CANCEL sent" << std::endl;
  }
  
  void DrillClientImpl::shutdownSocket(){
      m_io_service.stop();
      boost::system::error_code ignorederr;
      m_socket.shutdown(boost::asio::ip::tcp::socket::shutdown_both, ignorederr);
      m_bIsConnected=false;
      DRILL_LOG(LOG_TRACE) << "Socket shutdown" << std::endl;
  }
  
  // This COPIES the FieldMetadata definition for the record batch.  ColumnDefs held by this
  // class are used by the async callbacks.
  status_t DrillClientQueryResult::setupColumnDefs(exec::shared::QueryData* pQueryData) {
      bool hasSchemaChanged=false;
      bool isFirstIter=false;
      boost::lock_guard<boost::mutex> schLock(this->m_schemaMutex);
  
      isFirstIter=this->m_numBatches==1?true:false;
      std::map<std::string, Drill::FieldMetadata*> oldSchema;
      if(!m_columnDefs->empty()){
          for(std::vector<Drill::FieldMetadata*>::iterator it = this->m_columnDefs->begin(); it != this->m_columnDefs->end(); ++it){
              // the key is the field_name + type
              char type[256];
              sprintf(type, ":%d:%d",(*it)->getMinorType(), (*it)->getDataMode() );
              std::string k= (*it)->getName()+type;
              oldSchema[k]=*it;
              delete *it;
          }
      }
      m_columnDefs->clear();
      size_t numFields=pQueryData->def().field_size();
      if (numFields > 0){
          for(size_t i=0; i<numFields; i++){
              Drill::FieldMetadata* fmd= new Drill::FieldMetadata;
              fmd->set(pQueryData->def().field(i));
              this->m_columnDefs->push_back(fmd);
  
              //Look for changes in the vector and trigger a Schema change event if necessary.
              //If vectors are different, then call the schema change listener.
              char type[256];
              sprintf(type, ":%d:%d",fmd->getMinorType(), fmd->getDataMode() );
              std::string k= fmd->getName()+type;
              std::map<std::string, Drill::FieldMetadata*>::iterator iter=oldSchema.find(k);
              if(iter==oldSchema.end()){
                  // not found
                  hasSchemaChanged=true;
              }else{
                  oldSchema.erase(iter);
              }
          }
          if(oldSchema.size()>0){
              hasSchemaChanged=true;
              oldSchema.clear();
          }
      }
      this->m_bHasSchemaChanged=hasSchemaChanged&&!isFirstIter;
      if(this->m_bHasSchemaChanged){
          //invoke schema change Listener
          if(m_pSchemaListener!=NULL){
              m_pSchemaListener(this, m_columnDefs, NULL);
          }
      }
      return this->m_bHasSchemaChanged?QRY_SUCCESS_WITH_INFO:QRY_SUCCESS;
  }
  
  status_t DrillClientQueryResult::defaultQueryResultsListener(void* ctx,
          RecordBatch* b,
          DrillClientError* err) {
      //ctx; // unused, we already have the this pointer
      DRILL_LOG(LOG_TRACE) << "Query result listener called" << std::endl;
      //check if the query has been canceled. IF so then return FAILURE. Caller will send cancel to the server.
      if(this->m_bCancel){
          if(b!=NULL) delete b;
          return QRY_FAILURE;
      }
      if (!err) {
          // signal the cond var
          {
              if(b!=NULL){
  #ifdef DEBUG
                  DRILL_LOG(LOG_DEBUG)<<debugPrintQid(b->getQueryResult()->query_id())
                      << "Query result listener saved result to queue." << std::endl;
  #endif
                  boost::lock_guard<boost::mutex> cvLock(this->m_cvMutex);
                  this->m_recordBatches.push(b);
                  this->m_bHasData=true;
              }
          }
          m_cv.notify_one();
      }else{
          return QRY_FAILURE;
      }
      return QRY_SUCCESS;
  }
  
  RecordBatch*  DrillClientQueryResult::peekNext(){
      RecordBatch* pRecordBatch=NULL;
      boost::unique_lock<boost::mutex> cvLock(this->m_cvMutex);
      //if no more data, return NULL;
      if(!m_bIsQueryPending) return NULL;
      DRILL_LOG(LOG_TRACE) << "Synchronous read waiting for data." << std::endl;
      while(!this->m_bHasData && !m_bHasError && m_bIsQueryPending) {
          this->m_cv.wait(cvLock);
      }
      // READ but not remove first element from queue
      pRecordBatch = this->m_recordBatches.front();
      return pRecordBatch;
  }
  
  RecordBatch*  DrillClientQueryResult::getNext() {
      RecordBatch* pRecordBatch=NULL;
      boost::unique_lock<boost::mutex> cvLock(this->m_cvMutex);
      //if no more data, return NULL;
      if(!m_bIsQueryPending){
          DRILL_LOG(LOG_TRACE) << "Query is done." << std::endl;
          if(!m_recordBatches.empty()){
              DRILL_LOG(LOG_TRACE) << " But there is a Record batch left behind." << std::endl;
          }
          return NULL;
      }
  
      DRILL_LOG(LOG_TRACE) << "Synchronous read waiting for data." << std::endl;
      while(!this->m_bHasData && !m_bHasError && m_bIsQueryPending){
          this->m_cv.wait(cvLock);
      }
      // remove first element from queue
      pRecordBatch = this->m_recordBatches.front();
      this->m_recordBatches.pop();
      this->m_bHasData=!this->m_recordBatches.empty();
      // if vector is empty, set m_bHasDataPending to false;
      m_bIsQueryPending=!(this->m_recordBatches.empty()&&m_queryState==exec::shared::QueryResult_QueryState_COMPLETED);
      return pRecordBatch;
  }
  
  // Blocks until data is available
  void DrillClientQueryResult::waitForData() {
      boost::unique_lock<boost::mutex> cvLock(this->m_cvMutex);
      //if no more data, return NULL;
      if(!m_bIsQueryPending) return;
      while(!this->m_bHasData && !m_bHasError && m_bIsQueryPending) {
          this->m_cv.wait(cvLock);
      }
  }
  
  void DrillClientQueryResult::cancel() {
      this->m_bCancel=true;
  }
  
  void DrillClientQueryResult::signalError(DrillClientError* pErr){
      // Ignore return values from the listener.
      if(pErr!=NULL){
          if(m_pError!=NULL){
              delete m_pError; m_pError=NULL;
          }
          m_pError=pErr;
          pfnQueryResultsListener pResultsListener=this->m_pResultsListener;
          if(pResultsListener!=NULL){
              pResultsListener(this, NULL, pErr);
          }else{
              defaultQueryResultsListener(this, NULL, pErr);
          }
          {
              boost::lock_guard<boost::mutex> cvLock(this->m_cvMutex);
              m_bIsQueryPending=false;
              m_bHasData=false;
              m_bHasError=true;
          }
          //Signal the cv in case there is a client waiting for data already.
          m_cv.notify_one();
      }
      return;
  }
  
  void DrillClientQueryResult::signalComplete(){
      pfnQueryResultsListener pResultsListener=this->m_pResultsListener;
      if(pResultsListener!=NULL){
          pResultsListener(this, NULL, NULL);
      }else{
          defaultQueryResultsListener(this, NULL, NULL);
      }
      {
          boost::lock_guard<boost::mutex> cvLock(this->m_cvMutex);
          m_bIsQueryPending=false;
          m_bIsQueryPending=!(this->m_recordBatches.empty()&&m_queryState==exec::shared::QueryResult_QueryState_COMPLETED);
          m_bHasError=false;
      }
      //Signal the cv in case there is a client waiting for data already.
      m_cv.notify_one();
      return;
  }
  
  void DrillClientQueryResult::clearAndDestroy(){
      //free memory allocated for FieldMetadata objects saved in m_columnDefs;
      if(!m_columnDefs->empty()){
          for(std::vector<Drill::FieldMetadata*>::iterator it = m_columnDefs->begin(); it != m_columnDefs->end(); ++it){
              delete *it;
          }
          m_columnDefs->clear();
      }
      if(this->m_pQueryId!=NULL){
          DRILL_LOG(LOG_TRACE) << "Clearing state for Query Id - " << debugPrintQid(*this->m_pQueryId) << std::endl;
      }
      //Tell the parent to remove this from its lists
      m_pClient->clearMapEntries(this);
  
      //clear query id map entries.
      if(this->m_pQueryId!=NULL){
          delete this->m_pQueryId; this->m_pQueryId=NULL;
      }
      if(!m_recordBatches.empty()){
          // When multiple qwueries execute in parallel we sometimes get an empty record batch back from the server _after_
          // the last chunk has been received. We eventually delete it.
          DRILL_LOG(LOG_TRACE) << "Freeing Record batch(es) left behind "<< std::endl;
          RecordBatch* pR=NULL;
          while(!m_recordBatches.empty()){
              pR=m_recordBatches.front();
              m_recordBatches.pop();
              delete pR;
          }
      }
      if(m_pError!=NULL){
          delete m_pError; m_pError=NULL;
      }
  }
  
  char ZookeeperImpl::s_drillRoot[]="/drill/";
  char ZookeeperImpl::s_defaultCluster[]="drillbits1";
  
  ZookeeperImpl::ZookeeperImpl(){
      m_pDrillbits=new String_vector;
      m_bConnecting=true;
      memset(&m_id, 0, sizeof(m_id));
  }
  
  ZookeeperImpl::~ZookeeperImpl(){
      delete m_pDrillbits;
  }
  
  ZooLogLevel ZookeeperImpl::getZkLogLevel(){
      //typedef enum {ZOO_LOG_LEVEL_ERROR=1,
      //    ZOO_LOG_LEVEL_WARN=2,
      //    ZOO_LOG_LEVEL_INFO=3,
      //    ZOO_LOG_LEVEL_DEBUG=4
      //} ZooLogLevel;
      switch(DrillClientConfig::getLogLevel()){
          case LOG_TRACE:
          case LOG_DEBUG:
              return ZOO_LOG_LEVEL_DEBUG;
          case LOG_INFO:
              return ZOO_LOG_LEVEL_INFO;
          case LOG_WARNING:
              return ZOO_LOG_LEVEL_WARN;
          case LOG_ERROR:
          case LOG_FATAL:
          default:
              return ZOO_LOG_LEVEL_ERROR;
      }
      return ZOO_LOG_LEVEL_ERROR;
  }
  
  int ZookeeperImpl::connectToZookeeper(const char* connectStr, const char* pathToDrill){
      uint32_t waitTime=30000; // 10 seconds
      zoo_set_debug_level(getZkLogLevel());
      zoo_deterministic_conn_order(1); // enable deterministic order
      m_zh = zookeeper_init(connectStr, watcher, waitTime, 0, this, 0);
      if(!m_zh) {
          m_err = getMessage(ERR_CONN_ZKFAIL);
          return CONN_FAILURE;
      }else{
          m_err="";
          //Wait for the completion handler to signal successful connection
          boost::unique_lock<boost::mutex> bufferLock(this->m_cvMutex);
          boost::system_time const timeout=boost::get_system_time()+ boost::posix_time::milliseconds(waitTime);
          while(this->m_bConnecting) {
              if(!this->m_cv.timed_wait(bufferLock, timeout)){
                  m_err = getMessage(ERR_CONN_ZKTIMOUT);
                  return CONN_FAILURE;
              }
          }
      }
      if(m_state!=ZOO_CONNECTED_STATE){
          return CONN_FAILURE;
      }
      int rc = ZOK;
      char rootDir[MAX_CONNECT_STR+1];
      if(pathToDrill==NULL || strlen(pathToDrill)==0){
          strcpy(rootDir, (char*)s_drillRoot);
          strcat(rootDir, s_defaultCluster);
      }else{
          strncpy(rootDir, pathToDrill, MAX_CONNECT_STR); rootDir[MAX_CONNECT_STR]=0;
      }
      rc=zoo_get_children(m_zh, (char*)rootDir, 0, m_pDrillbits);
      if(rc!=ZOK){
          m_err=getMessage(ERR_CONN_ZKERR, rc);
          zookeeper_close(m_zh);
          return -1;
      }
  
      //Let's pick a random drillbit.
      if(m_pDrillbits && m_pDrillbits->count >0){
  
          std::vector<std::string> randomDrillbits;
          for(int i=0; i<m_pDrillbits->count; i++){
              randomDrillbits.push_back(m_pDrillbits->data[i]);
          }
          //Use the same random shuffle as the Java client instead of picking a drillbit at random.
          //Gives much better randomization when the size of the cluster is small.
          std::random_shuffle(randomDrillbits.begin(), randomDrillbits.end());
          const char * bit=randomDrillbits[0].c_str();
          std::string s;
  
          s=rootDir +  std::string("/") + bit;
          int buffer_len=MAX_CONNECT_STR;
          char buffer[MAX_CONNECT_STR+1];
          struct Stat stat;
          buffer[MAX_CONNECT_STR]=0;
          rc= zoo_get(m_zh, s.c_str(), 0, buffer,  &buffer_len, &stat);
          if(rc!=ZOK){
              m_err=getMessage(ERR_CONN_ZKDBITERR, rc);
              zookeeper_close(m_zh);
              return -1;
          }
          m_drillServiceInstance.ParseFromArray(buffer, buffer_len);
      }else{
          m_err=getMessage(ERR_CONN_ZKNODBIT);
          zookeeper_close(m_zh);
          return -1;
      }
      return 0;
  }
  
  void ZookeeperImpl::close(){
      zookeeper_close(m_zh);
  }
  
  void ZookeeperImpl::watcher(zhandle_t *zzh, int type, int state, const char *path, void* context) {
      //From cli.c
  
      /* Be careful using zh here rather than zzh - as this may be mt code
       * the client lib may call the watcher before zookeeper_init returns */
  
      ZookeeperImpl* self=(ZookeeperImpl*)context;
      self->m_state=state;
      if (type == ZOO_SESSION_EVENT) {
          if (state == ZOO_CONNECTED_STATE) {
          } else if (state == ZOO_AUTH_FAILED_STATE) {
              self->m_err= getMessage(ERR_CONN_ZKNOAUTH);
              zookeeper_close(zzh);
              self->m_zh=0;
          } else if (state == ZOO_EXPIRED_SESSION_STATE) {
              self->m_err= getMessage(ERR_CONN_ZKEXP);
              zookeeper_close(zzh);
              self->m_zh=0;
          }
      }
      // signal the cond var
      {
          if (state == ZOO_CONNECTED_STATE){
              DRILL_LOG(LOG_TRACE) << "Connected to Zookeeper." << std::endl;
          }
          boost::lock_guard<boost::mutex> bufferLock(self->m_cvMutex);
          self->m_bConnecting=false;
      }
      self->m_cv.notify_one();
  }
  
  void ZookeeperImpl:: debugPrint(){
      if(m_zh!=NULL && m_state==ZOO_CONNECTED_STATE){
          DRILL_LOG(LOG_TRACE) << m_drillServiceInstance.DebugString() << std::endl;
      }
  }
  
  } // namespace Drill
  