disttask: fix routine leak & refine logs & refactor failpoint

pkg/disttask/framework/integrationtests/BUILD.bazel

@@ -10,10 +10,11 @@ go_test(
         "framework_role_test.go",
         "framework_rollback_test.go",
         "framework_test.go",
+        "main_test.go",
     ],
     flaky = True,
     race = "off",
-    shard_count = 26,
+    shard_count = 22,
     deps = [
         "//pkg/disttask/framework/handle",
         "//pkg/disttask/framework/proto",
@@ -22,8 +23,10 @@ go_test(
         "//pkg/disttask/framework/taskexecutor",
         "//pkg/disttask/framework/testutil",
         "//pkg/testkit",
+        "//pkg/testkit/testsetup",
         "//pkg/util",
         "@com_github_pingcap_failpoint//:failpoint",
         "@com_github_stretchr_testify//require",
+        "@org_uber_go_goleak//:goleak",
     ],
 )

pkg/disttask/framework/integrationtests/framework_err_handling_test.go

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package integrationtests_test
+package integrationtests
 
 import (
 	"testing"

pkg/disttask/framework/integrationtests/framework_ha_test.go

@@ -12,17 +12,18 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package integrationtests_test
+package integrationtests
 
 import (
 	"context"
-	"sync"
+	"math"
 	"testing"
 
 	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/pkg/disttask/framework/proto"
 	"github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor"
 	"github.com/pingcap/tidb/pkg/disttask/framework/testutil"
+	"github.com/pingcap/tidb/pkg/testkit"
 	"github.com/stretchr/testify/require"
 )
 
@@ -33,78 +34,52 @@ func submitTaskAndCheckSuccessForHA(ctx context.Context, t *testing.T, taskKey s
 	})
 }
 
-func TestHABasic(t *testing.T) {
-	taskexecutor.TestContexts = sync.Map{}
-	ctx, ctrl, testContext, distContext := testutil.InitTestContext(t, 4)
-	defer ctrl.Finish()
-
-	testutil.RegisterTaskMeta(t, ctrl, testutil.GetMockHATestSchedulerExt(ctrl), testContext, nil)
-	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockCleanExecutor", "return()"))
-	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockStopManager", "4*return()"))
-	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockTiDBDown", "return(\":4000\")"))
-	submitTaskAndCheckSuccessForHA(ctx, t, "😊", testContext)
-	require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockTiDBDown"))
-	require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockStopManager"))
-	require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockCleanExecutor"))
-	distContext.Close()
+func TestHANodeRandomShutdown(t *testing.T) {
+	testkit.EnableFailPoint(t, "github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockTiDBShutdown", "return()")
+	c := testutil.NewDXFContextWithRandomNodes(t, 4, 15)
+	testutil.RegisterTaskMeta(t, c.MockCtrl, testutil.GetMockHATestSchedulerExt(c.MockCtrl), c.TestContext, nil)
+
+	// we keep [1, 10] nodes running, as we only have 10 subtask at stepOne
+	keepCount := int(math.Min(float64(c.NodeCount()-1), float64(c.Rand.Intn(10)+1)))
+	nodeNeedDown := c.GetRandNodeIDs(c.NodeCount() - keepCount)
+	t.Logf("started %d nodes, and we keep %d nodes, nodes that need shutdown: %v", c.NodeCount(), keepCount, nodeNeedDown)
+	taskexecutor.MockTiDBDown = func(execID string, _ *proto.Task) bool {
+		if _, ok := nodeNeedDown[execID]; ok {
+			c.AsyncShutdown(execID)
+			return true
+		}
+		return false
+	}
+	submitTaskAndCheckSuccessForHA(c.Ctx, t, "😊", c.TestContext)
 }
 
