//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
//  This source code is licensed under both the GPLv2 (found in the
//  COPYING file in the root directory) and Apache 2.0 License
//  (found in the LICENSE.Apache file in the root directory).
//
// Copyright (c) 2011 The LevelDB Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file. See the AUTHORS file for names of contributors.

#pragma once

#include <algorithm>
#include <atomic>
#include <deque>
#include <sstream>
#include <string>

#include "file/readahead_file_info.h"
#include "monitoring/statistics_impl.h"
#include "port/port.h"
#include "rocksdb/env.h"
#include "rocksdb/file_system.h"
#include "rocksdb/options.h"
#include "util/aligned_buffer.h"
#include "util/autovector.h"
#include "util/stop_watch.h"

namespace ROCKSDB_NAMESPACE {

#define DEFAULT_DECREMENT 8 * 1024

struct IOOptions;
class RandomAccessFileReader;

struct ReadaheadParams {
  ReadaheadParams() {}

  // The initial readahead size.
  size_t initial_readahead_size = 0;

  // The maximum readahead size.
  // If max_readahead_size > readahead_size, then readahead size will be doubled
  // on every IO until max_readahead_size is hit. Typically this is set as a
  // multiple of initial_readahead_size. initial_readahead_size should be
  // greater than equal to initial_readahead_size.
  size_t max_readahead_size = 0;

  // If true, Readahead is enabled implicitly by rocksdb
  // after doing sequential scans for num_file_reads_for_auto_readahead.
  bool implicit_auto_readahead = false;

  // TODO akanksha - Remove num_file_reads when BlockPrefetcher is refactored.
  uint64_t num_file_reads = 0;
  uint64_t num_file_reads_for_auto_readahead = 0;

  // Number of buffers to maintain that contains prefetched data. If num_buffers
  // > 1 then buffers will be filled asynchronously whenever they get emptied.
  size_t num_buffers = 1;
};

struct BufferInfo {
  void ClearBuffer() {
    buffer_.Clear();
    initial_end_offset_ = 0;
    async_req_len_ = 0;
  }

  AlignedBuffer buffer_;

  uint64_t offset_ = 0;

  // Below parameters are used in case of async read flow.
  // Length requested for in ReadAsync.
  size_t async_req_len_ = 0;

  // async_read_in_progress can be used as mutex. Callback can update the buffer
  // and its size but async_read_in_progress is only set by main thread.
  bool async_read_in_progress_ = false;

  // io_handle is allocated and used by underlying file system in case of
  // asynchronous reads.
  void* io_handle_ = nullptr;

  IOHandleDeleter del_fn_ = nullptr;

  // initial_end_offset is used to keep track of the end offset of the buffer
  // that was originally called. It's helpful in case of autotuning of readahead
  // size when callback is made to BlockBasedTableIterator.
  // initial end offset of this buffer which will be the starting
  // offset of next prefetch.
  //
  // For example - if end offset of previous buffer was 100 and because of
  // readahead_size optimization, end_offset was trimmed to 60. Then for next
  // prefetch call, start_offset should be intialized to 100 i.e  start_offset =
  // buf->initial_end_offset_.
  uint64_t initial_end_offset_ = 0;

  bool IsDataBlockInBuffer(uint64_t offset, size_t length) {
    assert(async_read_in_progress_ == false);
    return (offset >= offset_ &&
            offset + length <= offset_ + buffer_.CurrentSize());
  }

  bool IsOffsetInBuffer(uint64_t offset) {
    assert(async_read_in_progress_ == false);
    return (offset >= offset_ && offset < offset_ + buffer_.CurrentSize());
  }

  bool DoesBufferContainData() {
    assert(async_read_in_progress_ == false);
    return buffer_.CurrentSize() > 0;
  }

  bool IsBufferOutdated(uint64_t offset) {
    return (!async_read_in_progress_ && DoesBufferContainData() &&
            offset >= offset_ + buffer_.CurrentSize());
  }

