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 +}