br: fix flaky test br_txn (#56444)

close #54979
pingcap · Oct 11, 2024 · f773b6e · f773b6e
1 parent 93cad31
commit f773b6e
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 38 deletions.
diff --git a/br/pkg/rtree/BUILD.bazel b/br/pkg/rtree/BUILD.bazel
@@ -10,7 +10,6 @@ go_library(
     visibility = ["//visibility:public"],
     deps = [
         "//br/pkg/logutil",
-        "//pkg/kv",
         "//pkg/tablecodec",
         "//pkg/util/redact",
         "@com_github_google_btree//:btree",

diff --git a/br/pkg/rtree/rtree.go b/br/pkg/rtree/rtree.go
@@ -9,7 +9,6 @@ import (
 	backuppb "github.com/pingcap/kvproto/pkg/brpb"
 	"github.com/pingcap/log"
 	"github.com/pingcap/tidb/br/pkg/logutil"
-	"github.com/pingcap/tidb/pkg/kv"
 	"github.com/pingcap/tidb/pkg/tablecodec"
 	"github.com/pingcap/tidb/pkg/util/redact"
 	"github.com/pkg/errors"
@@ -153,8 +152,8 @@ func NeedsMerge(left, right *RangeStats, splitSizeBytes, splitKeyCount uint64) b
 	if leftKeys+rightKeys > splitKeyCount {
 		return false
 	}
-	tableID1, indexID1, isRecord1, err1 := tablecodec.DecodeKeyHead(kv.Key(left.StartKey))
-	tableID2, indexID2, isRecord2, err2 := tablecodec.DecodeKeyHead(kv.Key(right.StartKey))
+	tableID1, indexID1, isRecord1, err1 := tablecodec.DecodeKeyHead(left.StartKey)
+	tableID2, indexID2, isRecord2, err2 := tablecodec.DecodeKeyHead(right.StartKey)
 
 	// Failed to decode the file key head... can this happen?
 	if err1 != nil || err2 != nil {

diff --git a/br/pkg/task/restore_txn.go b/br/pkg/task/restore_txn.go
@@ -29,8 +29,7 @@ func RunRestoreTxn(c context.Context, g glue.Glue, cmdName string, cfg *Config)
 	defer cancel()
 
 	// Restore raw does not need domain.
-	needDomain := false
-	mgr, err := NewMgr(ctx, g, cfg.PD, cfg.TLS, GetKeepalive(cfg), cfg.CheckRequirements, needDomain, conn.NormalVersionChecker)
+	mgr, err := NewMgr(ctx, g, cfg.PD, cfg.TLS, GetKeepalive(cfg), cfg.CheckRequirements, false, conn.NormalVersionChecker)
 	if err != nil {
 		return errors.Trace(err)
 	}

diff --git a/br/tests/br_txn/client.go b/br/tests/br_txn/client.go
@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"hash/crc64"
 	"math/rand"
+	"sync"
 	"time"
 
 	"github.com/pingcap/errors"
@@ -78,20 +79,13 @@ func main() {
 
 func randGenWithDuration(client *txnkv.Client, startKey, endKey []byte,
 	maxLen int, concurrency int, duration int) error {
-	var err error
-	ok := make(chan struct{})
-	go func() {
-		err = randGen(client, startKey, endKey, maxLen, concurrency)
-		ok <- struct{}{}
-	}()
-	select {
-	case <-time.After(time.Second * time.Duration(duration)):
-	case <-ok:
-	}
-	return errors.Trace(err)
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*time.Duration(duration))
+	defer cancel()
+
+	return randGen(ctx, client, startKey, endKey, maxLen, concurrency)
 }
 
-func randGen(client *txnkv.Client, startKey, endKey []byte, maxLen int, concurrency int) error {
+func randGen(ctx context.Context, client *txnkv.Client, startKey, endKey []byte, maxLen int, concurrency int) error {
 	log.Info("Start rand-gen", zap.Int("maxlen", maxLen),
 		zap.String("startkey", hex.EncodeToString(startKey)), zap.String("endkey", hex.EncodeToString(endKey)))
 	log.Info("Rand-gen will keep running. Please Ctrl+C to stop manually.")
@@ -108,39 +102,57 @@ func randGen(client *txnkv.Client, startKey, endKey []byte, maxLen int, concurre
 	}
 
 	const batchSize = 32
+	const numBatch = 100
 
 	errCh := make(chan error, concurrency)
+	var wg sync.WaitGroup
+
 	for i := maxLen; i <= maxLen+concurrency; i++ {
+		wg.Add(1)
 		go func(i int) {
-			for {
-				txn, err := client.Begin()
-				if err != nil {
-					errCh <- errors.Trace(err)
-				}
-				for j := 0; j < batchSize; j++ {
-					key := randKey(startKey, endKey, i)
-					// append index to avoid write conflict
-					key = appendIndex(key, i)
-					value := randValue()
-					err = txn.Set(key, value)
+			defer wg.Done()
+			for k := 0; k < numBatch; k++ {
+				select {
+				case <-ctx.Done():
+					return
+				default:
+					txn, err := client.Begin()
+					if err != nil {
+						errCh <- errors.Trace(err)
+					}
+					for j := 0; j < batchSize; j++ {
+						key := randKey(startKey, endKey, i)
+						// append index to avoid write conflict
+						key = appendIndex(key, i)
+						value := randValue()
+						err = txn.Set(key, value)
+						if err != nil {
+							errCh <- errors.Trace(err)
+						}
+					}
+					err = txn.Commit(context.TODO())
 					if err != nil {
 						errCh <- errors.Trace(err)
 					}
-				}
-				err = txn.Commit(context.TODO())
-				if err != nil {
-					errCh <- errors.Trace(err)
 				}
 			}
 		}(i)
 	}
 
-	err := <-errCh
-	if err != nil {
-		return errors.Trace(err)
-	}
+	// Use a separate goroutine to wait for all workers to finish
+	done := make(chan struct{})
+	go func() {
+		wg.Wait()
+		close(done)
+	}()
 
-	return nil
+	select {
+	case <-done:
+		return nil
+	case err := <-errCh:
+		<-done
+		return err
+	}
 }
 
 func testRandKey(startKey, endKey []byte, maxLen int) {