Skip to content

mmaintegration: add mmaintegration datadriven tests #151510

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wenyihu6 wants to merge 17 commits into
base: master
Choose a base branch
from
Open

mmaintegration: add mmaintegration datadriven tests #151510

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
2019db5
mmaintegration: add allocator_op.go and allocator_sync.go
wenyihu6 Aug 2, 2025
c0c77a3
kvserver: plumb allocator sync
wenyihu6 Aug 5, 2025
9d98ffd
kvserver: use allocator sync to update store pool
wenyihu6 Aug 4, 2025
d2ad3ec
kvserver: register external changes with mma
wenyihu6 Aug 3, 2025
31184a6
mmaintegration: add MMAPreApply to allocator sync
wenyihu6 Aug 5, 2025
6d059c6
mmaintegration: integrate allocator sync to mma store rebalancer
wenyihu6 Aug 3, 2025
61c60e6
mmaintegration: add store load msg
wenyihu6 Aug 3, 2025
58077dc
mmaintegration: add InvalidSyncChangeID
wenyihu6 Aug 3, 2025
25b67dc
mmaintegration: remove pkg/kv/kvserver/allocator/mmaprototypehelpers
wenyihu6 Aug 3, 2025
766405a
mmaintegration: address minor code review comments
wenyihu6 Aug 6, 2025
1b86275
kvserver: update allocator sync on errors in AdminTransferLease
wenyihu6 Aug 6, 2025
aa5fe37
mmaintegration: add interface for store pool and mma
wenyihu6 Aug 6, 2025
c070516
mmaprototype: add MakePendingRangeChangeForTesting
wenyihu6 Aug 7, 2025
439ca95
mmaintegration: add mmaintegration datadriven tests
wenyihu6 Aug 7, 2025
02d255f
mmaintegration: add store load datadriven tests
wenyihu6 Aug 7, 2025
db61f25
fixup! mmaintegration: add mmaintegration datadriven tests
wenyihu6 Aug 7, 2025
39545ef
fixup! mmaprototype: add MakePendingRangeChangeForTesting
wenyihu6 Aug 7, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion pkg/BUILD.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -269,6 +269,7 @@ ALL_TESTS = [
"//pkg/kv/kvserver/logstore:logstore_test",
"//pkg/kv/kvserver/loqrecovery/loqrecoverypb:loqrecoverypb_test",
"//pkg/kv/kvserver/loqrecovery:loqrecovery_test",
"//pkg/kv/kvserver/mmaintegration:mmaintegration_test",
"//pkg/kv/kvserver/multiqueue:multiqueue_test",
"//pkg/kv/kvserver/print:print_test",
"//pkg/kv/kvserver/protectedts/ptcache:ptcache_test",
Expand Down Expand Up @@ -1462,7 +1463,6 @@ GO_TARGETS = [
"//pkg/kv/kvserver/allocator/load:load_test",
"//pkg/kv/kvserver/allocator/mmaprototype:mmaprototype",
"//pkg/kv/kvserver/allocator/mmaprototype:mmaprototype_test",
"//pkg/kv/kvserver/allocator/mmaprototypehelpers:mmaprototypehelpers",
"//pkg/kv/kvserver/allocator/plan:plan",
"//pkg/kv/kvserver/allocator/plan:plan_test",
"//pkg/kv/kvserver/allocator/storepool:storepool",
Expand Down Expand Up @@ -1557,6 +1557,8 @@ GO_TARGETS = [
"//pkg/kv/kvserver/loqrecovery/loqrecoverypb:loqrecoverypb_test",
"//pkg/kv/kvserver/loqrecovery:loqrecovery",
"//pkg/kv/kvserver/loqrecovery:loqrecovery_test",
"//pkg/kv/kvserver/mmaintegration:mmaintegration",
"//pkg/kv/kvserver/mmaintegration:mmaintegration_test",
"//pkg/kv/kvserver/multiqueue:multiqueue",
"//pkg/kv/kvserver/multiqueue:multiqueue_test",
"//pkg/kv/kvserver/print:print",
Expand Down
1 change: 1 addition & 0 deletions pkg/kv/kvserver/BUILD.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -168,6 +168,7 @@ go_library(
"//pkg/kv/kvserver/load",
"//pkg/kv/kvserver/lockspanset",
"//pkg/kv/kvserver/logstore",
"//pkg/kv/kvserver/mmaintegration",
"//pkg/kv/kvserver/multiqueue",
"//pkg/kv/kvserver/print",
"//pkg/kv/kvserver/raftentry",
Expand Down
38 changes: 31 additions & 7 deletions pkg/kv/kvserver/allocator/mmaprototype/cluster_state.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,9 +57,9 @@ func (rt ReplicaIDAndType) SafeFormat(w redact.SafePrinter, _ rune) {
default:
w.Print(rt.ReplicaID)
}
w.Printf(" type=%v", rt.ReplicaType.ReplicaType)
w.Printf(",type=%v", rt.ReplicaType.ReplicaType)
if rt.IsLeaseholder {
w.Print(" leaseholder=true")
w.Print(",leaseholder=true")
}
}

Expand Down Expand Up @@ -116,6 +116,14 @@ type ReplicaState struct {
// beneficial.
}

func (rs ReplicaState) String() string {
return redact.StringWithoutMarkers(rs)
}

func (rs ReplicaState) SafeFormat(w redact.SafePrinter, _ rune) {
w.Printf("%s,lagging=%v", rs.ReplicaIDAndType, rs.VoterIsLagging)
}

// ChangeID is a unique ID, in the context of this data-structure and when
// receiving updates about enactment having happened or having been rejected
// (by the component responsible for change enactment).
Expand All @@ -133,13 +141,13 @@ const (
func (s ReplicaChangeType) String() string {
switch s {
case AddLease:
return "AddLease"
return "add-lease"
case RemoveLease:
return "RemoveLease"
return "remove-lease"
case AddReplica:
return "AddReplica"
return "add-replica"
case RemoveReplica:
return "RemoveReplica"
return "remove-replica"
default:
panic("unknown ReplicaChangeType")
}
Expand Down Expand Up @@ -183,7 +191,7 @@ func (rc ReplicaChange) String() string {

// SafeFormat implements the redact.SafeFormatter interface.
func (rc ReplicaChange) SafeFormat(w redact.SafePrinter, _ rune) {
w.Printf("r%v type: %v target store %v (%v)->(%v)", rc.rangeID, rc.replicaChangeType, rc.target, rc.prev, rc.next)
w.Printf("r%v:%v,target=(%v), prev=(%v)->next=(%v)", rc.rangeID, rc.replicaChangeType, rc.target, rc.prev, rc.next)
}

// isRemoval returns true if the change is a removal of a replica.
Expand Down Expand Up @@ -375,6 +383,22 @@ type PendingRangeChange struct {
pendingReplicaChanges []*pendingReplicaChange
}

func MakePendingRangeChangeForTesting(
rangeID roachpb.RangeID, changes []ReplicaChange, changeIDs []ChangeID,
) PendingRangeChange {
prc := PendingRangeChange{
RangeID: rangeID,
pendingReplicaChanges: make([]*pendingReplicaChange, len(changes)),
}
for i, c := range changes {
prc.pendingReplicaChanges[i] = &pendingReplicaChange{
ChangeID: changeIDs[i],
ReplicaChange: c,
}
}
return prc
}

func (prc PendingRangeChange) String() string {
return redact.StringWithoutMarkers(prc)
}
Expand Down
9 changes: 9 additions & 0 deletions pkg/kv/kvserver/allocator/mmaprototype/load.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,15 @@ type RangeLoad struct {
RaftCPU LoadValue
}

func (rl RangeLoad) String() string {
return redact.StringWithoutMarkers(rl)
}

func (rl RangeLoad) SafeFormat(w redact.SafePrinter, _ rune) {
w.Printf("(cpu:%d,raft-cpu:%d,write-band:%d,byte-size:%d)",
rl.Load[CPURate], rl.RaftCPU, rl.Load[WriteBandwidth], rl.Load[ByteSize])
}

// storeLoad is the load information for a store. Roughly, this is the
// information we need each store to provide us periodically, i.e.,
// StoreLoadMsg is the input used to compute this.
Expand Down
24 changes: 12 additions & 12 deletions pkg/kv/kvserver/allocator/mmaprototype/testdata/cluster_state/rebalance_replica
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ store-id=1
ranges
----
range-id=1 load=[cpu:80, write-bandwidth:80, byte-size:80] raft-cpu=20
store-id=1 replica-id=1 type=VOTER_FULL leaseholder=true
store-id=1 replica-id=1,type=VOTER_FULL,leaseholder=true


get-load-info
Expand All @@ -39,16 +39,16 @@ make-pending-changes range-id=1
----
pending(2)
change-id=1 store-id=2 node-id=2 range-id=1 load-delta=[cpu:88, write-bandwidth:88, byte-size:88] start=0s
prev=(replica-id=none type=VOTER_FULL)
next=(replica-id=unknown type=VOTER_FULL leaseholder=true)
prev=(replica-id=none,type=VOTER_FULL,lagging=false)
next=(replica-id=unknown,type=VOTER_FULL,leaseholder=true)
change-id=2 store-id=1 node-id=1 range-id=1 load-delta=[cpu:-80, write-bandwidth:-80, byte-size:-80] start=0s
prev=(replica-id=1 type=VOTER_FULL leaseholder=true)
next=(replica-id=none type=VOTER_FULL)
prev=(replica-id=1,type=VOTER_FULL,leaseholder=true,lagging=false)
next=(replica-id=none,type=VOTER_FULL)

ranges
----
range-id=1 load=[cpu:80, write-bandwidth:80, byte-size:80] raft-cpu=20
store-id=2 replica-id=unknown type=VOTER_FULL leaseholder=true
store-id=2 replica-id=unknown,type=VOTER_FULL,leaseholder=true

get-load-info
----
Expand All @@ -75,11 +75,11 @@ get-pending-changes
----
pending(2)
change-id=1 store-id=2 node-id=2 range-id=1 load-delta=[cpu:88, write-bandwidth:88, byte-size:88] start=0s enacted=0s
prev=(replica-id=none type=VOTER_FULL)
next=(replica-id=unknown type=VOTER_FULL leaseholder=true)
prev=(replica-id=none,type=VOTER_FULL,lagging=false)
next=(replica-id=unknown,type=VOTER_FULL,leaseholder=true)
change-id=2 store-id=1 node-id=1 range-id=1 load-delta=[cpu:-80, write-bandwidth:-80, byte-size:-80] start=0s enacted=0s
prev=(replica-id=1 type=VOTER_FULL leaseholder=true)
next=(replica-id=none type=VOTER_FULL)
prev=(replica-id=1,type=VOTER_FULL,leaseholder=true,lagging=false)
next=(replica-id=none,type=VOTER_FULL)

store-load-msg
store-id=2 node-id=2 load=[80,80,80] capacity=[100,100,100] secondary-load=1 load-time=15s
Expand All @@ -89,8 +89,8 @@ get-pending-changes
----
pending(1)
change-id=2 store-id=1 node-id=1 range-id=1 load-delta=[cpu:-80, write-bandwidth:-80, byte-size:-80] start=0s enacted=0s
prev=(replica-id=1 type=VOTER_FULL leaseholder=true)
next=(replica-id=none type=VOTER_FULL)
prev=(replica-id=1,type=VOTER_FULL,leaseholder=true,lagging=false)
next=(replica-id=none,type=VOTER_FULL)

get-load-info
----
Expand Down
14 changes: 7 additions & 7 deletions pkg/kv/kvserver/allocator/mmaprototype/testdata/cluster_state/remove_replica
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,23 +33,23 @@ store-id=2 node-id=2 reported=[cpu:20, write-bandwidth:80, byte-size:80] adjuste
ranges
----
range-id=1 load=[cpu:80, write-bandwidth:80, byte-size:80] raft-cpu=20
store-id=1 replica-id=1 type=VOTER_FULL leaseholder=true
store-id=2 replica-id=2 type=VOTER_FULL
store-id=1 replica-id=1,type=VOTER_FULL,leaseholder=true
store-id=2 replica-id=2,type=VOTER_FULL

make-pending-changes range-id=1
remove-replica: remove-store-id=2
----
pending(1)
change-id=1 store-id=2 node-id=2 range-id=1 load-delta=[cpu:-20, write-bandwidth:-80, byte-size:-80] start=0s
prev=(replica-id=2 type=VOTER_FULL)
next=(replica-id=none type=VOTER_FULL)
prev=(replica-id=2,type=VOTER_FULL,lagging=false)
next=(replica-id=none,type=VOTER_FULL)

# We should see the change applied to the range state, with only the replica on
# store 1 remaining.
ranges
----
range-id=1 load=[cpu:80, write-bandwidth:80, byte-size:80] raft-cpu=20
store-id=1 replica-id=1 type=VOTER_FULL leaseholder=true
store-id=1 replica-id=1,type=VOTER_FULL,leaseholder=true

# The load info for s2 should also reflect the load delta [-20,-80,-80] being
# applied.
Expand All @@ -76,8 +76,8 @@ get-pending-changes
----
pending(1)
change-id=1 store-id=2 node-id=2 range-id=1 load-delta=[cpu:-20, write-bandwidth:-80, byte-size:-80] start=0s enacted=0s
prev=(replica-id=2 type=VOTER_FULL)
next=(replica-id=none type=VOTER_FULL)
prev=(replica-id=2,type=VOTER_FULL,lagging=false)
next=(replica-id=none,type=VOTER_FULL)

store-load-msg
store-id=2 node-id=2 load=[0,0,0] capacity=[100,100,100] secondary-load=0 load-time=15s
Expand Down
18 changes: 9 additions & 9 deletions pkg/kv/kvserver/allocator/mmaprototype/testdata/cluster_state/removed_ranges
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,18 +19,18 @@ store-id=1
ranges
----
range-id=1 load=[cpu:80, write-bandwidth:80, byte-size:80] raft-cpu=20
store-id=1 replica-id=1 type=VOTER_FULL leaseholder=true
store-id=1 replica-id=1,type=VOTER_FULL,leaseholder=true

make-pending-changes range-id=1
rebalance-replica: remove-store-id=1 add-store-id=2
----
pending(2)
change-id=1 store-id=2 node-id=2 range-id=1 load-delta=[cpu:88, write-bandwidth:88, byte-size:88] start=0s
prev=(replica-id=none type=VOTER_FULL)
next=(replica-id=unknown type=VOTER_FULL leaseholder=true)
prev=(replica-id=none,type=VOTER_FULL,lagging=false)
next=(replica-id=unknown,type=VOTER_FULL,leaseholder=true)
change-id=2 store-id=1 node-id=1 range-id=1 load-delta=[cpu:-80, write-bandwidth:-80, byte-size:-80] start=0s
prev=(replica-id=1 type=VOTER_FULL leaseholder=true)
next=(replica-id=none type=VOTER_FULL)
prev=(replica-id=1,type=VOTER_FULL,leaseholder=true,lagging=false)
next=(replica-id=none,type=VOTER_FULL)

# Advance time to 1s so that we can see pending changes for removed ranges are
# properly enacted.
Expand All @@ -47,11 +47,11 @@ get-pending-changes
----
pending(2)
change-id=1 store-id=2 node-id=2 range-id=1 load-delta=[cpu:88, write-bandwidth:88, byte-size:88] start=0s enacted=1s
prev=(replica-id=none type=VOTER_FULL)
next=(replica-id=unknown type=VOTER_FULL leaseholder=true)
prev=(replica-id=none,type=VOTER_FULL,lagging=false)
next=(replica-id=unknown,type=VOTER_FULL,leaseholder=true)
change-id=2 store-id=1 node-id=1 range-id=1 load-delta=[cpu:-80, write-bandwidth:-80, byte-size:-80] start=0s enacted=1s
prev=(replica-id=1 type=VOTER_FULL leaseholder=true)
next=(replica-id=none type=VOTER_FULL)
prev=(replica-id=1,type=VOTER_FULL,leaseholder=true,lagging=false)
next=(replica-id=none,type=VOTER_FULL)

ranges
----
20 changes: 10 additions & 10 deletions pkg/kv/kvserver/allocator/mmaprototype/testdata/cluster_state/transfer_lease
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,19 +37,19 @@ store-id=2 node-id=2 reported=[cpu:20, write-bandwidth:80, byte-size:80] adjuste
ranges
----
range-id=1 load=[cpu:80, write-bandwidth:80, byte-size:80] raft-cpu=20
store-id=1 replica-id=1 type=VOTER_FULL leaseholder=true
store-id=2 replica-id=2 type=VOTER_FULL
store-id=1 replica-id=1,type=VOTER_FULL,leaseholder=true
store-id=2 replica-id=2,type=VOTER_FULL

make-pending-changes range-id=1
transfer-lease: remove-store-id=1 add-store-id=2
----
pending(2)
change-id=1 store-id=1 node-id=1 range-id=1 load-delta=[cpu:-60, write-bandwidth:0, byte-size:0] start=0s
prev=(replica-id=1 type=VOTER_FULL leaseholder=true)
next=(replica-id=1 type=VOTER_FULL)
prev=(replica-id=1,type=VOTER_FULL,leaseholder=true,lagging=false)
next=(replica-id=1,type=VOTER_FULL)
change-id=2 store-id=2 node-id=2 range-id=1 load-delta=[cpu:66, write-bandwidth:0, byte-size:0] start=0s
prev=(replica-id=2 type=VOTER_FULL)
next=(replica-id=2 type=VOTER_FULL leaseholder=true)
prev=(replica-id=2,type=VOTER_FULL,lagging=false)
next=(replica-id=2,type=VOTER_FULL,leaseholder=true)

get-load-info
----
Expand All @@ -61,8 +61,8 @@ store-id=2 node-id=2 reported=[cpu:20, write-bandwidth:80, byte-size:80] adjuste
ranges
----
range-id=1 load=[cpu:80, write-bandwidth:80, byte-size:80] raft-cpu=20
store-id=1 replica-id=1 type=VOTER_FULL
store-id=2 replica-id=2 type=VOTER_FULL leaseholder=true
store-id=1 replica-id=1,type=VOTER_FULL
store-id=2 replica-id=2,type=VOTER_FULL,leaseholder=true

reject-pending-changes change-ids=(1,2)
----
Expand All @@ -78,8 +78,8 @@ store-id=2 node-id=2 reported=[cpu:20, write-bandwidth:80, byte-size:80] adjuste
ranges
----
range-id=1 load=[cpu:80, write-bandwidth:80, byte-size:80] raft-cpu=20
store-id=1 replica-id=1 type=VOTER_FULL leaseholder=true
store-id=2 replica-id=2 type=VOTER_FULL
store-id=1 replica-id=1,type=VOTER_FULL,leaseholder=true
store-id=2 replica-id=2,type=VOTER_FULL

# Set both existing stores again, this should have no effect.
set-store
Expand Down
23 changes: 0 additions & 23 deletions pkg/kv/kvserver/allocator/mmaprototypehelpers/BUILD.bazel
View file
Open in desktop

This file was deleted.

Loading
Loading