Repository: pingcap/tidb
File: pkg/server/tests/commontest/tidb_test.go (Go)

Code before:

	// We have these possible paths to retry:
	// 1. Auto-commit retry
	ts.RunTests(t, nil, func(dbt *testkit.DBTestKit) {
		db := dbt.GetDB()

		_, err := db.Exec("DROP TABLE IF EXISTS auto_retry_test")
		require.NoError(t, err)
		_, err = db.Exec("CREATE TABLE auto_retry_test (id INT PRIMARY KEY, val INT)")
		require.NoError(t, err)
		_, err = db.Exec("INSERT INTO auto_retry_test VALUES (1, 0)")
		require.NoError(t, err)

		// a big enough concurrency to trigger retries
		concurrency := 500
		// Usually the following retry-loop will succeed in the first try. However, if we are lucky
		// enough, it might need more times to trigger the retry.
		require.Eventually(t, func() bool {
			testResults = testResults[:0]
			var wg sync.WaitGroup
			for range concurrency {
				wg.Add(1)
				conn, err := db.Conn(context.Background())
				require.NoError(t, err)
				go func() {
					defer wg.Done()
					_, err := conn.QueryContext(context.Background(), "UPDATE auto_retry_test SET val = val + 1 WHERE id = 1")
					require.NoError(t, err)
				}()
			}
			wg.Wait()

			return len(testResults) > concurrency
		}, time.Second*10, time.Millisecond*100)

		nonRetryingCount := 0
		for _, res := range testResults {
			if !res.retrying {
				nonRetryingCount++
			}
		}
		require.Equal(t, concurrency, nonRetryingCount)
	})

	runExplicitTransactionRetry := func(db *sql.DB, isOptimistic bool) {
		_, err := db.Exec("DROP TABLE IF EXISTS retry_test")
		require.NoError(t, err)
		_, err = db.Exec("CREATE TABLE retry_test (id INT PRIMARY KEY, val INT)")
		require.NoError(t, err)
		_, err = db.Exec("INSERT INTO retry_test VALUES (1, 0)")
		require.NoError(t, err)

Review comment (suggestion):
note that we can use https://pkg.go.dev/sync#WaitGroup.Go now. Don't need to write wg.Add, go func(), defer wg.Done()

Code after:

	// We have these possible paths to retry:
	// 1. Auto-commit retry
	ts.RunTests(t, nil, func(dbt *testkit.DBTestKit) {
		db := dbt.GetDB()

		_, err := db.Exec("DROP TABLE IF EXISTS auto_retry_test")
		require.NoError(t, err)
		_, err = db.Exec("CREATE TABLE auto_retry_test (id INT PRIMARY KEY, val INT)")
		require.NoError(t, err)
		_, err = db.Exec("INSERT INTO auto_retry_test VALUES (1, 0)")
		require.NoError(t, err)

		// a big enough concurrency to trigger retries
		concurrency := 500
		// Usually the following retry-loop will succeed in the first try. However, if we are lucky
		// enough, it might need more times to trigger the retry.
		require.Eventually(t, func() bool {
			testResults = testResults[:0]
			var wg sync.WaitGroup
			for range concurrency {
				conn, err := db.Conn(context.Background())
				require.NoError(t, err)
				wg.Go(func() {
					_, err := conn.QueryContext(context.Background(), "UPDATE auto_retry_test SET val = val + 1 WHERE id = 1")
					require.NoError(t, err)
				})
			}
			wg.Wait()

			return len(testResults) > concurrency
		}, time.Second*10, time.Millisecond*100)

		nonRetryingCount := 0
		for _, res := range testResults {
			if !res.retrying {
				nonRetryingCount++
			}
		}
		require.Equal(t, concurrency, nonRetryingCount)
	})

	runExplicitTransactionRetry := func(db *sql.DB, isOptimistic bool) {
		_, err := db.Exec("DROP TABLE IF EXISTS retry_test")
		require.NoError(t, err)
		_, err = db.Exec("CREATE TABLE retry_test (id INT PRIMARY KEY, val INT)")
		require.NoError(t, err)
		_, err = db.Exec("INSERT INTO retry_test VALUES (1, 0)")
		require.NoError(t, err)

		step1T1Started := make(chan struct{})