diff --git a/pkg/cmd/roachtest/cluster_test.go b/pkg/cmd/roachtest/cluster_test.go index b0e31bac009e..401463fedd23 100644 --- a/pkg/cmd/roachtest/cluster_test.go +++ b/pkg/cmd/roachtest/cluster_test.go @@ -144,6 +144,10 @@ func (t testWrapper) Go(_ task.Func, _ ...task.Option) { panic("implement me") } +func (t testWrapper) NewGroup() task.Group { + panic("implement me") +} + var _ test2.Test = testWrapper{} // ArtifactsDir is part of the test.Test interface. diff --git a/pkg/cmd/roachtest/clusterstats/mock_test_generated_test.go b/pkg/cmd/roachtest/clusterstats/mock_test_generated_test.go index 35ee7ca580fe..f78dee987155 100644 --- a/pkg/cmd/roachtest/clusterstats/mock_test_generated_test.go +++ b/pkg/cmd/roachtest/clusterstats/mock_test_generated_test.go @@ -343,6 +343,20 @@ func (mr *MockTestMockRecorder) Name() *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Name", reflect.TypeOf((*MockTest)(nil).Name)) } +// NewGroup mocks base method. +func (m *MockTest) NewGroup() task.Group { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "NewGroup") + ret0, _ := ret[0].(task.Group) + return ret0 +} + +// NewGroup indicates an expected call of NewGroup. +func (mr *MockTestMockRecorder) NewGroup() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "NewGroup", reflect.TypeOf((*MockTest)(nil).NewGroup)) +} + // PerfArtifactsDir mocks base method. func (m *MockTest) PerfArtifactsDir() string { m.ctrl.T.Helper() diff --git a/pkg/cmd/roachtest/test/test_interface.go b/pkg/cmd/roachtest/test/test_interface.go index e5b3c9836b7d..302851aab810 100644 --- a/pkg/cmd/roachtest/test/test_interface.go +++ b/pkg/cmd/roachtest/test/test_interface.go @@ -83,6 +83,7 @@ type Test interface { Go(task.Func, ...task.Option) GoWithCancel(task.Func, ...task.Option) context.CancelFunc + NewGroup() task.Group // DeprecatedWorkload returns the path to the workload binary. // Don't use this, invoke `./cockroach workload` instead. diff --git a/pkg/cmd/roachtest/test_impl.go b/pkg/cmd/roachtest/test_impl.go index 9c99a275e470..0c6104bfa6c7 100644 --- a/pkg/cmd/roachtest/test_impl.go +++ b/pkg/cmd/roachtest/test_impl.go @@ -688,14 +688,21 @@ func (t *testImpl) IsBuildVersion(minVersion string) bool { return t.BuildVersion().AtLeast(vers) } -func panicHandler(_ context.Context, name string, l *logger.Logger, r interface{}) error { - return fmt.Errorf("test task %s panicked: %v", name, r) +// defaultTaskOptions returns the default options for a task started by the test. +func defaultTaskOptions() []task.Option { + return []task.Option{ + task.PanicHandler(func(_ context.Context, name string, l *logger.Logger, r interface{}) error { + return fmt.Errorf("test task %s panicked: %v", name, r) + }), + } } // GoWithCancel runs the given function in a goroutine and returns a // CancelFunc that can be used to cancel the function. func (t *testImpl) GoWithCancel(fn task.Func, opts ...task.Option) context.CancelFunc { - return t.taskManager.GoWithCancel(fn, task.PanicHandler(panicHandler), task.OptionList(opts...)) + return t.taskManager.GoWithCancel( + fn, task.OptionList(defaultTaskOptions()...), task.OptionList(opts...), + ) } // Go is like GoWithCancel but without a cancel function. @@ -703,6 +710,11 @@ func (t *testImpl) Go(fn task.Func, opts ...task.Option) { _ = t.GoWithCancel(fn, task.OptionList(opts...)) } +// NewGroup starts a new task group. +func (t *testImpl) NewGroup() task.Group { + return t.taskManager.NewGroup(defaultTaskOptions()...) +} + // TeamCityEscape escapes a string for use as in a key='' attribute // in TeamCity build output marker. // See https://www.jetbrains.com/help/teamcity/2023.05/service-messages.html#Escaped+Values