feat(mtr):hotfix mtr for schema embedding

to fixup the assert failed in Assertion failure: trx0sys.cc:680:trx->state.load(std::memory_order_relaxed) == TRX_STATE_NOT_STARTED
ha_pre_dd_shutdown() calls plugin_foreach(MYSQL_STORAGE_ENGINE_PLUGIN, ...).
Plugin registration order determines call order. InnoDB is a built-in plugin registered before Shannon Rapid,
so the actual sequence is:
```
ha_pre_dd_shutdown()
  → innodb_pre_dd_shutdown()       ← InnoDB fires FIRST
      → srv_pre_dd_shutdown()
          → trx_sys_before_pre_dd_shutdown_validate()   ← assertion here
  → rapid_pre_dd_shutdown()        ← Rapid fires SECOND
      → EmbeddingManager::shutdown()   ← too late
```

Author: ShannonBase

.github/workflows/nightly.yaml

@@ -109,7 +109,7 @@ jobs:
         sudo chown -R $USER:$USER /home/shannon-bin/
         cd /home/shannon-bin/mysql-test/
         sudo chmod -R u+rwx mysql-test-run.pl
-        ./mysql-test-run.pl --mysqld=--loose-rapid_schema_embedding=OFF --suite=main,innodb,federated,rpl,rpl_ndb,rpl_gtid,rpl_nogtid,funcs_1,funcs_2,information_schema,sys_vars,opt_trace,secondary_engine,ml \
+        ./mysql-test-run.pl --mysqld=--loose-rapid_schema_embedding=OFF --suite=main,innodb,innodb_zip,federated,rpl,rpl_ndb,rpl_gtid,rpl_nogtid,funcs_1,funcs_2,information_schema,sys_vars,opt_trace,secondary_engine,ml \
         --big-test --mysqld=--user=$USER --mysqld=--default-storage-engine=innodb --nowarnings --force --nocheck-testcases --retry=3 --parallel=$(nproc)
         # binlog,binlog_gtid,binlog_nogtid, due to disk space limition on github runner, when we have a fast git action runner, we can use the following command to run the test.
         #sudo chmod +x ./collections/default.push && ./collections/default.daily
@@ -171,7 +171,7 @@ jobs:
           sudo chown -R $USER:$USER /home/shannon-bin/
           cd /home/shannon-bin/mysql-test/
           sudo chmod -R u+rwx mysql-test-run.pl
-          ./mysql-test-run.pl --mysqld=--loose-rapid_schema_embedding=OFF --suite=main,innodb,binlog,binlog_gtid,binlog_nogtid,secondary_engine,ml \
+          ./mysql-test-run.pl --mysqld=--loose-rapid_schema_embedding=OFF --suite=main,innodb,innodb_zip,binlog,binlog_gtid,binlog_nogtid,secondary_engine,ml \
           --big-test --mysqld=--user=$USER --mysqld=--default-storage-engine=innodb --nowarnings --force --nocheck-testcases --skip-test-list=skip-tests.arm --retry=3 --parallel=$(nproc)
           # when we have a fast git action runner, we can use the following command to run the test.
           #sudo chmod +x ./collections/default.push && ./collections/default.daily

.github/workflows/pull-requests.yaml

@@ -125,7 +125,7 @@ jobs:
           export ASAN_OPTIONS="detect_leaks=1"
           export LSAN_OPTIONS="suppressions=/home/shannon-bin/mysql-test/lsan.supp"
           ./mysql-test-run.pl --mysqld=--loose-rapid_schema_embedding=OFF --sanitize --suite=main,innodb,binlog,binlog_gtid,binlog_nogtid,federated,rpl,rpl_gtid,rpl_nogtid,funcs_1,funcs_2,\
-          information_schema,secondary_engine,ml,sys_vars \
+          information_schema,secondary_engine,ml,sys_vars,innodb_zip \
           --mysqld=--user=$USER --mysqld=--default-storage-engine=innodb --nowarnings --force --nocheck-testcases --retry=3 --parallel=$(nproc)
         # when we have a fast git action runner, we can use the following command to run the test.
         # sudo chmod -R u+rwx mysql-test-run.pl && sudo chmod +x ./collections/default.push
@@ -190,7 +190,7 @@ jobs:
           sudo chmod -R u+rwx mysql-test-run.pl
           export ASAN_OPTIONS="detect_leaks=1"
           export LSAN_OPTIONS="suppressions=/home/shannon-bin/mysql-test/lsan.supp"
