From d18a1049a293d457b7974ebf9a2f597b602549fa Mon Sep 17 00:00:00 2001 From: Herko Lategan Date: Wed, 23 Oct 2024 16:43:48 +0100 Subject: [PATCH] roachtest: interface for go routines Previously, roachtests would use bare goroutines `go func...` to start tasks during a test to perform various background functions. This is problematic due to the nature of roachtest, which runs multiple tests in parallel using workers. If one such bare goroutine panics the whole test suite is interrupted and the process exits. This change is the first step and introduces an interface that both the traditional and mixedversion test frameworks can utilise. It intends to target both regular and MVT tests by supplying the test with the interface either through a helper or the test interface. In addition to providing the `Go` function, there are additional options that can be passed. The test framework will generally supply the panic and error handler, but users can supply additional options such as naming the task, or passing a custom logger. This makes logging more useful to determine what task panic or encountered an error. Informs: #118214 Epic: None Release note: None --- pkg/BUILD.bazel | 3 + .../roachtest/roachtestutil/task/BUILD.bazel | 29 ++++++ .../roachtest/roachtestutil/task/options.go | 91 +++++++++++++++++++ .../roachtest/roachtestutil/task/tasker.go | 24 +++++ 4 files changed, 147 insertions(+) create mode 100644 pkg/cmd/roachtest/roachtestutil/task/BUILD.bazel create mode 100644 pkg/cmd/roachtest/roachtestutil/task/options.go create mode 100644 pkg/cmd/roachtest/roachtestutil/task/tasker.go diff --git a/pkg/BUILD.bazel b/pkg/BUILD.bazel index a117b56d65eb..1a526dcf16c2 100644 --- a/pkg/BUILD.bazel +++ b/pkg/BUILD.bazel @@ -160,6 +160,7 @@ ALL_TESTS = [ "//pkg/cmd/roachtest/registry:registry_test", "//pkg/cmd/roachtest/roachtestflags:roachtestflags_test", "//pkg/cmd/roachtest/roachtestutil/mixedversion:mixedversion_test", + "//pkg/cmd/roachtest/roachtestutil/task:task_test", "//pkg/cmd/roachtest/roachtestutil:roachtestutil_test", "//pkg/cmd/roachtest/spec:spec_test", "//pkg/cmd/roachtest/tests:tests_test", @@ -1218,6 +1219,8 @@ GO_TARGETS = [ "//pkg/cmd/roachtest/roachtestutil/mixedversion:mixedversion", "//pkg/cmd/roachtest/roachtestutil/mixedversion:mixedversion_test", "//pkg/cmd/roachtest/roachtestutil/operations:operations", + "//pkg/cmd/roachtest/roachtestutil/task:task", + "//pkg/cmd/roachtest/roachtestutil/task:task_test", "//pkg/cmd/roachtest/roachtestutil:roachtestutil", "//pkg/cmd/roachtest/roachtestutil:roachtestutil_test", "//pkg/cmd/roachtest/spec:spec", diff --git a/pkg/cmd/roachtest/roachtestutil/task/BUILD.bazel b/pkg/cmd/roachtest/roachtestutil/task/BUILD.bazel new file mode 100644 index 000000000000..ceb09193ef8a --- /dev/null +++ b/pkg/cmd/roachtest/roachtestutil/task/BUILD.bazel @@ -0,0 +1,29 @@ +load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test") + +go_library( + name = "task", + srcs = [ + "manager.go", + "options.go", + "tasker.go", + "utils.go", + ], + importpath = "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/roachtestutil/task", + visibility = ["//visibility:public"], + deps = [ + "//pkg/roachprod/logger", + "//pkg/util/ctxgroup", + ], +) + +go_test( + name = "task_test", + srcs = ["manager_test.go"], + embed = [":task"], + deps = [ + "//pkg/roachprod/logger", + "//pkg/util/leaktest", + "@com_github_cockroachdb_errors//:errors", + "@com_github_stretchr_testify//require", + ], +) diff --git a/pkg/cmd/roachtest/roachtestutil/task/options.go b/pkg/cmd/roachtest/roachtestutil/task/options.go new file mode 100644 index 000000000000..b5bf07a9bded --- /dev/null +++ b/pkg/cmd/roachtest/roachtestutil/task/options.go @@ -0,0 +1,91 @@ +// Copyright 2024 The Cockroach Authors. +// +// Use of this software is governed by the CockroachDB Software License +// included in the /LICENSE file. + +package task + +import ( + "context" + + "github.com/cockroachdb/cockroach/pkg/roachprod/logger" +) + +type ( + // Options is a struct that contains the options that can be passed when + // starting a task. + Options struct { + Name string + L func() (*logger.Logger, error) + PanicHandler PanicHandlerFunc + ErrorHandler ErrorHandlerFunc + } + + // PanicHandlerFunc is a function that handles panics. If a panic is recovered + // during the execution of a task, the panic handler is called with the + // recovered value. The function has the option to either return an error or + // panic again. + PanicHandlerFunc func(context.Context, *logger.Logger, interface{}) error + + // ErrorHandlerFunc is a function that handles errors. If an error is returned + // from the execution of a task, or the task's panic handler, the error + // handler is called with the error. The error can be augmented or transformed + // by the error handler. + ErrorHandlerFunc func(context.Context, *logger.Logger, error) error +) + +type Option func(result *Options) + +// Name is an option that sets the name of the task. +func Name(name string) Option { + return func(result *Options) { + result.Name = name + } +} + +// LoggerFunc is an option that sets the logger function that will provide the +// task with a logger. Use Logger to provide a logger directly. +func LoggerFunc(loggerFn func() (*logger.Logger, error)) Option { + return func(result *Options) { + result.L = loggerFn + } +} + +// Logger is an option that sets the logger that will be used by the task. +func Logger(l *logger.Logger) Option { + return func(result *Options) { + result.L = func() (*logger.Logger, error) { + return l, nil + } + } +} + +// PanicHandler is an option that sets the panic handler that will be used by the task. +func PanicHandler(handler PanicHandlerFunc) Option { + return func(result *Options) { + result.PanicHandler = handler + } +} + +// ErrorHandler is an option that sets the error handler that will be used by the task. +func ErrorHandler(handler ErrorHandlerFunc) Option { + return func(result *Options) { + result.ErrorHandler = handler + } +} + +func OptionList(opts ...Option) Option { + return func(result *Options) { + for _, opt := range opts { + opt(result) + } + } +} + +func CombineOptions(opts ...Option) Options { + result := Options{} + for _, opt := range opts { + opt(&result) + } + return result +} diff --git a/pkg/cmd/roachtest/roachtestutil/task/tasker.go b/pkg/cmd/roachtest/roachtestutil/task/tasker.go new file mode 100644 index 000000000000..e5cf6c67311e --- /dev/null +++ b/pkg/cmd/roachtest/roachtestutil/task/tasker.go @@ -0,0 +1,24 @@ +// Copyright 2024 The Cockroach Authors. +// +// Use of this software is governed by the CockroachDB Software License +// included in the /LICENSE file. + +package task + +import ( + "context" + + "github.com/cockroachdb/cockroach/pkg/roachprod/logger" +) + +type Func func(context.Context, *logger.Logger) error + +// Tasker is an interface for executing tasks (goroutines). It is intended for +// use in tests, enabling the test framework to manage panics and errors. +type Tasker interface { + // Go runs the given function in a goroutine. + Go(fn Func, opts ...Option) + // GoWithCancel runs the given function in a goroutine and returns a + // CancelFunc that can be used to cancel the function. + GoWithCancel(fn Func, opts ...Option) context.CancelFunc +}