Add scheduler API to get CPU stats for current task

Summary: This is part 1 of the change to add CPU stats to the slow log report.

Differential Revision: D52791917

fbshipit-source-id: 2df865d

-----------------------------------------------------------------------------------------------

Add api to fetch wait stats for cpu scheduler

Summary: We want to fetch the wait stats for the worker threads and publish it to the slow logs.

Differential Revision: D53984553

fbshipit-source-id: 1efdceb

-----------------------------------------------------------------------------------------------

Add additional CPU scheduler APIs

Summary:
Introduce a few APIs to cover the scenarios for thrift plugin and child mysql tasks.
 - `tp_create_connection` takes additional parameter whether or not to acquire connection slot against `thread_pool_max_db_connections` and the resource shape quota. The child tasks would set this parameter to `false`.
 - `tp_is_scheduler_enabled` helps determine if the scheduler is enabled or not. If not, the code should fall back to the other APIs (mysql or thrift).
 - `tp_get_current_task_connection` returns connection of the current task, from which the tenant id could be obtained and used to enqueue child tasks.
 - `tp_get_tenant_id` looks up tenant id by db name and allows enqueueing a task without `THD` or `TpConn` which would need to be created/found on the task itself, and then attached.

Another small change is to get rid of the assert in `THD::get_connection_attr`. The scheduler could call this API when it is about to enqueue a new task for a `THD` which is done from a random thread. The `THD` in question doesn't have any current thread at that moment so it is safe to call this API and access the memory.

Differential Revision: D53980947

fbshipit-source-id: 2ef3d6f

include/mysql/service_cpu_scheduler.h

@@ -20,6 +20,8 @@
   @file include/mysql/service_cpu_scheduler.h
 */
 
+#include "my_inttypes.h"
+
 class THD;
 
 /**
@@ -43,15 +45,29 @@ using tp_tenant_id_handle = void *;
 */
 using tp_routine = void *(*)(void *);
 
+/**
+  CPU usage stats.
+*/
+struct tp_cpu_stats {
+  int64_t cpu_usage_ns;
+  int64_t delay_total_ns;
+};
+
 extern "C" struct cpu_scheduler_service_st {
   bool (*enqueue_task)(tp_routine routine, void *param,
                        tp_tenant_id_handle tenant_id, tp_scheduler_hint &hint);
-  tp_conn_handle (*create_connection)(THD *thd, const char *db);
+  tp_conn_handle (*create_connection)(THD *thd, const char *db,
+                                      bool acquire_conn_slot);
   void (*destroy_connection)(tp_conn_handle conn_handle);
   void (*attach_connection)(tp_conn_handle conn_handle);
   void (*detach_connection)(tp_conn_handle conn_handle);
   tp_tenant_id_handle (*get_connection_tenant_id)(tp_conn_handle conn_handle);
   void (*destroy_tenant_id)(tp_tenant_id_handle tenant_id);
+  bool (*get_current_task_cpu_stats)(tp_cpu_stats &cpu_stats);
+  int (*get_current_task_wait_stats)(char* buf_stats, size_t buf_len);
+  bool (*is_scheduler_enabled)();
+  tp_conn_handle (*get_current_task_connection)();
+  tp_tenant_id_handle (*get_tenant_id)(const char *db);
 } * cpu_scheduler_service;
 
 /**
@@ -61,8 +77,8 @@ extern "C" struct cpu_scheduler_service_st {
 
 #define tp_enqueue_task(_ROUTINE, _PARAM, _TENANT_ID, _HINT) \
   cpu_scheduler_service->enqueue_task(_ROUTINE, _PARAM, _TENANT_ID, _HINT)
-#define tp_create_connection(_THD, _DB) \
-  cpu_scheduler_service->create_connection(_THD, _DB)
+#define tp_create_connection(_THD, _DB, _ACQUIRE_SLOT) \
+  cpu_scheduler_service->create_connection(_THD, _DB, _ACQUIRE_SLOT)
 #define tp_destroy_connection(_CONN_HANDLE) \
   cpu_scheduler_service->destroy_connection(_CONN_HANDLE)
 #define tp_attach_connection(_CONN_HANDLE) \
@@ -73,6 +89,14 @@ extern "C" struct cpu_scheduler_service_st {
   cpu_scheduler_service->get_connection_tenant_id(_CONN_HANDLE)
 #define tp_destroy_tenant_id(_TENANT_ID) \
   cpu_scheduler_service->destroy_tenant_id(_TENANT_ID)
+#define tp_get_current_task_cpu_stats(_CPU_STATS) \
+  cpu_scheduler_service->get_current_task_cpu_stats(_CPU_STATS)
+#define tp_get_current_task_wait_stats(_BUF_, _BUF_LEN) \
+  cpu_scheduler_service->get_current_task_wait_stats(_BUF_, _BUF_LEN)
+#define tp_is_scheduler_enabled() cpu_scheduler_service->is_scheduler_enabled()
+#define tp_get_current_task_connection() \
+  cpu_scheduler_service->get_current_task_connection()
+#define tp_get_tenant_id(_DB) cpu_scheduler_service->get_tenant_id(_DB)
 
 #else
 
@@ -81,7 +105,8 @@ extern "C" struct cpu_scheduler_service_st {
 
   @param routine Routine to execute.
   @param param Routine parameter.
-  @param tenant_id Tenant id to associate task with.
+  @param tenant_id Tenant id to associate task with. Passing nullptr will
+                   enqueue the task in the system tenant.
   @param hint Scheduler hint to use, and returns new hint that can be used for
               subsequent enqueue to co-locate it together with previous.
   @return true if task was enqueued, false otherwise.
@@ -94,9 +119,11 @@ bool tp_enqueue_task(tp_routine routine, void *param,
 
   @param thd THD to attach to the connection.
   @param db Database name to assign the connection to.
+  @param acquire_conn_slot Account against thread_pool_max_db_connections.
   @return Connection handle, or nullptr on failure.
 */
