[release-16.0] Port two flaky test fixes #12603 and #12546

go/mysql/fakesqldb/server.go

@@ -25,6 +25,7 @@ import (
 	"regexp"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -67,7 +68,7 @@ type DB struct {
 	acceptWG sync.WaitGroup
 
 	// orderMatters is set when the query order matters.
-	orderMatters bool
+	orderMatters atomic.Bool
 
 	// Fields set at runtime.
 
@@ -77,16 +78,16 @@ type DB struct {
 	// Use SetName() to change.
 	name string
 	// isConnFail trigger a panic in the connection handler.
-	isConnFail bool
+	isConnFail atomic.Bool
 	// connDelay causes a sleep in the connection handler
 	connDelay time.Duration
 	// shouldClose, if true, tells ComQuery() to close the connection when
 	// processing the next query. This will trigger a MySQL client error with
 	// errno 2013 ("server lost").
-	shouldClose bool
-	// AllowAll: if set to true, ComQuery returns an empty result
+	shouldClose atomic.Bool
+	// allowAll: if set to true, ComQuery returns an empty result
 	// for all queries. This flag is used for benchmarking.
-	AllowAll bool
+	allowAll atomic.Bool
 
 	// Handler: interface that allows a caller to override the query handling
 	// implementation. By default it points to the DB itself
@@ -122,7 +123,7 @@ type DB struct {
 
 	// if fakesqldb is asked to serve queries or query patterns that it has not been explicitly told about it will
 	// error out by default. However if you set this flag then any unmatched query results in an empty result
-	neverFail bool
+	neverFail atomic.Bool
 }
 
 // QueryHandler is the interface used by the DB to simulate executed queries
@@ -216,12 +217,8 @@ func (db *DB) SetName(name string) *DB {
 }
 
 // OrderMatters sets the orderMatters flag.
-func (db *DB) OrderMatters() *DB {
-	db.mu.Lock()
-	defer db.mu.Unlock()
-
-	db.orderMatters = true
-	return db
+func (db *DB) OrderMatters() {
+	db.orderMatters.Store(true)
 }
 
 // Close closes the Listener and waits for it to stop accepting.
@@ -309,7 +306,7 @@ func (db *DB) NewConnection(c *mysql.Conn) {
 	db.mu.Lock()
 	defer db.mu.Unlock()
 
-	if db.isConnFail {
+	if db.isConnFail.Load() {
 		panic(fmt.Errorf("simulating a connection failure"))
 	}
 
@@ -346,11 +343,11 @@ func (db *DB) WarningCount(c *mysql.Conn) uint16 {
 
 // HandleQuery is the default implementation of the QueryHandler interface
 func (db *DB) HandleQuery(c *mysql.Conn, query string, callback func(*sqltypes.Result) error) error {
-	if db.AllowAll {
+	if db.allowAll.Load() {
 		return callback(&sqltypes.Result{})
 	}
 
-	if db.orderMatters {
+	if db.orderMatters.Load() {
 		result, err := db.comQueryOrdered(query)
 		if err != nil {
 			return err
@@ -363,7 +360,7 @@ func (db *DB) HandleQuery(c *mysql.Conn, query string, callback func(*sqltypes.R
 	db.queryCalled[key]++
 	db.querylog = append(db.querylog, key)
 	// Check if we should close the connection and provoke errno 2013.
-	if db.shouldClose {
+	if db.shouldClose.Load() {
 		c.Close()
 
 		//log error
@@ -412,7 +409,7 @@ func (db *DB) HandleQuery(c *mysql.Conn, query string, callback func(*sqltypes.R
 		}
 	}
 
-	if db.neverFail {
+	if db.neverFail.Load() {
 		return callback(&sqltypes.Result{})
 	}
 	// Nothing matched.
@@ -450,7 +447,7 @@ func (db *DB) comQueryOrdered(query string) (*sqltypes.Result, error) {
 	index := db.expectedExecuteFetchIndex
 
 	if index >= len(db.expectedExecuteFetch) {
-		if db.neverFail {
+		if db.neverFail.Load() {
 			return &sqltypes.Result{}, nil
 		}
 		db.t.Errorf("%v: got unexpected out of bound fetch: %v >= %v", db.name, index, len(db.expectedExecuteFetch))
@@ -465,15 +462,15 @@ func (db *DB) comQueryOrdered(query string) (*sqltypes.Result, error) {
 
 	if strings.HasSuffix(expected, "*") {
 		if !strings.HasPrefix(query, expected[0:len(expected)-1]) {
-			if db.neverFail {
+			if db.neverFail.Load() {
 				return &sqltypes.Result{}, nil
 			}
 			db.t.Errorf("%v: got unexpected query start (index=%v): %v != %v", db.name, index, query, expected)
 			return nil, errors.New("unexpected query")
 		}
 	} else {
 		if query != expected {
-			if db.neverFail {
+			if db.neverFail.Load() {
 				return &sqltypes.Result{}, nil
 			}
 			db.t.Errorf("%v: got unexpected query (index=%v): %v != %v", db.name, index, query, expected)
@@ -629,16 +626,12 @@ func (db *DB) ResetQueryLog() {
 
 // EnableConnFail makes connection to this fake DB fail.
 func (db *DB) EnableConnFail() {
-	db.mu.Lock()
-	defer db.mu.Unlock()
-	db.isConnFail = true
+	db.isConnFail.Store(true)
 }
 
 // DisableConnFail makes connection to this fake DB success.
 func (db *DB) DisableConnFail() {
-	db.mu.Lock()
-	defer db.mu.Unlock()
-	db.isConnFail = false
+	db.isConnFail.Store(false)
 }
 
 // SetConnDelay delays connections to this fake DB for the given duration
@@ -650,9 +643,7 @@ func (db *DB) SetConnDelay(d time.Duration) {
 
 // EnableShouldClose closes the connection when processing the next query.
 func (db *DB) EnableShouldClose() {
-	db.mu.Lock()
-	defer db.mu.Unlock()
-	db.shouldClose = true
+	db.shouldClose.Store(true)
 }
 
 //
@@ -757,8 +748,12 @@ func (db *DB) VerifyAllExecutedOrFail() {
 	}
 }
 
+func (db *DB) SetAllowAll(allowAll bool) {
+	db.allowAll.Store(allowAll)
+}
+
 func (db *DB) SetNeverFail(neverFail bool) {
-	db.neverFail = neverFail
+	db.neverFail.Store(neverFail)
 }
 
 func (db *DB) MockQueriesForTable(table string, result *sqltypes.Result) {

go/vt/log/log.go

@@ -22,6 +22,10 @@ limitations under the License.
 package log
 
 import (
+	"fmt"
+	"strconv"
+	"sync/atomic"
+
 	"github.com/golang/glog"
 	"github.com/spf13/pflag"
 )
@@ -78,5 +82,32 @@ var (
 // calls this function, or call this function directly before parsing
 // command-line arguments.
 func RegisterFlags(fs *pflag.FlagSet) {
-	fs.Uint64Var(&glog.MaxSize, "log_rotate_max_size", glog.MaxSize, "size in bytes at which logs are rotated (glog.MaxSize)")
+	flagVal := logRotateMaxSize{
+		val: fmt.Sprintf("%d", atomic.LoadUint64(&glog.MaxSize)),
+	}
+	fs.Var(&flagVal, "log_rotate_max_size", "size in bytes at which logs are rotated (glog.MaxSize)")
+}
+
+// logRotateMaxSize implements pflag.Value and is used to
+// try and provide thread-safe access to glog.MaxSize.
+type logRotateMaxSize struct {
+	val string
+}
+
+func (lrms *logRotateMaxSize) Set(s string) error {
+	maxSize, err := strconv.ParseUint(s, 10, 64)
+	if err != nil {
+		return err
+	}
+	atomic.StoreUint64(&glog.MaxSize, maxSize)
+	lrms.val = s
+	return nil
+}
+
+func (lrms *logRotateMaxSize) String() string {
+	return lrms.val
+}
+
+func (lrms *logRotateMaxSize) Type() string {
+	return "uint64"
 }

go/vt/vtctl/vdiff2_test.go

@@ -35,7 +35,7 @@ var (
 )
 
 func TestVDiff2Unsharded(t *testing.T) {
-	env := newTestVDiffEnv([]string{"0"}, []string{"0"}, "", nil)
+	env := newTestVDiffEnv(t, []string{"0"}, []string{"0"}, "", nil)
 	defer env.close()
 
 	UUID := uuid.New().String()
@@ -275,7 +275,7 @@ func TestVDiff2Unsharded(t *testing.T) {
 }
 
 func TestVDiff2Sharded(t *testing.T) {
-	env := newTestVDiffEnv([]string{"-40", "40-"}, []string{"-80", "80-"}, "", map[string]string{
+	env := newTestVDiffEnv(t, []string{"-40", "40-"}, []string{"-80", "80-"}, "", map[string]string{
 		"-80": "MySQL56/0e45e704-7cb9-11ed-a1eb-0242ac120002:1-890",
 		"80-": "MySQL56/1497ddb0-7cb9-11ed-a1eb-0242ac120002:1-891",
 	})

go/vt/vtctl/vdiff_env_test.go

@@ -19,7 +19,9 @@ package vtctl
 import (
 	"context"
 	"fmt"
+	"math/rand"
 	"sync"
+	"testing"
 
 	"vitess.io/vitess/go/sqltypes"
 	"vitess.io/vitess/go/vt/grpcclient"
@@ -64,25 +66,10 @@ type testVDiffEnv struct {
 	tablets map[int]*testVDiffTablet
 }
 
-// vdiffEnv has to be a global for RegisterDialer to work.
-var vdiffEnv *testVDiffEnv
-
-func init() {
-	tabletconn.RegisterDialer("VDiffTest", func(tablet *topodatapb.Tablet, failFast grpcclient.FailFast) (queryservice.QueryService, error) {
-		vdiffEnv.mu.Lock()
-		defer vdiffEnv.mu.Unlock()
-		if qs, ok := vdiffEnv.tablets[int(tablet.Alias.Uid)]; ok {
-			return qs, nil
-		}
-		return nil, fmt.Errorf("tablet %d not found", tablet.Alias.Uid)
-	})
-}
-
 //----------------------------------------------
 // testVDiffEnv
 
-func newTestVDiffEnv(sourceShards, targetShards []string, query string, positions map[string]string) *testVDiffEnv {
-	tabletconntest.SetProtocol("go.vt.vtctl.vdiff_env_test", "VDiffTest")
+func newTestVDiffEnv(t testing.TB, sourceShards, targetShards []string, query string, positions map[string]string) *testVDiffEnv {
 	env := &testVDiffEnv{
 		workflow:   "vdiffTest",
 		tablets:    make(map[int]*testVDiffTablet),
@@ -94,6 +81,18 @@ func newTestVDiffEnv(sourceShards, targetShards []string, query string, position
 	}
 	env.wr = wrangler.NewTestWrangler(env.cmdlog, env.topoServ, env.tmc)
 
+	// Generate a unique dialer name.
+	dialerName := fmt.Sprintf("VDiffTest-%s-%d", t.Name(), rand.Intn(1000000000))
+	tabletconn.RegisterDialer(dialerName, func(tablet *topodatapb.Tablet, failFast grpcclient.FailFast) (queryservice.QueryService, error) {
+		env.mu.Lock()
+		defer env.mu.Unlock()
+		if qs, ok := env.tablets[int(tablet.Alias.Uid)]; ok {
+			return qs, nil
+		}
+		return nil, fmt.Errorf("tablet %d not found", tablet.Alias.Uid)
+	})
+	tabletconntest.SetProtocol("go.vt.vtctl.vdiff_env_test", dialerName)
+
 	tabletID := 100
 	for _, shard := range sourceShards {
 		_ = env.addTablet(tabletID, "source", shard, topodatapb.TabletType_PRIMARY)

go/vt/vttablet/tabletserver/bench_test.go

@@ -68,7 +68,7 @@ func BenchmarkExecuteVarBinary(b *testing.B) {
 	}
 
 	target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY}
-	db.AllowAll = true
+	db.SetAllowAll(true)
 	for i := 0; i < b.N; i++ {
 		if _, err := tsv.Execute(context.Background(), &target, benchQuery, bv, 0, 0, nil); err != nil {
 			panic(err)
@@ -93,7 +93,7 @@ func BenchmarkExecuteExpression(b *testing.B) {
 	}
 
 	target := querypb.Target{TabletType: topodatapb.TabletType_PRIMARY}
-	db.AllowAll = true
+	db.SetAllowAll(true)
 	for i := 0; i < b.N; i++ {
 		if _, err := tsv.Execute(context.Background(), &target, benchQuery, bv, 0, 0, nil); err != nil {
 			panic(err)

go/vt/wrangler/vdiff_env_test.go

@@ -19,7 +19,9 @@ package wrangler
 import (
 	"context"
 	"fmt"
+	"math/rand"
 	"sync"
+	"testing"
 
 	"vitess.io/vitess/go/sqltypes"
 	"vitess.io/vitess/go/vt/grpcclient"
@@ -62,25 +64,10 @@ type testVDiffEnv struct {
 	tablets map[int]*testVDiffTablet
 }
 
-// vdiffEnv has to be a global for RegisterDialer to work.
-var vdiffEnv *testVDiffEnv
-
-func init() {
-	tabletconn.RegisterDialer("VDiffTest", func(tablet *topodatapb.Tablet, failFast grpcclient.FailFast) (queryservice.QueryService, error) {
-		vdiffEnv.mu.Lock()
-		defer vdiffEnv.mu.Unlock()
-		if qs, ok := vdiffEnv.tablets[int(tablet.Alias.Uid)]; ok {
-			return qs, nil
-		}
-		return nil, fmt.Errorf("tablet %d not found", tablet.Alias.Uid)
-	})
-}
-
 //----------------------------------------------
 // testVDiffEnv
 
-func newTestVDiffEnv(sourceShards, targetShards []string, query string, positions map[string]string) *testVDiffEnv {
-	tabletconntest.SetProtocol("go.vt.wrangler.vdiff_env_test", "VDiffTest")
+func newTestVDiffEnv(t testing.TB, sourceShards, targetShards []string, query string, positions map[string]string) *testVDiffEnv {
 	env := &testVDiffEnv{
 		workflow:   "vdiffTest",
 		tablets:    make(map[int]*testVDiffTablet),
@@ -91,6 +78,18 @@ func newTestVDiffEnv(sourceShards, targetShards []string, query string, position
 	}
 	env.wr = New(logutil.NewConsoleLogger(), env.topoServ, env.tmc)
 
+	// Generate a unique dialer name.
+	dialerName := fmt.Sprintf("VDiffTest-%s-%d", t.Name(), rand.Intn(1000000000))
+	tabletconn.RegisterDialer(dialerName, func(tablet *topodatapb.Tablet, failFast grpcclient.FailFast) (queryservice.QueryService, error) {
+		env.mu.Lock()
+		defer env.mu.Unlock()
+		if qs, ok := env.tablets[int(tablet.Alias.Uid)]; ok {
+			return qs, nil
+		}
+		return nil, fmt.Errorf("tablet %d not found", tablet.Alias.Uid)
+	})
+	tabletconntest.SetProtocol("go.vt.wrangler.vdiff_env_test", dialerName)
+
 	tabletID := 100
 	for _, shard := range sourceShards {
 		_ = env.addTablet(tabletID, "source", shard, topodatapb.TabletType_PRIMARY)
@@ -167,7 +166,6 @@ func newTestVDiffEnv(sourceShards, targetShards []string, query string, position
 
 		tabletID += 10
 	}
-	vdiffEnv = env
 	return env
 }