disk-io-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_DISK_IO_MGR_H
#define IMPALA_RUNTIME_DISK_IO_MGR_H

#include <list>
#include <vector>
#include <boost/scoped_ptr.hpp>
#include <boost/unordered_set.hpp>
#include <boost/thread/mutex.hpp>
#include <boost/thread/condition_variable.hpp>
#include <boost/thread/thread.hpp>

#include "common/atomic.h"
#include "common/hdfs.h"
#include "common/object-pool.h"
#include "common/status.h"
#include "runtime/thread-resource-mgr.h"
#include "util/bit-util.h"
#include "util/error-util.h"
#include "util/internal-queue.h"
#include "util/runtime-profile.h"
#include "util/thread.h"

namespace impala {

class MemTracker;

//
//
//
//
//
//
//
//
//
//
//
//
//
//
//
//
class DiskIoMgr {
 public:
  class RequestContext;
  class ScanRange;

  class BufferDescriptor {
   public:
    ScanRange* scan_range() { return scan_range_; }
    char* buffer() { return buffer_; }
    int64_t buffer_len() { return buffer_len_; }
    int64_t len() { return len_; }
    bool eosr() { return eosr_; }

    int64_t scan_range_offset() const { return scan_range_offset_; }

    void SetMemTracker(MemTracker* tracker);

    void Return();

   private:
    friend class DiskIoMgr;
    BufferDescriptor(DiskIoMgr* io_mgr);

    void Reset(RequestContext* reader, ScanRange* range, char* buffer,
        int64_t buffer_len);

    DiskIoMgr* io_mgr_;

    RequestContext* reader_;

    MemTracker* mem_tracker_;

    ScanRange* scan_range_;

    char* buffer_;

    int64_t buffer_len_;

    int64_t len_;

    bool eosr_;

    Status status_;

    int64_t scan_range_offset_;
  };

  struct RequestType {
    enum type {
      READ,
      WRITE,
    };
  };

  class RequestRange : public InternalQueue<RequestRange>::Node {
   public:
    hdfsFS fs() const { return fs_; }
    const char* file() const { return file_.c_str(); }
    int64_t offset() const { return offset_; }
    int64_t len() const { return len_; }
    int disk_id() const { return disk_id_; }
    RequestType::type request_type() const { return request_type_; }

   protected:
    hdfsFS fs_;

    std::string file_;

    int64_t offset_;

    int64_t len_;

    int disk_id_;

    RequestType::type request_type_;
  };

  class ScanRange : public RequestRange {
   public:

    const static int64_t NEVER_CACHE = -1;

    ScanRange(int initial_capacity = -1);

    virtual ~ScanRange();

    void Reset(hdfsFS fs, const char* file, int64_t len, int64_t offset, int disk_id,
        bool try_cache, bool expected_local, int64_t mtime, void* metadata = NULL);

    void* meta_data() const { return meta_data_; }
    bool try_cache() const { return try_cache_; }
    bool expected_local() const { return expected_local_; }
    int ready_buffers_capacity() const { return ready_buffers_capacity_; }

    Status GetNext(BufferDescriptor** buffer);

    void Cancel(const Status& status);

    std::string DebugString() const;

    int64_t mtime() const { return mtime_; }

   private:
    friend class DiskIoMgr;

    void InitInternal(DiskIoMgr* io_mgr, RequestContext* reader);

    bool EnqueueBuffer(BufferDescriptor* buffer);

    void CleanupQueuedBuffers();

    bool Validate();

    int64_t MaxReadChunkSize() const;

    Status Open();

    void Close();

    Status Read(char* buffer, int64_t* bytes_read, bool* eosr);

    Status ReadFromCache(bool* read_succeeded);

    void* meta_data_;

    bool try_cache_;

    bool expected_local_;

    DiskIoMgr* io_mgr_;

    RequestContext* reader_;

    union {
      FILE* local_file_;
      hdfsFile hdfs_file_;
    };

    struct hadoopRzBuffer* cached_buffer_;

    boost::mutex lock_;

    int bytes_read_;

    Status status_;

    bool eosr_queued_;

    bool eosr_returned_;

    bool blocked_on_queue_;

    boost::condition_variable buffer_ready_cv_;
    std::list<BufferDescriptor*> ready_buffers_;

    int ready_buffers_capacity_;

    boost::mutex hdfs_lock_;

    bool is_cancelled_;

    int64_t mtime_;
  };

  class WriteRange : public RequestRange {
   public:
    typedef boost::function<void (const Status&)> WriteDoneCallback;
    WriteRange(const std::string& file, int64_t file_offset, int disk_id,
        WriteDoneCallback callback);

    void SetData(const uint8_t* buffer, int64_t len);

   private:
    friend class DiskIoMgr;

    const uint8_t* data_;

    WriteDoneCallback callback_;
  };

