MyRocks: merge the key packing buffer allocation, annotate m_pack_buffer lifetime (#1446)

Summary:
All the key packing buffers in the MyRocks handler have exactly the same
lifetime between alloc_key_buffers and free_key_buffers. Thus allocate a single
large buffer and point the key packing buffers to offsets inside it, reducing
number of calls to the heap allocator on table opening/close code path. Annotate
the padding bytes for Valgrind as inaccessible. Add additional handling for
m_pack_buffer: since its use it's fully contained in Rdb_key_def::pack_record,
annotate it as such.

Pull Request resolved: #1446

Differential Revision: D56051395

fbshipit-source-id: f503d56

Author: laurynas-biveinis

include/memory_debugging.h

@@ -41,11 +41,27 @@
 #include <valgrind/memcheck.h>
 
 #define MEM_UNDEFINED(a, len) VALGRIND_MAKE_MEM_UNDEFINED(a, len)
-#define MEM_DEFINED_IF_ADDRESSABLE(a, len) \
-  VALGRIND_MAKE_MEM_DEFINED_IF_ADDRESSABLE(a, len)
+// #define MEM_DEFINED_IF_ADDRESSABLE(a, len) \
+//   VALGRIND_MAKE_MEM_DEFINED_IF_ADDRESSABLE(a, len)
 #define MEM_NOACCESS(a, len) VALGRIND_MAKE_MEM_NOACCESS(a, len)
 #define MEM_CHECK_ADDRESSABLE(a, len) VALGRIND_CHECK_MEM_IS_ADDRESSABLE(a, len)
 
+#elif defined(HAVE_ASAN)
+
+#include <cassert>
+#include <sanitizer/asan_interface.h>
+
+#define MEM_MALLOCLIKE_BLOCK(p1, p2, p3, p4) ASAN_UNPOISON_MEMORY_REGION(p1, p2)
+#define MEM_FREELIKE_BLOCK(p1, p2, p3, p4) ASAN_POISON_MEMORY_REGION(p1, p2)
+
+#define MEM_UNDEFINED(a, len) ASAN_UNPOISON_MEMORY_REGION(a, len)
+#define MEM_NOACCESS(a, len) ASAN_POISON_MEMORY_REGION(a, len)
+
+// In the case of error, this will immediatelly terminate the process instead of
+// printing an error and continuing, which is more common for ASan errors.
+// Change it to log and continue if that becomes an issue.
+#define MEM_CHECK_ADDRESSABLE(a, len) assert(!__asan_region_is_poisoned(a, len))
+
 #else /* HAVE_VALGRIND */
 
 #define MEM_MALLOCLIKE_BLOCK(p1, p2, p3, p4) \
@@ -55,13 +71,18 @@
   do {                             \
   } while (0)
 #define MEM_UNDEFINED(a, len) ((void)0)
-#define MEM_DEFINED_IF_ADDRESSABLE(a, len) ((void)0)
+// #define MEM_DEFINED_IF_ADDRESSABLE(a, len) ((void)0)
 #define MEM_NOACCESS(a, len) ((void)0)
 #define MEM_CHECK_ADDRESSABLE(a, len) ((void)0)
 
-#endif
+#endif /* HAVE_VALGRIND */
+
+// Not used in the current sources. If it starts being used after a rebase,
+// somehow implement it (could be an unconditional unpoison?) for the HAVE_ASAN
+// case above.
+#pragma GCC poison MEM_DEFINED_IF_ADDRESSABLE
 
-#if !defined(NDEBUG) || defined(HAVE_VALGRIND)
+#if !defined(NDEBUG) || defined(HAVE_VALGRIND) || defined(HAVE_ASAN)
 
 /**
   Put bad content in memory to be sure it will segfault if dereferenced.

storage/rocksdb/ha_rocksdb.cc

@@ -9364,15 +9364,22 @@ int ha_rocksdb::convert_record_from_storage_format(
                                            : rc;
 }
 
+static constexpr std::size_t get_next_aligned(std::size_t &offset) {
+  constexpr auto bits_below_alignment = alignof(std::max_align_t) - 1;
+  constexpr auto bits_above_alignment = ~bits_below_alignment;
+
+  const auto padding =
+      ((offset + bits_below_alignment) & bits_above_alignment) - offset;
+  offset += padding;
+
+  return padding;
+}
+
 int ha_rocksdb::alloc_key_buffers(const TABLE &table_arg,
                                   const Rdb_tbl_def &tbl_def_arg) {
   DBUG_ENTER_FUNC();
 
   std::shared_ptr<Rdb_key_def> *const kd_arr = tbl_def_arg.m_key_descr_arr;
-
-  uint max_packed_sk_len = 0;
-  uint pack_key_len = 0;
-
   m_pk_descr = kd_arr[pk_index(table_arg, tbl_def_arg)];
 
   // move this into get_table_handler() ??
@@ -9381,12 +9388,11 @@ int ha_rocksdb::alloc_key_buffers(const TABLE &table_arg,
     return rtn;
   }
 
-  pack_key_len = m_pk_descr->max_storage_fmt_length();
-  m_pk_packed_tuple = reinterpret_cast<uchar *>(
-      my_malloc(PSI_NOT_INSTRUMENTED, pack_key_len, MYF(0)));
+  const auto pack_key_len = m_pk_descr->max_storage_fmt_length();
+  auto buf_size = static_cast<std::size_t>(pack_key_len);
 
   /* Sometimes, we may use m_sk_packed_tuple for storing packed PK */
-  max_packed_sk_len = pack_key_len;
+  auto max_packed_sk_len = pack_key_len;
   for (uint i = 0; i < table_arg.s->keys; i++) {
     /* Primary key was processed above */
     if (i == table_arg.s->primary_key) continue;
@@ -9403,48 +9409,57 @@ int ha_rocksdb::alloc_key_buffers(const TABLE &table_arg,
     }
   }
 
-  m_sk_packed_tuple = reinterpret_cast<uchar *>(
-      my_malloc(PSI_NOT_INSTRUMENTED, max_packed_sk_len, MYF(0)));
-  m_sk_packed_tuple_old = reinterpret_cast<uchar *>(
-      my_malloc(PSI_NOT_INSTRUMENTED, max_packed_sk_len, MYF(0)));
-  m_sk_packed_tuple_updated = reinterpret_cast<uchar *>(
-      my_malloc(PSI_NOT_INSTRUMENTED, max_packed_sk_len, MYF(0)));
-  m_end_key_packed_tuple = reinterpret_cast<uchar *>(
-      my_malloc(PSI_NOT_INSTRUMENTED, max_packed_sk_len, MYF(0)));
-  m_pack_buffer = reinterpret_cast<uchar *>(
-      my_malloc(PSI_NOT_INSTRUMENTED, max_packed_sk_len, MYF(0)));
+  const auto pad1 [[maybe_unused]] = get_next_aligned(buf_size);
+  const auto m_sk_packed_tuple_offset = buf_size;
+  buf_size += max_packed_sk_len;
+
+  const auto pad2 [[maybe_unused]] = get_next_aligned(buf_size);
+  const auto m_sk_packed_tuple_old_offset = buf_size;
+  buf_size += max_packed_sk_len;
+
+  const auto pad3 [[maybe_unused]] = get_next_aligned(buf_size);
+  const auto m_sk_packed_tuple_updated_offset = buf_size;
+  buf_size += max_packed_sk_len;
 
-  if (m_pk_packed_tuple == nullptr || m_sk_packed_tuple == nullptr ||
-      m_sk_packed_tuple_old == nullptr ||
-      m_sk_packed_tuple_updated == nullptr ||
-      m_end_key_packed_tuple == nullptr || m_pack_buffer == nullptr) {
-    // One or more of the above allocations failed.  Clean up and exit
+  const auto pad4 [[maybe_unused]] = get_next_aligned(buf_size);
+  const auto m_end_key_packed_tuple_offset = buf_size;
+  buf_size += max_packed_sk_len;
+
+  const auto pad5 [[maybe_unused]] = get_next_aligned(buf_size);
+  const auto m_pack_buffer_offset = buf_size;
+  buf_size += max_packed_sk_len;
+
+  buffers.reset(static_cast<uchar *>(
+      my_malloc(PSI_NOT_INSTRUMENTED, buf_size, MYF(0))));
+  if (buffers == nullptr) {
     free_key_buffers();
 
     DBUG_RETURN(HA_ERR_OUT_OF_MEM);
   }
 
+  m_pk_packed_tuple = buffers.get();
+  MEM_NOACCESS(m_pk_packed_tuple + pack_key_len, pad1);
+  m_sk_packed_tuple = buffers.get() + m_sk_packed_tuple_offset;
+  MEM_NOACCESS(m_sk_packed_tuple + max_packed_sk_len, pad2);
+  m_sk_packed_tuple_old = buffers.get() + m_sk_packed_tuple_old_offset;
+  MEM_NOACCESS(m_sk_packed_tuple_old + max_packed_sk_len, pad3);
+  m_sk_packed_tuple_updated = buffers.get() + m_sk_packed_tuple_updated_offset;
+  MEM_NOACCESS(m_sk_packed_tuple_updated + max_packed_sk_len, pad4);
+  m_end_key_packed_tuple = buffers.get() + m_end_key_packed_tuple_offset;
+  MEM_NOACCESS(m_end_key_packed_tuple + max_packed_sk_len, pad5);
+  m_pack_buffer = buffers.get() + m_pack_buffer_offset;
+
   DBUG_RETURN(HA_EXIT_SUCCESS);
 }
 
 void ha_rocksdb::free_key_buffers() {
-  my_free(m_pk_packed_tuple);
   m_pk_packed_tuple = nullptr;
-
-  my_free(m_sk_packed_tuple);
   m_sk_packed_tuple = nullptr;
-
-  my_free(m_sk_packed_tuple_old);
   m_sk_packed_tuple_old = nullptr;
-
-  my_free(m_sk_packed_tuple_updated);
   m_sk_packed_tuple_updated = nullptr;
-
-  my_free(m_end_key_packed_tuple);
   m_end_key_packed_tuple = nullptr;
-
-  my_free(m_pack_buffer);
   m_pack_buffer = nullptr;
+  buffers.reset();
 
   release_blob_buffer();
 }

storage/rocksdb/ha_rocksdb.h

@@ -194,8 +194,14 @@ class ha_rocksdb : public my_core::handler, public blob_buffer {
   */
   mutable bool m_pk_can_be_decoded;
 
+  // The common buffer for m_pk_packed_tuple, m_sk_packed_tuple,
+  // m_sk_packed_tuple_old, m_sk_packed_tuple_updated, m_end_key_packed_tuple,
+  // & m_pack_buffer.
+  unique_ptr_my_free<uchar[]> buffers;
+
   uchar *m_pk_packed_tuple; /* Buffer for storing PK in StorageFormat */
-  // ^^ todo: change it to 'char*'? TODO: ^ can we join this with last_rowkey?
+  // TODO: ^ can we join this with last_rowkey?
+  // ^^ todo: change it to 'char[]'?
 
   /*
     Temporary buffers for storing the key part of the Key/Value pair

storage/rocksdb/rdb_datadic.cc

@@ -1320,6 +1320,8 @@ uint Rdb_key_def::pack_record(const TABLE *const tbl, uchar *const pack_buffer,
   assert_IMP(should_store_row_debug_checksums,
              (m_index_type == INDEX_TYPE_SECONDARY));
 
+  MEM_UNDEFINED(pack_buffer, max_storage_fmt_length());
+
   uchar *tuple = packed_tuple;
   size_t unpack_start_pos = size_t(-1);
   size_t unpack_len_pos = size_t(-1);
@@ -1471,6 +1473,8 @@ uint Rdb_key_def::pack_record(const TABLE *const tbl, uchar *const pack_buffer,
 
   assert(is_storage_available(tuple - packed_tuple, 0));
 
+  MEM_NOACCESS(pack_buffer, max_storage_fmt_length());
+
   return tuple - packed_tuple;
 }