diff --git a/pkg/cmd/roachtest/BUILD.bazel b/pkg/cmd/roachtest/BUILD.bazel index efae87705d1b..6b4dbb327764 100644 --- a/pkg/cmd/roachtest/BUILD.bazel +++ b/pkg/cmd/roachtest/BUILD.bazel @@ -33,6 +33,7 @@ go_library( "//pkg/cmd/roachtest/roachtestflags", "//pkg/cmd/roachtest/roachtestutil", "//pkg/cmd/roachtest/roachtestutil/operations", + "//pkg/cmd/roachtest/roachtestutil/task", "//pkg/cmd/roachtest/spec", "//pkg/cmd/roachtest/test", "//pkg/cmd/roachtest/tests", @@ -107,6 +108,7 @@ go_test( "//pkg/cmd/roachtest/option", "//pkg/cmd/roachtest/registry", "//pkg/cmd/roachtest/roachtestflags", + "//pkg/cmd/roachtest/roachtestutil/task", "//pkg/cmd/roachtest/spec", "//pkg/cmd/roachtest/test", "//pkg/cmd/roachtest/testselector", diff --git a/pkg/cmd/roachtest/cluster_test.go b/pkg/cmd/roachtest/cluster_test.go index 0c55ef3a8387..881bd2696560 100644 --- a/pkg/cmd/roachtest/cluster_test.go +++ b/pkg/cmd/roachtest/cluster_test.go @@ -6,6 +6,7 @@ package main import ( + "context" "fmt" "strconv" "strings" @@ -13,6 +14,7 @@ import ( "time" "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/option" + "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/roachtestutil/task" "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/spec" test2 "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/test" "github.com/cockroachdb/cockroach/pkg/roachprod/logger" @@ -121,6 +123,14 @@ func (t testWrapper) IsDebug() bool { return false } +func (t testWrapper) GoWithCancel(_ task.Func, _ ...task.Option) context.CancelFunc { + panic("implement me") +} + +func (t testWrapper) Go(_ task.Func, _ ...task.Option) { + panic("implement me") +} + var _ test2.Test = testWrapper{} // ArtifactsDir is part of the test.Test interface. diff --git a/pkg/cmd/roachtest/clusterstats/BUILD.bazel b/pkg/cmd/roachtest/clusterstats/BUILD.bazel index 4014e81d9978..a1d75d077666 100644 --- a/pkg/cmd/roachtest/clusterstats/BUILD.bazel +++ b/pkg/cmd/roachtest/clusterstats/BUILD.bazel @@ -45,6 +45,7 @@ go_test( "//pkg/cmd/roachtest/cluster", "//pkg/cmd/roachtest/option", "//pkg/cmd/roachtest/registry", + "//pkg/cmd/roachtest/roachtestutil/task", "//pkg/cmd/roachtest/spec", "//pkg/cmd/roachtest/test", "//pkg/roachprod", diff --git a/pkg/cmd/roachtest/clusterstats/mocks_generated_test_test.go b/pkg/cmd/roachtest/clusterstats/mocks_generated_test_test.go index b871281c7e96..f34c372d8e44 100644 --- a/pkg/cmd/roachtest/clusterstats/mocks_generated_test_test.go +++ b/pkg/cmd/roachtest/clusterstats/mocks_generated_test_test.go @@ -1,8 +1,3 @@ -// Copyright 2024 The Cockroach Authors. -// -// Use of this software is governed by the CockroachDB Software License -// included in the /LICENSE file. - // Code generated by MockGen. DO NOT EDIT. // Source: github.com/cockroachdb/cockroach/pkg/cmd/roachtest/test (interfaces: Test) @@ -10,8 +5,10 @@ package clusterstats import ( + context "context" reflect "reflect" + task "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/roachtestutil/task" logger "github.com/cockroachdb/cockroach/pkg/roachprod/logger" version "github.com/cockroachdb/cockroach/pkg/util/version" gomock "github.com/golang/mock/gomock" @@ -202,6 +199,23 @@ func (mr *MockTestMockRecorder) Fatalf(arg0 interface{}, arg1 ...interface{}) *g return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Fatalf", reflect.TypeOf((*MockTest)(nil).Fatalf), varargs...) } +// Go mocks base method. +func (m *MockTest) Go(arg0 task.Func, arg1 ...task.Option) { + m.ctrl.T.Helper() + varargs := []interface{}{arg0} + for _, a := range arg1 { + varargs = append(varargs, a) + } + m.ctrl.Call(m, "Go", varargs...) +} + +// Go indicates an expected call of Go. +func (mr *MockTestMockRecorder) Go(arg0 interface{}, arg1 ...interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]interface{}{arg0}, arg1...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Go", reflect.TypeOf((*MockTest)(nil).Go), varargs...) +} + // GoCoverArtifactsDir mocks base method. func (m *MockTest) GoCoverArtifactsDir() string { m.ctrl.T.Helper() @@ -216,6 +230,25 @@ func (mr *MockTestMockRecorder) GoCoverArtifactsDir() *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GoCoverArtifactsDir", reflect.TypeOf((*MockTest)(nil).GoCoverArtifactsDir)) } +// GoWithCancel mocks base method. +func (m *MockTest) GoWithCancel(arg0 task.Func, arg1 ...task.Option) context.CancelFunc { + m.ctrl.T.Helper() + varargs := []interface{}{arg0} + for _, a := range arg1 { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "GoWithCancel", varargs...) + ret0, _ := ret[0].(context.CancelFunc) + return ret0 +} + +// GoWithCancel indicates an expected call of GoWithCancel. +func (mr *MockTestMockRecorder) GoWithCancel(arg0 interface{}, arg1 ...interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]interface{}{arg0}, arg1...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GoWithCancel", reflect.TypeOf((*MockTest)(nil).GoWithCancel), varargs...) +} + // Helper mocks base method. func (m *MockTest) Helper() { m.ctrl.T.Helper() diff --git a/pkg/cmd/roachtest/test/BUILD.bazel b/pkg/cmd/roachtest/test/BUILD.bazel index a8758568e176..6269f778e499 100644 --- a/pkg/cmd/roachtest/test/BUILD.bazel +++ b/pkg/cmd/roachtest/test/BUILD.bazel @@ -6,6 +6,7 @@ go_library( importpath = "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/test", visibility = ["//visibility:public"], deps = [ + "//pkg/cmd/roachtest/roachtestutil/task", "//pkg/roachprod/logger", "//pkg/util/version", ], diff --git a/pkg/cmd/roachtest/test/test_interface.go b/pkg/cmd/roachtest/test/test_interface.go index 9dff50c838b2..1875c9be126e 100644 --- a/pkg/cmd/roachtest/test/test_interface.go +++ b/pkg/cmd/roachtest/test/test_interface.go @@ -6,6 +6,9 @@ package test import ( + "context" + + "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/roachtestutil/task" "github.com/cockroachdb/cockroach/pkg/roachprod/logger" "github.com/cockroachdb/cockroach/pkg/util/version" ) @@ -77,6 +80,9 @@ type Test interface { WorkerProgress(float64) IsDebug() bool + Go(task.Func, ...task.Option) + GoWithCancel(task.Func, ...task.Option) context.CancelFunc + // DeprecatedWorkload returns the path to the workload binary. // Don't use this, invoke `./cockroach workload` instead. DeprecatedWorkload() string diff --git a/pkg/cmd/roachtest/test_impl.go b/pkg/cmd/roachtest/test_impl.go index 05da8257f8d5..720df93a5db6 100644 --- a/pkg/cmd/roachtest/test_impl.go +++ b/pkg/cmd/roachtest/test_impl.go @@ -11,12 +11,14 @@ import ( "io" "math/rand" "os" + "runtime/debug" "strings" "sync" "time" "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/registry" "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/roachtestflags" + "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/roachtestutil/task" "github.com/cockroachdb/cockroach/pkg/roachprod/logger" "github.com/cockroachdb/cockroach/pkg/testutils/skip" "github.com/cockroachdb/cockroach/pkg/util/syncutil" @@ -72,6 +74,9 @@ type testImpl struct { // l is the logger that the test will use for its output. l *logger.Logger + // taskManager manages tasks (goroutines) for tests. + taskManager task.Manager + runner string // runnerID is the test's main goroutine ID. runnerID int64 @@ -583,6 +588,22 @@ func (t *testImpl) IsBuildVersion(minVersion string) bool { return t.BuildVersion().AtLeast(vers) } +func panicHandler(_ context.Context, l *logger.Logger, r interface{}) error { + l.Printf("panic stack trace:\n%s", string(debug.Stack())) + return fmt.Errorf("panic (stack trace above): %v", 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...)) +} + +// Go is like GoWithCancel but without a cancel function. +func (t *testImpl) Go(fn task.Func, opts ...task.Option) { + _ = t.GoWithCancel(fn, task.OptionList(opts...)) +} + // 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 diff --git a/pkg/cmd/roachtest/test_runner.go b/pkg/cmd/roachtest/test_runner.go index 8d26b3ed4dd9..fcf920d3db0b 100644 --- a/pkg/cmd/roachtest/test_runner.go +++ b/pkg/cmd/roachtest/test_runner.go @@ -32,6 +32,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/registry" "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/roachtestflags" "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/roachtestutil" + "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/roachtestutil/task" "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/spec" "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/test" "github.com/cockroachdb/cockroach/pkg/roachprod/config" @@ -836,6 +837,7 @@ func (r *testRunner) runWorker( debug: clustersOpt.debugMode.IsDebug(), goCoverEnabled: topt.goCoverEnabled, exportOpenmetrics: topt.exportOpenMetrics, + taskManager: task.NewManager(ctx, testL), } github := newGithubIssues(r.config.disableIssue, c, vmCreateOpts) @@ -1281,6 +1283,26 @@ func (r *testRunner) runTest( s.Run(runCtx, t, c) }() + // Monitor the task manager for completed events, or failure events and log + // them. A failure will call t.Errorf which cancels the test's context. + go func() { + for { + select { + case event := <-t.taskManager.CompletedEvents(): + if event.Err == nil { + t.L().Printf("task finished: %s", event.Name) + continue + } else if event.ExpectedCancel { + t.L().Printf("task canceled by test: %s", event.Name) + continue + } + t.Errorf("task `%s` returned error: %v", event.Name, event.Err) + case <-runCtx.Done(): + return + } + } + }() + var timedOut bool timeout := testTimeout(t.spec) @@ -1557,6 +1579,10 @@ func (r *testRunner) teardownTest( getGoCoverArtifacts(ctx, c, t) } + // Terminate tasks to ensure that any stray tasks are cleaned up. + t.L().Printf("terminating tasks") + t.taskManager.Terminate(t.L()) + return "", nil }