-          ./mysql-test-run.pl --mysqld=--loose-rapid_schema_embedding=OFF --sanitize --suite=main,innodb,secondary_engine,ml,sys_vars \
+          ./mysql-test-run.pl --mysqld=--loose-rapid_schema_embedding=OFF --sanitize --suite=main,innodb,innodb_zip,secondary_engine,ml,sys_vars \
           --mysqld=--user=$USER --mysqld=--default-storage-engine=innodb --nowarnings --force --nocheck-testcases --retry=3 --parallel=$(nproc)
         # when we have a fast git action runner, we can use the following command to run the test.
         # sudo chmod -R u+rwx mysql-test-run.pl && sudo chmod +x ./collections/default.push

ml/ml_retrieve_schema_metadata.cpp

@@ -71,11 +71,6 @@ std::atomic<embedding_state_t> EmbeddingManager::m_state{embedding_state_t::EMBE
 std::condition_variable EmbeddingManager::m_manager_cv;
 std::mutex EmbeddingManager::m_manager_mutex;
 
-std::condition_variable EmbeddingManager::m_fully_stopped_cv;
-std::mutex EmbeddingManager::m_fully_stopped_mutex;
-bool EmbeddingManager::m_fully_stopped{true};
-std::atomic<bool> EmbeddingManager::m_shutdown_initiated{false};
-
 struct ScopedInternalTHD {
   THD *thd{nullptr};
 
@@ -662,7 +657,6 @@ static void *embedding_table_worker_func(void *arg) {
   }
 
   DBUG_PRINT("ml", ("ML TableWorker [%s]: exiting", ctx->key.c_str()));
-  mgr->on_thread_exiting();
   return nullptr;
 }
 
@@ -824,12 +818,11 @@ static void *embedding_manager_func(void *arg) {
   }
 
   DBUG_PRINT("ml", ("ML EmbeddingManager: event loop finished, exiting."));
-  mgr->on_thread_exiting();
   return nullptr;
 }
 
 void EmbeddingManager::start_impl() {
-  if (!m_initialized.load() || m_shutdown_initiated.load(std::memory_order_acquire)) return;
+  if (!m_initialized.load()) return;
   embedding_state_t expected = embedding_state_t::EMBEDDING_STATE_EXIT;
   if (!m_state.compare_exchange_strong(expected, embedding_state_t::EMBEDDING_STATE_RUN, std::memory_order_acq_rel,
                                        std::memory_order_acquire)) {
@@ -840,13 +833,6 @@ void EmbeddingManager::start_impl() {
     }
   }
 
-  {
-    std::lock_guard<std::mutex> lk(m_fully_stopped_mutex);
-    m_fully_stopped = false;
-  }
-
-  m_active_thread_count.fetch_add(1, std::memory_order_relaxed);
-
   {
     std::lock_guard<std::mutex> lk(m_embedder_mutex);
     m_embedder = std::make_unique<ML_embedding_row>();
@@ -862,15 +848,36 @@ void EmbeddingManager::start_impl() {
   if (my_thread_create(&m_manager_thread, &attr, embedding_manager_func, this) != 0) {
     sql_print_error("[EmbeddingManager] start: failed to create coordinator thread");
     my_thread_attr_destroy(&attr);
-    m_active_thread_count.fetch_sub(1, std::memory_order_relaxed);
     m_state.store(embedding_state_t::EMBEDDING_STATE_STOP, std::memory_order_release);
     return;
   }
   my_thread_attr_destroy(&attr);
 }
 
+void EmbeddingManager::shutdown_impl() {
+  initiate_shutdown_impl();
+
+  if (m_manager_thread.thread != 0) {
+    my_thread_join(&m_manager_thread, nullptr);
+    m_manager_thread = {};
+  }
+
+  {
+    std::lock_guard<std::mutex> lk(m_table_workers_mutex);
+    for (auto &[key, ctx] : m_table_workers) {
+      if (ctx->thread.thread != 0) {
+        my_thread_join(&ctx->thread, nullptr);
+        ctx->thread = {};
+      }
+    }
+    m_table_workers.clear();
+  }
+
+  m_state.store(embedding_state_t::EMBEDDING_STATE_EXIT, std::memory_order_release);
+  m_initialized.store(false);
+}
+
 void EmbeddingManager::initiate_shutdown_impl() {
-  m_shutdown_initiated.store(true, std::memory_order_release);
   embedding_state_t expected = embedding_state_t::EMBEDDING_STATE_RUN;
   if (!m_state.compare_exchange_strong(expected, embedding_state_t::EMBEDDING_STATE_STOP, std::memory_order_acq_rel,
                                        std::memory_order_acquire)) {
@@ -907,55 +914,6 @@ void EmbeddingManager::initiate_shutdown_impl() {
   m_manager_cv.notify_all();
 }
 
-void EmbeddingManager::on_thread_exiting() {
-  int prev = m_active_thread_count.fetch_sub(1, std::memory_order_acq_rel);
-  if (prev == 1) {
-    std::lock_guard<std::mutex> lk(m_fully_stopped_mutex);
-    m_fully_stopped = true;
-    m_fully_stopped_cv.notify_all();
-    DBUG_PRINT("ml", ("ML EmbeddingManager: all threads exited — shutdown gate open."));
-  }
-}
-
-void EmbeddingManager::initiate_shutdown() {
-  auto *mgr = instance();
-  if (mgr && EmbeddingManager::is_running()) mgr->initiate_shutdown_impl();
-}
-
-bool EmbeddingManager::wait_until_fully_stopped(std::chrono::milliseconds timeout) {
-  std::unique_lock<std::mutex> lk(m_fully_stopped_mutex);
-  return m_fully_stopped_cv.wait_for(lk, timeout, [] { return m_fully_stopped; });
-}
-
-void EmbeddingManager::shutdown_impl() {
-  initiate_shutdown_impl();
-
-  if (m_manager_thread.thread != 0) {
-    my_thread_join(&m_manager_thread, nullptr);
-    m_manager_thread = {};
-  }
-
-  {
-    std::lock_guard<std::mutex> lk(m_table_workers_mutex);
-    for (auto &[key, ctx] : m_table_workers) {
-      if (ctx->thread.thread != 0) {
-        my_thread_join(&ctx->thread, nullptr);
-        ctx->thread = {};
-      }
-    }
-    m_table_workers.clear();
-  }
-
-  {
-    std::lock_guard<std::mutex> lk(m_fully_stopped_mutex);
-    m_fully_stopped = true;
-  }
-  m_fully_stopped_cv.notify_all();
-
-  m_state.store(embedding_state_t::EMBEDDING_STATE_EXIT, std::memory_order_release);
-  m_initialized.store(false);
-}
-
 /**
  * When the pool is below MAX_WORKER_THREADS, spawn a new dedicated worker for
  * this (schema, table) key.  Once the pool is full, return the existing worker
@@ -977,13 +935,11 @@ TableWorkerContext *EmbeddingManager::get_or_create_worker(const std::string &ke
     my_thread_attr_t attr;
     my_thread_attr_init(&attr);
     my_thread_attr_setdetachstate(&attr, MY_THREAD_CREATE_JOINABLE);
-    m_active_thread_count.fetch_add(1, std::memory_order_relaxed);
     int rc = my_thread_create(&ctx->thread, &attr, embedding_table_worker_func, ctx.get());
     my_thread_attr_destroy(&attr);
 
     if (rc != 0) {
       DBUG_PRINT("ml", ("ML EmbeddingManager: failed to spawn worker for %s", key.c_str()));
-      m_active_thread_count.fetch_sub(1, std::memory_order_relaxed);
       return nullptr;
     }
 

ml/ml_retrieve_schema_metadata.h

@@ -152,10 +152,6 @@ class EmbeddingManager {
     instance()->shutdown_impl();
   }
 
-  static void initiate_shutdown();
-
-  static bool wait_until_fully_stopped(std::chrono::milliseconds timeout = std::chrono::seconds(30));
-
   bool initialized() const { return m_initialized.load(); }
 
   static inline bool is_running() noexcept {
@@ -192,10 +188,6 @@ class EmbeddingManager {
 
   std::atomic<THD *> m_current_thd{nullptr};  // THD inside open_and_lock_tables
 
-  static std::condition_variable m_fully_stopped_cv;
-  static std::mutex m_fully_stopped_mutex;
-  static bool m_fully_stopped;
-
   TableWorkerContext *get_or_create_worker(const std::string &key);
   void consume_results(THD *thd, TABLE *schema_embedding_table_ptr);
   void on_thread_exiting();
@@ -212,11 +204,9 @@ class EmbeddingManager {
 
   my_thread_handle m_manager_thread{};
   std::atomic<bool> m_initialized{false};
-  std::atomic<int> m_active_thread_count{0};
 
   static std::once_flag s_once;
   static EmbeddingManager *s_instance;
-  static std::atomic<bool> m_shutdown_initiated;
 };
 
 void shannon_ml_on_ddl_event(const DDLEvent &event);

mysql-test/suite/innodb_zip/r/16k.result

@@ -21,6 +21,8 @@ AND	t.name LIKE 'mysql%'
         AND     t.name NOT LIKE 'mysql/ndb_binlog_index'
 	ORDER BY t.name, i.index_id;
 table_name	n_cols	table_flags	index_name	root_page	type	n_fields	merge_threshold
+mysql/schema_embeddings	10	161	PRIMARY	41	3	9	50
+mysql/schema_embeddings	10	161	unique_schema_table	42	2	3	50
 CREATE TABLE t1 (a INT KEY, b TEXT) ROW_FORMAT=REDUNDANT ENGINE=innodb;
 CREATE TABLE t2 (a INT KEY, b TEXT) ROW_FORMAT=COMPACT ENGINE=innodb;
 CREATE TABLE t3 (a INT KEY, b TEXT) ROW_FORMAT=COMPRESSED ENGINE=innodb;

mysql-test/suite/innodb_zip/r/4k.result

@@ -21,6 +21,8 @@ AND	t.name LIKE 'mysql%'
         AND     t.name NOT LIKE 'mysql/ndb_binlog_index'
 	ORDER BY t.name, i.index_id;
 table_name	n_cols	table_flags	index_name	root_page	type	n_fields	merge_threshold
+mysql/schema_embeddings	10	161	PRIMARY	138	3	9	50
+mysql/schema_embeddings	10	161	unique_schema_table	139	2	3	50
 CREATE TABLE t1 (a INT KEY, b TEXT) ROW_FORMAT=REDUNDANT ENGINE=innodb;
 CREATE TABLE t2 (a INT KEY, b TEXT) ROW_FORMAT=COMPACT ENGINE=innodb;
 CREATE TABLE t3 (a INT KEY, b TEXT) ROW_FORMAT=COMPRESSED ENGINE=innodb;

mysql-test/suite/innodb_zip/r/8k.result

@@ -21,6 +21,8 @@ AND	t.name LIKE 'mysql%'
         AND     t.name NOT LIKE 'mysql/ndb_binlog_index'
 	ORDER BY t.name, i.index_id;
 table_name	n_cols	table_flags	index_name	root_page	type	n_fields	merge_threshold
+mysql/schema_embeddings	10	161	PRIMARY	73	3	9	50
+mysql/schema_embeddings	10	161	unique_schema_table	74	2	3	50
 CREATE TABLE t1 (a INT KEY, b TEXT) ROW_FORMAT=REDUNDANT ENGINE=innodb;
 CREATE TABLE t2 (a INT KEY, b TEXT) ROW_FORMAT=COMPACT ENGINE=innodb;
 CREATE TABLE t3 (a INT KEY, b TEXT) ROW_FORMAT=COMPRESSED ENGINE=innodb;

sql/mysqld.cc

@@ -993,6 +993,7 @@ MySQL clients support the protocol:
 #include "sql/server_component/persistent_dynamic_loader_imp.h"
 #include "sql/srv_session.h"
 
+#include "ml/ml_retrieve_schema_metadata.h"  // ml::Schema_manager
 using mysql::binlog::event::enum_binlog_checksum_alg;
 using std::max;
 using std::min;
@@ -2713,6 +2714,8 @@ static void clean_up(bool print_message) {
   authentication_policy::deinit();
   denit_command_maps();
 
+  ShannonBase::ML::EmbeddingManager::shutdown();
+
   ha_pre_dd_shutdown();
   dd::shutdown();
 

storage/rapid_engine/handler/ha_shannon_rapid.cc

@@ -1943,18 +1943,9 @@ static handler *rapid_create_handler(handlerton *hton, TABLE_SHARE *table_share,
 
 static void rapid_pre_dd_shutdown(handlerton *) {
   auto *mgr = ShannonBase::ML::EmbeddingManager::instance();
-  if ((!mgr || !mgr->initialized()) || !ShannonBase::ML::EmbeddingManager::is_running())
-    return;  // already stopped or never started
-
-  ShannonBase::ML::EmbeddingManager::initiate_shutdown();
-
-  constexpr auto kTimeout = std::chrono::seconds(60);
-  if (!ShannonBase::ML::EmbeddingManager::wait_until_fully_stopped(kTimeout)) {
-    sql_print_warning(
-        "[EmbeddingManager] shannon_ml_pre_dd_shutdown: "
-        "threads did not stop within 60 s — proceeding anyway.");
-  }
+  if ((!mgr || !mgr->initialized())) return;
 
+  ShannonBase::ML::EmbeddingManager::shutdown();
   DBUG_PRINT("ml", ("ML EmbeddingManager: shannon_ml_pre_dd_shutdown — all threads stopped."));
 }
 
@@ -1963,7 +1954,7 @@ static void rapid_pre_dd_shutdown(handlerton *) {
 @retval 0 always */
 static int rapid_shutdown(handlerton *, ha_panic_function) {
   DBUG_TRACE;
-  // embedding worker thread.
+  // embedding worker thread shut down. Idempotent operation.
   ShannonBase::ML::EmbeddingManager::shutdown();
 
   // recovery worker