buffered-block-mgr.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_BUFFERED_BLOCK_MGR
#define IMPALA_RUNTIME_BUFFERED_BLOCK_MGR

#include <boost/shared_ptr.hpp>

#include "runtime/disk-io-mgr.h"
#include "runtime/tmp-file-mgr.h"

#include <openssl/aes.h>
#include <openssl/sha.h>

namespace impala {

class RuntimeState;

//
//
//
//
//
//
//
//
class BufferedBlockMgr {
 private:
  struct BufferDescriptor;

 public:
  struct Client;

  //
  //
  //
  class Block : public InternalQueue<Block>::Node {
   public:
    Status Pin(bool* pinned, Block* release_block = NULL, bool unpin = true);

    Status Unpin();

    Status Delete();

    void AddRow() { ++num_rows_; }
    int num_rows() const { return num_rows_; }

    template <typename T> T* Allocate(int size) {
      DCHECK_GE(BytesRemaining(), size);
      uint8_t* current_location = buffer_desc_->buffer + valid_data_len_;
      valid_data_len_ += size;
      return reinterpret_cast<T*>(current_location);
    }

    int BytesRemaining() const {
      DCHECK_NOTNULL(buffer_desc_);
      return buffer_desc_->len - valid_data_len_;
    }

    void ReturnAllocation(int size) {
      DCHECK_GE(valid_data_len_, size);
      valid_data_len_ -= size;
    }

    uint8_t* buffer() const {
      DCHECK_NOTNULL(buffer_desc_);
      return buffer_desc_->buffer;
    }

    int64_t valid_data_len() const { return valid_data_len_; }

    int64_t buffer_len() const {
      DCHECK(is_pinned());
      return buffer_desc_->len;
    }

    bool is_max_size() const {
      DCHECK(is_pinned());
      return buffer_desc_->len == block_mgr_->max_block_size();
    }

    bool is_pinned() const { return is_pinned_; }

    std::string DebugString() const;

   private:
    friend class BufferedBlockMgr;

    Block(BufferedBlockMgr* block_mgr);

    void Init();

    bool Validate() const;

    BufferDescriptor* buffer_desc_;

    BufferedBlockMgr* block_mgr_;

    Client* client_;

    DiskIoMgr::WriteRange* write_range_;

    int64_t valid_data_len_;

    int num_rows_;

    boost::scoped_array<uint8_t> encrypted_write_buffer_;

    uint8_t key_[32];

    uint8_t iv_[AES_BLOCK_SIZE];

    uint8_t hash_[SHA256_DIGEST_LENGTH];


    bool is_pinned_;

    bool in_write_;

    bool is_deleted_;

    boost::condition_variable write_complete_cv_;

    bool client_local_;
  }; // class Block

  static Status Create(RuntimeState* state, MemTracker* parent,
      RuntimeProfile* profile, int64_t mem_limit, int64_t buffer_size,
      boost::shared_ptr<BufferedBlockMgr>* block_mgr);

  ~BufferedBlockMgr();

  Status RegisterClient(int num_reserved_buffers, MemTracker* tracker,
      RuntimeState* state, Client** client);

  void ClearReservations(Client* client);

  bool TryAcquireTmpReservation(Client* client, int num_buffers);

  void ClearTmpReservation(Client* client);

  static int GetNumReservedBlocks() {
    return TmpFileMgr::num_tmp_devices();
  }

  Status GetNewBlock(Client* client, Block* unpin_block, Block** block, int64_t len = -1);

  void Cancel();

  std::string DebugString(Client* client = NULL);

  bool ConsumeMemory(Client* client, int64_t size);

  void ReleaseMemory(Client* client, int64_t size);

  int64_t available_buffers(Client* client) const;

  Status MemLimitTooLowError(Client* client);

  int available_allocated_buffers() const { return all_io_buffers_.size(); }
  int num_free_buffers() const { return free_io_buffers_.size(); }

