locate: clamp bucket boundaries to region boundaries (#1822)

 

Signed-off-by: ekexium <[email protected]>

Author: ekexium

internal/locate/region_cache.go

@@ -1124,7 +1124,7 @@ func (l *KeyLocation) LocateBucket(key []byte) *Bucket {
 	bucket := l.locateBucket(key)
 	// Return the bucket when locateBucket can locate the key
 	if bucket != nil {
-		return bucket
+		return l.clampBucketToRegion(bucket)
 	}
 	// Case one returns nil too.
 	if !l.Contains(key) {
@@ -1187,6 +1187,46 @@ func (b *Bucket) Contains(key []byte) bool {
 	return contains(b.StartKey, b.EndKey, key)
 }
 
+// clampBucketToRegion ensures bucket boundaries don't exceed region boundaries.
+// This is a defensive measure against stale bucket metadata that can have keys
+// outside the current region boundary.
+func (l *KeyLocation) clampBucketToRegion(bucket *Bucket) *Bucket {
+	startKey := bucket.StartKey
+	endKey := bucket.EndKey
+
+	// Clamp start: max(bucket.StartKey, region.StartKey)
+	if bytes.Compare(startKey, l.StartKey) < 0 {
+		startKey = l.StartKey
+	}
+
+	// Clamp end: min(bucket.EndKey, region.EndKey)
+	// Note: empty endKey means infinity, so clamp if bucket is unbounded but region is not
+	if len(l.EndKey) > 0 && (len(endKey) == 0 || bytes.Compare(endKey, l.EndKey) > 0) {
+		endKey = l.EndKey
+	}
+
+	// Ensure clamped bucket is valid (start < end)
+	if len(endKey) > 0 && bytes.Compare(startKey, endKey) >= 0 {
+		logutil.BgLogger().Warn("Clamped bucket invalid, using region boundaries",
+			zap.Uint64("regionID", l.Region.GetID()),
+			zap.String("bucketStart", redact.Key(bucket.StartKey)),
+			zap.String("bucketEnd", redact.Key(bucket.EndKey)),
+			zap.String("regionStart", redact.Key(l.StartKey)),
+			zap.String("regionEnd", redact.Key(l.EndKey)))
+		startKey = l.StartKey
+		endKey = l.EndKey
+	}
+
+	// If no clamping needed, return original bucket
+	if bytes.Equal(startKey, bucket.StartKey) && bytes.Equal(endKey, bucket.EndKey) {
+		return bucket
+	}
+
+	metrics.TiKVBucketClampedCounter.Inc()
+
+	return &Bucket{StartKey: startKey, EndKey: endKey}
+}
+
 // LocateKeyRange lists region and range that key in [start_key,end_key).
 // Regions without leader won't be returned.
 func (c *RegionCache) LocateKeyRange(bo *retry.Backoffer, startKey, endKey []byte) ([]*KeyLocation, error) {

internal/locate/region_cache_test.go

@@ -1918,6 +1918,184 @@ func (s *testRegionCacheSuite) TestLocateBucket() {
 	}
 }
 
+// TestBucketClampingToRegion tests that bucket boundaries exceeding region
+// boundaries are clamped correctly.
+func TestBucketClampingToRegion(t *testing.T) {
+	tests := []struct {
+		name            string
+		regionStart     []byte
+		regionEnd       []byte
+		bucketStart     []byte
+		bucketEnd       []byte
+		wantStart       []byte
+		wantEnd         []byte
+		shouldBeClamped bool
+	}{
+		{
+			name:            "bucket within region - no clamping",
+			regionStart:     []byte("a"),
+			regionEnd:       []byte("z"),
+			bucketStart:     []byte("f"),
+			bucketEnd:       []byte("m"),
+			wantStart:       []byte("f"),
+			wantEnd:         []byte("m"),
+			shouldBeClamped: false,
+		},
+		{
+			name:            "bucket start before region - clamp start",
+			regionStart:     []byte("f"),
+			regionEnd:       []byte("z"),
+			bucketStart:     []byte("a"), // Before region start
+			bucketEnd:       []byte("m"),
+			wantStart:       []byte("f"), // Clamped to region start
+			wantEnd:         []byte("m"),
+			shouldBeClamped: true,
+		},
+		{
+			name:            "bucket end after region - clamp end",
+			regionStart:     []byte("a"),
+			regionEnd:       []byte("m"),
+			bucketStart:     []byte("f"),
+			bucketEnd:       []byte("z"), // After region end
+			wantStart:       []byte("f"),
+			wantEnd:         []byte("m"), // Clamped to region end
+			shouldBeClamped: true,
+		},
+		{
+			name:            "bucket exceeds both boundaries - clamp both",
+			regionStart:     []byte("f"),
+			regionEnd:       []byte("m"),
+			bucketStart:     []byte("a"), // Before region start
+			bucketEnd:       []byte("z"), // After region end
+			wantStart:       []byte("f"), // Clamped to region start
+			wantEnd:         []byte("m"), // Clamped to region end
+			shouldBeClamped: true,
+		},
+		{
+			name:            "empty region start - bucket start not clamped",
+			regionStart:     []byte{}, // Beginning of keyspace
+			regionEnd:       []byte("m"),
+			bucketStart:     []byte("a"),
+			bucketEnd:       []byte("z"),
+			wantStart:       []byte("a"),
+			wantEnd:         []byte("m"), // Only end clamped
+			shouldBeClamped: true,
+		},
+		{
+			name:            "empty region end - bucket end not clamped",
+			regionStart:     []byte("f"),
+			regionEnd:       []byte{}, // End of keyspace (infinity)
+			bucketStart:     []byte("a"),
+			bucketEnd:       []byte("z"),
+			wantStart:       []byte("f"), // Only start clamped
+			wantEnd:         []byte("z"),
+			shouldBeClamped: true,
+		},
+		{
+			name:            "bucket completely before region - fallback to region",
+			regionStart:     []byte("m"),
+			regionEnd:       []byte("z"),
+			bucketStart:     []byte("a"),
+			bucketEnd:       []byte("f"), // Ends before region starts
+			wantStart:       []byte("m"), // Falls back to region boundaries
+			wantEnd:         []byte("z"),
+			shouldBeClamped: true,
+		},
+		{
+			name:            "bucket with empty startKey - clamp to region start",
+			regionStart:     []byte("f"),
+			regionEnd:       []byte("z"),
+			bucketStart:     []byte{}, // Beginning of keyspace
+			bucketEnd:       []byte("m"),
+			wantStart:       []byte("f"), // Clamped to region start
+			wantEnd:         []byte("m"),
+			shouldBeClamped: true,
+		},
+		{
+			name:            "bucket with empty endKey - clamp to region end",
+			regionStart:     []byte("a"),
+			regionEnd:       []byte("m"),
+			bucketStart:     []byte("f"),
+			bucketEnd:       []byte{}, // End of keyspace (infinity)
+			wantStart:       []byte("f"),
+			wantEnd:         []byte("m"), // Clamped to region end
+			shouldBeClamped: true,
+		},
+		{
+			name:            "bucket with empty startKey and region with empty startKey - no start clamping",
+			regionStart:     []byte{}, // Beginning of keyspace
+			regionEnd:       []byte("z"),
+			bucketStart:     []byte{}, // Beginning of keyspace
+			bucketEnd:       []byte("m"),
+			wantStart:       []byte{}, // No clamping needed
+			wantEnd:         []byte("m"),
+			shouldBeClamped: false,
+		},
+		{
+			name:            "bucket with empty endKey and region with empty endKey - no end clamping",
+			regionStart:     []byte("a"),
+			regionEnd:       []byte{}, // End of keyspace (infinity)
+			bucketStart:     []byte("f"),
+			bucketEnd:       []byte{}, // End of keyspace (infinity)
+			wantStart:       []byte("f"),
+			wantEnd:         []byte{}, // No clamping needed
+			shouldBeClamped: false,
+		},
+		{
+			name:            "bucket spanning full keyspace - clamp to region",
+			regionStart:     []byte("f"),
+			regionEnd:       []byte("m"),
+			bucketStart:     []byte{},    // Beginning of keyspace
+			bucketEnd:       []byte{},    // End of keyspace (infinity)
+			wantStart:       []byte("f"), // Clamped to region start
+			wantEnd:         []byte("m"), // Clamped to region end
+			shouldBeClamped: true,
+		},
+		{
+			name:            "infinity region - no clamping needed",
+			regionStart:     []byte{}, // Beginning of keyspace
+			regionEnd:       []byte{}, // End of keyspace (infinity)
+			bucketStart:     []byte("a"),
+			bucketEnd:       []byte("z"),
+			wantStart:       []byte("a"),
+			wantEnd:         []byte("z"),
+			shouldBeClamped: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			loc := &KeyLocation{
+				Region:   RegionVerID{id: 1},
+				StartKey: tt.regionStart,
+				EndKey:   tt.regionEnd,
+			}
+			bucket := &Bucket{
+				StartKey: tt.bucketStart,
+				EndKey:   tt.bucketEnd,
+			}
+
+			result := loc.clampBucketToRegion(bucket)
+
+			if !bytes.Equal(result.StartKey, tt.wantStart) {
+				t.Errorf("StartKey = %q, want %q", result.StartKey, tt.wantStart)
+			}
+			if !bytes.Equal(result.EndKey, tt.wantEnd) {
+				t.Errorf("EndKey = %q, want %q", result.EndKey, tt.wantEnd)
+			}
+
+			// Verify that clamped bucket is different from original if clamping was expected
+			isSameBucket := result == bucket
+			if tt.shouldBeClamped && isSameBucket {
+				t.Error("Expected bucket to be clamped (new object), but got same object")
+			}
+			if !tt.shouldBeClamped && !isSameBucket {
+				t.Error("Expected no clamping (same object), but got new object")
+			}
+		})
+	}
+}
+
 func (s *testRegionCacheSuite) TestRemoveIntersectingRegions() {
 	// Split at "b", "c", "d", "e"
 	regions := s.cluster.AllocIDs(4)

metrics/metrics.go

@@ -121,6 +121,7 @@ var (
 	TiKVValidateReadTSFromPDCount                  prometheus.Counter
 	TiKVLowResolutionTSOUpdateIntervalSecondsGauge prometheus.Gauge
 	TiKVStaleRegionFromPDCounter                   prometheus.Counter
+	TiKVBucketClampedCounter                       prometheus.Counter
 	TiKVPipelinedFlushThrottleSecondsHistogram     prometheus.Histogram
 	TiKVTxnWriteConflictCounter                    prometheus.Counter
 	TiKVAsyncSendReqCounter                        *prometheus.CounterVec
@@ -893,6 +894,14 @@ func initMetrics(namespace, subsystem string, constLabels prometheus.Labels) {
 			Help:        "Counter of stale region from PD",
 			ConstLabels: constLabels,
 		})
+	TiKVBucketClampedCounter = prometheus.NewCounter(
+		prometheus.CounterOpts{
+			Namespace:   namespace,
+			Subsystem:   subsystem,
+			Name:        "bucket_clamped",
+			Help:        "Counter of bucket boundaries clamped to region boundaries",
+			ConstLabels: constLabels,
+		})
 	TiKVPipelinedFlushThrottleSecondsHistogram = prometheus.NewHistogram(
 		prometheus.HistogramOpts{
 			Namespace:   namespace,
@@ -1035,6 +1044,7 @@ func RegisterMetrics() {
 	prometheus.MustRegister(TiKVValidateReadTSFromPDCount)
 	prometheus.MustRegister(TiKVLowResolutionTSOUpdateIntervalSecondsGauge)
 	prometheus.MustRegister(TiKVStaleRegionFromPDCounter)
+	prometheus.MustRegister(TiKVBucketClampedCounter)
 	prometheus.MustRegister(TiKVPipelinedFlushThrottleSecondsHistogram)
 	prometheus.MustRegister(TiKVTxnWriteConflictCounter)
 	prometheus.MustRegister(TiKVAsyncSendReqCounter)