-tp_conn_handle tp_create_connection(THD *thd, const char *db);
+tp_conn_handle tp_create_connection(THD *thd, const char *db,
+                                    bool acquire_conn_slot);
 
 /**
   Destroy connection object.
@@ -133,4 +160,55 @@ tp_tenant_id_handle tp_get_connection_tenant_id(tp_conn_handle conn_handle);
 */
 void tp_destroy_tenant_id(tp_tenant_id_handle tenant_id);
 
+/**
+  Get CPU stats for current task.
+
+  @param cpu_stats Stats struct to populate.
+  @return true  if stats are populated,
+          false if current thread doesn't have a CPU scheduler task.
+*/
+bool tp_get_current_task_cpu_stats(tp_cpu_stats &cpu_stats);
+
+/**
+  Get wait stats for current task.
+
+  @param buf buffer to store the wait stats.
+  @param len buffer len.
+  @return expected buffer size.
+*/
+int tp_get_current_task_wait_stats(char* buf_stats, size_t buf_len);
+
+/**
+  Is CPU scheduler enabled? This helps decide whether to use scheduler APIs or
+  fall back to the regular OS threads. If the scheduler is disabled in the
+  middle of a sequence of API calls the remaining APIs will still succeed.
+
+  @return true if enabled, false otherwise.
+*/
+bool tp_is_scheduler_enabled();
+
+/**
+  Get connection for current task. This connection can be used to obtain its
+  tenant and then enqueue child tasks.
+
+  @return Connection if current thread is running CPU scheduler task,
+          nullptr if current thread is not a CPU scheduler task.
+*/
+tp_conn_handle tp_get_current_task_connection();
+
+/**
+  Get tenant for db name. This allows altering the sequence of enqueueing a
+  task. Instead of creating THD and connection first and getting the tenant
+  from the connection, this API could be used to get the tenant. After
+  enqueueing the task, THD and connection can be created on the task itself.
+
+  @param db Database name of the future connection. Note: this API does not
+            acquire connection slot.
+  @return Tenant id handle or nullptr if tenant for db doesn't exist.
+          Note: tenant is created on demand when a first connection to db is
+          created. So this API will fail until first successful
+          tp_create_connection call for this db name.
+*/
+tp_tenant_id_handle tp_get_tenant_id(const char *db);
+
 #endif

