From 0d4be10a8af30cacc038c589792e772cb42683fc Mon Sep 17 00:00:00 2001
From: Dmitry Shulga <dmitry.shulga@mariadb.com>
Date: Fri, 1 Sep 2023 19:25:33 +0700
Subject: [PATCH] MDEV-14959: Control over memory allocated for SP/PS

This patch adds support for controlling of memory allocation
done by SP/PS that could happen on second and following executions.
As soon as SP or PS has been executed the first time its memory root
is marked as read only since no further memory allocation should
be performed on it. In case such allocation takes place it leads to
the assert hit for invariant that force no new memory allocations
takes place as soon as the SP/PS has been marked as read only.

The feature for control of memory allocation made on behalf SP/PS
is turned on when both debug build is on and the cmake option
-DWITH_PROTECT_STATEMENT_MEMROOT is set.

The reason for introduction of the new cmake option
  -DWITH_PROTECT_STATEMENT_MEMROOT
to control memory allocation of second and following executions of
SP/PS is that for the current server implementation there are too many
places where such memory allocation takes place. As soon as all such
incorrect allocations be fixed the cmake option
 -DWITH_PROTECT_STATEMENT_MEMROOT
can be removed and control of memory allocation made on second and
following executions can be turned on only for debug build. Before
every incorrect memory allocation be fixed it makes sense to guard
the checking of memory allocation on read only memory by extra cmake
option else we would get a lot of failing test on buildbot.

Moreover, fixing of all incorrect memory allocations could take pretty
long period of time, so for introducing the feature without necessary
to wait until all places throughout the source code be fixed it makes
sense to add the new cmake option.
---
 CMakeLists.txt     |  9 +++++++
 include/my_alloc.h |  4 +++
 mysys/my_alloc.c   |  7 +++++
 sql/sp_head.cc     | 67 ++++++++++++++++++++++++++++++++++++++++++++++
 sql/sp_head.h      | 37 +++++++++++++++++++++++++
 sql/sql_prepare.cc | 44 ++++++++++++++++++++++++++++++
 6 files changed, 168 insertions(+)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 7e5f3e157c3..9605d786c93 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -193,6 +193,15 @@ ENDIF()
 
 OPTION(NOT_FOR_DISTRIBUTION "Allow linking with GPLv2-incompatible system libraries. Only set it you never plan to distribute the resulting binaries" OFF)
 
+#
+#  Enable protection of statement's memory root after first SP/PS execution.
+#  Can be switched on only for debug build.
+#
+OPTION(WITH_PROTECT_STATEMENT_MEMROOT "Enable protection of statement's memory root after first SP/PS execution. Turned into account only for debug build" OFF)
+IF (CMAKE_BUILD_TYPE MATCHES "Debug" AND WITH_PROTECT_STATEMENT_MEMROOT)
+  ADD_DEFINITIONS(-DPROTECT_STATEMENT_MEMROOT)
+ENDIF()
+
 INCLUDE(check_compiler_flag)
 INCLUDE(check_linker_flag)
 
diff --git a/include/my_alloc.h b/include/my_alloc.h
index 9b0aad26956..6399ce67667 100644
--- a/include/my_alloc.h
+++ b/include/my_alloc.h
@@ -51,6 +51,10 @@ typedef struct st_mem_root
   */
   unsigned int first_block_usage;
 
+#ifdef PROTECT_STATEMENT_MEMROOT
+  int read_only;
+#endif
+
   void (*error_handler)(void);
   const char *name;
 } MEM_ROOT;
