client-cache.h

Go to the documentation of this file.

// Copyright 2012 Cloudera Inc.
//
// Licensed 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 IMPALA_RUNTIME_CLIENT_CACHE_H
#define IMPALA_RUNTIME_CLIENT_CACHE_H

#include <vector>
#include <list>
#include <string>
#include <boost/unordered_map.hpp>
#include <boost/thread/mutex.hpp>
#include <boost/bind.hpp>

#include "util/metrics.h"
#include "rpc/thrift-client.h"
#include "rpc/thrift-util.h"

#include "common/status.h"

namespace impala {

typedef void* ClientKey;

//
//
//
//
class ClientCacheHelper {
 public:
  typedef boost::function<ThriftClientImpl* (const TNetworkAddress& address,
                                             ClientKey* client_key)> ClientFactory;

  //
  Status GetClient(const TNetworkAddress& address, ClientFactory factory_method,
      ClientKey* client_key);

  //
  Status ReopenClient(ClientFactory factory_method, ClientKey* client_key);

  void ReleaseClient(ClientKey* client_key);

  void CloseConnections(const TNetworkAddress& address);

  std::string DebugString();

  void TestShutdown();

  void InitMetrics(MetricGroup* metrics, const std::string& key_prefix);

 private:
  template <class T> friend class ClientCache;
  ClientCacheHelper(uint32_t num_tries, uint64_t wait_ms, int32_t send_timeout_ms,
      int32_t recv_timeout_ms)
      : num_tries_(num_tries),
        wait_ms_(wait_ms),
        send_timeout_ms_(send_timeout_ms),
        recv_timeout_ms_(recv_timeout_ms),
        metrics_enabled_(false) { }


  //
  struct PerHostCache {
    boost::mutex lock;

    std::list<ClientKey> clients;
  };

  boost::mutex cache_lock_;

  typedef boost::unordered_map<
      TNetworkAddress, boost::shared_ptr<PerHostCache> > PerHostCacheMap;
  PerHostCacheMap per_host_caches_;

  boost::mutex client_map_lock_;

  typedef std::map<ClientKey, boost::shared_ptr<ThriftClientImpl> > ClientMap;
  ClientMap client_map_;

  const uint32_t num_tries_;

  const uint64_t wait_ms_;

  const int32_t send_timeout_ms_;

  const int32_t recv_timeout_ms_;

  bool metrics_enabled_;

  IntGauge* clients_in_use_metric_;

  IntGauge* total_clients_metric_;

  Status CreateClient(const TNetworkAddress& address, ClientFactory factory_method,
      ClientKey* client_key);
};

template<class T>
class ClientCache;

template<class T>
class ClientConnection {
 public:
  ClientConnection(ClientCache<T>* client_cache, TNetworkAddress address, Status* status)
    : client_cache_(client_cache), client_(NULL) {
    *status = client_cache_->GetClient(address, &client_);
    if (status->ok()) DCHECK(client_ != NULL);
  }

  ~ClientConnection() {
    if (client_ != NULL) {
      client_cache_->ReleaseClient(&client_);
    }
  }

  Status Reopen() {
    return client_cache_->ReopenClient(&client_);
  }

  T* operator->() const { return client_; }

  //
  //
  template <class F, class Request, class Response>
  Status DoRpc(const F& f, const Request& request, Response* response) {
    DCHECK(response != NULL);
    try {
      (client_->*f)(*response, request);
    } catch (const apache::thrift::TException& e) {
      if (IsTimeoutTException(e)) return Status(TErrorCode::RPC_TIMEOUT);

      // Client may have unexpectedly been closed, so re-open and retry.
      // TODO: ThriftClient should return proper error codes.
      RETURN_IF_ERROR(Reopen());
      try {
        (client_->*f)(*response, request);
      } catch (apache::thrift::TException& e) {
        // By this point the RPC really has failed.
        return Status(TErrorCode::RPC_GENERAL_ERROR, e.what());
      }
    }
    return Status::OK;
  }

 private:
  ClientCache<T>* client_cache_;
  T* client_;
};

template<class T>
class ClientCache {
 public:
  typedef ThriftClient<T> Client;

  ClientCache(const std::string& service_name = "") : client_cache_helper_(1, 0, 0, 0) {
    client_factory_ = boost::bind<ThriftClientImpl*>(
        boost::mem_fn(&ClientCache::MakeClient), this, _1, _2, service_name);
  }

  ClientCache(uint32_t num_tries, uint64_t wait_ms, int32_t send_timeout_ms = 0,
      int32_t recv_timeout_ms = 0, const std::string& service_name = "")
      : client_cache_helper_(num_tries, wait_ms, send_timeout_ms, recv_timeout_ms) {
    client_factory_ =
        boost::bind<ThriftClientImpl*>(
            boost::mem_fn(&ClientCache::MakeClient), this, _1, _2, service_name);
  }

  void CloseConnections(const TNetworkAddress& address) {
    return client_cache_helper_.CloseConnections(address);
  }

  std::string DebugString() {
    return client_cache_helper_.DebugString();
  }

  void TestShutdown() {
    return client_cache_helper_.TestShutdown();
  }

  void InitMetrics(MetricGroup* metrics, const std::string& key_prefix) {
    client_cache_helper_.InitMetrics(metrics, key_prefix);
  }

 private:
  friend class ClientConnection<T>;

  ClientCacheHelper client_cache_helper_;

  ClientCacheHelper::ClientFactory client_factory_;

  Status GetClient(const TNetworkAddress& address, T** iface) {
    return client_cache_helper_.GetClient(address, client_factory_,
        reinterpret_cast<ClientKey*>(iface));
  }

  Status ReopenClient(T** client) {
    return client_cache_helper_.ReopenClient(client_factory_,
        reinterpret_cast<ClientKey*>(client));
  }

  void ReleaseClient(T** client) {
    return client_cache_helper_.ReleaseClient(reinterpret_cast<ClientKey*>(client));
  }

  ThriftClientImpl* MakeClient(const TNetworkAddress& address, ClientKey* client_key,
      const std::string service_name) {
    Client* client = new Client(address.hostname, address.port, service_name);
    *client_key = reinterpret_cast<ClientKey>(client->iface());
    return client;
  }

};


class ImpalaInternalServiceClient;
typedef ClientCache<ImpalaInternalServiceClient> ImpalaInternalServiceClientCache;
typedef ClientConnection<ImpalaInternalServiceClient> ImpalaInternalServiceConnection;

class CatalogServiceClient;
typedef ClientCache<CatalogServiceClient> CatalogServiceClientCache;
typedef ClientConnection<CatalogServiceClient> CatalogServiceConnection;
}

#endif