Merge #104512

104512: upgrade: improve delete dropped udf upgrade job efficiency r=chengxiong-ruan a=chengxiong-ruan

Informs: https://github.com/cockroachlabs/support/issues/2364

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.

Co-authored-by: Chengxiong Ruan <[email protected]>

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
 }