  bool IsBufferOutdatedWithAsyncProgress(uint64_t offset) {
    return (async_read_in_progress_ && io_handle_ != nullptr &&
            offset >= offset_ + async_req_len_);
  }

  bool IsOffsetInBufferWithAsyncProgress(uint64_t offset) {
    return (async_read_in_progress_ && offset >= offset_ &&
            offset < offset_ + async_req_len_);
  }

  size_t CurrentSize() { return buffer_.CurrentSize(); }
};

enum class FilePrefetchBufferUsage {
  kTableOpenPrefetchTail,
  kUserScanPrefetch,
  kUnknown,
};

// Implementation:
// FilePrefetchBuffer maintains a dequeu of free buffers (free_bufs_) with no
// data and bufs_ which contains the prefetched data. Whenever a buffer is
// consumed or is outdated (w.r.t. to requested offset), that buffer is cleared
// and returned to free_bufs_.
//
// If a buffer is available in free_bufs_, it's moved to bufs_ and is sent for
// prefetching.
// num_buffers_ defines how many buffers FilePrefetchBuffer can maintain at a
// time that contains prefetched data with num_buffers_ == bufs_.size() +
// free_bufs_.size().
//
// If num_buffers_ == 1, it's a sequential read flow. Read API will be called on
// that one buffer whenever the data is requested and is not in the buffer.
// If num_buffers_ > 1, then the data is prefetched asynchronosuly in the
// buffers whenever the data is consumed from the buffers and that buffer is
// freed.
// If num_buffers > 1, then requested data can be overlapping between 2 buffers.
// To return the continuous buffer, overlap_buf_ is used. The requested data is
// copied from 2 buffers to the overlap_buf_ and overlap_buf_ is returned to
// the caller.

// FilePrefetchBuffer is a smart buffer to store and read data from a file.
class FilePrefetchBuffer {
 public:
  // Constructor.
  //
  // All arguments are optional.
  // ReadaheadParams  : Parameters to control the readahead behavior.
  // enable           : controls whether reading from the buffer is enabled.
  //                    If false, TryReadFromCache() always return false, and we
  //                    only take stats for the minimum offset if
  //                    track_min_offset = true.
  //                    See below NOTE about mmap reads.
  // track_min_offset : Track the minimum offset ever read and collect stats on
  //                    it. Used for adaptable readahead of the file
  //                    footer/metadata.
  //
  // A user can construct a FilePrefetchBuffer without any arguments, but use
  // `Prefetch` to load data into the buffer.
  // NOTE: FilePrefetchBuffer is incompatible with prefetching from
  // RandomAccessFileReaders using mmap reads, so it is common to use
  // `!use_mmap_reads` for the `enable` parameter.
  FilePrefetchBuffer(
      const ReadaheadParams& readahead_params = {}, bool enable = true,
      bool track_min_offset = false, FileSystem* fs = nullptr,
      SystemClock* clock = nullptr, Statistics* stats = nullptr,
      const std::function<void(bool, uint64_t&, uint64_t&)>& cb = nullptr,
      FilePrefetchBufferUsage usage = FilePrefetchBufferUsage::kUnknown)
      : readahead_size_(readahead_params.initial_readahead_size),
        initial_auto_readahead_size_(readahead_params.initial_readahead_size),
        max_readahead_size_(readahead_params.max_readahead_size),
        min_offset_read_(std::numeric_limits<size_t>::max()),
        enable_(enable),
        track_min_offset_(track_min_offset),
        implicit_auto_readahead_(readahead_params.implicit_auto_readahead),
        prev_offset_(0),
        prev_len_(0),
        num_file_reads_for_auto_readahead_(
            readahead_params.num_file_reads_for_auto_readahead),
        num_file_reads_(readahead_params.num_file_reads),
        explicit_prefetch_submitted_(false),
        fs_(fs),
        clock_(clock),
        stats_(stats),
        usage_(usage),
        readaheadsize_cb_(cb),
        num_buffers_(readahead_params.num_buffers) {
    assert((num_file_reads_ >= num_file_reads_for_auto_readahead_ + 1) ||
           (num_file_reads_ == 0));

    // If num_buffers_ > 1, data is asynchronously filled in the
    // queue. As result, data can be overlapping in two buffers. It copies the
    // data to overlap_buf_ in order to to return continuous buffer.
    if (num_buffers_ > 1) {
      overlap_buf_ = new BufferInfo();
    }

    free_bufs_.resize(num_buffers_);
    for (uint32_t i = 0; i < num_buffers_; i++) {
      free_bufs_[i] = new BufferInfo();
    }
  }

