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

owner/manager.go

@@ -154,7 +154,9 @@ func (m *ownerManager) CampaignOwner() error {
 		return errors.Trace(err)
 	}
 	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
 }
 
@@ -194,9 +196,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 {

owner/manager_test.go

@@ -36,6 +36,23 @@ import (
 
 const testLease = 5 * time.Millisecond
 
+type testInfo struct {
+	cluster *integration.ClusterV3
+	client  *clientv3.Client
+}
+
+func newTestInfo(t *testing.T) *testInfo {
+	cluster := integration.NewClusterV3(t, &integration.ClusterConfig{Size: 1})
+	return &testInfo{
+		cluster: cluster,
+		client:  cluster.Client(0),
+	}
+}
+
+func (ti *testInfo) Close(t *testing.T) {
+	ti.cluster.Terminate(t)
+}
+
 func TestSingle(t *testing.T) {
 	if runtime.GOOS == "windows" {
 		t.Skip("integration.NewClusterV3 will create file contains a colon which is not allowed on Windows")
@@ -213,3 +230,20 @@ func deleteLeader(cli *clientv3.Client, prefixKey string) error {
 	_, err = cli.Delete(context.Background(), string(resp.Kvs[0].Key))
 	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)
+	defer tInfo.Close(t)
+	ownerMgr := owner.NewOwnerManager(context.Background(), tInfo.client, "ddl", "1", "/owner/key")
+	defer ownerMgr.Cancel()
+	for i := 0; i < 10; i++ {
+		err := ownerMgr.CampaignOwner()
+		require.NoError(t, err)
+		ownerMgr.CampaignCancel()
+	}
+}