diff --git a/sql/sql_lex.cc b/sql/sql_lex.cc
index e411cc488..24549bdaa 100644
--- a/sql/sql_lex.cc
+++ b/sql/sql_lex.cc
@@ -4398,23 +4398,13 @@ void Query_block::include_chain_in_global(Query_block **start) {
   *start = this;
 }
 
-// stonedb8 start TODO: a better way.
-//
-// follows: from MySQL 5.7.36
-//void st_select_lex::set_join(JOIN *join_arg)
-//{
-//  master_unit()->thd->lock_query_plan();
-//  join= join_arg;
-//  master_unit()->thd->unlock_query_plan();
-//}
-//
-void Query_block::set_join(JOIN *join_arg)
+void Query_block::set_join(THD *thd, JOIN *join_arg)
 {
+  thd->lock_query_plan();
   join= join_arg;
+  thd->unlock_query_plan();
 }
 
-// stonedb8 end
-
 /**
    Helper function which handles the "ON conditions" part of
    Query_block::get_optimizable_conditions().
diff --git a/sql/sql_lex.h b/sql/sql_lex.h
index 7f6bb9d46..c3e997536 100644
--- a/sql/sql_lex.h
+++ b/sql/sql_lex.h
@@ -1797,7 +1797,7 @@ class Query_block {
    The function sets the pointer only after acquiring THD::LOCK_query_plan
    mutex. This is needed to avoid races when EXPLAIN FOR CONNECTION is used.
 */
-  void set_join(JOIN *join_arg);
+  void set_join(THD *thd, JOIN *join_arg);
 
   /**
     Does permanent transformations which are local to a query block (which do
diff --git a/storage/tianmu/core/engine_execute.cpp b/storage/tianmu/core/engine_execute.cpp
index b2d90b4c3..0f07fbb43 100644
--- a/storage/tianmu/core/engine_execute.cpp
+++ b/storage/tianmu/core/engine_execute.cpp
@@ -361,7 +361,7 @@ int optimize_select(THD *thd, ulong select_options, Query_result *result,
         }       
         if (!(join = new JOIN(thd, select_lex)))
             return true; /* purecov: inspected */
-        select_lex->set_join(join); // stonedb8 TODO
+        select_lex->set_join(thd, join);
         
 	}
     join->best_rowcount = 2;
@@ -564,7 +564,7 @@ int Query_expression::optimize_for_tianmu(THD *thd) {
                     cleanup(thd, 0); // stonedb8
                     return true;
                 }
-                sl->set_join(join);
+                sl->set_join(thd, join);
             }           
             if (is_optimized())
                 sl->join->reset();
@@ -651,7 +651,7 @@ int Query_expression::optimize_after_tianmu(THD *thd)
                 cleanup(thd, 0);  // stonedb8
                 return true;
             }
-            sl->set_join(join);
+            sl->set_join(thd, join);
         }           
         int res = sl->join->optimize(2);
         if (res) {