  ~FilePrefetchBuffer() {
    // Abort any pending async read request before destroying the class object.
    if (fs_ != nullptr) {
      std::vector<void*> handles;
      for (auto& buf : bufs_) {
        if (buf->async_read_in_progress_ && buf->io_handle_ != nullptr) {
          handles.emplace_back(buf->io_handle_);
        }
      }
      if (!handles.empty()) {
        StopWatch sw(clock_, stats_, ASYNC_PREFETCH_ABORT_MICROS);
        Status s = fs_->AbortIO(handles);
        assert(s.ok());
      }

      for (auto& buf : bufs_) {
        if (buf->io_handle_ != nullptr) {
          DestroyAndClearIOHandle(buf);
          buf->ClearBuffer();
        }
        buf->async_read_in_progress_ = false;
      }
    }

    // Prefetch buffer bytes discarded.
    uint64_t bytes_discarded = 0;
    // Iterated over buffers.
    for (auto& buf : bufs_) {
      if (buf->DoesBufferContainData()) {
        // If last read was from this block and some bytes are still unconsumed.
        if (prev_offset_ >= buf->offset_ &&
            prev_offset_ + prev_len_ < buf->offset_ + buf->CurrentSize()) {
          bytes_discarded +=
              buf->CurrentSize() - (prev_offset_ + prev_len_ - buf->offset_);
        }
        // If last read was from previous blocks and this block is unconsumed.
        else if (prev_offset_ < buf->offset_ &&
                 prev_offset_ + prev_len_ <= buf->offset_) {
          bytes_discarded += buf->CurrentSize();
        }
      }
    }

    RecordInHistogram(stats_, PREFETCHED_BYTES_DISCARDED, bytes_discarded);

    for (auto& buf : bufs_) {
      delete buf;
      buf = nullptr;
    }

    for (auto& buf : free_bufs_) {
      delete buf;
      buf = nullptr;
    }

    if (overlap_buf_ != nullptr) {
      delete overlap_buf_;
      overlap_buf_ = nullptr;
    }
  }

  bool Enabled() const { return enable_; }

  // Called externally by user to only load data into the buffer from a file
  // with num_buffers_ should be set to default(1).
  //
  // opts                  : the IO options to use.
  // reader                : the file reader.
  // offset                : the file offset to start reading from.
  // n                     : the number of bytes to read.
  //
  Status Prefetch(const IOOptions& opts, RandomAccessFileReader* reader,
                  uint64_t offset, size_t n);

  // Request for reading the data from a file asynchronously.
  // If data already exists in the buffer, result will be updated.
  // reader                : the file reader.
  // offset                : the file offset to start reading from.
  // n                     : the number of bytes to read.
  // result                : if data already exists in the buffer, result will
  //                         be updated with the data.
  //
  // If data already exist in the buffer, it will return Status::OK, otherwise
  // it will send asynchronous request and return Status::TryAgain.
  Status PrefetchAsync(const IOOptions& opts, RandomAccessFileReader* reader,
                       uint64_t offset, size_t n, Slice* result);

