mariadb/storage/rocksdb/rdb_index_merge.cc

/*
   Copyright (c) 2016, Facebook, Inc.

   This program is free software; you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation; version 2 of the License.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program; if not, write to the Free Software
   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA */

#include <my_config.h>

/* This C++ file's header file */
#include "./rdb_index_merge.h"

/* MyRocks header files */
#include "./ha_rocksdb.h"
#include "./rdb_datadic.h"

namespace myrocks {

Rdb_index_merge::Rdb_index_merge(const ulonglong merge_buf_size,
                                 const ulonglong merge_combine_read_size,
                                 const rocksdb::Comparator* comparator) :
  m_merge_buf_size(merge_buf_size),
  m_merge_combine_read_size(merge_combine_read_size),
  m_comparator(comparator),
  m_rec_buf_unsorted(nullptr),
  m_output_buf(nullptr)
{
}

Rdb_index_merge::~Rdb_index_merge()
{
  /*
    Close tmp file, we don't need to worry about deletion, mysql handles it.
  */
  my_close(m_merge_file.fd, MYF(MY_WME));

  /* There should be no records left in the offset tree */
  DBUG_ASSERT(m_offset_tree.empty());

  /* There should be no pointers left on the merge heap */
  DBUG_ASSERT(m_merge_min_heap.empty());
}

int Rdb_index_merge::init()
{
  /*
    Create a temporary merge file on disk to store sorted chunks during
    inplace index creation.
  */
  if (merge_file_create())
  {
    return HA_ERR_INTERNAL_ERROR;
  }

  /*
    Then, allocate buffer to store unsorted records before they are written
    to disk. They will be written to disk sorted. A sorted tree is used to
    keep track of the offset of each record within the unsorted buffer.
  */
  m_rec_buf_unsorted= std::make_shared<merge_buf_info>(m_merge_buf_size);

  /*
    Allocate output buffer that will contain sorted block that is written to
    disk.
  */
  m_output_buf= std::make_shared<merge_buf_info>(m_merge_buf_size);

  return 0;
}

/**
  Create a merge file in the given location.
*/
int Rdb_index_merge::merge_file_create()
{
  DBUG_ASSERT(m_merge_file.fd == -1);

  int fd = mysql_tmpfile("myrocks");

  if (fd < 0)
  {
    return HA_ERR_INTERNAL_ERROR;
  }

  m_merge_file.fd = fd;
  m_merge_file.num_sort_buffers = 0;

  return 0;
}

/**
  Add record to offset tree (and unsorted merge buffer) in preparation for
  writing out to disk in sorted chunks.

  If buffer in memory is full, write the buffer out to disk sorted using the
  offset tree, and clear the tree. (Happens in merge_buf_write)
*/
int Rdb_index_merge::add(const rocksdb::Slice& key,
                         const rocksdb::Slice& val)
{
  /* Adding a record after heap is already created results in error */
  DBUG_ASSERT(m_merge_min_heap.empty());

  /*
    Check if sort buffer is going to be out of space, if so write it
    out to disk in sorted order using offset tree.
  */
  uint total_offset= RDB_MERGE_CHUNK_LEN + m_rec_buf_unsorted->curr_offset +
                     RDB_MERGE_KEY_DELIMITER + RDB_MERGE_VAL_DELIMITER +
                     key.size() + val.size();
  if (total_offset >= m_rec_buf_unsorted->total_size)
  {
    if (merge_buf_write())
    {
      // NO_LINT_DEBUG
      sql_print_error("Error writing sort buffer to disk.");
      return HA_ERR_INTERNAL_ERROR;
    }
  }

  ulonglong rec_offset= m_rec_buf_unsorted->curr_offset;

  /*
    Store key and value in temporary unsorted in memory buffer pointed to by
    offset tree.
  */
  m_rec_buf_unsorted->store_key_value(key, val);

  /* Find sort order of the new record */
  m_offset_tree.emplace(m_rec_buf_unsorted->block.get() + rec_offset,
                        m_comparator);

  return 0;
}

/**
  Sort + write merge buffer chunk out to disk.
*/
int Rdb_index_merge::merge_buf_write()
{
  DBUG_ASSERT(m_merge_file.fd != -1);
  DBUG_ASSERT(m_rec_buf_unsorted != nullptr);
  DBUG_ASSERT(m_output_buf != nullptr);
  DBUG_ASSERT(!m_offset_tree.empty());

  /* Write actual chunk size to first 8 bytes of the merge buffer */
  merge_store_uint64(m_output_buf->block.get(),
                     m_rec_buf_unsorted->curr_offset + RDB_MERGE_CHUNK_LEN);
  m_output_buf->curr_offset += RDB_MERGE_CHUNK_LEN;

  /*
    Iterate through the offset tree.  Should be ordered by the secondary key
    at this point.
  */
  for (auto& rec : m_offset_tree)
  {
    DBUG_ASSERT(m_output_buf->curr_offset <= m_merge_buf_size);

    /* Read record from offset (should never fail) */
    rocksdb::Slice key;
    rocksdb::Slice val;
    merge_read_rec(rec.block, &key, &val);

    /* Store key and value into sorted output buffer */
    m_output_buf->store_key_value(key, val);
  }

  DBUG_ASSERT(m_output_buf->curr_offset <= m_output_buf->total_size);

  /*
    Write output buffer to disk.

    Need to position cursor to the chunk it needs to be at on filesystem
    then write into the respective merge buffer.
  */
  if (my_seek(m_merge_file.fd, m_merge_file.num_sort_buffers * m_merge_buf_size,
      SEEK_SET, MYF(0)) == MY_FILEPOS_ERROR)
  {
    // NO_LINT_DEBUG
    sql_print_error("Error seeking to location in merge file on disk.");
    return HA_ERR_INTERNAL_ERROR;
  }

  if (my_write(m_merge_file.fd, m_output_buf->block.get(),
        m_output_buf->total_size, MYF(MY_WME | MY_NABP)))
  {
    // NO_LINT_DEBUG
    sql_print_error("Error writing sorted merge buffer to disk.");
    return HA_ERR_INTERNAL_ERROR;
  }

  /* Increment merge file offset to track number of merge buffers written */
  m_merge_file.num_sort_buffers += 1;

  /* Reset everything for next run */
  merge_reset();

  return 0;
}

/**
  Prepare n-way merge of n sorted buffers on disk, using a heap sorted by
  secondary key records.
*/
int Rdb_index_merge::merge_heap_prepare()
{
  DBUG_ASSERT(m_merge_min_heap.empty());

  /*
    If the offset tree is not empty, there are still some records that need to
    be written to disk. Write them out now.
  */
  if (!m_offset_tree.empty() && merge_buf_write())
  {
    return HA_ERR_INTERNAL_ERROR;
  }

  DBUG_ASSERT(m_merge_file.num_sort_buffers > 0);

  /*
    For an n-way merge, we need to read chunks of each merge file
    simultaneously.
  */
  ulonglong chunk_size= m_merge_combine_read_size/
    m_merge_file.num_sort_buffers;
  if (chunk_size >= m_merge_buf_size)
  {
    chunk_size= m_merge_buf_size;
  }

  /* Allocate buffers for each chunk */
  for (ulonglong i = 0; i < m_merge_file.num_sort_buffers; i++)
  {
    auto entry= std::make_shared<merge_heap_entry>(m_comparator);

    /*
      Read chunk_size bytes from each chunk on disk, and place inside
      respective chunk buffer.
    */
    size_t total_size=
      entry->prepare(m_merge_file.fd, i * m_merge_buf_size, chunk_size);

    if (total_size == (size_t) - 1)
    {
      return HA_ERR_INTERNAL_ERROR;
    }

    /* Can reach this condition if an index was added on table w/ no rows */
    if (total_size - RDB_MERGE_CHUNK_LEN == 0)
    {
      break;
    }

    /* Read the first record from each buffer to initially populate the heap */
    if (entry->read_rec(&entry->key, &entry->val))
    {
      // NO_LINT_DEBUG
      sql_print_error("Chunk size is too small to process merge.");
      return HA_ERR_INTERNAL_ERROR;
    }

    m_merge_min_heap.push(std::move(entry));
  }

  return 0;
}

/**
  Create and/or iterate through keys in the merge heap.
*/
int Rdb_index_merge::next(rocksdb::Slice* key, rocksdb::Slice* val)
{
  /*
    If table fits in one sort buffer, we can optimize by writing
    the sort buffer directly through to the sstfilewriter instead of
    needing to create tmp files/heap to merge the sort buffers.

    If there are no sort buffer records (alters on empty tables),
    also exit here.
  */
  if (m_merge_file.num_sort_buffers == 0)
  {
    if (m_offset_tree.empty())
    {
      return -1;
    }

    auto rec= m_offset_tree.begin();

    /* Read record from offset */
    merge_read_rec(rec->block, key, val);

    m_offset_tree.erase(rec);
    return 0;
  }

  int res;

  /*
    If heap and heap chunk info are empty, we must be beginning the merge phase
    of the external sort. Populate the heap with initial values from each
    disk chunk.
  */
  if (m_merge_min_heap.empty())
  {
    if ((res= merge_heap_prepare()))
    {
      // NO_LINT_DEBUG
      sql_print_error("Error during preparation of heap.");
      return res;
    }

    /*
      Return the first top record without popping, as we haven't put this
      inside the SST file yet.
    */
    merge_heap_top(key, val);
    return 0;
  }

  DBUG_ASSERT(!m_merge_min_heap.empty());
  return merge_heap_pop_and_get_next(key, val);
}

/**
  Get current top record from the heap.
*/
void Rdb_index_merge::merge_heap_top(rocksdb::Slice* key,
                                     rocksdb::Slice* val)
{
  DBUG_ASSERT(!m_merge_min_heap.empty());

  const std::shared_ptr<merge_heap_entry>& entry= m_merge_min_heap.top();
  *key= entry->key;
  *val= entry->val;
}

/**
  Pops the top record, and uses it to read next record from the
  corresponding sort buffer and push onto the heap.

  Returns -1 when there are no more records in the heap.
*/
int Rdb_index_merge::merge_heap_pop_and_get_next(rocksdb::Slice* key,
                                                 rocksdb::Slice* val)
{
  /*
    Make a new reference to shared ptr so it doesn't get destroyed
    during pop(). We are going to push this entry back onto the heap.
  */
  const std::shared_ptr<merge_heap_entry> entry= m_merge_min_heap.top();
  m_merge_min_heap.pop();

  /*
    We are finished w/ current chunk if:
    current_offset + disk_offset == total_size

    Return without adding entry back onto heap.
    If heap is also empty, we must be finished with merge.
  */
  if (entry->chunk_info->is_chunk_finished())
  {
    if (m_merge_min_heap.empty())
    {
      return -1;
    }

    merge_heap_top(key, val);
    return 0;
  }

  /*
    Make sure we haven't reached the end of the chunk.
  */
  DBUG_ASSERT(!entry->chunk_info->is_chunk_finished());

  /*
    If merge_read_rec fails, it means the either the chunk was cut off
    or we've reached the end of the respective chunk.
  */
  if (entry->read_rec(&entry->key, &entry->val))
  {
    if (entry->read_next_chunk_from_disk(m_merge_file.fd))
    {
      return HA_ERR_INTERNAL_ERROR;
    }

    /* Try reading record again, should never fail. */
    if (entry->read_rec(&entry->key, &entry->val))
    {
      return HA_ERR_INTERNAL_ERROR;
    }
  }

  /* Push entry back on to the heap w/ updated buffer + offset ptr */
  m_merge_min_heap.push(std::move(entry));

  /* Return the current top record on heap */
  merge_heap_top(key, val);
  return 0;
}

int Rdb_index_merge::merge_heap_entry::read_next_chunk_from_disk(File fd)
{
  if (chunk_info->read_next_chunk_from_disk(fd))
  {
    return 1;
  }

  block= chunk_info->block.get();
  return 0;
}

int Rdb_index_merge::merge_buf_info::read_next_chunk_from_disk(File fd)
{
  disk_curr_offset += curr_offset;

  if (my_seek(fd, disk_curr_offset, SEEK_SET, MYF(0)) == MY_FILEPOS_ERROR)
  {
    // NO_LINT_DEBUG
    sql_print_error("Error seeking to location in merge file on disk.");
    return 1;
  }

  /* Overwrite the old block */
  size_t bytes_read= my_read(fd, block.get(), block_len, MYF(MY_WME));
  if (bytes_read == (size_t) -1)
  {
    // NO_LINT_DEBUG
    sql_print_error("Error reading merge file from disk.");
    return 1;
  }

  curr_offset= 0;
  return 0;
}

/**
  Get records from offset within sort buffer and compare them.
  Sort by least to greatest.
*/
int Rdb_index_merge::merge_record_compare(const uchar* a_block,
    const uchar* b_block,
    const rocksdb::Comparator* const comparator)
{
  return comparator->Compare(as_slice(a_block), as_slice(b_block));
}

/**
  Given an offset in a merge sort buffer, read out the keys + values.
  After this, block will point to the next record in the buffer.
**/
void Rdb_index_merge::merge_read_rec(const uchar* block,
                                    rocksdb::Slice* key,
                                    rocksdb::Slice* val)
{
  /* Read key at block offset into key slice and the value into value slice*/
  read_slice(key, block);
  read_slice(val, block + RDB_MERGE_REC_DELIMITER + key->size());
}

void Rdb_index_merge::read_slice(rocksdb::Slice* slice, const uchar* block_ptr)
{
  uint64 slice_len;
  merge_read_uint64(&block_ptr, &slice_len);

  *slice= rocksdb::Slice(reinterpret_cast<const char*>(block_ptr), slice_len);
}

int Rdb_index_merge::merge_heap_entry::read_rec(rocksdb::Slice *key,
                                                rocksdb::Slice *val)
{
  const uchar* block_ptr= block;

  /* Read key at block offset into key slice and the value into value slice*/
  if (read_slice(key, &block_ptr) != 0 || read_slice(val, &block_ptr) != 0)
  {
    return 1;
  }

  chunk_info->curr_offset += (uintptr_t) block_ptr - (uintptr_t) block;
  block += (uintptr_t) block_ptr - (uintptr_t) block;

  return 0;
}

int Rdb_index_merge::merge_heap_entry::read_slice(rocksdb::Slice* slice,
                                                  const uchar** block_ptr)
{
  if (!chunk_info->has_space(RDB_MERGE_REC_DELIMITER))
  {
    return 1;
  }

  uint64 slice_len;
  merge_read_uint64(block_ptr, &slice_len);
  if (!chunk_info->has_space(RDB_MERGE_REC_DELIMITER + slice_len))
  {
    return 1;
  }

  *slice= rocksdb::Slice(reinterpret_cast<const char*>(*block_ptr), slice_len);
  *block_ptr += slice_len;
  return 0;
}

size_t Rdb_index_merge::merge_heap_entry::prepare(File fd, ulonglong f_offset,
                                                  ulonglong chunk_size)
{
  chunk_info= std::make_shared<merge_buf_info>(chunk_size);
  size_t res = chunk_info->prepare(fd, f_offset);
  if (res != (size_t) - 1)
  {
    block= chunk_info->block.get() + RDB_MERGE_CHUNK_LEN;
  }

  return res;
}

size_t Rdb_index_merge::merge_buf_info::prepare(File fd, ulonglong f_offset)
{
  disk_start_offset= f_offset;
  disk_curr_offset= f_offset;

  /*
    Need to position cursor to the chunk it needs to be at on filesystem
    then read 'chunk_size' bytes into the respective chunk buffer.
  */
  if (my_seek(fd, f_offset, SEEK_SET, MYF(0)) == MY_FILEPOS_ERROR)
  {
    // NO_LINT_DEBUG
    sql_print_error("Error seeking to location in merge file on disk.");
    return (size_t) - 1;
  }

  size_t bytes_read= my_read(fd, block.get(), total_size, MYF(MY_WME));
  if (bytes_read == (size_t) - 1)
  {
    // NO_LINT_DEBUG
    sql_print_error("Error reading merge file from disk.");
    return (size_t) - 1;
  }

  /*
    Read the first 8 bytes of each chunk, this gives us the actual
    size of each chunk.
  */
  const uchar *block_ptr= block.get();
  merge_read_uint64(&block_ptr, &total_size);
  curr_offset += RDB_MERGE_CHUNK_LEN;
  return total_size;
}

/* Store key and value w/ their respective delimiters at the given offset */
void Rdb_index_merge::merge_buf_info::store_key_value(
    const rocksdb::Slice& key, const rocksdb::Slice& val)
{
  store_slice(key);
  store_slice(val);
}

void Rdb_index_merge::merge_buf_info::store_slice(const rocksdb::Slice& slice)
{
  /* Store length delimiter */
  merge_store_uint64(&block[curr_offset], slice.size());

  /* Store slice data */
  memcpy(&block[curr_offset + RDB_MERGE_REC_DELIMITER], slice.data(),
      slice.size());

  curr_offset += slice.size() + RDB_MERGE_REC_DELIMITER;
}


void Rdb_index_merge::merge_reset()
{
  /*
    Either error, or all values in the sort buffer have been written to disk,
    so we need to clear the offset tree.
  */
  m_offset_tree.clear();

  /* Reset sort buffer block */
  if (m_rec_buf_unsorted && m_rec_buf_unsorted->block)
  {
    m_rec_buf_unsorted->curr_offset= 0;
  }

  /* Reset output buf */
  if (m_output_buf && m_output_buf->block)
  {
    m_output_buf->curr_offset= 0;
  }
}

}  // namespace myrocks