Merge pull request #152460 from iskettaneh/backport25.2-144303

release-25.2: storepool: lock individual storepool details instead of all details

Author: iskettaneh

pkg/kv/kvserver/allocator/allocatorimpl/allocator_test.go

@@ -480,60 +480,68 @@ func mockStorePool(
 	decommissionedStoreIDs []roachpb.StoreID,
 	suspectedStoreIDs []roachpb.StoreID,
 ) {
-	storePool.DetailsMu.Lock()
-	defer storePool.DetailsMu.Unlock()
-
 	liveNodeSet := map[roachpb.NodeID]livenesspb.NodeLivenessStatus{}
-	storePool.DetailsMu.StoreDetails = map[roachpb.StoreID]*storepool.StoreDetail{}
 	for _, storeID := range aliveStoreIDs {
 		liveNodeSet[roachpb.NodeID(storeID)] = livenesspb.NodeLivenessStatus_LIVE
-		detail := storePool.GetStoreDetailLocked(storeID)
+		detail := storePool.GetStoreDetail(storeID)
+		detail.Lock()
 		detail.Desc = &roachpb.StoreDescriptor{
 			StoreID: storeID,
 			Node:    roachpb.NodeDescriptor{NodeID: roachpb.NodeID(storeID)},
 		}
+		detail.Unlock()
 	}
 	for _, storeID := range unavailableStoreIDs {
 		liveNodeSet[roachpb.NodeID(storeID)] = livenesspb.NodeLivenessStatus_UNAVAILABLE
-		detail := storePool.GetStoreDetailLocked(storeID)
+		detail := storePool.GetStoreDetail(storeID)
+		detail.Lock()
 		detail.Desc = &roachpb.StoreDescriptor{
 			StoreID: storeID,
 			Node:    roachpb.NodeDescriptor{NodeID: roachpb.NodeID(storeID)},
 		}
+		detail.Unlock()
 	}
 	for _, storeID := range deadStoreIDs {
 		liveNodeSet[roachpb.NodeID(storeID)] = livenesspb.NodeLivenessStatus_DEAD
-		detail := storePool.GetStoreDetailLocked(storeID)
+		detail := storePool.GetStoreDetail(storeID)
+		detail.Lock()
 		detail.Desc = &roachpb.StoreDescriptor{
 			StoreID: storeID,
 			Node:    roachpb.NodeDescriptor{NodeID: roachpb.NodeID(storeID)},
 		}
+		detail.Unlock()
 	}
 	for _, storeID := range decommissioningStoreIDs {
 		liveNodeSet[roachpb.NodeID(storeID)] = livenesspb.NodeLivenessStatus_DECOMMISSIONING
-		detail := storePool.GetStoreDetailLocked(storeID)
+		detail := storePool.GetStoreDetail(storeID)
+		detail.Lock()
 		detail.Desc = &roachpb.StoreDescriptor{
 			StoreID: storeID,
 			Node:    roachpb.NodeDescriptor{NodeID: roachpb.NodeID(storeID)},
 		}
+		detail.Unlock()
 	}
 	for _, storeID := range decommissionedStoreIDs {
 		liveNodeSet[roachpb.NodeID(storeID)] = livenesspb.NodeLivenessStatus_DECOMMISSIONED
-		detail := storePool.GetStoreDetailLocked(storeID)
+		detail := storePool.GetStoreDetail(storeID)
+		detail.Lock()
 		detail.Desc = &roachpb.StoreDescriptor{
 			StoreID: storeID,
 			Node:    roachpb.NodeDescriptor{NodeID: roachpb.NodeID(storeID)},
 		}
+		detail.Unlock()
 	}
 
 	for _, storeID := range suspectedStoreIDs {
 		liveNodeSet[roachpb.NodeID(storeID)] = livenesspb.NodeLivenessStatus_LIVE
-		detail := storePool.GetStoreDetailLocked(storeID)
+		detail := storePool.GetStoreDetail(storeID)
+		detail.Lock()
 		detail.LastUnavailable = storePool.Clock().Now()
 		detail.Desc = &roachpb.StoreDescriptor{
 			StoreID: storeID,
 			Node:    roachpb.NodeDescriptor{NodeID: roachpb.NodeID(storeID)},
 		}
+		detail.Unlock()
 	}
 
 	// Set the node liveness function using the set we constructed.
@@ -1380,16 +1388,20 @@ func TestAllocatorRebalanceDeadNodes(t *testing.T) {
 	}
 
 	// Initialize 8 stores: where store 6 is the target for rebalancing.
-	sp.DetailsMu.Lock()
-	sp.GetStoreDetailLocked(1).Desc.Capacity = ranges(100)
-	sp.GetStoreDetailLocked(2).Desc.Capacity = ranges(100)
-	sp.GetStoreDetailLocked(3).Desc.Capacity = ranges(100)
-	sp.GetStoreDetailLocked(4).Desc.Capacity = ranges(100)
-	sp.GetStoreDetailLocked(5).Desc.Capacity = ranges(100)
-	sp.GetStoreDetailLocked(6).Desc.Capacity = ranges(0)
-	sp.GetStoreDetailLocked(7).Desc.Capacity = ranges(100)
-	sp.GetStoreDetailLocked(8).Desc.Capacity = ranges(100)
-	sp.DetailsMu.Unlock()
+	updateDescCapacity := func(storeID roachpb.StoreID, capacity roachpb.StoreCapacity) {
+		sd := sp.GetStoreDetail(storeID)
+		sd.Lock()
+		sd.Desc.Capacity = capacity
+		sd.Unlock()
+	}
+	updateDescCapacity(1, ranges(100))
+	updateDescCapacity(2, ranges(100))
+	updateDescCapacity(3, ranges(100))
+	updateDescCapacity(4, ranges(100))
+	updateDescCapacity(5, ranges(100))
+	updateDescCapacity(6, ranges(0))
+	updateDescCapacity(7, ranges(100))
+	updateDescCapacity(8, ranges(100))
 
 	// Each test case should describe a repair situation which has a lower
 	// priority than the previous test case.

pkg/kv/kvserver/allocator/storepool/override_store_pool.go

@@ -114,38 +114,31 @@ func (o *OverrideStorePool) DecommissioningReplicas(
 func (o *OverrideStorePool) GetStoreList(
 	filter StoreFilter,
 ) (StoreList, int, ThrottledStoreReasons) {
-	o.sp.DetailsMu.Lock()
-	defer o.sp.DetailsMu.Unlock()
-
 	var storeIDs roachpb.StoreIDSlice
-	for storeID := range o.sp.DetailsMu.StoreDetails {
+	o.sp.Details.StoreDetails.Range(func(storeID roachpb.StoreID, _ *StoreDetailMu) bool {
 		storeIDs = append(storeIDs, storeID)
-	}
-	return o.sp.getStoreListFromIDsLocked(storeIDs, o.overrideNodeLivenessFn, filter)
+		return true
+	})
+	return o.sp.getStoreListFromIDs(storeIDs, o.overrideNodeLivenessFn, filter)
 }
 
 // GetStoreListFromIDs implements the AllocatorStorePool interface.
 func (o *OverrideStorePool) GetStoreListFromIDs(
 	storeIDs roachpb.StoreIDSlice, filter StoreFilter,
 ) (StoreList, int, ThrottledStoreReasons) {
-	o.sp.DetailsMu.Lock()
-	defer o.sp.DetailsMu.Unlock()
-	return o.sp.getStoreListFromIDsLocked(storeIDs, o.overrideNodeLivenessFn, filter)
+	return o.sp.getStoreListFromIDs(storeIDs, o.overrideNodeLivenessFn, filter)
 }
 
 // GetStoreListForTargets implements the AllocatorStorePool interface.
 func (o *OverrideStorePool) GetStoreListForTargets(
 	candidates []roachpb.ReplicationTarget, filter StoreFilter,
 ) (StoreList, int, ThrottledStoreReasons) {
-	o.sp.DetailsMu.Lock()
-	defer o.sp.DetailsMu.Unlock()
-
 	storeIDs := make(roachpb.StoreIDSlice, 0, len(candidates))
 	for _, tgt := range candidates {
 		storeIDs = append(storeIDs, tgt.StoreID)
 	}
 
-	return o.sp.getStoreListFromIDsLocked(storeIDs, o.overrideNodeLivenessFn, filter)
+	return o.sp.getStoreListFromIDs(storeIDs, o.overrideNodeLivenessFn, filter)
 }
 
 // LiveAndDeadReplicas implements the AllocatorStorePool interface.

pkg/kv/kvserver/allocator/storepool/override_store_pool_test.go

@@ -330,9 +330,11 @@ func TestOverrideStorePoolGetStoreList(t *testing.T) {
 	livenessOverrides[decommissioningStore.Node.NodeID] = livenesspb.NodeLivenessStatus_DECOMMISSIONING
 
 	// Set suspectedStore as suspected.
-	testStorePool.DetailsMu.Lock()
-	testStorePool.DetailsMu.StoreDetails[suspectedStore.StoreID].LastUnavailable = testStorePool.clock.Now()
-	testStorePool.DetailsMu.Unlock()
+	val, ok := testStorePool.Details.StoreDetails.Load(suspectedStore.StoreID)
+	require.True(t, ok)
+	val.Lock()
+	val.LastUnavailable = testStorePool.clock.Now()
+	val.Unlock()
 
 	// No filter or limited set of store IDs.
 	require.NoError(t, verifyStoreList(