  // Tries returning the data for a file read from this buffer if that data is
  // in the buffer.
  // It handles tracking the minimum read offset if track_min_offset = true.
  // It also does the exponential readahead when readahead_size is set as part
  // of the constructor.
  //
  // opts                  : the IO options to use.
  // reader                : the file reader.
  // offset                : the file offset.
  // n                     : the number of bytes.
  // result                : output buffer to put the data into.
  // s                     : output status.
  // for_compaction        : true if cache read is done for compaction read.
  bool TryReadFromCache(const IOOptions& opts, RandomAccessFileReader* reader,
                        uint64_t offset, size_t n, Slice* result, Status* s,
                        bool for_compaction = false);

  // The minimum `offset` ever passed to TryReadFromCache(). This will nly be
  // tracked if track_min_offset = true.
  size_t min_offset_read() const { return min_offset_read_; }

  size_t GetPrefetchOffset() const { return bufs_.front()->offset_; }

  // Called in case of implicit auto prefetching.
  void UpdateReadPattern(const uint64_t& offset, const size_t& len,
                         bool decrease_readaheadsize) {
    if (decrease_readaheadsize) {
      DecreaseReadAheadIfEligible(offset, len);
    }
    prev_offset_ = offset;
    prev_len_ = len;
    explicit_prefetch_submitted_ = false;
  }

  void GetReadaheadState(ReadaheadFileInfo::ReadaheadInfo* readahead_info) {
    readahead_info->readahead_size = readahead_size_;
    readahead_info->num_file_reads = num_file_reads_;
  }

  void DecreaseReadAheadIfEligible(uint64_t offset, size_t size,
                                   size_t value = DEFAULT_DECREMENT) {
    if (bufs_.empty()) {
      return;
    }

    // Decrease the readahead_size if
    // - its enabled internally by RocksDB (implicit_auto_readahead_) and,
    // - readahead_size is greater than 0 and,
    // - this block would have called prefetch API if not found in cache for
    //   which conditions are:
    //   - few/no bytes are in buffer and,
    //   - block is sequential with the previous read and,
    //   - num_file_reads_ + 1 (including this read) >
    //   num_file_reads_for_auto_readahead_

    size_t curr_size = bufs_.front()->async_read_in_progress_
                           ? bufs_.front()->async_req_len_
                           : bufs_.front()->CurrentSize();
    if (implicit_auto_readahead_ && readahead_size_ > 0) {
      if ((offset + size > bufs_.front()->offset_ + curr_size) &&
          IsBlockSequential(offset) &&
          (num_file_reads_ + 1 > num_file_reads_for_auto_readahead_)) {
        readahead_size_ =
            std::max(initial_auto_readahead_size_,
                     (readahead_size_ >= value ? readahead_size_ - value : 0));
      }
    }
  }

  // Callback function passed to underlying FS in case of asynchronous reads.
  void PrefetchAsyncCallback(const FSReadRequest& req, void* cb_arg);

  void TEST_GetBufferOffsetandSize(
      std::vector<std::pair<uint64_t, size_t>>& buffer_info) {
    for (size_t i = 0; i < bufs_.size(); i++) {
      buffer_info[i].first = bufs_[i]->offset_;
      buffer_info[i].second = bufs_[i]->async_read_in_progress_
                                  ? bufs_[i]->async_req_len_
                                  : bufs_[i]->CurrentSize();
    }
  }

 private:
  // Calculates roundoff offset and length to be prefetched based on alignment
  // and data present in buffer_. It also allocates new buffer or refit tail if
  // required.
  void PrepareBufferForRead(BufferInfo* buf, size_t alignment, uint64_t offset,
                            size_t roundup_len, bool refit_tail,
                            uint64_t& aligned_useful_len);

  void AbortOutdatedIO(uint64_t offset);

  void AbortAllIOs();

  void ClearOutdatedData(uint64_t offset, size_t len);

  // It calls Poll API to check for any pending asynchronous request.
  void PollIfNeeded(uint64_t offset, size_t len);