-func TestHAManyNodes(t *testing.T) {
-	taskexecutor.TestContexts = sync.Map{}
-	ctx, ctrl, testContext, distContext := testutil.InitTestContext(t, 30)
-	defer ctrl.Finish()
-
-	testutil.RegisterTaskMeta(t, ctrl, testutil.GetMockHATestSchedulerExt(ctrl), testContext, nil)
-	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockCleanExecutor", "return()"))
-	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockStopManager", "30*return()"))
-	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockTiDBDown", "return(\":4000\")"))
-	submitTaskAndCheckSuccessForHA(ctx, t, "😊", testContext)
-	require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockTiDBDown"))
-	require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockStopManager"))
-	require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockCleanExecutor"))
-	distContext.Close()
-}
-
-func TestHAFailInDifferentStage(t *testing.T) {
-	taskexecutor.TestContexts = sync.Map{}
-	ctx, ctrl, testContext, distContext := testutil.InitTestContext(t, 6)
-	defer ctrl.Finish()
-
-	testutil.RegisterTaskMeta(t, ctrl, testutil.GetMockHATestSchedulerExt(ctrl), testContext, nil)
-	// stage1 : server num from 6 to 3.
-	// stage2 : server num from 3 to 2.
-	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockCleanExecutor", "return()"))
-	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockStopManager", "6*return()"))
-	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockTiDBDown", "return(\":4000\")"))
-	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockTiDBDown2", "return()"))
-
-	submitTaskAndCheckSuccessForHA(ctx, t, "😊", testContext)
-	require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockTiDBDown"))
-	require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockTiDBDown2"))
-	require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockStopManager"))
-	require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockCleanExecutor"))
-	distContext.Close()
-}
-
-func TestHAFailInDifferentStageManyNodes(t *testing.T) {
-	taskexecutor.TestContexts = sync.Map{}
-	ctx, ctrl, testContext, distContext := testutil.InitTestContext(t, 30)
-	defer ctrl.Finish()
-
-	testutil.RegisterTaskMeta(t, ctrl, testutil.GetMockHATestSchedulerExt(ctrl), testContext, nil)
-	// stage1 : server num from 30 to 27.
-	// stage2 : server num from 27 to 26.
-	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockCleanExecutor", "return()"))
-	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockStopManager", "30*return()"))
-	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockTiDBDown", "return(\":4000\")"))
-	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockTiDBDown2", "return()"))
-
-	submitTaskAndCheckSuccessForHA(ctx, t, "😊", testContext)
-	require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockTiDBDown"))
-	require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockTiDBDown2"))
-	require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockStopManager"))
-	require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockCleanExecutor"))
-	distContext.Close()
+func TestHARandomShutdownInDifferentStep(t *testing.T) {
+	testkit.EnableFailPoint(t, "github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockTiDBShutdown", "return()")
+	c := testutil.NewDXFContextWithRandomNodes(t, 6, 15)
+
+	testutil.RegisterTaskMeta(t, c.MockCtrl, testutil.GetMockHATestSchedulerExt(c.MockCtrl), c.TestContext, nil)
+	// they might overlap, but will leave at least 2 nodes running
+	nodeNeedDownAtStepOne := c.GetRandNodeIDs(c.NodeCount()/2 - 1)
+	nodeNeedDownAtStepTwo := c.GetRandNodeIDs(c.NodeCount()/2 - 1)
+	t.Logf("started %d nodes, shutdown nodes at step 1: %v, shutdown nodes at step 2: %v",
+		c.NodeCount(), nodeNeedDownAtStepOne, nodeNeedDownAtStepTwo)
+	taskexecutor.MockTiDBDown = func(execID string, task *proto.Task) bool {
+		var targetNodes map[string]struct{}
+		switch task.Step {
+		case proto.StepOne:
+			targetNodes = nodeNeedDownAtStepOne
+		case proto.StepTwo:
+			targetNodes = nodeNeedDownAtStepTwo
+		default:
+			return false
+		}
+		if _, ok := targetNodes[execID]; ok {
+			c.AsyncShutdown(execID)
+			return true
+		}
+		return false
+	}
+	submitTaskAndCheckSuccessForHA(c.Ctx, t, "😊", c.TestContext)
 }
 
 func TestHAReplacedButRunning(t *testing.T) {

pkg/disttask/framework/integrationtests/framework_pause_and_resume_test.go

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package integrationtests_test
+package integrationtests
 
 import (
 	"context"

pkg/disttask/framework/integrationtests/framework_role_test.go

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package integrationtests_test
+package integrationtests
 
 import (
 	"context"

pkg/disttask/framework/integrationtests/framework_rollback_test.go

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package integrationtests_test
+package integrationtests
 
 import (
 	"testing"

pkg/disttask/framework/integrationtests/framework_test.go

@@ -12,21 +12,19 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package integrationtests_test
+package integrationtests
 
 import (
 	"context"
 	"fmt"
 	"math/rand"
-	"sync"
 	"testing"
 	"time"
 
 	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/pkg/disttask/framework/proto"
 	"github.com/pingcap/tidb/pkg/disttask/framework/scheduler"
 	"github.com/pingcap/tidb/pkg/disttask/framework/storage"
-	"github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor"
 	"github.com/pingcap/tidb/pkg/disttask/framework/testutil"
 	"github.com/pingcap/tidb/pkg/testkit"
 	"github.com/pingcap/tidb/pkg/util"
@@ -170,43 +168,6 @@ func TestOwnerChangeWhenSchedule(t *testing.T) {
 	require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/disttask/framework/scheduler/mockOwnerChange"))
 }
 
-func TestTaskExecutorDownBasic(t *testing.T) {
-	taskexecutor.TestContexts = sync.Map{}
-	ctx, ctrl, testContext, distContext := testutil.InitTestContext(t, 4)
-	defer ctrl.Finish()
-
-	testutil.RegisterTaskMeta(t, ctrl, testutil.GetMockBasicSchedulerExt(ctrl), testContext, nil)
-	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockCleanExecutor", "return()"))
-	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockStopManager", "4*return(\":4000\")"))
-	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockTiDBDown", "return(\":4000\")"))
-	submitTaskAndCheckSuccess(ctx, t, "😊", testContext, map[proto.Step]int{
-		proto.StepOne: 3,
-		proto.StepTwo: 1,
-	})
-	require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockCleanExecutor"))
-	require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockTiDBDown"))
-	require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockStopManager"))
-	distContext.Close()
-}
-
-func TestTaskExecutorDownManyNodes(t *testing.T) {
-	taskexecutor.TestContexts = sync.Map{}
-	ctx, ctrl, testContext, distContext := testutil.InitTestContext(t, 30)
-	defer ctrl.Finish()
-	testutil.RegisterTaskMeta(t, ctrl, testutil.GetMockBasicSchedulerExt(ctrl), testContext, nil)
-	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockCleanExecutor", "return()"))
-	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockStopManager", "30*return(\":4000\")"))
-	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockTiDBDown", "return(\":4000\")"))
-	submitTaskAndCheckSuccess(ctx, t, "😊", testContext, map[proto.Step]int{
-		proto.StepOne: 3,
-		proto.StepTwo: 1,
-	})
-	require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockCleanExecutor"))
-	require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockTiDBDown"))
-	require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/mockStopManager"))
-	distContext.Close()
-}
-
 func TestGC(t *testing.T) {
 	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/disttask/framework/storage/subtaskHistoryKeepSeconds", "return(1)"))
 	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/disttask/framework/scheduler/historySubtaskTableGcInterval", "return(1)"))

pkg/disttask/framework/integrationtests/main_test.go

@@ -0,0 +1,36 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package integrationtests
+
+import (
+	"testing"
+
+	"github.com/pingcap/tidb/pkg/testkit/testsetup"
+	"go.uber.org/goleak"
+)
+
+func TestMain(m *testing.M) {
+	testsetup.SetupForCommonTest()
+
+	opts := []goleak.Option{
+		goleak.IgnoreTopFunction("github.com/golang/glog.(*fileSink).flushDaemon"),
+		goleak.IgnoreTopFunction("github.com/bazelbuild/rules_go/go/tools/bzltestutil.RegisterTimeoutHandler.func1"),
+		goleak.IgnoreTopFunction("github.com/lestrrat-go/httprc.runFetchWorker"),
+		goleak.IgnoreTopFunction("go.etcd.io/etcd/client/pkg/v3/logutil.(*MergeLogger).outputLoop"),
+		goleak.IgnoreTopFunction("go.opencensus.io/stats/view.(*worker).start"),
+		goleak.IgnoreTopFunction("syscall.syscall"),
+	}
+	goleak.VerifyTestMain(m, opts...)
+}

pkg/disttask/framework/scheduler/scheduler.go

@@ -226,12 +226,10 @@ func (s *BaseScheduler) scheduleTask() {
 				s.logger.Info("schedule task meet err, reschedule it", zap.Error(err))
 			}
 
-			failpoint.Inject("mockOwnerChange", func(val failpoint.Value) {
-				if val.(bool) {
-					s.logger.Info("mockOwnerChange called")
-					MockOwnerChange()
-					time.Sleep(time.Second)
-				}
+			failpoint.Inject("mockOwnerChange", func() {
+				s.logger.Info("mockOwnerChange called")
+				MockOwnerChange()
+				time.Sleep(time.Second)
 			})
 		}
 	}

pkg/disttask/framework/scheduler/scheduler_manager.go

@@ -176,6 +176,12 @@ func (sm *Manager) Start() {
 	sm.initialized = true
 }
 
+// Cancel cancels the scheduler manager.
+// used in test to simulate tidb node shutdown.
+func (sm *Manager) Cancel() {
+	sm.cancel()
+}
+
 // Stop the schedulerManager.
 func (sm *Manager) Stop() {
 	sm.cancel()