Skip to content

Commit 3565057

Browse files
pav-kvpav-kv
committed
replica_rac2: squash datadriven command parsing
Epic: none Release note: none
1 parent de909f7 commit 3565057

File tree

4 files changed

+64
-94
lines changed

4 files changed

+64
-94
lines changed

‎pkg/kv/kvserver/kvflowcontrol/replica_rac2/BUILD.bazel‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@ go_test(
5656
"//pkg/roachpb",
5757
"//pkg/server/serverpb",
5858
"//pkg/testutils/datapathutils",
59+
"//pkg/testutils/dd",
5960
"//pkg/util/admission/admissionpb",
6061
"//pkg/util/leaktest",
6162
"//pkg/util/log",

‎pkg/kv/kvserver/kvflowcontrol/replica_rac2/admission_test.go‎

Lines changed: 11 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import (
1212

1313
"github.com/cockroachdb/cockroach/pkg/raft/raftpb"
1414
"github.com/cockroachdb/cockroach/pkg/testutils/datapathutils"
15+
"github.com/cockroachdb/cockroach/pkg/testutils/dd"
1516
"github.com/cockroachdb/cockroach/pkg/util/leaktest"
1617
"github.com/cockroachdb/cockroach/pkg/util/log"
1718
"github.com/cockroachdb/datadriven"
@@ -47,15 +48,11 @@ func TestLowPriOverrideState(t *testing.T) {
4748
//
4849
// Provides side-channel info for the the interval [5, 10] for the given
4950
// leader-term and with the specific override.
50-
var leaderTerm uint64
51-
d.ScanArgs(t, "leader-term", &leaderTerm)
52-
var first, last uint64
53-
d.ScanArgs(t, "first", &first)
54-
d.ScanArgs(t, "last", &last)
55-
var lowPriOverride bool
56-
if d.HasArg("low-pri") {
57-
lowPriOverride = true
58-
}
51+
leaderTerm := dd.ScanArg[uint64](t, d, "leader-term")
52+
first := dd.ScanArg[uint64](t, d, "first")
53+
last := dd.ScanArg[uint64](t, d, "last")
54+
lowPriOverride := d.HasArg("low-pri")
55+
5956
notStaleTerm := lpos.sideChannelForLowPriOverride(leaderTerm, first, last, lowPriOverride)
6057
return fmt.Sprintf("not-stale-term: %t\n%s", notStaleTerm, lposString())
6158

@@ -64,9 +61,8 @@ func TestLowPriOverrideState(t *testing.T) {
6461
// get-effective-priority index=4 pri=HighPri
6562
// Gets the effective priority for index 4, where the original
6663
// priority is HighPri
67-
var index uint64
68-
d.ScanArgs(t, "index", &index)
69-
pri := readPriority(t, d)
64+
index := dd.ScanArg[uint64](t, d, "index")
65+
pri := parsePriority(t, dd.ScanArg[string](t, d, "pri"))
7066
effectivePri := lpos.getEffectivePriority(index, pri)
7167
return fmt.Sprintf("pri: %s\n%s", effectivePri, lposString())
7268

@@ -76,10 +72,8 @@ func TestLowPriOverrideState(t *testing.T) {
7672
})
7773
}
7874

79-
func readPriority(t *testing.T, d *datadriven.TestData) raftpb.Priority {
80-
var priStr string
81-
d.ScanArgs(t, "pri", &priStr)
82-
switch priStr {
75+
func parsePriority(t *testing.T, str string) raftpb.Priority {
76+
switch str {
8377
case "LowPri":
8478
return raftpb.LowPri
8579
case "NormalPri":
@@ -89,7 +83,7 @@ func readPriority(t *testing.T, d *datadriven.TestData) raftpb.Priority {
8983
case "HighPri":
9084
return raftpb.HighPri
9185
default:
92-
t.Fatalf("unknown pri %s", priStr)
86+
t.Fatalf("unknown pri %s", str)
9387
}
9488
return 0
9589
}

‎pkg/kv/kvserver/kvflowcontrol/replica_rac2/close_scheduler_test.go‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import (
1515
"time"
1616

1717
"github.com/cockroachdb/cockroach/pkg/roachpb"
18+
"github.com/cockroachdb/cockroach/pkg/testutils/dd"
1819
"github.com/cockroachdb/cockroach/pkg/util/leaktest"
1920
"github.com/cockroachdb/cockroach/pkg/util/log"
2021
"github.com/cockroachdb/cockroach/pkg/util/stop"
@@ -119,10 +120,7 @@ func TestStreamCloseScheduler(t *testing.T) {
119120
return buf.String()
120121

121122
case "tick":
122-
var durationStr string
123-
d.ScanArgs(t, "duration", &durationStr)
124-
duration, err := time.ParseDuration(durationStr)
125-
require.NoError(t, err)
123+
duration := dd.ScanArg[time.Duration](t, d, "duration")
126124
clock.Advance(duration)
127125
// Delay to allow the channel selects to fire.
128126
time.Sleep(20 * time.Millisecond)

‎pkg/kv/kvserver/kvflowcontrol/replica_rac2/processor_test.go‎

Lines changed: 50 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ import (
2424
"github.com/cockroachdb/cockroach/pkg/roachpb"
2525
"github.com/cockroachdb/cockroach/pkg/server/serverpb"
2626
"github.com/cockroachdb/cockroach/pkg/testutils/datapathutils"
27+
"github.com/cockroachdb/cockroach/pkg/testutils/dd"
2728
"github.com/cockroachdb/cockroach/pkg/util/admission/admissionpb"
2829
"github.com/cockroachdb/cockroach/pkg/util/leaktest"
2930
"github.com/cockroachdb/cockroach/pkg/util/log"
@@ -330,51 +331,43 @@ func TestProcessorBasic(t *testing.T) {
330331
return builderStr()
331332

332333
case "init-raft":
333-
var mark rac2.LogMark
334-
d.ScanArgs(t, "log-term", &mark.Term)
335-
d.ScanArgs(t, "log-index", &mark.Index)
336-
r.initRaft(mark)
334+
r.initRaft(rac2.LogMark{
335+
Term: dd.ScanArg[uint64](t, d, "log-term"),
336+
Index: dd.ScanArg[uint64](t, d, "log-index"),
337+
})
337338
unlockFunc := LockRaftMuAndReplicaMu(&muAsserter)
338339
p.InitRaftLocked(ctx, r.raftNode, r.raftNode.mark)
339340
unlockFunc()
340341
return builderStr()
341342

342343
case "set-raft-state":
343-
if d.HasArg("leader") {
344-
var leaderID int
345-
d.ScanArgs(t, "leader", &leaderID)
346-
r.raftNode.isLeader = leaderID == replicaID
347-
r.raftNode.leader = roachpb.ReplicaID(leaderID)
344+
if id, ok := dd.ScanArgOpt[roachpb.ReplicaID](t, d, "leader"); ok {
345+
r.raftNode.isLeader = id == replicaID
346+
r.raftNode.leader = id
348347
}
349-
if d.HasArg("next-unstable-index") {
350-
var nextUnstableIndex uint64
351-
d.ScanArgs(t, "next-unstable-index", &nextUnstableIndex)
352-
r.raftNode.nextUnstableIndex = nextUnstableIndex
348+
if idx, ok := dd.ScanArgOpt[uint64](t, d, "next-unstable-index"); ok {
349+
r.raftNode.nextUnstableIndex = idx
353350
}
354-
if d.HasArg("term") {
355-
var term uint64
356-
d.ScanArgs(t, "term", &term)
351+
if term, ok := dd.ScanArgOpt[uint64](t, d, "term"); ok {
357352
r.raftNode.term = term
358353
}
359-
if d.HasArg("leaseholder") {
360-
var leaseholder int
361-
d.ScanArgs(t, "leaseholder", &leaseholder)
362-
r.leaseholder = roachpb.ReplicaID(leaseholder)
354+
if lh, ok := dd.ScanArgOpt[roachpb.ReplicaID](t, d, "leaseholder"); ok {
355+
r.leaseholder = lh
363356
}
364-
if d.HasArg("log-term") {
365-
var mark rac2.LogMark
366-
d.ScanArgs(t, "log-term", &mark.Term)
367-
d.ScanArgs(t, "log-index", &mark.Index)
368-
r.raftNode.setMark(t, mark)
357+
if term, ok := dd.ScanArgOpt[uint64](t, d, "log-term"); ok {
358+
r.raftNode.setMark(t, rac2.LogMark{
359+
Term: term,
360+
Index: dd.ScanArg[uint64](t, d, "log-index"),
361+
})
369362
}
370363
r.raftNode.print()
371364
return builderStr()
372365

373366
case "synced-log":
374-
var mark rac2.LogMark
375-
d.ScanArgs(t, "term", &mark.Term)
376-
d.ScanArgs(t, "index", &mark.Index)
377-
p.SyncedLogStorage(ctx, mark)
367+
p.SyncedLogStorage(ctx, rac2.LogMark{
368+
Term: dd.ScanArg[uint64](t, d, "term"),
369+
Index: dd.ScanArg[uint64](t, d, "index"),
370+
})
378371
printLogTracker()
379372
return builderStr()
380373

@@ -396,13 +389,11 @@ func TestProcessorBasic(t *testing.T) {
396389
// unused by processorImpl, and simply passed down to RangeController
397390
// (which we've mocked out in this test).
398391
var event rac2.RaftEvent
399-
if d.HasArg("entries") {
400-
var arg string
401-
d.ScanArgs(t, "entries", &arg)
392+
if arg, ok := dd.ScanArgOpt[string](t, d, "entries"); ok {
402393
event.Entries = createEntries(t, parseEntryInfos(t, arg))
403394
}
404395
if len(event.Entries) > 0 {
405-
d.ScanArgs(t, "leader-term", &event.Term)
396+
event.Term = dd.ScanArg[uint64](t, d, "leader-term")
406397
}
407398
fmt.Fprintf(&b, "HandleRaftReady:\n")
408399
var state RaftNodeBasicState
@@ -427,23 +418,22 @@ func TestProcessorBasic(t *testing.T) {
427418
return builderStr()
428419

429420
case "enqueue-piggybacked-admitted":
430-
var from, to uint64
431-
d.ScanArgs(t, "from", &from)
432-
d.ScanArgs(t, "to", &to)
433-
require.Equal(t, p.opts.ReplicaID, roachpb.ReplicaID(to))
434-
435-
var term, index, pri int
436-
d.ScanArgs(t, "term", &term)
437-
d.ScanArgs(t, "index", &index)
438-
d.ScanArgs(t, "pri", &pri)
439-
require.Less(t, pri, int(raftpb.NumPriorities))
421+
from := dd.ScanArg[roachpb.ReplicaID](t, d, "from")
422+
to := dd.ScanArg[roachpb.ReplicaID](t, d, "to")
423+
require.Equal(t, p.opts.ReplicaID, to)
424+
425+
term := dd.ScanArg[uint64](t, d, "term")
426+
index := dd.ScanArg[uint64](t, d, "index")
427+
pri := dd.ScanArg[raftpb.Priority](t, d, "pri")
428+
require.Less(t, pri, raftpb.NumPriorities)
429+
440430
as := kvflowcontrolpb.AdmittedState{
441-
Term: uint64(term),
431+
Term: term,
442432
Admitted: make([]uint64, raftpb.NumPriorities),
443433
}
444-
as.Admitted[pri] = uint64(index)
434+
as.Admitted[pri] = index
445435

446-
p.EnqueuePiggybackedAdmittedAtLeader(roachpb.ReplicaID(from), as)
436+
p.EnqueuePiggybackedAdmittedAtLeader(from, as)
447437
return builderStr()
448438

449439
case "process-piggybacked-admitted":
@@ -453,34 +443,25 @@ func TestProcessorBasic(t *testing.T) {
453443
return builderStr()
454444

455445
case "side-channel":
456-
var leaderTerm uint64
457-
d.ScanArgs(t, "leader-term", &leaderTerm)
458-
var first, last uint64
459-
d.ScanArgs(t, "first", &first)
460-
d.ScanArgs(t, "last", &last)
461-
var lowPriOverride bool
462-
if d.HasArg("low-pri") {
463-
lowPriOverride = true
464-
}
465446
info := SideChannelInfoUsingRaftMessageRequest{
466-
LeaderTerm: leaderTerm,
467-
First: first,
468-
Last: last,
469-
LowPriOverride: lowPriOverride,
447+
LeaderTerm: dd.ScanArg[uint64](t, d, "leader-term"),
448+
First: dd.ScanArg[uint64](t, d, "first"),
449+
Last: dd.ScanArg[uint64](t, d, "last"),
450+
LowPriOverride: d.HasArg("low-pri"),
470451
}
471452
unlockFunc := LockRaftMu(&muAsserter)
472453
p.SideChannelForPriorityOverrideAtFollowerRaftMuLocked(info)
473454
unlockFunc()
474455
return builderStr()
475456

476457
case "admitted-log-entry":
477-
var cb EntryForAdmissionCallbackState
478-
d.ScanArgs(t, "leader-term", &cb.Mark.Term)
479-
d.ScanArgs(t, "index", &cb.Mark.Index)
480-
var pri int
481-
d.ScanArgs(t, "pri", &pri)
482-
cb.Priority = raftpb.Priority(pri)
483-
p.AdmittedLogEntry(ctx, cb)
458+
p.AdmittedLogEntry(ctx, EntryForAdmissionCallbackState{
459+
Mark: rac2.LogMark{
460+
Term: dd.ScanArg[uint64](t, d, "leader-term"),
461+
Index: dd.ScanArg[uint64](t, d, "index"),
462+
},
463+
Priority: dd.ScanArg[raftpb.Priority](t, d, "pri"),
464+
})
484465
printLogTracker()
485466
return builderStr()
486467

@@ -500,9 +481,7 @@ func TestProcessorBasic(t *testing.T) {
500481
rc := rcFactory.rcs[len(rcFactory.rcs)-1]
501482
d.ScanArgs(t, "waited", &rc.waited)
502483
rc.waitForEvalErr = nil
503-
if d.HasArg("err") {
504-
var errStr string
505-
d.ScanArgs(t, "err", &errStr)
484+
if errStr, ok := dd.ScanArgOpt[string](t, d, "err"); ok {
506485
rc.waitForEvalErr = errors.Errorf("%s", errStr)
507486
}
508487
return builderStr()
@@ -525,8 +504,7 @@ func TestProcessorBasic(t *testing.T) {
525504
}
526505

527506
func parseAdmissionPriority(t *testing.T, td *datadriven.TestData) admissionpb.WorkPriority {
528-
var priStr string
529-
td.ScanArgs(t, "pri", &priStr)
507+
priStr := dd.ScanArg[string](t, td, "pri")
530508
for k, v := range admissionpb.WorkPriorityDict {
531509
if v == priStr {
532510
return k
@@ -537,8 +515,7 @@ func parseAdmissionPriority(t *testing.T, td *datadriven.TestData) admissionpb.W
537515
}
538516

539517
func parseRangeDescriptor(t *testing.T, td *datadriven.TestData) roachpb.RangeDescriptor {
540-
var replicaStr string
541-
td.ScanArgs(t, "replicas", &replicaStr)
518+
replicaStr := dd.ScanArg[string](t, td, "replicas")
542519
parts := strings.Split(replicaStr, ",")
543520
var desc roachpb.RangeDescriptor
544521
for _, part := range parts {

0 commit comments

Comments
 (0)