diff --git a/pkg/roachprod/failureinjection/failures/BUILD.bazel b/pkg/roachprod/failureinjection/failures/BUILD.bazel index 9d7209f997c..ba95ecc0bcc 100644 --- a/pkg/roachprod/failureinjection/failures/BUILD.bazel +++ b/pkg/roachprod/failureinjection/failures/BUILD.bazel @@ -5,6 +5,7 @@ go_library( srcs = [ "failure.go", "network_partition.go", + "registry.go", ], importpath = "github.com/cockroachdb/cockroach/pkg/roachprod/failureinjection/failures", visibility = ["//visibility:public"], diff --git a/pkg/roachprod/failureinjection/failures/network_partition.go b/pkg/roachprod/failureinjection/failures/network_partition.go index 6da68b3d5cf..676befda564 100644 --- a/pkg/roachprod/failureinjection/failures/network_partition.go +++ b/pkg/roachprod/failureinjection/failures/network_partition.go @@ -65,6 +65,10 @@ func MakeIPTablesPartitionNode( const IPTablesNetworkPartitionName = "iptables-network-partition" +func registerIPTablesPartitionNode(r *FailureRegistry) { + r.add(IPTablesNetworkPartitionName, NetworkPartitionArgs{}, MakeIPTablesPartitionNode) +} + func (f *IPTablesPartitionNode) Description() string { return IPTablesNetworkPartitionName } diff --git a/pkg/roachprod/failureinjection/failures/registry.go b/pkg/roachprod/failureinjection/failures/registry.go new file mode 100644 index 00000000000..68fc086f629 --- /dev/null +++ b/pkg/roachprod/failureinjection/failures/registry.go @@ -0,0 +1,72 @@ +// Copyright 2025 The Cockroach Authors. +// +// Use of this software is governed by the CockroachDB Software License +// included in the /LICENSE file. + +package failures + +import ( + "fmt" + "regexp" + + "github.com/cockroachdb/cockroach/pkg/roachprod/logger" +) + +type failureSpec struct { + makeFailureFunc func(clusterName string, l *logger.Logger, secure bool) (FailureMode, error) + args FailureArgs +} +type FailureRegistry struct { + failures map[string]failureSpec +} + +func NewFailureRegistry() *FailureRegistry { + return &FailureRegistry{ + failures: make(map[string]failureSpec), + } +} + +func (r *FailureRegistry) Register() { + registerIPTablesPartitionNode(r) +} + +func (r *FailureRegistry) add( + failureName string, + args FailureArgs, + makeFailureFunc func(clusterName string, l *logger.Logger, secure bool) (FailureMode, error), +) { + if _, ok := r.failures[failureName]; ok { + panic(fmt.Sprintf("failure %s already exists", failureName)) + } + r.failures[failureName] = failureSpec{ + makeFailureFunc: makeFailureFunc, + args: args, + } +} + +func (r *FailureRegistry) List(regex string) []string { + var filter *regexp.Regexp + if regex == "" { + filter = regexp.MustCompile(`.`) + } else { + filter = regexp.MustCompile(regex) + } + + var matches []string + for name := range r.failures { + if filter.MatchString(name) { + matches = append(matches, name) + } + } + return matches +} + +func (r *FailureRegistry) GetFailure( + clusterName, failureName string, l *logger.Logger, secure bool, +) (FailureMode, error) { + spec, ok := r.failures[failureName] + if !ok { + return nil, fmt.Errorf("unknown failure %s", failureName) + } + return spec.makeFailureFunc(clusterName, l, secure) +}