include/mysql/services.h.pp

@@ -237,29 +237,46 @@
                                 void *service_callbacks_ctx);
 #include <mysql/service_cpu_scheduler.h>
 
+#include "my_inttypes.h"
 class THD;
 using tp_scheduler_hint = void *;
 using tp_conn_handle = void *;
 using tp_tenant_id_handle = void *;
 using tp_routine = void *(*)(void *);
+struct tp_cpu_stats {
+  int64_t cpu_usage_ns;
+  int64_t delay_total_ns;
+};
 extern "C" struct cpu_scheduler_service_st {
   bool (*enqueue_task)(tp_routine routine, void *param,
                        tp_tenant_id_handle tenant_id, tp_scheduler_hint &hint);
-  tp_conn_handle (*create_connection)(THD *thd, const char *db);
+  tp_conn_handle (*create_connection)(THD *thd, const char *db,
+                                      bool acquire_conn_slot);
   void (*destroy_connection)(tp_conn_handle conn_handle);
   void (*attach_connection)(tp_conn_handle conn_handle);
   void (*detach_connection)(tp_conn_handle conn_handle);
   tp_tenant_id_handle (*get_connection_tenant_id)(tp_conn_handle conn_handle);
   void (*destroy_tenant_id)(tp_tenant_id_handle tenant_id);
+  bool (*get_current_task_cpu_stats)(tp_cpu_stats &cpu_stats);
+  int (*get_current_task_wait_stats)(char* buf_stats, size_t buf_len);
+  bool (*is_scheduler_enabled)();
+  tp_conn_handle (*get_current_task_connection)();
+  tp_tenant_id_handle (*get_tenant_id)(const char *db);
 } * cpu_scheduler_service;
 bool tp_enqueue_task(tp_routine routine, void *param,
                      tp_tenant_id_handle tenant_id, tp_scheduler_hint &hint);
-tp_conn_handle tp_create_connection(THD *thd, const char *db);
+tp_conn_handle tp_create_connection(THD *thd, const char *db,
+                                    bool acquire_conn_slot);
 void tp_destroy_connection(tp_conn_handle conn_handle);
 void tp_attach_connection(tp_conn_handle conn_handle);
 void tp_detach_connection(tp_conn_handle conn_handle);
 tp_tenant_id_handle tp_get_connection_tenant_id(tp_conn_handle conn_handle);
 void tp_destroy_tenant_id(tp_tenant_id_handle tenant_id);