  Status PrefetchInternal(const IOOptions& opts, RandomAccessFileReader* reader,
                          uint64_t offset, size_t length, size_t readahead_size,
                          bool& copy_to_third_buffer);

  Status Read(BufferInfo* buf, const IOOptions& opts,
              RandomAccessFileReader* reader, uint64_t read_len,
              uint64_t aligned_useful_len, uint64_t start_offset);

  Status ReadAsync(BufferInfo* buf, const IOOptions& opts,
                   RandomAccessFileReader* reader, uint64_t read_len,
                   uint64_t start_offset);

  // Copy the data from src to overlap_buf_.
  void CopyDataToBuffer(BufferInfo* src, uint64_t& offset, size_t& length);

  bool IsBlockSequential(const size_t& offset) {
    return (prev_len_ == 0 || (prev_offset_ + prev_len_ == offset));
  }

  // Called in case of implicit auto prefetching.
  void ResetValues() {
    num_file_reads_ = 1;
    readahead_size_ = initial_auto_readahead_size_;
  }

  // Called in case of implicit auto prefetching.
  bool IsEligibleForPrefetch(uint64_t offset, size_t n) {
    // Prefetch only if this read is sequential otherwise reset readahead_size_
    // to initial value.
    if (!IsBlockSequential(offset)) {
      UpdateReadPattern(offset, n, false /*decrease_readaheadsize*/);
      ResetValues();
      return false;
    }
    num_file_reads_++;

    // Since async request was submitted in last call directly by calling
    // PrefetchAsync, it skips num_file_reads_ check as this call is to poll the
    // data submitted in previous call.
    if (explicit_prefetch_submitted_) {
      return true;
    }
    if (num_file_reads_ <= num_file_reads_for_auto_readahead_) {
      UpdateReadPattern(offset, n, false /*decrease_readaheadsize*/);
      return false;
    }
    return true;
  }

  bool IsEligibleForFurtherPrefetching() {
    if (free_bufs_.empty()) {
      return false;
    }
    // Readahead size can be 0 because of trimming.
    if (readahead_size_ == 0) {
      return false;
    }
    return true;
  }

  void DestroyAndClearIOHandle(BufferInfo* buf) {
    if (buf->io_handle_ != nullptr && buf->del_fn_ != nullptr) {
      buf->del_fn_(buf->io_handle_);
      buf->io_handle_ = nullptr;
      buf->del_fn_ = nullptr;
    }
    buf->async_read_in_progress_ = false;
  }

  Status HandleOverlappingData(const IOOptions& opts,
                               RandomAccessFileReader* reader, uint64_t offset,
                               size_t length, size_t readahead_size,
                               bool& copy_to_third_buffer, uint64_t& tmp_offset,
                               size_t& tmp_length);

  bool TryReadFromCacheUntracked(const IOOptions& opts,
                                 RandomAccessFileReader* reader,
                                 uint64_t offset, size_t n, Slice* result,
                                 Status* s,
                                 bool for_compaction = false);

  void ReadAheadSizeTuning(BufferInfo* buf, bool read_curr_block,
                           bool refit_tail, uint64_t prev_buf_end_offset,
                           size_t alignment, size_t length,
                           size_t readahead_size, uint64_t& offset,
                           uint64_t& end_offset, size_t& read_len,
                           uint64_t& aligned_useful_len);

  void UpdateStats(bool found_in_buffer, size_t length_found) {
    if (found_in_buffer) {
      RecordTick(stats_, PREFETCH_HITS);
    }
    if (length_found > 0) {
      RecordTick(stats_, PREFETCH_BYTES_USEFUL, length_found);
    }
  }

  void UpdateReadAheadTrimmedStat(size_t initial_length,
                                  size_t updated_length) {
    if (initial_length != updated_length) {
      RecordTick(stats_, READAHEAD_TRIMMED);
    }
  }

