From 571865166e9c2c3428f31804058fcb898f4ef854 Mon Sep 17 00:00:00 2001
From: Hongqin-Li <ihongqinli@gmail.com>
Date: Thu, 13 Apr 2023 05:56:17 +0000
Subject: [PATCH] Adjust timeout of drop database rpc according to total tablet
 count

---
 src/rootserver/ddl_task/ob_ddl_retry_task.cpp |  4 ++-
 src/share/ob_ddl_common.cpp                   | 31 +++++++++++++++++++
 src/share/ob_ddl_common.h                     |  1 +
 3 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/src/rootserver/ddl_task/ob_ddl_retry_task.cpp b/src/rootserver/ddl_task/ob_ddl_retry_task.cpp
index 8991a8682..4bba74cda 100644
--- a/src/rootserver/ddl_task/ob_ddl_retry_task.cpp
+++ b/src/rootserver/ddl_task/ob_ddl_retry_task.cpp
@@ -355,11 +355,13 @@ int ObDDLRetryTask::drop_schema(const ObDDLTaskStatus next_task_status)
     obrpc::ObCommonRpcProxy common_rpc_proxy = root_service_->get_common_rpc_proxy().to(GCTX.self_addr()).timeout(GCONF._ob_ddl_timeout);
     switch(task_type_) {
       case ObDDLType::DDL_DROP_DATABASE: {
+        int64_t timeout_us = 0;
         obrpc::ObDropDatabaseRes drop_database_res;
         obrpc::ObDropDatabaseArg *arg = static_cast<obrpc::ObDropDatabaseArg *>(ddl_arg_);
         arg->is_add_to_scheduler_ = false;
         arg->task_id_ = task_id_;
-        if (OB_FAIL(common_rpc_proxy.drop_database(*arg, drop_database_res))) {
+        ObDDLUtil::get_ddl_rpc_timeout_for_database(tenant_id_, object_id_, timeout_us);
+        if (OB_FAIL(common_rpc_proxy.timeout(timeout_us).drop_database(*arg, drop_database_res))) {
           LOG_WARN("fail to drop database", K(ret));
         } else {
           affected_rows_ = drop_database_res.affected_row_;
diff --git a/src/share/ob_ddl_common.cpp b/src/share/ob_ddl_common.cpp
index 09b13dad7..4852b0d57 100644
--- a/src/share/ob_ddl_common.cpp
+++ b/src/share/ob_ddl_common.cpp
@@ -1010,6 +1010,37 @@ int ObDDLUtil::get_ddl_rpc_timeout(const int64_t tenant_id, const int64_t table_
   }
   return ret;
 }
+
+void ObDDLUtil::get_ddl_rpc_timeout_for_database(const int64_t tenant_id, const int64_t database_id, int64_t &ddl_rpc_timeout_us)
+{
+  int ret = OB_SUCCESS;
+  int64_t tablet_count = 0;
+  share::schema::ObSchemaGetterGuard schema_guard;
+  ObArray<uint64_t> table_ids;
+  int64_t total_tablet_cnt = 0;
+  if (OB_INVALID_ID == tenant_id || OB_INVALID_ID == database_id) {
+    ret = OB_INVALID_ARGUMENT;
+    LOG_WARN("invalid arguments", K(ret), K(tenant_id), K(database_id));
+  } else if (OB_FAIL(share::schema::ObMultiVersionSchemaService::get_instance().get_tenant_schema_guard(
+      tenant_id, schema_guard))) {
+    LOG_WARN("get tenant schema guard failed", K(ret), K(tenant_id));
+  } else if (OB_FAIL(schema_guard.get_table_ids_in_database(tenant_id,
+                                                            database_id,
+                                                            table_ids))) {
+    LOG_WARN("failed to get table ids in database", K(ret));
+  }
+  for (int64_t i = 0; i < table_ids.count(); i++) {
+    int64_t tablet_count = 0;
+    if (OB_SUCCESS != get_tablet_count(tenant_id, table_ids[i], tablet_count)) {
+      tablet_count = 0;
+    }
+    total_tablet_cnt += tablet_count;
+  }
+  (void)get_ddl_rpc_timeout(total_tablet_cnt, ddl_rpc_timeout_us);
+  ddl_rpc_timeout_us = max(ddl_rpc_timeout_us, GCONF._ob_ddl_timeout);
+  return;
+}
+
 int ObDDLUtil::get_ddl_tx_timeout(const int64_t tablet_count, int64_t &ddl_tx_timeout_us)
 {
   int ret = OB_SUCCESS;
diff --git a/src/share/ob_ddl_common.h b/src/share/ob_ddl_common.h
index f51ddeefb..d75a63d5f 100644
--- a/src/share/ob_ddl_common.h
+++ b/src/share/ob_ddl_common.h
@@ -326,6 +326,7 @@ public:
   static int get_ddl_rpc_timeout(const int64_t tenant_id, const int64_t table_id, int64_t &ddl_rpc_timeout_us);
   static int get_ddl_tx_timeout(const int64_t tablet_count, int64_t &ddl_tx_timeout_us);
   static int get_ddl_tx_timeout(const int64_t tenant_id, const int64_t table_id, int64_t &ddl_tx_timeout_us);
+  static void get_ddl_rpc_timeout_for_database(const int64_t tenant_id, const int64_t database_id, int64_t &ddl_rpc_timeout_us);
   static int64_t get_default_ddl_rpc_timeout();
   static int64_t get_default_ddl_tx_timeout();