  int num_pinned_buffers(Client* client) const;
  int num_reserved_buffers_remaining(Client* client) const;
  MemTracker* get_tracker(Client* client) const;
  int64_t max_block_size() const { return max_block_size_; }
  int64_t bytes_allocated() const;
  RuntimeProfile* profile() { return profile_.get(); }
  int writes_issued() const { return writes_issued_; }

 private:
  friend struct Client;

  struct BufferDescriptor : public InternalQueue<BufferDescriptor>::Node {
    uint8_t* buffer;

    int64_t len;

    Block* block;

    std::list<BufferDescriptor*>::iterator all_buffers_it;

    BufferDescriptor(uint8_t* buf, int64_t len)
      : buffer(buf), len(len), block(NULL) {
    }
  };

  BufferedBlockMgr(RuntimeState* state, int64_t block_size);

  void Init(DiskIoMgr* io_mgr, RuntimeProfile* profile,
      MemTracker* parent_tracker, int64_t mem_limit);

  Status InitTmpFiles();

  Status PinBlock(Block* block, bool* pinned, Block* src, bool unpin);
  Status UnpinBlock(Block* block);
  Status DeleteBlock(Block* block);

  Status DeleteOrUnpinBlock(Block* block, bool unpin);

  Status TransferBuffer(Block* dst, Block* src, bool unpin);

  int64_t remaining_unreserved_buffers() const;

  Status FindBufferForBlock(Block* block, bool* in_mem);

  Status FindBuffer(boost::unique_lock<boost::mutex>& lock,
      BufferDescriptor** buffer);

  Status WriteUnpinnedBlocks();

  Status WriteUnpinnedBlock(Block* block);

  void WriteComplete(Block* block, const Status& write_status);

  void ReturnUnusedBlock(Block* block);

  Block* GetUnusedBlock(Client* client);

  bool Validate() const;
  std::string DebugInternal() const;

  const int64_t max_block_size_;

  const int block_write_threshold_;

  const bool disable_spill_;

  const TUniqueId query_id_;

  ObjectPool obj_pool_;

  boost::scoped_ptr<MemTracker> mem_tracker_;

  boost::mutex lock_;

  bool initialized_;

  int unfullfilled_reserved_buffers_;

  int total_pinned_buffers_;

  int non_local_outstanding_writes_;

  boost::condition_variable buffer_available_cv_;

  InternalQueue<Block> unpinned_blocks_;

  InternalQueue<Block> unused_blocks_;

  InternalQueue<BufferDescriptor> free_io_buffers_;

  std::list<BufferDescriptor*> all_io_buffers_;

  boost::ptr_vector<TmpFileMgr::File> tmp_files_;

  int next_block_index_;

  DiskIoMgr* io_mgr_;
  DiskIoMgr::RequestContext* io_request_context_;

  bool is_cancelled_;

  boost::scoped_ptr<RuntimeProfile> profile_;

  RuntimeProfile::Counter* mem_limit_counter_;
  RuntimeProfile::Counter* block_size_counter_;

  RuntimeProfile::Counter* created_block_counter_;

  RuntimeProfile::Counter* recycled_blocks_counter_;

  RuntimeProfile::Counter* buffered_pin_counter_;

  RuntimeProfile::Counter* disk_read_timer_;

  RuntimeProfile::Counter* buffer_wait_timer_;

  RuntimeProfile::Counter* bytes_written_counter_;

  RuntimeProfile::Counter* outstanding_writes_counter_;

  RuntimeProfile::Counter* encryption_timer_;

  RuntimeProfile::Counter* integrity_check_timer_;

  int writes_issued_;

  static SpinLock static_block_mgrs_lock_;

  typedef boost::unordered_map<TUniqueId, boost::weak_ptr<BufferedBlockMgr> >
      BlockMgrsMap;
  static BlockMgrsMap query_to_block_mgrs_;

  Status Encrypt(Block* block, uint8_t** outbuf);

  void EncryptDone(Block* block);

  Status Decrypt(Block* block);

  void SetHash(Block* block);

  Status VerifyHash(Block* block);

  const bool encryption_;

  const bool check_integrity_;
}; // class BufferedBlockMgr

} // namespace impala.

#endif