diff --git a/mysys/my_alloc.c b/mysys/my_alloc.c
index a6f38dcb145..4a5e48c9e80 100644
--- a/mysys/my_alloc.c
+++ b/mysys/my_alloc.c
@@ -74,6 +74,9 @@ void init_alloc_root(MEM_ROOT *mem_root, const char *name, size_t block_size,
   mem_root->first_block_usage= 0;
   mem_root->total_alloc= 0;
   mem_root->name= name;
+#ifdef PROTECT_STATEMENT_MEMROOT
+  mem_root->read_only= 0;
+#endif
 
 #if !(defined(HAVE_valgrind) && defined(EXTRA_DEBUG))
   if (pre_alloc_size)
@@ -218,6 +221,10 @@ void *alloc_root(MEM_ROOT *mem_root, size_t length)
   DBUG_PRINT("enter",("root: %p  name: %s", mem_root, mem_root->name));
   DBUG_ASSERT(alloc_root_inited(mem_root));
 
+#ifdef PROTECT_STATEMENT_MEMROOT
+  DBUG_ASSERT(mem_root->read_only == 0);
+#endif
+
   DBUG_EXECUTE_IF("simulate_out_of_memory",
                   {
                     /* Avoid reusing an already allocated block */
diff --git a/sql/sp_head.cc b/sql/sp_head.cc
index 8801d8880b1..b1a697ada2a 100644
--- a/sql/sp_head.cc
+++ b/sql/sp_head.cc
@@ -493,6 +493,9 @@ sp_head::sp_head(MEM_ROOT *mem_root_arg, sp_package *parent,
   :Query_arena(NULL, STMT_INITIALIZED_FOR_SP),
    Database_qualified_name(&null_clex_str, &null_clex_str),
    main_mem_root(*mem_root_arg),
+#ifdef PROTECT_STATEMENT_MEMROOT
+   executed_counter(0),
+#endif
    m_parent(parent),
    m_handler(sph),
    m_flags(0),
@@ -738,6 +741,10 @@ sp_head::init(LEX *lex)
   */
   lex->trg_table_fields.empty();
 
+#ifdef PROTECT_STATEMENT_MEMROOT
+  executed_counter= 0;
+#endif
+
   DBUG_VOID_RETURN;
 }
 
@@ -1364,6 +1371,11 @@ sp_head::execute(THD *thd, bool merge_da_on_success)
 #endif /* WITH_WSREP */
     err_status= i->execute(thd, &ip);
 
+#ifdef PROTECT_STATEMENT_MEMROOT
+    if (!err_status)
+      i->mark_as_run();
+#endif
+
 #ifdef WITH_WSREP
     if (WSREP(thd))
     {
@@ -1460,6 +1472,16 @@ sp_head::execute(THD *thd, bool merge_da_on_success)
     /* Reset sp_rcontext::end_partial_result_set flag. */
     ctx->end_partial_result_set= FALSE;
 
+#ifdef PROTECT_STATEMENT_MEMROOT
+    if (thd->is_error())
+    {
+      // Don't count a call ended with an error as normal run
+      executed_counter= 0;
+      main_mem_root.read_only= 0;
+      reset_instrs_executed_counter();
+    }
+#endif
+
   } while (!err_status && likely(!thd->killed) &&
            likely(!thd->is_fatal_error) &&
            !thd->spcont->pause_state);
@@ -1571,6 +1593,20 @@ sp_head::execute(THD *thd, bool merge_da_on_success)
 
     err_status|= mysql_change_db(thd, (LEX_CSTRING*)&saved_cur_db_name, TRUE) != 0;
   }
+
+#ifdef PROTECT_STATEMENT_MEMROOT
+  if (!err_status)
+  {
+    if (!main_mem_root.read_only &&
+        has_all_instrs_executed())
+    {
+      main_mem_root.read_only= 1;
+    }
+    ++executed_counter;
+    DBUG_PRINT("info", ("execute counter: %lu", executed_counter));
+  }
+#endif
+
   m_flags&= ~IS_INVOKED;
   if (m_parent)
     m_parent->m_invoked_subroutine_count--;
@@ -3214,6 +3250,37 @@ void sp_head::add_mark_lead(uint ip, List<sp_instr> *leads)
     leads->push_front(i);
 }
 
+#ifdef PROTECT_STATEMENT_MEMROOT
+
+int sp_head::has_all_instrs_executed()
+{
+  sp_instr *ip;
+  uint count= 0;
+
+  for (uint i= 0; i < m_instr.elements; ++i)
+  {
+    get_dynamic(&m_instr, (uchar*)&ip, i);
+    if (ip->has_been_run())
+      ++count;
+  }
+
+  return count == m_instr.elements;
+}
+
+
+void sp_head::reset_instrs_executed_counter()
+{
+  sp_instr *ip;
+
+  for (uint i= 0; i < m_instr.elements; ++i)
+  {
+    get_dynamic(&m_instr, (uchar*)&ip, i);
+    ip->mark_as_not_run();
+  }
+}
+
+#endif
+
 void
 sp_head::opt_mark()
 {
diff --git a/sql/sp_head.h b/sql/sp_head.h
index 06f5d5234ae..693a5e78703 100644
--- a/sql/sp_head.h
+++ b/sql/sp_head.h
@@ -132,6 +132,16 @@ class sp_head :private Query_arena,
 
 protected:
   MEM_ROOT main_mem_root;
+#ifdef PROTECT_STATEMENT_MEMROOT
+  /*
+    The following data member is wholly for debugging purpose.
+    It can be used for possible crash analysis to determine how many times
+    the stored routine was executed before the mem_root marked read_only
+    was requested for a memory chunk. Additionally, a value of this data
+    member is output to the log with DBUG_PRINT.
+  */
+  ulong executed_counter;
+#endif
 public:
   /** Possible values of m_flags */
   enum {
@@ -801,6 +811,11 @@ public:
     return ip;
   }
 
+#ifdef PROTECT_STATEMENT_MEMROOT
+  int has_all_instrs_executed();
+  void reset_instrs_executed_counter();
+#endif
+
   /* Add tables used by routine to the table list. */
   bool add_used_tables_to_table_list(THD *thd,
                                      TABLE_LIST ***query_tables_last_ptr,
@@ -1084,6 +1099,9 @@ public:
   /// Should give each a name or type code for debugging purposes?
   sp_instr(uint ip, sp_pcontext *ctx)
     :Query_arena(0, STMT_INITIALIZED_FOR_SP), marked(0), m_ip(ip), m_ctx(ctx)
+#ifdef PROTECT_STATEMENT_MEMROOT
+  , m_has_been_run(false)
+#endif
   {}
 
   virtual ~sp_instr()
@@ -1173,6 +1191,25 @@ public:
     m_ip= dst;
   }
 
+#ifdef PROTECT_STATEMENT_MEMROOT
+  bool has_been_run() const
+  {
+    return m_has_been_run;
+  }
+
+  void mark_as_run()
+  {
+    m_has_been_run= true;
+  }
+
+  void mark_as_not_run()
+  {
+    m_has_been_run= false;
+  }
+
+private:
+  bool m_has_been_run;
+#endif
 }; // class sp_instr : public Sql_alloc
 
 
diff --git a/sql/sql_prepare.cc b/sql/sql_prepare.cc
index 9c52aea2931..63ce4d4c26d 100644
--- a/sql/sql_prepare.cc
+++ b/sql/sql_prepare.cc
@@ -171,6 +171,16 @@ public:
   Server_side_cursor *cursor;
   uchar *packet;
   uchar *packet_end;
+#ifdef PROTECT_STATEMENT_MEMROOT
+  /*
+    The following data member is wholly for debugging purpose.
+    It can be used for possible crash analysis to determine how many times
+    the stored routine was executed before the mem_root marked read_only
+    was requested for a memory chunk. Additionally, a value of this data
+    member is output to the log with DBUG_PRINT.
+  */
+  ulong executed_counter;
+#endif
   uint param_count;
   uint last_errno;
   uint flags;
@@ -3995,6 +4005,9 @@ Prepared_statement::Prepared_statement(THD *thd_arg)
   cursor(0),
   packet(0),
   packet_end(0),
+#ifdef PROTECT_STATEMENT_MEMROOT
+  executed_counter(0),
+#endif
   param_count(0),
   last_errno(0),
   flags((uint) IS_IN_USE),
@@ -4070,8 +4083,13 @@ void Prepared_statement::setup_set_params()
 Prepared_statement::~Prepared_statement()
 {
   DBUG_ENTER("Prepared_statement::~Prepared_statement");
+#ifdef PROTECT_STATEMENT_MEMROOT
+  DBUG_PRINT("enter",("stmt: %p  cursor: %p executed_counter: %lu",
+                      this, cursor, executed_counter));
+#else
   DBUG_PRINT("enter",("stmt: %p  cursor: %p",
                       this, cursor));
+#endif
   delete cursor;
   /*
     We have to call free on the items even if cleanup is called as some items,
@@ -4237,6 +4255,10 @@ bool Prepared_statement::prepare(const char *packet, uint packet_len)
   }
   lex->set_trg_event_type_for_tables();
 
+#ifdef PROTECT_STATEMENT_MEMROOT
+  executed_counter= 0;
+#endif
+
   /*
     While doing context analysis of the query (in check_prepared_statement)
     we allocate a lot of additional memory: for open tables, JOINs, derived
@@ -4506,9 +4528,31 @@ reexecute:
     error= reprepare();
 
     if (likely(!error))                         /* Success */
+    {
+#ifdef PROTECT_STATEMENT_MEMROOT
+      // There was reprepare so the counter of runs should be reset
+      executed_counter= 0;
+      mem_root->read_only= 0;
+#endif
       goto reexecute;
+    }
   }
   reset_stmt_params(this);
+#ifdef PROTECT_STATEMENT_MEMROOT
+  if (!error)
+  {
+    mem_root->read_only= 1;
+    ++executed_counter;
+
+    DBUG_PRINT("info", ("execute counter: %lu", executed_counter));
+  }
+  else
+  {
+    // Error on call shouldn't be counted as a normal run
+    executed_counter= 0;
+    mem_root->read_only= 0;
+  }
+#endif
 
   return error;
 }