resource-broker.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 RESOURCE_BROKER_H_
#define RESOURCE_BROKER_H_

#include <boost/unordered_map.hpp>
#include <boost/scoped_ptr.hpp>
#include <boost/uuid/uuid.hpp>

#include "runtime/client-cache.h"
#include "util/collection-metrics.h"
#include "util/promise.h"
#include "gen-cpp/LlamaAMService.h"
#include "gen-cpp/LlamaNotificationService.h"
#include "gen-cpp/ResourceBrokerService_types.h"

namespace impala {

class QueryResourceMgr;
class Status;
class MetricGroup;
class Scheduler;
class ResourceBrokerNotificationServiceClient;

class ResourceBroker {
 public:
  ResourceBroker(const std::vector<TNetworkAddress>& llama_addresses,
      const TNetworkAddress& llama_callback_address, MetricGroup* metrics);

  Status Init();

  void Close();

  Status Reserve(const TResourceBrokerReservationRequest& request,
      TResourceBrokerReservationResponse* response);

  Status Expand(const TResourceBrokerExpansionRequest& request,
      TResourceBrokerExpansionResponse* response);

  void ClearRequests(const TUniqueId& reservation_id, bool include_reservation);

  Status Release(const TResourceBrokerReleaseRequest& request,
      TResourceBrokerReleaseResponse* response);

  void AMNotification(const llama::TLlamaAMNotificationRequest& request,
      llama::TLlamaAMNotificationResponse& response);

  void NMNotification(const llama::TLlamaNMNotificationRequest& request,
      llama::TLlamaNMNotificationResponse& response);

  const std::vector<std::string>& llama_nodes() { return llama_nodes_; }

  Status RefreshLlamaNodes();

  void set_scheduler(Scheduler* scheduler) { scheduler_ = scheduler; };

  bool GetQueryResourceMgr(const TUniqueId& query_id,
      const TUniqueId& reservation_id, const TNetworkAddress& local_resource_address,
      QueryResourceMgr** res_mgr);

  void UnregisterQueryResourceMgr(const TUniqueId& query_id);

 private:
  typedef std::map<TNetworkAddress, llama::TAllocatedResource> ResourceMap;

  bool has_standby_llama() { return llama_addresses_.size() > 1; }

  Status RegisterWithLlama();

  template <typename LlamaReqType, typename LlamaRespType>
  Status LlamaRpc(LlamaReqType* request, LlamaRespType* response,
      StatsMetric<double>* rpc_time_metric);

  template <typename LlamaReqType, typename LlamaRespType>
  void SendLlamaRpc(ClientConnection<llama::LlamaAMServiceClient>* llama_client,
      const LlamaReqType& request, LlamaRespType* response);

  template <typename LlamaReqType, typename LlamaRespType>
  Status ReRegisterWithLlama(const LlamaReqType& request, LlamaRespType* response);

  bool LlamaHasRestarted(const llama::TStatus& status) const;

  void CreateLlamaReservationRequest(const TResourceBrokerReservationRequest& src,
      llama::TLlamaAMReservationRequest& dest);

  void CreateLlamaReleaseRequest(const TResourceBrokerReleaseRequest& src,
      llama::TLlamaAMReleaseRequest& dest);

  class PendingRequest;
  bool WaitForNotification(int64_t timeout, ResourceMap* resources, bool* timed_out,
      PendingRequest* reservation);

  std::vector<TNetworkAddress> llama_addresses_;

  int active_llama_addr_idx_;

  TNetworkAddress llama_callback_address_;

  MetricGroup* metrics_;

  Scheduler* scheduler_;

  StringProperty* active_llama_metric_;

  StringProperty* active_llama_handle_metric_;

  StatsMetric<double>* reservation_rpc_time_metric_;

  StatsMetric<double>* reservation_response_time_metric_;

  IntCounter* reservation_requests_total_metric_;

  IntCounter* reservation_requests_fulfilled_metric_;

  IntCounter* reservation_requests_failed_metric_;

