Merge #158003

158003: sql/catalog/lease: skip dropped regions when counting leases r=spilchen a=spilchen

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

Co-authored-by: Matt Spilchen <[email protected]>

Author: craig[bot]

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