Skip to content

Commit

Permalink
add test impl sol
Browse files Browse the repository at this point in the history
  • Loading branch information
herkolategan committed Oct 24, 2024
1 parent da04eaf commit 669f912
Show file tree
Hide file tree
Showing 6 changed files with 66 additions and 0 deletions.
2 changes: 2 additions & 0 deletions pkg/cmd/roachtest/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down
10 changes: 10 additions & 0 deletions pkg/cmd/roachtest/cluster_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,13 +6,15 @@
package main

import (
"context"
"fmt"
"strconv"
"strings"
"testing"
"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"
Expand Down Expand Up @@ -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.
Expand Down
1 change: 1 addition & 0 deletions pkg/cmd/roachtest/test/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -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",
],
Expand Down
6 changes: 6 additions & 0 deletions pkg/cmd/roachtest/test/test_interface.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)
Expand Down Expand Up @@ -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
Expand Down
21 changes: 21 additions & 0 deletions pkg/cmd/roachtest/test_impl.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 <value> in a key='<value>' attribute
// in TeamCity build output marker.
// See https://www.jetbrains.com/help/teamcity/2023.05/service-messages.html#Escaped+Values
Expand Down
26 changes: 26 additions & 0 deletions pkg/cmd/roachtest/test_runner.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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
}

Expand Down

0 comments on commit 669f912

Please sign in to comment.