  DiskIoMgr(int num_disks, int threads_per_disk, int min_buffer_size,
      int max_buffer_size);

  DiskIoMgr();

  ~DiskIoMgr();

  Status Init(MemTracker* process_mem_tracker);

  Status RegisterContext(RequestContext** request_context,
      MemTracker* reader_mem_tracker = NULL);

  void UnregisterContext(RequestContext* context);

  void CancelContext(RequestContext* context, bool wait_for_disks_completion = false);

  Status AddScanRanges(RequestContext* reader, const std::vector<ScanRange*>& ranges,
      bool schedule_immediately = false);

  Status AddWriteRange(RequestContext* writer, WriteRange* write_range);

  Status GetNextRange(RequestContext* reader, ScanRange** range);

  Status Read(RequestContext* reader, ScanRange* range, BufferDescriptor** buffer);

  int AssignQueue(const char* file, int disk_id, bool expected_local);

  Status context_status(RequestContext* context) const;

  int num_unstarted_ranges(RequestContext* reader) const;

  void set_bytes_read_counter(RequestContext*, RuntimeProfile::Counter*);
  void set_read_timer(RequestContext*, RuntimeProfile::Counter*);
  void set_active_read_thread_counter(RequestContext*, RuntimeProfile::Counter*);
  void set_disks_access_bitmap(RequestContext*, RuntimeProfile::Counter*);

  int64_t queue_size(RequestContext* reader) const;
  int64_t bytes_read_local(RequestContext* reader) const;
  int64_t bytes_read_short_circuit(RequestContext* reader) const;
  int64_t bytes_read_dn_cache(RequestContext* reader) const;
  int num_remote_ranges(RequestContext* reader) const;
  int64_t unexpected_remote_bytes(RequestContext* reader) const;

  int64_t GetReadThroughput();

  int max_read_buffer_size() const { return max_buffer_size_; }

  int num_total_disks() const { return disk_queues_.size(); }

  int num_remote_disks() const { return REMOTE_NUM_DISKS; }

  int num_local_disks() const { return num_total_disks() - num_remote_disks(); }

  int RemoteS3DiskId() const { return num_local_disks() + REMOTE_S3_DISK_OFFSET; }

  int num_allocated_buffers() const { return num_allocated_buffers_; }

  int num_buffers_in_readers() const { return num_buffers_in_readers_; }

  std::string DebugString();

  bool Validate() const;

  static const int DEFAULT_QUEUE_CAPACITY;

  enum {
    REMOTE_S3_DISK_OFFSET = 0,
    REMOTE_NUM_DISKS
  };

 private:
  friend class BufferDescriptor;
  struct DiskQueue;
  class RequestContextCache;

  friend class DiskIoMgrTest_Buffers_Test;

  ObjectPool pool_;

  MemTracker* process_mem_tracker_;

  const int num_threads_per_disk_;

  const int max_buffer_size_;

  const int min_buffer_size_;

  ThreadGroup disk_thread_group_;

  struct hadoopRzOptions* cached_read_options_;

  volatile bool shut_down_;

  RuntimeProfile::Counter total_bytes_read_counter_;

  RuntimeProfile::Counter read_timer_;

  boost::scoped_ptr<RequestContextCache> request_context_cache_;

  boost::mutex free_buffers_lock_;

  //
  std::vector<std::list<char*> > free_buffers_;

  std::list<BufferDescriptor*> free_buffer_descs_;

  AtomicInt<int> num_allocated_buffers_;

  AtomicInt<int> num_buffers_in_readers_;

  std::vector<DiskQueue*> disk_queues_;

  int free_buffers_idx(int64_t buffer_size);

  BufferDescriptor* GetBufferDesc(
      RequestContext* reader, ScanRange* range, char* buffer, int64_t buffer_size);

  void ReturnBufferDesc(BufferDescriptor* desc);

  void ReturnBuffer(BufferDescriptor* buffer);

  char* GetFreeBuffer(int64_t* buffer_size);

  void GcIoBuffers();

  void ReturnFreeBuffer(char* buffer, int64_t buffer_size);

  void ReturnFreeBuffer(BufferDescriptor* desc);

  void WorkLoop(DiskQueue* queue);

  bool GetNextRequestRange(DiskQueue* disk_queue, RequestRange** range,
      RequestContext** request_context);

  void HandleReadFinished(DiskQueue*, RequestContext*, BufferDescriptor*);

  void HandleWriteFinished(RequestContext* writer, WriteRange* write_range,
      const Status& write_status);

  Status ValidateScanRange(ScanRange* range);

  void Write(RequestContext* writer_context, WriteRange* write_range);

  Status WriteRangeHelper(FILE* file_handle, WriteRange* write_range);

  void ReadRange(DiskQueue* disk_queue, RequestContext* reader,
      ScanRange* range);
};

}

#endif