From b20d5ac3b86753180d9f66cd3ce0e30b7fd45513 Mon Sep 17 00:00:00 2001
From: Chengxiong Ruan <chengxiongruan@gmail.com>
Date: Wed, 7 Jun 2023 10:37:05 -0400
Subject: [PATCH] upgrade: improve delete dropped udf upgrade job efficiency

Release note (performance improvement): this commit makes the
`delete descriptors of dropped functions` upgrade job more efficient.
It used to look at every single id until the max descriptor id which
was very inefficient when the max descriptor id really large, in which
case the upgrade job took very long even there was no function
descriptor at all. This commit changed it to actually query upper bound
id of each batch.
---
 ...delete_descriptors_of_dropped_functions.go | 58 +++++++++++++++----
 1 file changed, 46 insertions(+), 12 deletions(-)

diff --git a/pkg/upgrade/upgrades/delete_descriptors_of_dropped_functions.go b/pkg/upgrade/upgrades/delete_descriptors_of_dropped_functions.go
index b53de459ffab..1907e7e80ea4 100644
--- a/pkg/upgrade/upgrades/delete_descriptors_of_dropped_functions.go
+++ b/pkg/upgrade/upgrades/delete_descriptors_of_dropped_functions.go
@@ -15,10 +15,22 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/clusterversion"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descs"
+	"github.com/cockroachdb/cockroach/pkg/sql/isql"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/upgrade"
 )
 
+const getMaxIDInBatch = `
+WITH id_batch AS (
+    SELECT id
+    FROM system.descriptor
+    WHERE id > $1
+    ORDER BY id
+    LIMIT $2
+)
+SELECT max(id) from id_batch;
+`
+
 const deleteDroppedFunctionQuery = `
 WITH to_json AS (
     SELECT 
@@ -30,6 +42,8 @@ WITH to_json AS (
       ) AS d
     FROM 
       system.descriptor
+    WHERE id > $1
+    AND id <= $2
 ),
 to_delete AS (
     SELECT id
@@ -40,15 +54,14 @@ to_delete AS (
       AND d->'function'->>'state' = 'DROP'
 )
 SELECT crdb_internal.unsafe_delete_descriptor(id, false)
-FROM to_delete
-WHERE id >= $1
-AND id < $2;
+FROM to_delete;
 `
 
 func deleteDescriptorsOfDroppedFunctions(
 	ctx context.Context, _ clusterversion.ClusterVersion, d upgrade.TenantDeps,
 ) error {
-	if err := d.DB.DescsTxn(ctx, func(ctx context.Context, txn descs.Txn) error {
+	var maxID int64
+	if err := d.DB.Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
 		row, err := txn.QueryRow(
 			ctx,
 			"upgrade-delete-dropped-function-descriptors-get-max-descriptor-id",
@@ -60,25 +73,46 @@ func deleteDescriptorsOfDroppedFunctions(
 			return err
 		}
 
-		maxID := int(tree.MustBeDInt(row[0]))
-		const batchSize = 50
+		maxID = int64(tree.MustBeDInt(row[0]))
+		return nil
+	}); err != nil {
+		return err
+	}
+
+	const batchSize = 50
+	for curID := int64(0); curID < maxID; {
+		if err := d.DB.DescsTxn(ctx, func(ctx context.Context, txn descs.Txn) error {
+
+			row, err := txn.QueryRow(
+				ctx,
+				"upgrade-delete-dropped-function-descriptors-get-batch-max-id",
+				txn.KV(),
+				getMaxIDInBatch,
+				curID,
+				batchSize,
+			)
+			if err != nil {
+				return err
+			}
+			batchMaxID := int64(tree.MustBeDInt(row[0]))
 
-		for curID := 1; curID <= maxID; curID += batchSize {
-			_, err := txn.Exec(
+			_, err = txn.Exec(
 				ctx,
 				"upgrade-delete-dropped-function-descriptors", /* opName */
 				txn.KV(),
 				deleteDroppedFunctionQuery,
 				curID,
-				curID+batchSize,
+				batchMaxID,
 			)
 			if err != nil {
 				return err
 			}
+			curID = batchMaxID
+
+			return nil
+		}); err != nil {
+			return err
 		}
-		return nil
-	}); err != nil {
-		return err
 	}
 	return nil
 }