@@ -15,6 +15,7 @@ import (
15
15
"slices"
16
16
"sort"
17
17
"sync/atomic"
18
+ "time"
18
19
19
20
"github.com/cockroachdb/cockroach/pkg/keys"
20
21
"github.com/cockroachdb/cockroach/pkg/kv/kvnemesis/kvnemesisutil"
@@ -28,6 +29,7 @@ import (
28
29
"github.com/cockroachdb/cockroach/pkg/util/encoding"
29
30
"github.com/cockroachdb/cockroach/pkg/util/hlc"
30
31
"github.com/cockroachdb/cockroach/pkg/util/syncutil"
32
+ "github.com/cockroachdb/cockroach/pkg/util/timeutil"
31
33
"github.com/cockroachdb/errors"
32
34
"golang.org/x/exp/maps"
33
35
)
@@ -645,7 +647,7 @@ func GeneratorDataSpan() roachpb.Span {
645
647
646
648
// GetReplicasFn is a function that returns the current voting and non-voting
647
649
// replicas, respectively, for the range containing a key.
648
- type GetReplicasFn func (roachpb.Key ) ([]roachpb.ReplicationTarget , []roachpb.ReplicationTarget )
650
+ type GetReplicasFn func (context. Context , roachpb.Key ) ([]roachpb.ReplicationTarget , []roachpb.ReplicationTarget )
649
651
650
652
// Generator incrementally constructs KV traffic designed to maximally test edge
651
653
// cases.
@@ -676,7 +678,9 @@ type Generator struct {
676
678
}
677
679
678
680
// MakeGenerator constructs a Generator.
679
- func MakeGenerator (config GeneratorConfig , replicasFn GetReplicasFn ) (* Generator , error ) {
681
+ func MakeGenerator (
682
+ config GeneratorConfig , replicasFn GetReplicasFn , mode TestMode ,
683
+ ) (* Generator , error ) {
680
684
if config .NumNodes <= 0 {
681
685
return nil , errors .Errorf (`NumNodes must be positive got: %d` , config .NumNodes )
682
686
}
@@ -693,7 +697,11 @@ func MakeGenerator(config GeneratorConfig, replicasFn GetReplicasFn) (*Generator
693
697
}
694
698
for i := 1 ; i <= config .NumNodes ; i ++ {
695
699
for j := 1 ; j <= config .NumNodes ; j ++ {
696
- if i == j {
700
+ // In liveness mode, we don't allow adding partitions between the two
701
+ // protected nodes (node 1 and node 2), so we don't include those
702
+ // connections in the set of healthy connections at all.
703
+ protectedConn := (i == 1 && j == 2 ) || (i == 2 && j == 1 )
704
+ if i == j || (mode == Liveness && protectedConn ) {
697
705
continue
698
706
}
699
707
conn := connection {from : i , to : j }
@@ -708,6 +716,7 @@ func MakeGenerator(config GeneratorConfig, replicasFn GetReplicasFn) (*Generator
708
716
currentSplits : make (map [string ]struct {}),
709
717
historicalSplits : make (map [string ]struct {}),
710
718
partitions : p ,
719
+ mode : mode ,
711
720
}
712
721
return g , nil
713
722
}
@@ -745,6 +754,10 @@ type generator struct {
745
754
// partitions contains the sets of healthy and partitioned connections
746
755
// between nodes.
747
756
partitions
757
+
758
+ // mode is the test mode (e.g. Liveness or Safety). The generator needs it in
759
+ // order to set a timeout for range lookups under safety mode.
760
+ mode TestMode
748
761
}
749
762
750
763
type connection struct {
@@ -777,26 +790,45 @@ func (g *generator) RandStep(rng *rand.Rand) Step {
777
790
}
778
791
779
792
key := randKey (rng )
780
- voters , nonVoters := g .replicasFn (roachpb .Key (key ))
793
+ var voters , nonVoters []roachpb.ReplicationTarget
794
+ if g .mode == Safety {
795
+ if err := timeutil .RunWithTimeout (context .Background (), "getting replicas" , 3 * time .Second ,
796
+ func (ctx context.Context ) error {
797
+ voters , nonVoters = g .replicasFn (ctx , roachpb .Key (key ))
798
+ return nil
799
+ }); err != nil {
800
+ voters , nonVoters = []roachpb.ReplicationTarget {}, []roachpb.ReplicationTarget {}
801
+ }
802
+ } else {
803
+ voters , nonVoters = g .replicasFn (context .Background (), roachpb .Key (key ))
804
+ }
781
805
numVoters , numNonVoters := len (voters ), len (nonVoters )
782
806
numReplicas := numVoters + numNonVoters
783
807
if numReplicas < g .Config .NumNodes {
784
- addVoterFn := makeAddReplicaFn (key , voters , false /* atomicSwap */ , true /* voter */ )
785
- addOpGen (& allowed , addVoterFn , g .Config .Ops .ChangeReplicas .AddVotingReplica )
786
- addNonVoterFn := makeAddReplicaFn (key , nonVoters , false /* atomicSwap */ , false /* voter */ )
787
- addOpGen (& allowed , addNonVoterFn , g .Config .Ops .ChangeReplicas .AddNonVotingReplica )
808
+ if len (voters ) > 0 {
809
+ addVoterFn := makeAddReplicaFn (key , voters , false /* atomicSwap */ , true /* voter */ )
810
+ addOpGen (& allowed , addVoterFn , g .Config .Ops .ChangeReplicas .AddVotingReplica )
811
+ }
812
+ if len (nonVoters ) > 0 {
813
+ addNonVoterFn := makeAddReplicaFn (key , nonVoters , false /* atomicSwap */ , false /* voter */ )
814
+ addOpGen (& allowed , addNonVoterFn , g .Config .Ops .ChangeReplicas .AddNonVotingReplica )
815
+ }
788
816
}
789
817
if numReplicas == g .Config .NumReplicas && numReplicas < g .Config .NumNodes {
790
- atomicSwapVoterFn := makeAddReplicaFn (key , voters , true /* atomicSwap */ , true /* voter */ )
791
- addOpGen (& allowed , atomicSwapVoterFn , g .Config .Ops .ChangeReplicas .AtomicSwapVotingReplica )
818
+ if len (voters ) > 0 {
819
+ atomicSwapVoterFn := makeAddReplicaFn (key , voters , true /* atomicSwap */ , true /* voter */ )
820
+ addOpGen (& allowed , atomicSwapVoterFn , g .Config .Ops .ChangeReplicas .AtomicSwapVotingReplica )
821
+ }
792
822
if numNonVoters > 0 {
793
823
atomicSwapNonVoterFn := makeAddReplicaFn (key , nonVoters , true /* atomicSwap */ , false /* voter */ )
794
824
addOpGen (& allowed , atomicSwapNonVoterFn , g .Config .Ops .ChangeReplicas .AtomicSwapNonVotingReplica )
795
825
}
796
826
}
797
827
if numReplicas > g .Config .NumReplicas {
798
- removeVoterFn := makeRemoveReplicaFn (key , voters , true /* voter */ )
799
- addOpGen (& allowed , removeVoterFn , g .Config .Ops .ChangeReplicas .RemoveVotingReplica )
828
+ if len (voters ) > 0 {
829
+ removeVoterFn := makeRemoveReplicaFn (key , voters , true /* voter */ )
830
+ addOpGen (& allowed , removeVoterFn , g .Config .Ops .ChangeReplicas .RemoveVotingReplica )
831
+ }
800
832
if numNonVoters > 0 {
801
833
removeNonVoterFn := makeRemoveReplicaFn (key , nonVoters , false /* voter */ )
802
834
addOpGen (& allowed , removeNonVoterFn , g .Config .Ops .ChangeReplicas .RemoveNonVotingReplica )
@@ -810,8 +842,10 @@ func (g *generator) RandStep(rng *rand.Rand) Step {
810
842
promoteNonVoterFn := makePromoteReplicaFn (key , nonVoters )
811
843
addOpGen (& allowed , promoteNonVoterFn , g .Config .Ops .ChangeReplicas .PromoteReplica )
812
844
}
813
- transferLeaseFn := makeTransferLeaseFn (key , append (voters , nonVoters ... ))
814
- addOpGen (& allowed , transferLeaseFn , g .Config .Ops .ChangeLease .TransferLease )
845
+ if numVoters > 0 {
846
+ transferLeaseFn := makeTransferLeaseFn (key , append (voters , nonVoters ... ))
847
+ addOpGen (& allowed , transferLeaseFn , g .Config .Ops .ChangeLease .TransferLease )
848
+ }
815
849
816
850
addOpGen (& allowed , setLeaseType , g .Config .Ops .ChangeSetting .SetLeaseType )
817
851
addOpGen (& allowed , toggleGlobalReads , g .Config .Ops .ChangeZone .ToggleGlobalReads )
0 commit comments