owner: fix data race on ownerManager.campaignCancel (#56362)

close #56053
pingcap · Sep 27, 2024 · afdd5c2 · afdd5c2
1 parent bad2ecd
commit afdd5c2
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 5 deletions.
diff --git a/pkg/owner/BUILD.bazel b/pkg/owner/BUILD.bazel
@@ -37,7 +37,7 @@ go_test(
     ],
     embed = [":owner"],
     flaky = True,
-    shard_count = 8,
+    shard_count = 9,
     deps = [
         "//pkg/ddl",
         "//pkg/infoschema",

diff --git a/pkg/owner/manager.go b/pkg/owner/manager.go
@@ -196,7 +196,9 @@ func (m *ownerManager) CampaignOwner(withTTL ...int) error {
 	}
 	m.sessionLease.Store(int64(session.Lease()))
 	m.wg.Add(1)
-	go m.campaignLoop(session)
+	var campaignContext context.Context
+	campaignContext, m.campaignCancel = context.WithCancel(m.ctx)
+	go m.campaignLoop(campaignContext, session)
 	return nil
 }
 
@@ -241,9 +243,7 @@ func (m *ownerManager) CampaignCancel() {
 	m.wg.Wait()
 }
 
-func (m *ownerManager) campaignLoop(etcdSession *concurrency.Session) {
-	var campaignContext context.Context
-	campaignContext, m.campaignCancel = context.WithCancel(m.ctx)
+func (m *ownerManager) campaignLoop(campaignContext context.Context, etcdSession *concurrency.Session) {
 	defer func() {
 		m.campaignCancel()
 		if r := recover(); r != nil {

diff --git a/pkg/owner/manager_test.go b/pkg/owner/manager_test.go
@@ -439,6 +439,23 @@ func deleteLeader(cli *clientv3.Client, prefixKey string) error {
 	return errors.Trace(err)
 }
 
+func TestImmediatelyCancel(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skip("integration.NewClusterV3 will create file contains a colon which is not allowed on Windows")
+	}
+	integration.BeforeTestExternal(t)
+
+	tInfo := newTestInfo(t)
+	d := tInfo.ddl
+	defer tInfo.Close(t)
+	ownerManager := d.OwnerManager()
+	for i := 0; i < 10; i++ {
+		err := ownerManager.CampaignOwner()
+		require.NoError(t, err)
+		ownerManager.CampaignCancel()
+	}
+}
+
 func TestAcquireDistributedLock(t *testing.T) {
 	const addrFmt = "http://127.0.0.1:%d"
 	cfg := embed.NewConfig()