+bool tp_get_current_task_cpu_stats(tp_cpu_stats &cpu_stats);
+int tp_get_current_task_wait_stats(char* buf_stats, size_t buf_len);
+bool tp_is_scheduler_enabled();
+tp_conn_handle tp_get_current_task_connection();
+tp_tenant_id_handle tp_get_tenant_id(const char *db);
 #include <mysql/service_locking.h>
 #include "my_inttypes.h"
 enum enum_locking_service_lock_type {

sql/conn_handler/connection_handler_manager.h

@@ -56,12 +56,18 @@ struct THD_event_functions {
 struct Cpu_scheduler_functions {
   bool (*enqueue_task)(tp_routine routine, void *param,
                        tp_tenant_id_handle tenant_id, tp_scheduler_hint &hint);
-  tp_conn_handle (*create_connection)(THD *thd, const char *db);
+  tp_conn_handle (*create_connection)(THD *thd, const char *db,
+                                      bool acquire_conn_slot);
   void (*destroy_connection)(tp_conn_handle conn_handle);
   void (*attach_connection)(tp_conn_handle conn_handle);
   void (*detach_connection)(tp_conn_handle conn_handle);
   tp_tenant_id_handle (*get_connection_tenant_id)(tp_conn_handle conn_handle);
   void (*destroy_tenant_id)(tp_tenant_id_handle tenant_id);
+  bool (*get_current_task_cpu_stats)(tp_cpu_stats &cpu_stats);
+  int (*get_current_task_wait_stats)(char* buf_stats, size_t buf_len);
+  bool (*is_scheduler_enabled)();
+  tp_conn_handle (*get_current_task_connection)();
+  tp_tenant_id_handle (*get_tenant_id)(const char *db);
 };
 
 /**

sql/sql_class.cc

@@ -4961,9 +4961,10 @@ const std::string &THD::get_query_attr(const std::string &qattr_key) {
 const std::string &THD::get_connection_attr(const std::string &cattr_key) {
   /*
    * pointers (to db name, table names etc) are only valid until
-   * the end of the current query
+   * the end of the current query. However cannot assert that it
+   * is only called by current THD as CPU scheduler calls it for
+   * any THD before it has a worker thread assigned.
    */
-  assert(this == current_thd);
 
   bool found_ignored = false;
   return found_connection_attr(cattr_key, &found_ignored);

sql/sql_plugin_services.h

@@ -256,7 +256,12 @@ static struct my_thread_scheduler_service my_thread_scheduler_handler = {
 static struct cpu_scheduler_service_st cpu_scheduler_handler = {
     tp_enqueue_task,      tp_create_connection, tp_destroy_connection,
     tp_attach_connection, tp_detach_connection, tp_get_connection_tenant_id,
-    tp_destroy_tenant_id};
+    tp_destroy_tenant_id,
+    tp_get_current_task_cpu_stats,
+    tp_get_current_task_wait_stats,
+    tp_is_scheduler_enabled,
+    tp_get_current_task_connection,
+    tp_get_tenant_id};
 
 static struct my_plugin_log_service my_plugin_log_handler = {
     my_plugin_log_message};

sql/sql_thd_api.cc

@@ -800,9 +800,10 @@ bool tp_enqueue_task(tp_routine routine, void *param,
   return false;
 }
 
-tp_conn_handle tp_create_connection(THD *thd, const char *db) {
+tp_conn_handle tp_create_connection(THD *thd, const char *db,
+                                    bool acquire_conn_slot) {
   MYSQL_CALLBACK_RETURN(Connection_handler_manager::cpu_scheduler_functions,
-                        create_connection, (thd, db));
+                        create_connection, (thd, db, acquire_conn_slot));
   return nullptr;
 }
 
@@ -831,3 +832,33 @@ void tp_destroy_tenant_id(tp_tenant_id_handle tenant_id) {
   MYSQL_CALLBACK(Connection_handler_manager::cpu_scheduler_functions,
                  destroy_tenant_id, (tenant_id));
 }
+
+bool tp_get_current_task_cpu_stats(tp_cpu_stats &cpu_stats) {
+  MYSQL_CALLBACK_RETURN(Connection_handler_manager::cpu_scheduler_functions,
+                        get_current_task_cpu_stats, (cpu_stats));
+  return false;
+}
+
+int tp_get_current_task_wait_stats(char* buf_stats, size_t buf_len) {
+  MYSQL_CALLBACK_RETURN(Connection_handler_manager::cpu_scheduler_functions,
+                        get_current_task_wait_stats, (buf_stats, buf_len));
+  return -1;
+}
+
+bool tp_is_scheduler_enabled() {
+  MYSQL_CALLBACK_RETURN(Connection_handler_manager::cpu_scheduler_functions,
+                        is_scheduler_enabled, ());
+  return false;
+}
+
+tp_conn_handle tp_get_current_task_connection() {
+  MYSQL_CALLBACK_RETURN(Connection_handler_manager::cpu_scheduler_functions,
+                        get_current_task_connection, ());
+  return nullptr;
+}
+
+tp_tenant_id_handle tp_get_tenant_id(const char *db) {
+  MYSQL_CALLBACK_RETURN(Connection_handler_manager::cpu_scheduler_functions,
+                        get_tenant_id, (db));
+  return nullptr;
+}