  Status PrefetchRemBuffers(const IOOptions& opts,
                            RandomAccessFileReader* reader,
                            uint64_t end_offset1, size_t alignment,
                            size_t readahead_size);

  // *** BEGIN APIs related to allocating and freeing buffers ***
  bool IsBufferQueueEmpty() { return bufs_.empty(); }

  BufferInfo* GetFirstBuffer() { return bufs_.front(); }

  BufferInfo* GetLastBuffer() { return bufs_.back(); }

  size_t NumBuffersAllocated() { return bufs_.size(); }

  void AllocateBuffer() {
    assert(!free_bufs_.empty());
    BufferInfo* buf = free_bufs_.front();
    free_bufs_.pop_front();
    bufs_.emplace_back(buf);
  }

  void AllocateBufferIfEmpty() {
    if (bufs_.empty()) {
      AllocateBuffer();
    }
  }

  void FreeFrontBuffer() {
    BufferInfo* buf = bufs_.front();
    buf->ClearBuffer();
    bufs_.pop_front();
    free_bufs_.emplace_back(buf);
  }

  void FreeLastBuffer() {
    BufferInfo* buf = bufs_.back();
    buf->ClearBuffer();
    bufs_.pop_back();
    free_bufs_.emplace_back(buf);
  }

  void FreeAllBuffers() {
    while (!bufs_.empty()) {
      BufferInfo* buf = bufs_.front();
      buf->ClearBuffer();
      bufs_.pop_front();
      free_bufs_.emplace_back(buf);
    }
  }

  void FreeEmptyBuffers() {
    if (bufs_.empty()) {
      return;
    }

    std::deque<BufferInfo*> tmp_buf;
    while (!bufs_.empty()) {
      BufferInfo* buf = bufs_.front();
      bufs_.pop_front();
      if (buf->async_read_in_progress_ || buf->DoesBufferContainData()) {
        tmp_buf.emplace_back(buf);
      } else {
        free_bufs_.emplace_back(buf);
      }
    }
    bufs_ = tmp_buf;
  }

  // *** END APIs related to allocating and freeing buffers ***

  std::deque<BufferInfo*> bufs_;
  std::deque<BufferInfo*> free_bufs_;
  BufferInfo* overlap_buf_ = nullptr;

  size_t readahead_size_;
  size_t initial_auto_readahead_size_;
  // FilePrefetchBuffer object won't be created from Iterator flow if
  // max_readahead_size_ = 0.
  size_t max_readahead_size_;

  // The minimum `offset` ever passed to TryReadFromCache().
  size_t min_offset_read_;
  // if false, TryReadFromCache() always return false, and we only take stats
  // for track_min_offset_ if track_min_offset_ = true
  bool enable_;
  // If true, track minimum `offset` ever passed to TryReadFromCache(), which
  // can be fetched from min_offset_read().
  bool track_min_offset_;

  // implicit_auto_readahead is enabled by rocksdb internally after 2
  // sequential IOs.
  bool implicit_auto_readahead_;
  uint64_t prev_offset_;
  size_t prev_len_;
  // num_file_reads_ and num_file_reads_for_auto_readahead_ is only used when
  // implicit_auto_readahead_ is set.
  uint64_t num_file_reads_for_auto_readahead_;
  uint64_t num_file_reads_;

  // If explicit_prefetch_submitted_ is set then it indicates RocksDB called
  // PrefetchAsync to submit request. It needs to call TryReadFromCache to
  // poll the submitted request without checking if data is sequential and
  // num_file_reads_.
  bool explicit_prefetch_submitted_;

  FileSystem* fs_;
  SystemClock* clock_;
  Statistics* stats_;

  FilePrefetchBufferUsage usage_;

  std::function<void(bool, uint64_t&, uint64_t&)> readaheadsize_cb_;

  // num_buffers_ is the number of buffers maintained by FilePrefetchBuffer to
  // prefetch the data at a time.
  size_t num_buffers_;
};
}  // namespace ROCKSDB_NAMESPACE