sql/catalog/lease: skip dropped regions when counting leases

countLeasesByRegion now skips regions that were just dropped
instead of reporting “failed to count leases”. This was found in
the TestMultiRegionTenantRegions test. We did have pre-existing
logic to handle missing regions in handleRegionLivenessErrors, but
it wasn't handled correctly by all the callers.

Resolves: #157122
Release note: none

Author: spilchen

pkg/sql/catalog/lease/BUILD.bazel

@@ -77,6 +77,7 @@ go_test(
     name = "lease_test",
     size = "large",
     srcs = [
+        "count_test.go",
         "helpers_test.go",
         "ie_writer_test.go",
         "kv_writer_test.go",

pkg/sql/catalog/lease/count.go

@@ -208,9 +208,13 @@ func countLeasesByRegion(
 		} else {
 			err = queryRegionRows(ctx)
 		}
-		if err := handleRegionLivenessErrors(ctx, prober, region, err); err != nil {
+		skipRegion, err := handleRegionLivenessErrors(ctx, prober, region, err)
+		if err != nil {
 			return err
 		}
+		if skipRegion {
+			return nil
+		}
 		if values == nil {
 			return errors.New("failed to count leases")
 		}
@@ -231,27 +235,27 @@ func getCountLeaseColumns() string {
 }
 
 // handleRegionLivenessErrors handles errors that are linked to region liveness
-// timeouts.
+// timeouts. Return true if the region should be skipped.
 func handleRegionLivenessErrors(
 	ctx context.Context, prober regionliveness.Prober, region string, err error,
-) error {
+) (bool, error) {
 	if err != nil {
 		if regionliveness.IsQueryTimeoutErr(err) {
 			// Probe and mark the region potentially.
 			probeErr := prober.ProbeLiveness(ctx, region)
 			if probeErr != nil {
 				err = errors.WithSecondaryError(err, probeErr)
-				return err
+				return false, err
 			}
-			return errors.Wrapf(err, "count-lease timed out reading from a region")
+			return false, errors.Wrapf(err, "count-lease timed out reading from a region")
 		} else if regionliveness.IsMissingRegionEnumErr(err) {
 			// Skip this region because we were unable to find region in
 			// type descriptor. Since the database regions are cached, they
 			// may be stale and have dropped regions.
 			log.Dev.Infof(ctx, "count-lease skipping region %s due to error: %v", region, err)
-			return nil
+			return true, nil
 		}
-		return err
+		return false, err
 	}
-	return err
+	return false, err
 }

pkg/sql/catalog/lease/count_test.go

@@ -0,0 +1,26 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package lease
+
+import (
+	"context"
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/sql/types"
+	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/stretchr/testify/require"
+)
+
+func TestHandleRegionLivenessErrorsSkipsMissingRegion(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	skip, err := handleRegionLivenessErrors(ctx, nil /* prober */, "us-east2", types.EnumValueNotYetPublicError)
+	require.NoError(t, err)
+	require.True(t, skip)
+}

pkg/sql/catalog/lease/lease.go

@@ -425,7 +425,11 @@ func (m *Manager) WaitForInitialVersion(
 						sessionIDs, err = getSessionsHoldingDescriptor(ctx, txn, schemaID, region)
 					}
 					if err != nil {
-						return handleRegionLivenessErrors(ctx, prober, region, err)
+						skipRegion, handledErr := handleRegionLivenessErrors(ctx, prober, region, err)
+						if skipRegion {
+							return nil
+						}
+						return handledErr
 					}
 					sessionsPerRegion[region] = sessionIDs
 					expectedSessions += len(sessionIDs)
@@ -469,9 +473,13 @@ func (m *Manager) WaitForInitialVersion(
 						} else {
 							regionCount, err = countDescriptorsHeldBySessionIDs(ctx, txn, descIDsForSchema, region, sessionIDs)
 						}
-						if err := handleRegionLivenessErrors(ctx, prober, region, err); err != nil {
+						skipRegion, err := handleRegionLivenessErrors(ctx, prober, region, err)
+						if err != nil {
 							return err
 						}
+						if skipRegion {
+							return nil
+						}
 						count += regionCount
 						return nil
 					})
@@ -620,7 +628,13 @@ func (m *Manager) WaitForNewVersion(
 					regionStaleSessionCount, err = countSessionsHoldingStaleDescriptor(ctx, txn, desc, region)
 				}
 				if err != nil {
-					return handleRegionLivenessErrors(ctx, prober, region, err)
+					skipRegion, handledErr := handleRegionLivenessErrors(ctx, prober, region, err)
+					if handledErr != nil {
+						return handledErr
+					}
+					if skipRegion {
+						continue
+					}
 				}
 
 				staleSessionCount += regionStaleSessionCount