  IntCounter* reservation_requests_rejected_metric_;

  IntCounter* reservation_requests_timedout_metric_;

  StatsMetric<double>* expansion_rpc_time_metric_;

  StatsMetric<double>* expansion_response_time_metric_;

  IntCounter* expansion_requests_total_metric_;

  IntCounter* expansion_requests_fulfilled_metric_;

  IntCounter* expansion_requests_failed_metric_;

  IntCounter* expansion_requests_rejected_metric_;

  IntCounter* expansion_requests_timedout_metric_;

  UIntGauge* allocated_memory_metric_;

  UIntGauge* allocated_vcpus_metric_;

  IntCounter* requests_released_metric_;

  boost::uuids::uuid llama_client_id_;

  boost::shared_ptr<llama::LlamaNotificationServiceIf> llama_callback_thrift_iface_;
  boost::scoped_ptr<ThriftServer> llama_callback_server_;

  boost::scoped_ptr<ClientCache<llama::LlamaAMServiceClient> > llama_client_cache_;

  boost::mutex llama_registration_lock_;

  llama::TUniqueId llama_handle_;

  std::vector<std::string> llama_nodes_;

  //
  class PendingRequest {
   public:
    PendingRequest(const llama::TUniqueId& reservation_id,
        const llama::TUniqueId& request_id, bool is_expansion)
        : reservation_id_(reservation_id), request_id_(request_id),
          is_expansion_(is_expansion) {
      DCHECK(is_expansion || reservation_id == request_id);
    }

    Promise<bool>* promise() { return &promise_; }

    void GetResources(ResourceMap* resources);

    void SetResources(const std::vector<llama::TAllocatedResource>& resources);

    const llama::TUniqueId& request_id() const { return request_id_; }
    const llama::TUniqueId& reservation_id() const { return reservation_id_; }

    bool is_expansion() const { return is_expansion_; }

   private:
    Promise<bool> promise_;

    std::vector<llama::TAllocatedResource> allocated_resources_;

    llama::TUniqueId reservation_id_;

    llama::TUniqueId request_id_;

    bool is_expansion_;
  };

  boost::mutex pending_requests_lock_;

  typedef boost::unordered_map<llama::TUniqueId, PendingRequest*> PendingRequestMap;
  PendingRequestMap pending_requests_;

  class AllocatedRequest {
   public:
    AllocatedRequest(const llama::TUniqueId& reservation_id,
        uint64_t memory_mb, uint32_t vcpus, bool is_expansion)
        : reservation_id_(reservation_id), memory_mb_(memory_mb), vcpus_(vcpus),
          is_expansion_(is_expansion) { }

    const llama::TUniqueId reservation_id() const { return reservation_id_; }
    uint64_t memory_mb() const { return memory_mb_; }
    uint32_t vcpus() const { return vcpus_; }
    bool is_expansion() const { return is_expansion_; }

   private:
    llama::TUniqueId reservation_id_;

    uint64_t memory_mb_;

    uint32_t vcpus_;

    bool is_expansion_;
  };

  boost::mutex allocated_requests_lock_;

  typedef boost::unordered_map<llama::TUniqueId, std::vector<AllocatedRequest> >
      AllocatedRequestMap;
  AllocatedRequestMap allocated_requests_;

  boost::mutex query_resource_mgrs_lock_;
  typedef boost::unordered_map<TUniqueId, std::pair<int32_t, QueryResourceMgr*> >
      QueryResourceMgrsMap;

  QueryResourceMgrsMap query_resource_mgrs_;
};

std::ostream& operator<<(std::ostream& os,
    const TResourceBrokerReservationRequest& request);

std::ostream& operator<<(std::ostream& os,
    const TResourceBrokerReservationResponse& reservation);

std::ostream& operator<<(std::ostream& os,
    const TResourceBrokerExpansionRequest& request);

std::ostream& operator<<(std::ostream& os,
    const TResourceBrokerExpansionResponse& expansion);
}

#endif