|
| 1 | +// Copyright 2023 The Cockroach Authors. |
| 2 | +// |
| 3 | +// Use of this software is governed by the Business Source License |
| 4 | +// included in the file licenses/BSL.txt. |
| 5 | +// |
| 6 | +// As of the Change Date specified in that file, in accordance with |
| 7 | +// the Business Source License, use of this software will be governed |
| 8 | +// by the Apache License, Version 2.0, included in the file |
| 9 | +// licenses/APL.txt. |
| 10 | + |
| 11 | +package mixedversion |
| 12 | + |
| 13 | +import ( |
| 14 | + "context" |
| 15 | + gosql "database/sql" |
| 16 | + "fmt" |
| 17 | + "math/rand" |
| 18 | + "path" |
| 19 | + "strings" |
| 20 | + "sync/atomic" |
| 21 | + |
| 22 | + "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/option" |
| 23 | + "github.com/cockroachdb/cockroach/pkg/roachprod/logger" |
| 24 | +) |
| 25 | + |
| 26 | +func (h *Helper) RandomNode(prng *rand.Rand, nodes option.NodeListOption) int { |
| 27 | + return nodes[prng.Intn(len(nodes))] |
| 28 | +} |
| 29 | + |
| 30 | +// RandomDB returns a (nodeID, connection) tuple for a randomly picked |
| 31 | +// cockroach node according to the parameters passed. |
| 32 | +func (h *Helper) RandomDB(prng *rand.Rand, nodes option.NodeListOption) (int, *gosql.DB) { |
| 33 | + node := h.RandomNode(prng, nodes) |
| 34 | + return node, h.Connect(node) |
| 35 | +} |
| 36 | + |
| 37 | +// QueryRow performs `db.QueryRowContext` on a randomly picked |
| 38 | +// database node. The query and the node picked are logged in the logs |
| 39 | +// of the step that calls this function. |
| 40 | +func (h *Helper) QueryRow(rng *rand.Rand, query string, args ...interface{}) *gosql.Row { |
| 41 | + node, db := h.RandomDB(rng, h.runner.crdbNodes) |
| 42 | + h.stepLogger.Printf("running SQL statement:\n%s\nArgs: %v\nNode: %d", query, args, node) |
| 43 | + return db.QueryRowContext(h.ctx, query, args...) |
| 44 | +} |
| 45 | + |
| 46 | +// Exec performs `db.ExecContext` on a randomly picked database node. |
| 47 | +// The query and the node picked are logged in the logs of the step |
| 48 | +// that calls this function. |
| 49 | +func (h *Helper) Exec(rng *rand.Rand, query string, args ...interface{}) error { |
| 50 | + node, db := h.RandomDB(rng, h.runner.crdbNodes) |
| 51 | + h.stepLogger.Printf("running SQL statement:\n%s\nArgs: %v\nNode: %d", query, args, node) |
| 52 | + _, err := db.ExecContext(h.ctx, query, args...) |
| 53 | + return err |
| 54 | +} |
| 55 | + |
| 56 | +func (h *Helper) Connect(node int) *gosql.DB { |
| 57 | + return h.runner.conn(node) |
| 58 | +} |
| 59 | + |
| 60 | +// SetContext should be called by steps that need access to the test |
| 61 | +// context, as that is only visible to them. |
| 62 | +func (h *Helper) SetContext(c *Context) { |
| 63 | + h.testContext = c |
| 64 | +} |
| 65 | + |
| 66 | +// Context returns the test context associated with a certain step. It |
| 67 | +// is made available for user-functions (see runHookStep). |
| 68 | +func (h *Helper) Context() *Context { |
| 69 | + return h.testContext |
| 70 | +} |
| 71 | + |
| 72 | +// Background allows test authors to create functions that run in the |
| 73 | +// background in mixed-version hooks. |
| 74 | +func (h *Helper) Background( |
| 75 | + name string, fn func(context.Context, *logger.Logger) error, |
| 76 | +) context.CancelFunc { |
| 77 | + return h.runner.background.Start(name, func(ctx context.Context) error { |
| 78 | + bgLogger, err := h.loggerFor(name) |
| 79 | + if err != nil { |
| 80 | + return fmt.Errorf("failed to create logger for background function %q: %w", name, err) |
| 81 | + } |
| 82 | + |
| 83 | + err = panicAsError(bgLogger, func() error { return fn(ctx, bgLogger) }) |
| 84 | + if err != nil { |
| 85 | + if isContextCanceled(ctx) { |
| 86 | + return err |
| 87 | + } |
| 88 | + |
| 89 | + desc := fmt.Sprintf("error in background function %s: %s", name, err) |
| 90 | + return h.runner.testFailure(desc, bgLogger) |
| 91 | + } |
| 92 | + |
| 93 | + return nil |
| 94 | + }) |
| 95 | +} |
| 96 | + |
| 97 | +// BackgroundCommand has the same semantics of `Background()`; the |
| 98 | +// command passed will run and the test will fail if the command is |
| 99 | +// not successful. |
| 100 | +func (h *Helper) BackgroundCommand(cmd string, nodes option.NodeListOption) context.CancelFunc { |
| 101 | + desc := fmt.Sprintf("run command: %q", cmd) |
| 102 | + return h.Background(desc, func(ctx context.Context, l *logger.Logger) error { |
| 103 | + l.Printf("running command `%s` on nodes %v in the background", cmd, nodes) |
| 104 | + return h.runner.cluster.RunE(ctx, nodes, cmd) |
| 105 | + }) |
| 106 | +} |
| 107 | + |
| 108 | +// ExpectDeath alerts the testing infrastructure that a node is |
| 109 | +// expected to die. Regular restarts as part of the mixedversion |
| 110 | +// testing are already taken into account. This function should only |
| 111 | +// be used by tests that perform their own node restarts or chaos |
| 112 | +// events. |
| 113 | +func (h *Helper) ExpectDeath() { |
| 114 | + h.ExpectDeaths(1) |
| 115 | +} |
| 116 | + |
| 117 | +// ExpectDeaths is the general version of `ExpectDeath()`. |
| 118 | +func (h *Helper) ExpectDeaths(n int) { |
| 119 | + h.runner.monitor.ExpectDeaths(n) |
| 120 | +} |
| 121 | + |
| 122 | +// loggerFor creates a logger instance to be used by background |
| 123 | +// functions (created by calling `Background` on the helper |
| 124 | +// instance). It is similar to the logger instances created for |
| 125 | +// mixed-version steps, but with the `background_` prefix. |
| 126 | +func (h *Helper) loggerFor(name string) (*logger.Logger, error) { |
| 127 | + atomic.AddInt64(&h.bgCount, 1) |
| 128 | + |
| 129 | + fileName := invalidChars.ReplaceAllString(strings.ToLower(name), "") |
| 130 | + fileName = fmt.Sprintf("background_%s_%d", fileName, h.bgCount) |
| 131 | + fileName = path.Join(logPrefix, fileName) |
| 132 | + |
| 133 | + return prefixedLogger(h.runner.logger, fileName) |
| 134 | +} |
0 commit comments