text: // Copyright (c) 2017, Baidu.com, Inc. All Rights Reserved
  
  // 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.
  
  #include "runtime/disk_io_mgr.h"
  
  #include <sched.h>
  #include <sys/stat.h>
  #include <boost/bind.hpp>
  #include <boost/thread/thread.hpp>
  
  #include <gtest/gtest.h>
  
  #include "util/cpu_info.h"
  #include "util/disk_info.h"
  #include "util/logging.h"
  
  using std::string;
  using std::stringstream;
  using std::vector;
  using std::list;
  
  using boost::lock_guard;
  using boost::unique_lock;
  using boost::mutex;
  using boost::mem_fn;
  using boost::condition_variable;
  using boost::scoped_ptr;
  using boost::thread;
  using boost::thread_group;
  
  namespace palo {
  
  const int MIN_BUFFER_SIZE = 512;
  const int MAX_BUFFER_SIZE = 1024;
  const int LARGE_MEM_LIMIT = 1024 * 1024 * 1024;
  
  class DiskIoMgrTest : public testing::Test {
  public:
      void write_validate_callback(int num_writes, DiskIoMgr::WriteRange** written_range,
              DiskIoMgr* io_mgr, DiskIoMgr::RequestContext* reader, int32_t* data,
              Status expected_status, const Status& status) {
          if (expected_status.code() == TStatusCode::CANCELLED) {
              EXPECT_TRUE(status.ok() || status.is_cancelled());
          } else {
              EXPECT_TRUE(status.code() == expected_status.code());
          }
          if (status.ok()) {
              DiskIoMgr::ScanRange* scan_range = _pool->add(new DiskIoMgr::ScanRange());
              scan_range->reset(NULL, (*written_range)->file(), (*written_range)->len(),
                      (*written_range)->offset(), 0, false, false, DiskIoMgr::ScanRange::NEVER_CACHE);
              validate_sync_read(io_mgr, reader, scan_range, reinterpret_cast<const char*>(data),
                      sizeof(int32_t));
          }
  
          {
              lock_guard<mutex> l(_written_mutex);
              ++_num_ranges_written;
              if (_num_ranges_written == num_writes) {
                  _writes_done.notify_one();
              }
          }
      }
  
      void write_complete_callback(int num_writes, const Status& status) {
          EXPECT_TRUE(status.ok());
          {
              lock_guard<mutex> l(_written_mutex);
              ++_num_ranges_written;
              if (_num_ranges_written == num_writes) {
                  _writes_done.notify_all();
              }
          }
      }
  
  protected:
      void CreateTempFile(const char* filename, const char* data) {
          FILE* file = fopen(filename, "w");
          EXPECT_TRUE(file != NULL);
          fwrite(data, 1, strlen(data), file);
          fclose(file);
      }
  
      int CreateTempFile(const char* filename, int file_size) {
          FILE* file = fopen(filename, "w");
          EXPECT_TRUE(file != NULL);
          int success = fclose(file);
          if (success != 0) {
              LOG(ERROR) << "Error closing file " << filename;
              return success;
          }
          return truncate(filename, file_size);
      }
  
      // Validates that buffer[i] is \0 or expected[i]
      static void validate_empty_or_correct(const char* expected, const char* buffer, int len) {
          for (int i = 0; i < len; ++i) {
              if (buffer[i] != '\0') {
                  EXPECT_EQ(expected[i], buffer[i]) << (int)expected[i] << " != " << (int)buffer[i];
              }
          }
      }
  
      static void validate_sync_read(DiskIoMgr* io_mgr, DiskIoMgr::RequestContext* reader,
              DiskIoMgr::ScanRange* range, const char* expected, int expected_len = -1) {
          DiskIoMgr::BufferDescriptor* buffer;
          Status status = io_mgr->read(reader, range, &buffer);
          ASSERT_TRUE(status.ok());
          ASSERT_TRUE(buffer != NULL);
          EXPECT_EQ(buffer->len(), range->len());
          if (expected_len < 0) {
              expected_len = strlen(expected);
          }
          int cmp = memcmp(buffer->buffer(), expected, expected_len);
          EXPECT_TRUE(cmp == 0);
          buffer->return_buffer();
      }
  
      static void validate_scan_range(DiskIoMgr::ScanRange* range, const char* expected,
              int expected_len, const Status& expected_status) {
          char result[expected_len + 1];
          memset(result, 0, expected_len + 1);
  
          while (true) {
              DiskIoMgr::BufferDescriptor* buffer = NULL;
              Status status = range->get_next(&buffer);
              ASSERT_TRUE(status.ok() || status.code() == expected_status.code());
              if (buffer == NULL || !status.ok()) {
                  if (buffer != NULL) buffer->return_buffer();
                  break;
              }
              ASSERT_LE(buffer->len(), expected_len);
              memcpy(result + range->offset() + buffer->scan_range_offset(),
                      buffer->buffer(), buffer->len());
              buffer->return_buffer();
          }
          validate_empty_or_correct(expected, result, expected_len);
      }
  
      // Continues pulling scan ranges from the io mgr until they are all done.
      // Updates num_ranges_processed with the number of ranges seen by this thread.
      static void scan_range_thread(DiskIoMgr* io_mgr, DiskIoMgr::RequestContext* reader,
              const char* expected_result, int expected_len, const Status& expected_status,
              int max_ranges, AtomicInt<int>* num_ranges_processed) {
          int num_ranges = 0;
          while (max_ranges == 0 || num_ranges < max_ranges) {
              DiskIoMgr::ScanRange* range;
              Status status = io_mgr->get_next_range(reader, &range);
              ASSERT_TRUE(status.ok() || status.code() == expected_status.code());
              if (range == NULL) break;
              validate_scan_range(range, expected_result, expected_len, expected_status);
              ++(*num_ranges_processed);
              ++num_ranges;
          }
      }
  
      DiskIoMgr::ScanRange* init_range(int num_buffers, const char* file_path, int offset,
              int len, int disk_id, int64_t mtime, void* meta_data = NULL, bool is_cached = false) {
          DiskIoMgr::ScanRange* range = _pool->add(new DiskIoMgr::ScanRange(num_buffers));
          range->reset(NULL, file_path, len, offset, disk_id, is_cached, true, mtime, meta_data);
          EXPECT_EQ(mtime, range->mtime());
          return range;
      }
  
      scoped_ptr<ObjectPool> _pool;
  
      mutex _written_mutex;
      condition_variable _writes_done;
      int _num_ranges_written;
  };
  
  // Test a single writer with multiple disks and threads per disk. Each WriteRange
  // writes random 4-byte integers, and upon completion, the written data is validated
  // by reading the data back via a separate IoMgr instance. All writes are expected to
  // complete successfully.
  TEST_F(DiskIoMgrTest, SingleWriter) {
      // MemTracker mem_tracker(LARGE_MEM_LIMIT);
      MemTracker mem_tracker(LARGE_MEM_LIMIT);
      _num_ranges_written = 0;
      string tmp_file = "/tmp/disk_io_mgr_test.txt";
      int num_ranges = 100;
      int64_t file_size = 1024 * 1024;
      int64_t cur_offset = 0;
      int success = CreateTempFile(tmp_file.c_str(), file_size);
      if (success != 0) {
          LOG(ERROR) << "Error creating temp file " << tmp_file.c_str() << " of size " <<
              file_size;
          EXPECT_TRUE(false);
      }
  
      scoped_ptr<DiskIoMgr> read_io_mgr(new DiskIoMgr(1, 1, 1, 10));
      // MemTracker reader_mem_tracker(LARGE_MEM_LIMIT);
      MemTracker reader_mem_tracker(LARGE_MEM_LIMIT);
      Status status = read_io_mgr->init(&reader_mem_tracker);
      ASSERT_TRUE(status.ok());
      DiskIoMgr::RequestContext* reader;
      status = read_io_mgr->register_context(&reader, &reader_mem_tracker);
      ASSERT_TRUE(status.ok());
      for (int num_threads_per_disk = 1; num_threads_per_disk <= 5; ++num_threads_per_disk) {
          for (int num_disks = 1; num_disks <= 5; num_disks += 2) {
              _pool.reset(new ObjectPool);
              DiskIoMgr io_mgr(num_disks, num_threads_per_disk, 1, 10);
              status = io_mgr.init(&mem_tracker);
              ASSERT_TRUE(status.ok());
              DiskIoMgr::RequestContext* writer;
              io_mgr.register_context(&writer, &mem_tracker);
              for (int i = 0; i < num_ranges; ++i) {
                  int32_t* data = _pool->add(new int32_t);
                  *data = rand();
                  DiskIoMgr::WriteRange** new_range = _pool->add(new DiskIoMgr::WriteRange*);
                  DiskIoMgr::WriteRange::WriteDoneCallback callback =
                      bind(mem_fn(&DiskIoMgrTest::write_validate_callback), this, num_ranges,
                              new_range, read_io_mgr.get(), reader, data, Status::OK, _1);
                  *new_range = _pool->add(new DiskIoMgr::WriteRange(tmp_file, cur_offset,
                              num_ranges % num_disks, callback));
                  (*new_range)->set_data(reinterpret_cast<uint8_t*>(data), sizeof(int32_t));
                  Status add_status = io_mgr.add_write_range(writer, *new_range);
                  EXPECT_TRUE(add_status.ok());
                  cur_offset += sizeof(int32_t);
              }
  
              {
                  unique_lock<mutex> lock(_written_mutex);
                  while (_num_ranges_written < num_ranges) {
                      _writes_done.wait(lock);
                  }
              }
              _num_ranges_written = 0;
              io_mgr.unregister_context(writer);
          }
      }
  
      read_io_mgr->unregister_context(reader);
      read_io_mgr.reset();
  }
  
  // Perform invalid writes (e.g. non-existent file, negative offset) and validate
  // that an error status is returned via the write callback.
  TEST_F(DiskIoMgrTest, InvalidWrite) {
      MemTracker mem_tracker(LARGE_MEM_LIMIT);
      _num_ranges_written = 0;
      string tmp_file = "/tmp/non-existent.txt";
      DiskIoMgr io_mgr(1, 1, 1, 10);
      Status status = io_mgr.init(&mem_tracker);
      ASSERT_TRUE(status.ok());
      DiskIoMgr::RequestContext* writer;
      status = io_mgr.register_context(&writer);
      _pool.reset(new ObjectPool);
      int32_t* data = _pool->add(new int32_t);
      *data = rand();
  
      // Write to a non-existent file.
      DiskIoMgr::WriteRange** new_range = _pool->add(new DiskIoMgr::WriteRange*);
      DiskIoMgr::WriteRange::WriteDoneCallback callback =
          bind(mem_fn(&DiskIoMgrTest::write_validate_callback), this, 2,
                  new_range, (DiskIoMgr*)NULL, (DiskIoMgr::RequestContext*)NULL,
                  data, Status("Test Failure"), _1);
      *new_range = _pool->add(new DiskIoMgr::WriteRange(tmp_file, rand(), 0, callback));
  
      (*new_range)->set_data(reinterpret_cast<uint8_t*>(data), sizeof(int32_t));
      status = io_mgr.add_write_range(writer, *new_range);
      EXPECT_TRUE(status.ok());
  
      // Write to a bad location in a file that exists.
      tmp_file = "/tmp/disk_io_mgr_test.txt";
      int success = CreateTempFile(tmp_file.c_str(), 100);
      if (success != 0) {
          LOG(ERROR) << "Error creating temp file " << tmp_file.c_str() << " of size 100";
          EXPECT_TRUE(false);
      }
  
      new_range = _pool->add(new DiskIoMgr::WriteRange*);
      callback = bind(mem_fn(&DiskIoMgrTest::write_validate_callback), this, 2,
              new_range, (DiskIoMgr*)NULL, (DiskIoMgr::RequestContext*)NULL,
              data, Status("Test Failure"), _1);
  
      *new_range = _pool->add(new DiskIoMgr::WriteRange(tmp_file, -1, 0, callback));
      (*new_range)->set_data(reinterpret_cast<uint8_t*>(data), sizeof(int32_t));
      status = io_mgr.add_write_range(writer, *new_range);
      EXPECT_TRUE(status.ok());
  
      {
          unique_lock<mutex> lock(_written_mutex);
          while (_num_ranges_written < 2) {
              _writes_done.wait(lock);
          }
      }
      _num_ranges_written = 0;
      io_mgr.unregister_context(writer);
  }
  
  // Issue a number of writes, cancel the writer context and issue more writes.
  // add_write_range() is expected to succeed before the cancel and fail after it.
  // The writes themselves may finish with status cancelled or ok.
  TEST_F(DiskIoMgrTest, SingleWriterCancel) {
      MemTracker mem_tracker(LARGE_MEM_LIMIT);
      _num_ranges_written = 0;
      string tmp_file = "/tmp/disk_io_mgr_test.txt";
      int num_ranges = 100;
      int num_ranges_before_cancel = 25;
      int64_t file_size = 1024 * 1024;
      int64_t cur_offset = 0;
      int success = CreateTempFile(tmp_file.c_str(), file_size);
      if (success != 0) {
          LOG(ERROR) << "Error creating temp file " << tmp_file.c_str() << " of size " <<
              file_size;
          EXPECT_TRUE(false);
      }
  
      scoped_ptr<DiskIoMgr> read_io_mgr(new DiskIoMgr(1, 1, 1, 10));
      MemTracker reader_mem_tracker(LARGE_MEM_LIMIT);
      Status status = read_io_mgr->init(&reader_mem_tracker);
      ASSERT_TRUE(status.ok());
      DiskIoMgr::RequestContext* reader;
      status = read_io_mgr->register_context(&reader, &reader_mem_tracker);
      ASSERT_TRUE(status.ok());
      for (int num_threads_per_disk = 1; num_threads_per_disk <= 5; ++num_threads_per_disk) {
          for (int num_disks = 1; num_disks <= 5; num_disks += 2) {
              _pool.reset(new ObjectPool);
              DiskIoMgr io_mgr(num_disks, num_threads_per_disk, 1, 10);
              status = io_mgr.init(&mem_tracker);
              DiskIoMgr::RequestContext* writer;
              io_mgr.register_context(&writer, &mem_tracker);
              Status validate_status = Status::OK;
              for (int i = 0; i < num_ranges; ++i) {
                  if (i == num_ranges_before_cancel) {
                      io_mgr.cancel_context(writer);
                      validate_status = Status::CANCELLED;
                  }
                  int32_t* data = _pool->add(new int32_t);
                  *data = rand();
                  DiskIoMgr::WriteRange** new_range = _pool->add(new DiskIoMgr::WriteRange*);
                  DiskIoMgr::WriteRange::WriteDoneCallback callback =
                      bind(mem_fn(&DiskIoMgrTest::write_validate_callback), this,
                              num_ranges_before_cancel, new_range, read_io_mgr.get(), reader, data,
                              Status::CANCELLED, _1);
                  *new_range = _pool->add(new DiskIoMgr::WriteRange(tmp_file, cur_offset,
                              num_ranges % num_disks, callback));
                  (*new_range)->set_data(reinterpret_cast<uint8_t*>(data), sizeof(int32_t));
                  cur_offset += sizeof(int32_t);
                  Status add_status = io_mgr.add_write_range(writer, *new_range);
                  EXPECT_TRUE(add_status.code() == validate_status.code());
              }
  
              {
                  unique_lock<mutex> lock(_written_mutex);
                  while (_num_ranges_written < num_ranges_before_cancel) {
                      _writes_done.wait(lock);
                  }
              }
              _num_ranges_written = 0;
              io_mgr.unregister_context(writer);
          }
      }
  
      read_io_mgr->unregister_context(reader);
      read_io_mgr.reset();
  }
  
  // Basic test with a single reader, testing multiple threads, disks and a different
  // number of buffers.
  TEST_F(DiskIoMgrTest, SingleReader) {
      MemTracker mem_tracker(LARGE_MEM_LIMIT);
      const char* tmp_file = "/tmp/disk_io_mgr_test.txt";
      const char* data = "abcdefghijklm";
      int len = strlen(data);
      CreateTempFile(tmp_file, data);
  
      // Get mtime for file
      struct stat stat_val;
      stat(tmp_file, &stat_val);
  
      int64_t iters = 0;
      for (int num_threads_per_disk = 1; num_threads_per_disk <= 5; ++num_threads_per_disk) {
          for (int num_disks = 1; num_disks <= 5; num_disks += 2) {
              for (int num_buffers = 1; num_buffers <= 5; ++num_buffers) {
                  for (int num_read_threads = 1; num_read_threads <= 5; ++num_read_threads) {
                      _pool.reset(new ObjectPool);
                      LOG(INFO) << "Starting test with num_threads_per_disk=" << num_threads_per_disk
                          << " num_disk=" << num_disks << " num_buffers=" << num_buffers
                          << " num_read_threads=" << num_read_threads;
  
                      if (++iters % 5000 == 0) {
                          LOG(ERROR) << "Starting iteration " << iters;
                      }
                      DiskIoMgr io_mgr(num_disks, num_threads_per_disk, 1, 1);
  
                      Status status = io_mgr.init(&mem_tracker);
                      ASSERT_TRUE(status.ok());
                      MemTracker reader_mem_tracker(LARGE_MEM_LIMIT);
                      DiskIoMgr::RequestContext* reader;
                      status = io_mgr.register_context(&reader, &reader_mem_tracker);
                      ASSERT_TRUE(status.ok());
  
                      vector<DiskIoMgr::ScanRange*> ranges;
                      for (int i = 0; i < len; ++i) {
                          int disk_id = i % num_disks;
                          ranges.push_back(init_range(num_buffers, tmp_file, 0, len, disk_id,
                                      stat_val.st_mtime));
                      }
                      status = io_mgr.add_scan_ranges(reader, ranges);
                      ASSERT_TRUE(status.ok());
  
                      AtomicInt<int> num_ranges_processed;
                      thread_group threads;
                      for (int i = 0; i < num_read_threads; ++i) {
                          threads.add_thread(new thread(scan_range_thread, &io_mgr, reader, data,
                                      len, Status::OK, 0, &num_ranges_processed));
                      }
                      threads.join_all();
  
                      EXPECT_EQ(num_ranges_processed, ranges.size());
                      io_mgr.unregister_context(reader);
                      EXPECT_EQ(reader_mem_tracker.consumption(), 0);
                  }
              }
          }
      }
      EXPECT_EQ(mem_tracker.consumption(), 0);
  }
  
  // This test issues adding additional scan ranges while there are some still in flight.
  TEST_F(DiskIoMgrTest, AddScanRangeTest) {
      MemTracker mem_tracker(LARGE_MEM_LIMIT);
      const char* tmp_file = "/tmp/disk_io_mgr_test.txt";
      const char* data = "abcdefghijklm";
      int len = strlen(data);
      CreateTempFile(tmp_file, data);
  
      // Get mtime for file
      struct stat stat_val;
      stat(tmp_file, &stat_val);
  
      int64_t iters = 0;
      for (int num_threads_per_disk = 1; num_threads_per_disk <= 5; ++num_threads_per_disk) {
          for (int num_disks = 1; num_disks <= 5; num_disks += 2) {
              for (int num_buffers = 1; num_buffers <= 5; ++num_buffers) {
                  _pool.reset(new ObjectPool);
                  LOG(INFO) << "Starting test with num_threads_per_disk=" << num_threads_per_disk
                      << " num_disk=" << num_disks << " num_buffers=" << num_buffers;
  
                  if (++iters % 5000 == 0) LOG(ERROR) << "Starting iteration " << iters;
                  DiskIoMgr io_mgr(num_disks, num_threads_per_disk, 1, 1);
  
                  Status status = io_mgr.init(&mem_tracker);
                  ASSERT_TRUE(status.ok());
                  MemTracker reader_mem_tracker(LARGE_MEM_LIMIT);
                  DiskIoMgr::RequestContext* reader;
                  status = io_mgr.register_context(&reader, &reader_mem_tracker);
                  ASSERT_TRUE(status.ok());
  
                  vector<DiskIoMgr::ScanRange*> ranges_first_half;
                  vector<DiskIoMgr::ScanRange*> ranges_second_half;
                  for (int i = 0; i < len; ++i) {
                      int disk_id = i % num_disks;
                      if (i > len / 2) {
                          ranges_second_half.push_back(
                                  init_range(num_buffers, tmp_file, i, 1, disk_id,
                                      stat_val.st_mtime));
                      } else {
                          ranges_first_half.push_back(init_range(num_buffers, tmp_file, i, 1, disk_id,
                                      stat_val.st_mtime));
                      }
                  }
                  AtomicInt<int> num_ranges_processed;
  
                  // Issue first half the scan ranges.
                  status = io_mgr.add_scan_ranges(reader, ranges_first_half);
                  ASSERT_TRUE(status.ok());
  
                  // Read a couple of them
                  scan_range_thread(&io_mgr, reader, data, strlen(data), Status::OK, 2,
                          &num_ranges_processed);
  
                  // Issue second half
                  status = io_mgr.add_scan_ranges(reader, ranges_second_half);
                  ASSERT_TRUE(status.ok());
  
                  // Start up some threads and then cancel
                  thread_group threads;
                  for (int i = 0; i < 3; ++i) {
                      threads.add_thread(new thread(scan_range_thread, &io_mgr, reader, data,
                                  strlen(data), Status::CANCELLED, 0, &num_ranges_processed));
                  }
  
                  threads.join_all();
                  EXPECT_EQ(num_ranges_processed, len);
                  io_mgr.unregister_context(reader);
                  EXPECT_EQ(reader_mem_tracker.consumption(), 0);
              }
          }
      }
      EXPECT_EQ(mem_tracker.consumption(), 0);
  }
  
  // Test to make sure that sync reads and async reads work together
  // Note: this test is constructed so the number of buffers is greater than the
  // number of scan ranges.
  TEST_F(DiskIoMgrTest, SyncReadTest) {
      MemTracker mem_tracker(LARGE_MEM_LIMIT);
      const char* tmp_file = "/tmp/disk_io_mgr_test.txt";
      const char* data = "abcdefghijklm";
      int len = strlen(data);
      CreateTempFile(tmp_file, data);
  
      // Get mtime for file
      struct stat stat_val;
      stat(tmp_file, &stat_val);
  
      int64_t iters = 0;
      for (int num_threads_per_disk = 1; num_threads_per_disk <= 5; ++num_threads_per_disk) {
          for (int num_disks = 1; num_disks <= 5; num_disks += 2) {
              for (int num_buffers = 1; num_buffers <= 5; ++num_buffers) {
                  _pool.reset(new ObjectPool);
                  LOG(INFO) << "Starting SyncReadTest test with num_threads_per_disk=" << num_threads_per_disk
                      << " num_disk=" << num_disks << " num_buffers=" << num_buffers;
  
                  if (++iters % 5000 == 0) {
                      LOG(ERROR) << "Starting iteration " << iters;
                  }
                  DiskIoMgr io_mgr(
                          num_disks, num_threads_per_disk, MIN_BUFFER_SIZE, MAX_BUFFER_SIZE);
  
                  Status status = io_mgr.init(&mem_tracker);
                  ASSERT_TRUE(status.ok());
                  MemTracker reader_mem_tracker(LARGE_MEM_LIMIT);
                  DiskIoMgr::RequestContext* reader;
                  status = io_mgr.register_context(&reader, &reader_mem_tracker);
                  ASSERT_TRUE(status.ok());
  
                  DiskIoMgr::ScanRange* complete_range = init_range(1, tmp_file, 0, strlen(data), 0,
                          stat_val.st_mtime);
  
                  // Issue some reads before the async ones are issued
                  validate_sync_read(&io_mgr, reader, complete_range, data);
                  validate_sync_read(&io_mgr, reader, complete_range, data);
  
                  vector<DiskIoMgr::ScanRange*> ranges;
                  for (int i = 0; i < len; ++i) {
                      int disk_id = i % num_disks;
                      ranges.push_back(init_range(num_buffers, tmp_file, 0, len, disk_id,
                                  stat_val.st_mtime));
                  }
                  status = io_mgr.add_scan_ranges(reader, ranges);
                  ASSERT_TRUE(status.ok());
  
                  AtomicInt<int> num_ranges_processed;
                  thread_group threads;
                  for (int i = 0; i < 5; ++i) {
                      threads.add_thread(new thread(scan_range_thread, &io_mgr, reader, data,
                                  strlen(data), Status::OK, 0, &num_ranges_processed));
                  }
  
                  // Issue some more sync ranges
                  for (int i = 0; i < 5; ++i) {
                      sched_yield();
                      validate_sync_read(&io_mgr, reader, complete_range, data);
                  }
  
                  threads.join_all();
  
                  validate_sync_read(&io_mgr, reader, complete_range, data);
                  validate_sync_read(&io_mgr, reader, complete_range, data);
  
                  EXPECT_EQ(num_ranges_processed, ranges.size());
                  io_mgr.unregister_context(reader);
                  EXPECT_EQ(reader_mem_tracker.consumption(), 0);
              }
          }
      }
      EXPECT_EQ(mem_tracker.consumption(), 0);
  }
  
  // Tests a single reader cancelling half way through scan ranges.
  TEST_F(DiskIoMgrTest, SingleReaderCancel) {
      MemTracker mem_tracker(LARGE_MEM_LIMIT);
      const char* tmp_file = "/tmp/disk_io_mgr_test.txt";
      const char* data = "abcdefghijklm";
      int len = strlen(data);
      CreateTempFile(tmp_file, data);
  
      // Get mtime for file
      struct stat stat_val;
      stat(tmp_file, &stat_val);
  
      int64_t iters = 0;
      for (int num_threads_per_disk = 1; num_threads_per_disk <= 5; ++num_threads_per_disk) {
          for (int num_disks = 1; num_disks <= 5; num_disks += 2) {
              for (int num_buffers = 1; num_buffers <= 5; ++num_buffers) {
                  _pool.reset(new ObjectPool);
                  LOG(INFO) << "Starting test with num_threads_per_disk=" << num_threads_per_disk
                      << " num_disk=" << num_disks << " num_buffers=" << num_buffers;
  
                  if (++iters % 5000 == 0) LOG(ERROR) << "Starting iteration " << iters;
                  DiskIoMgr io_mgr(num_disks, num_threads_per_disk, 1, 1);
  
                  Status status = io_mgr.init(&mem_tracker);
                  ASSERT_TRUE(status.ok());
                  MemTracker reader_mem_tracker(LARGE_MEM_LIMIT);
                  DiskIoMgr::RequestContext* reader;
                  status = io_mgr.register_context(&reader, &reader_mem_tracker);
                  ASSERT_TRUE(status.ok());
  
                  vector<DiskIoMgr::ScanRange*> ranges;
                  for (int i = 0; i < len; ++i) {
                      int disk_id = i % num_disks;
                      ranges.push_back(init_range(num_buffers, tmp_file, 0, len, disk_id,
                                  stat_val.st_mtime));
                  }
                  status = io_mgr.add_scan_ranges(reader, ranges);
                  ASSERT_TRUE(status.ok());
  
                  AtomicInt<int> num_ranges_processed;
                  int num_succesful_ranges = ranges.size() / 2;
                  // Read half the ranges
                  for (int i = 0; i < num_succesful_ranges; ++i) {
                      scan_range_thread(&io_mgr, reader, data, strlen(data), Status::OK, 1,
                              &num_ranges_processed);
                  }
                  EXPECT_EQ(num_ranges_processed, num_succesful_ranges);
  
                  // Start up some threads and then cancel
                  thread_group threads;
                  for (int i = 0; i < 3; ++i) {
                      threads.add_thread(new thread(scan_range_thread, &io_mgr, reader, data,
                                  strlen(data), Status::CANCELLED, 0, &num_ranges_processed));
                  }
  
                  io_mgr.cancel_context(reader);
                  sched_yield();
  
                  threads.join_all();
                  EXPECT_TRUE(io_mgr.context_status(reader).is_cancelled());
                  io_mgr.unregister_context(reader);
                  EXPECT_EQ(reader_mem_tracker.consumption(), 0);
              }
          }
      }
      EXPECT_EQ(mem_tracker.consumption(), 0);
  }
  
  // Test when the reader goes over the mem limit
  TEST_F(DiskIoMgrTest, MemTrackers) {
      const char* tmp_file = "/tmp/disk_io_mgr_test.txt";
      const char* data = "abcdefghijklm";
      int len = strlen(data);
      CreateTempFile(tmp_file, data);
  
      // Get mtime for file
      struct stat stat_val;
      stat(tmp_file, &stat_val);
  
      const int num_buffers = 25;
      // Give the reader more buffers than the limit
      const int mem_limit_num_buffers = 2;
  
      int64_t iters = 0;
      {
          _pool.reset(new ObjectPool);
          if (++iters % 1000 == 0) {
              LOG(ERROR) << "Starting iteration " << iters;
          }
  
          MemTracker mem_tracker(mem_limit_num_buffers * MAX_BUFFER_SIZE);
          DiskIoMgr io_mgr(1, 1, MIN_BUFFER_SIZE, MAX_BUFFER_SIZE);
  
          Status status = io_mgr.init(&mem_tracker);
          ASSERT_TRUE(status.ok());
          MemTracker reader_mem_tracker(LARGE_MEM_LIMIT);
          DiskIoMgr::RequestContext* reader;
          status = io_mgr.register_context(&reader, &reader_mem_tracker);
          ASSERT_TRUE(status.ok());
  
          vector<DiskIoMgr::ScanRange*> ranges;
          for (int i = 0; i < num_buffers; ++i) {
              ranges.push_back(init_range(num_buffers, tmp_file, 0, len, 0,
                          stat_val.st_mtime));
          }
          status = io_mgr.add_scan_ranges(reader, ranges);
          ASSERT_TRUE(status.ok());
  
          // Don't return buffers to force memory pressure
          vector<DiskIoMgr::BufferDescriptor*> buffers;
  
          AtomicInt<int> num_ranges_processed;
          scan_range_thread(&io_mgr, reader, data, strlen(data), Status::MEM_LIMIT_EXCEEDED,
                  1, &num_ranges_processed);
  
          char result[strlen(data) + 1];
          // Keep reading new ranges without returning buffers. This forces us
          // to go over the limit eventually.
          while (true) {
              memset(result, 0, strlen(data) + 1);
              DiskIoMgr::ScanRange* range = NULL;
              status = io_mgr.get_next_range(reader, &range);
              ASSERT_TRUE(status.ok() || status.is_mem_limit_exceeded());
              if (range == NULL) break;
  
              while (true) {
                  DiskIoMgr::BufferDescriptor* buffer = NULL;
                  Status status = range->get_next(&buffer);
                  ASSERT_TRUE(status.ok() || status.is_mem_limit_exceeded());
                  if (buffer == NULL) break;
                  memcpy(result + range->offset() + buffer->scan_range_offset(),
                          buffer->buffer(), buffer->len());
                  buffers.push_back(buffer);
              }
              validate_empty_or_correct(data, result, strlen(data));
          }
  
          for (int i = 0; i < buffers.size(); ++i) {
              buffers[i]->return_buffer();
          }
  
          EXPECT_TRUE(io_mgr.context_status(reader).is_mem_limit_exceeded());
          io_mgr.unregister_context(reader);
          EXPECT_EQ(reader_mem_tracker.consumption(), 0);
      }
  }
  #if 0
  // Test when some scan ranges are marked as being cached.
  // Since these files are not in HDFS, the cached path always fails so this
  // only tests the fallback mechanism.
  // TODO: we can fake the cached read path without HDFS
  TEST_F(DiskIoMgrTest, CachedReads) {
      MemTracker mem_tracker(LARGE_MEM_LIMIT);
      const char* tmp_file = "/tmp/disk_io_mgr_test.txt";
      const char* data = "abcdefghijklm";
      int len = strlen(data);
      CreateTempFile(tmp_file, data);
  
      // Get mtime for file
      struct stat stat_val;
      stat(tmp_file, &stat_val);
  
      const int num_disks = 2;
      const int num_buffers = 3;
  
      int64_t iters = 0;
      {
          _pool.reset(new ObjectPool);
          if (++iters % 5000 == 0) LOG(ERROR) << "Starting iteration " << iters;
          DiskIoMgr io_mgr(num_disks, 1, MIN_BUFFER_SIZE, MAX_BUFFER_SIZE);
  
          Status status = io_mgr.init(&mem_tracker);
          ASSERT_TRUE(status.ok());
          MemTracker reader_mem_tracker;
          DiskIoMgr::RequestContext* reader;
          status = io_mgr.register_context(&reader, &reader_mem_tracker);
          ASSERT_TRUE(status.ok());
  
          DiskIoMgr::ScanRange* complete_range =
              init_range(1, tmp_file, 0, strlen(data), 0, stat_val.st_mtime, NULL, true);
  
          // Issue some reads before the async ones are issued
          validate_sync_read(&io_mgr, reader, complete_range, data);
          validate_sync_read(&io_mgr, reader, complete_range, data);
  
          vector<DiskIoMgr::ScanRange*> ranges;
          for (int i = 0; i < len; ++i) {
              int disk_id = i % num_disks;
              ranges.push_back(init_range(num_buffers, tmp_file, 0, len, disk_id,
                          stat_val.st_mtime, NULL, true));
          }
          status = io_mgr.add_scan_ranges(reader, ranges);
          ASSERT_TRUE(status.ok());
  
          AtomicInt<int> num_ranges_processed;
          thread_group threads;
          for (int i = 0; i < 5; ++i) {
              threads.add_thread(new thread(scan_range_thread, &io_mgr, reader, data,
                          strlen(data), Status::OK, 0, &num_ranges_processed));
          }
  
          // Issue some more sync ranges
          for (int i = 0; i < 5; ++i) {
              sched_yield();
              validate_sync_read(&io_mgr, reader, complete_range, data);
          }
  
          threads.join_all();
  
          validate_sync_read(&io_mgr, reader, complete_range, data);
          validate_sync_read(&io_mgr, reader, complete_range, data);
  
          EXPECT_EQ(num_ranges_processed, ranges.size());
          io_mgr.unregister_context(reader);
          EXPECT_EQ(reader_mem_tracker.consumption(), 0);
      }
      EXPECT_EQ(mem_tracker.consumption(), 0);
  }
  #endif // end #if 0
  
  TEST_F(DiskIoMgrTest, MultipleReaderWriter) {
      MemTracker mem_tracker(LARGE_MEM_LIMIT);
      const int ITERATIONS = 1;
      const char* data = "abcdefghijklmnopqrstuvwxyz";
      const int num_contexts = 5;
      const int file_size = 4 * 1024;
      const int num_writes_queued = 5;
      const int num_reads_queued = 5;
  
      string file_name = "/tmp/disk_io_mgr_test.txt";
      int success = CreateTempFile(file_name.c_str(), file_size);
      if (success != 0) {
          LOG(ERROR) << "Error creating temp file " << file_name.c_str() << " of size " <<
              file_size;
          ASSERT_TRUE(false);
      }
  
      // Get mtime for file
      struct stat stat_val;
      stat(file_name.c_str(), &stat_val);
  
      int64_t iters = 0;
      vector<DiskIoMgr::RequestContext*> contexts(num_contexts);
      Status status;
      for (int iteration = 0; iteration < ITERATIONS; ++iteration) {
          for (int threads_per_disk = 1; threads_per_disk <= 5; ++threads_per_disk) {
              for (int num_disks = 1; num_disks <= 5; num_disks += 2) {
                  DiskIoMgr io_mgr(num_disks, threads_per_disk, MIN_BUFFER_SIZE, MAX_BUFFER_SIZE);
                  io_mgr.init(&mem_tracker);
                  for (int file_index = 0; file_index < num_contexts; ++file_index) {
                      status = io_mgr.register_context(&contexts[file_index]);
                      ASSERT_TRUE(status.ok());
                  }
                  _pool.reset(new ObjectPool);
                  int read_offset = 0;
                  int write_offset = 0;
                  while (read_offset < file_size) {
                      for (int context_index = 0; context_index < num_contexts; ++context_index) {
                          if (++iters % 5000 == 0) {
                              LOG(ERROR) << "Starting iteration " << iters;
                          }
                          AtomicInt<int> num_ranges_processed;
                          thread_group threads;
                          vector<DiskIoMgr::ScanRange*> ranges;
                          int num_scan_ranges = std::min<int>(num_reads_queued, write_offset - read_offset);
                          for (int i = 0; i < num_scan_ranges; ++i) {
                              ranges.push_back(init_range(1, file_name.c_str(), read_offset, 1,
                                          i % num_disks, stat_val.st_mtime));
                              threads.add_thread(new thread(scan_range_thread, &io_mgr,
                                          contexts[context_index],
                                          reinterpret_cast<const char*>(data + (read_offset % strlen(data))), 1,
                                          Status::OK, num_scan_ranges, &num_ranges_processed));
                              ++read_offset;
                          }
  
                          _num_ranges_written = 0;
                          int num_write_ranges = std::min<int>(num_writes_queued, file_size - write_offset);
                          for (int i = 0; i < num_write_ranges; ++i) {
                              DiskIoMgr::WriteRange::WriteDoneCallback callback =
                                  bind(mem_fn(&DiskIoMgrTest::write_complete_callback),
                                          this, num_write_ranges, _1);
                              DiskIoMgr::WriteRange* new_range = _pool->add(
                                      new DiskIoMgr::WriteRange(file_name,
                                          write_offset, i % num_disks, callback));
                              new_range->set_data(reinterpret_cast<const uint8_t*>
                                      (data + (write_offset % strlen(data))), 1);
                              status = io_mgr.add_write_range(contexts[context_index], new_range);
                              ++write_offset;
                          }
  
                          {
                              unique_lock<mutex> lock(_written_mutex);
                              while (_num_ranges_written < num_write_ranges) {
                                  _writes_done.wait(lock);
                              }
                          }
  
                          threads.join_all();
                      } // for (int context_index
                  } // while (read_offset < file_size)
  
  
                  for (int file_index = 0; file_index < num_contexts; ++file_index) {
                      io_mgr.unregister_context(contexts[file_index]);
                  }
              } // for (int num_disks
          } // for (int threads_per_disk
      } // for (int iteration
  }
  
  // This test will test multiple concurrent reads each reading a different file.
  TEST_F(DiskIoMgrTest, MultipleReader) {
      MemTracker mem_tracker(LARGE_MEM_LIMIT);
      const int NUM_READERS = 5;
      const int DATA_LEN = 50;
      const int ITERATIONS = 25;
      const int NUM_THREADS_PER_READER = 3;
  
      vector<string> file_names;
      vector<int64_t> mtimes;
      vector<string> data;
      vector<DiskIoMgr::RequestContext*> readers;
      vector<char*> results;
  
      file_names.resize(NUM_READERS);
      readers.resize(NUM_READERS);
      mtimes.resize(NUM_READERS);
      data.resize(NUM_READERS);
      results.resize(NUM_READERS);
  
      // Initialize data for each reader.  The data will be
      // 'abcd...' for reader one, 'bcde...' for reader two (wrapping around at 'z')
      for (int i = 0; i < NUM_READERS; ++i) {
          char buf[DATA_LEN];
          for (int j = 0; j < DATA_LEN; ++j) {
              int c = (j + i) % 26;
              buf[j] = 'a' + c;
          }
          data[i] = string(buf, DATA_LEN);
  
          stringstream ss;
          ss << "/tmp/disk_io_mgr_test" << i << ".txt";
          file_names[i] = ss.str();
          CreateTempFile(ss.str().c_str(), data[i].c_str());
  
          // Get mtime for file
          struct stat stat_val;
          stat(file_names[i].c_str(), &stat_val);
          mtimes[i] = stat_val.st_mtime;
  
          results[i] = new char[DATA_LEN + 1];
          memset(results[i], 0, DATA_LEN + 1);
      }
  
      // This exercises concurrency, run the test multiple times
      int64_t iters = 0;
      for (int iteration = 0; iteration < ITERATIONS; ++iteration) {
          for (int threads_per_disk = 1; threads_per_disk <= 5; ++threads_per_disk) {
              for (int num_disks = 1; num_disks <= 5; num_disks += 2) {
                  for (int num_buffers = 1; num_buffers <= 5; ++num_buffers) {
                      _pool.reset(new ObjectPool);
                      LOG(INFO) << "Starting test with num_threads_per_disk=" << threads_per_disk
                          << " num_disk=" << num_disks << " num_buffers=" << num_buffers;
                      if (++iters % 2500 == 0) LOG(ERROR) << "Starting iteration " << iters;
  
                      DiskIoMgr io_mgr(num_disks, threads_per_disk, MIN_BUFFER_SIZE, MAX_BUFFER_SIZE);
                      Status status = io_mgr.init(&mem_tracker);
                      ASSERT_TRUE(status.ok());
  
                      for (int i = 0; i < NUM_READERS; ++i) {
                          status = io_mgr.register_context(&readers[i], NULL);
                          ASSERT_TRUE(status.ok());
  
                          vector<DiskIoMgr::ScanRange*> ranges;
                          for (int j = 0; j < DATA_LEN; ++j) {
                              int disk_id = j % num_disks;
                              ranges.push_back(
                                      init_range(num_buffers,file_names[i].c_str(), j, 1, disk_id,
                                          mtimes[i]));
                          }
                          status = io_mgr.add_scan_ranges(readers[i], ranges);
                          ASSERT_TRUE(status.ok());
                      }
  
                      AtomicInt<int> num_ranges_processed;
                      thread_group threads;
                      for (int i = 0; i < NUM_READERS; ++i) {
                          for (int j = 0; j < NUM_THREADS_PER_READER; ++j) {
                              threads.add_thread(new thread(scan_range_thread, &io_mgr, readers[i],
                                          data[i].c_str(), data[i].size(), Status::OK, 0,
                                          &num_ranges_processed));
                          }
                      }
                      threads.join_all();
                      EXPECT_EQ(num_ranges_processed, DATA_LEN * NUM_READERS);
                      for (int i = 0; i < NUM_READERS; ++i) {
                          io_mgr.unregister_context(readers[i]);
                      }
                  }
              }
          }
      }
      EXPECT_EQ(mem_tracker.consumption(), 0);
  }
  
  #if 0
  // Stress test for multiple clients with cancellation
  // TODO: the stress app should be expanded to include sync reads and adding scan
  // ranges in the middle.
  TEST_F(DiskIoMgrTest, StressTest) {
    // Run the test with 5 disks, 5 threads per disk, 10 clients and with cancellation
    DiskIoMgrStress test(5, 5, 10, true);
    test.Run(2); // In seconds
  }
  #endif
  
  TEST_F(DiskIoMgrTest, Buffers) {
      // Test default min/max buffer size
      int min_buffer_size = 1024;
      int max_buffer_size = 8 * 1024 * 1024; // 8 MB
      MemTracker mem_tracker(max_buffer_size * 2);
  
      DiskIoMgr io_mgr(1, 1, min_buffer_size, max_buffer_size);
      Status status = io_mgr.init(&mem_tracker);
      ASSERT_TRUE(status.ok());
      ASSERT_EQ(mem_tracker.consumption(), 0);
  
      // buffer length should be rounded up to min buffer size
      int64_t buffer_len = 1;
      char* buf = io_mgr.get_free_buffer(&buffer_len);
      EXPECT_EQ(buffer_len, min_buffer_size);
      EXPECT_EQ(io_mgr._num_allocated_buffers, 1);
      io_mgr.return_free_buffer(buf, buffer_len);
      EXPECT_EQ(mem_tracker.consumption(), min_buffer_size);
  
      // reuse buffer
      buffer_len = min_buffer_size;
      buf = io_mgr.get_free_buffer(&buffer_len);
      EXPECT_EQ(buffer_len, min_buffer_size);
      EXPECT_EQ(io_mgr._num_allocated_buffers, 1);
      io_mgr.return_free_buffer(buf, buffer_len);
      EXPECT_EQ(mem_tracker.consumption(), min_buffer_size);
  
      // bump up to next buffer size
      buffer_len = min_buffer_size + 1;
      buf = io_mgr.get_free_buffer(&buffer_len);
      EXPECT_EQ(buffer_len, min_buffer_size * 2);
      EXPECT_EQ(io_mgr._num_allocated_buffers, 2);
      EXPECT_EQ(mem_tracker.consumption(), min_buffer_size * 3);
  
      // gc unused buffer
      io_mgr.gc_io_buffers();
      EXPECT_EQ(io_mgr._num_allocated_buffers, 1);
      EXPECT_EQ(mem_tracker.consumption(), min_buffer_size * 2);
  
      io_mgr.return_free_buffer(buf, buffer_len);
  
      // max buffer size
      buffer_len = max_buffer_size;
      buf = io_mgr.get_free_buffer(&buffer_len);
      EXPECT_EQ(buffer_len, max_buffer_size);
      EXPECT_EQ(io_mgr._num_allocated_buffers, 2);
      io_mgr.return_free_buffer(buf, buffer_len);
      EXPECT_EQ(mem_tracker.consumption(), min_buffer_size * 2 + max_buffer_size);
  
      // gc buffers
      io_mgr.gc_io_buffers();
      EXPECT_EQ(io_mgr._num_allocated_buffers, 0);
      EXPECT_EQ(mem_tracker.consumption(), 0);
  }
  
  // IMPALA-2366: handle partial read where range goes past end of file.
  TEST_F(DiskIoMgrTest, PartialRead) {
      MemTracker mem_tracker(LARGE_MEM_LIMIT);
      const char* tmp_file = "/tmp/disk_io_mgr_test.txt";
      const char* data = "the quick brown fox jumped over the lazy dog";
      int len = strlen(data);
      int read_len = len + 1000; // Read past end of file.
      CreateTempFile(tmp_file, data);
  
      // Get mtime for file
      struct stat stat_val;
      stat(tmp_file, &stat_val);
  
      _pool.reset(new ObjectPool);
      scoped_ptr<DiskIoMgr> io_mgr(new DiskIoMgr(1, 1, read_len, read_len));
  
      Status status = io_mgr->init(&mem_tracker);
      ASSERT_TRUE(status.ok());
      MemTracker reader_mem_tracker(LARGE_MEM_LIMIT);
      DiskIoMgr::RequestContext* reader;
      status = io_mgr->register_context(&reader, &reader_mem_tracker);
      ASSERT_TRUE(status.ok());
  
      // We should not read past the end of file.
      DiskIoMgr::ScanRange* range = init_range(1, tmp_file, 0, read_len, 0, stat_val.st_mtime);
      DiskIoMgr::BufferDescriptor* buffer;
      status = io_mgr->read(reader, range, &buffer);
      ASSERT_TRUE(status.ok());
      ASSERT_TRUE(buffer->eosr());
      ASSERT_EQ(len, buffer->len());
      ASSERT_TRUE(memcmp(buffer->buffer(), data, len) == 0);
      buffer->return_buffer();
  
      io_mgr->unregister_context(reader);
      _pool.reset();
      io_mgr.reset();
      EXPECT_EQ(reader_mem_tracker.consumption(), 0);
      EXPECT_EQ(mem_tracker.consumption(), 0);
  }
  
  } // end namespace palo
  
  int main(int argc, char** argv) {
      // std::string conffile = std::string(getenv("PALO_HOME")) + "/conf/be.conf";
      // if (!palo::config::init(conffile.c_str(), false)) {
      //     fprintf(stderr, "error read config file. \n");
      //     return -1;
      // }
      palo::config::query_scratch_dirs = "/tmp";
      palo::config::max_free_io_buffers = 128;
      palo::config::disable_mem_pools = false;
      palo::init_glog("be-test");
      ::testing::InitGoogleTest(&argc, argv);
  
      palo::CpuInfo::init();
      palo::DiskInfo::init();
  
      return RUN_ALL_TESTS();
  }
  
  