disttask: add cleanup routine for framework

Makefile

@@ -385,7 +385,7 @@ mock_lightning: tools/bin/mockgen
 
 gen_mock: tools/bin/mockgen
 	tools/bin/mockgen -package mock github.com/pingcap/tidb/disttask/framework/scheduler TaskTable,Pool,Scheduler,Extension > disttask/framework/mock/scheduler_mock.go
-	tools/bin/mockgen -package mock github.com/pingcap/tidb/disttask/framework/dispatcher Dispatcher > disttask/framework/mock/dispatcher_mock.go
+	tools/bin/mockgen -package mock github.com/pingcap/tidb/disttask/framework/dispatcher Dispatcher,CleanUpRoutine > disttask/framework/mock/dispatcher_mock.go
 	tools/bin/mockgen -package execute github.com/pingcap/tidb/disttask/framework/scheduler/execute SubtaskExecutor > disttask/framework/mock/execute/execute_mock.go
 	tools/bin/mockgen -package mock github.com/pingcap/tidb/disttask/importinto MiniTaskExecutor > disttask/importinto/mock/import_mock.go
 	tools/bin/mockgen -package mock github.com/pingcap/tidb/disttask/framework/planner LogicalPlan,PipelineSpec > disttask/framework/mock/plan_mock.go

ddl/BUILD.bazel

@@ -12,6 +12,7 @@ go_library(
     name = "ddl",
     srcs = [
         "backfilling.go",
+        "backfilling_clean_s3.go",
         "backfilling_dispatcher.go",
         "backfilling_dist_scheduler.go",
         "backfilling_import_cloud.go",

ddl/backfilling_clean_s3.go

@@ -0,0 +1,83 @@
+// Copyright 2023 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 ddl
+
+import (
+	"context"
+	"encoding/json"
+	"strconv"
+
+	"github.com/pingcap/tidb/br/pkg/lightning/backend/external"
+	"github.com/pingcap/tidb/br/pkg/storage"
+	"github.com/pingcap/tidb/disttask/framework/dispatcher"
+	"github.com/pingcap/tidb/disttask/framework/proto"
+	"github.com/pingcap/tidb/parser/ast"
+	"github.com/pingcap/tidb/util/logutil"
+	"go.uber.org/zap"
+)
+
+var _ dispatcher.CleanUpRoutine = (*BackfillCleanUpS3)(nil)
+
+// BackfillCleanUpS3 implements dispatcher.CleanUpRoutine.
+type BackfillCleanUpS3 struct {
+}
+
+func newBackfillCleanUpS3() dispatcher.CleanUpRoutine {
+	return &BackfillCleanUpS3{}
+}
+
+// CleanUp implements the CleanUpRoutine.CleanUp interface.
+func (*BackfillCleanUpS3) CleanUp(ctx context.Context, task *proto.Task) error {
+	var gTaskMeta BackfillGlobalMeta
+	if err := json.Unmarshal(task.Meta, &gTaskMeta); err != nil {
+		return err
+	}
+	// Not use cloud storage, no need to cleanUp.
+	if len(gTaskMeta.CloudStorageURI) == 0 {
+		return nil
+	}
+	backend, err := storage.ParseBackend(gTaskMeta.CloudStorageURI, nil)
+	if err != nil {
+		logutil.Logger(ctx).Warn("failed to parse cloud storage uri", zap.Error(err))
+		return err
+	}
+	extStore, err := storage.NewWithDefaultOpt(ctx, backend)
+	if err != nil {
+		logutil.Logger(ctx).Warn("failed to create cloud storage", zap.Error(err))
+		return err
+	}
+	prefix := strconv.Itoa(int(gTaskMeta.Job.ID))
+	err = external.CleanUpFiles(ctx, extStore, prefix)
+	if err != nil {
+		logutil.Logger(ctx).Warn("cannot cleanup cloud storage files", zap.Error(err))
+		return err
+	}
+	redactCloudStorageURI(ctx, task, &gTaskMeta)
+	return nil
+}
+
+func redactCloudStorageURI(
+	ctx context.Context,
+	gTask *proto.Task,
+	origin *BackfillGlobalMeta,
+) {
+	origin.CloudStorageURI = ast.RedactURL(origin.CloudStorageURI)
+	metaBytes, err := json.Marshal(origin)
+	if err != nil {
+		logutil.Logger(ctx).Warn("failed to marshal task meta", zap.Error(err))
+		return
+	}
+	gTask.Meta = metaBytes
+}

ddl/backfilling_dispatcher.go

@@ -21,7 +21,6 @@ import (
 	"encoding/json"
 	"math"
 	"sort"
-	"strconv"
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/br/pkg/lightning/backend/external"
@@ -34,7 +33,6 @@ import (
 	"github.com/pingcap/tidb/domain/infosync"
 	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/meta"
-	"github.com/pingcap/tidb/parser/ast"
 	"github.com/pingcap/tidb/parser/model"
 	"github.com/pingcap/tidb/store/helper"
 	"github.com/pingcap/tidb/table"
@@ -81,13 +79,6 @@ func (h *backfillingDispatcherExt) OnNextSubtasksBatch(
 
 	job := &gTaskMeta.Job
 	useExtStore := len(gTaskMeta.CloudStorageURI) > 0
-	defer func() {
-		// Only redact when the task is complete.
-		if len(taskMeta) == 0 && useExtStore {
-			cleanupCloudStorageFiles(ctx, &gTaskMeta)
-			redactCloudStorageURI(ctx, gTask, &gTaskMeta)
-		}
-	}()
 
 	tblInfo, err := getTblInfo(h.d, job)
 	if err != nil {
@@ -182,7 +173,7 @@ func (h *backfillingDispatcherExt) GetEligibleInstances(ctx context.Context, _ *
 	return serverInfos, nil
 }
 
-// IsRetryableErr implements TaskFlowHandle.IsRetryableErr interface.
+// IsRetryableErr implements dispatcher.Extension.IsRetryableErr interface.
 func (*backfillingDispatcherExt) IsRetryableErr(error) bool {
 	return true
 }
@@ -494,36 +485,3 @@ func getSummaryFromLastStep(
 	}
 	return minKey, maxKey, totalKVSize, allDataFiles, allStatFiles, nil
 }
-
-func cleanupCloudStorageFiles(ctx context.Context, gTaskMeta *BackfillGlobalMeta) {
-	backend, err := storage.ParseBackend(gTaskMeta.CloudStorageURI, nil)
-	if err != nil {
-		logutil.Logger(ctx).Warn("failed to parse cloud storage uri", zap.Error(err))
-		return
-	}
-	extStore, err := storage.NewWithDefaultOpt(ctx, backend)
-	if err != nil {
-		logutil.Logger(ctx).Warn("failed to create cloud storage", zap.Error(err))
-		return
-	}
-	prefix := strconv.Itoa(int(gTaskMeta.Job.ID))
-	err = external.CleanUpFiles(ctx, extStore, prefix)
-	if err != nil {
-		logutil.Logger(ctx).Warn("cannot cleanup cloud storage files", zap.Error(err))
-		return
-	}
-}
-
-func redactCloudStorageURI(
-	ctx context.Context,
-	gTask *proto.Task,
-	origin *BackfillGlobalMeta,
-) {
-	origin.CloudStorageURI = ast.RedactURL(origin.CloudStorageURI)
-	metaBytes, err := json.Marshal(origin)
-	if err != nil {
-		logutil.Logger(ctx).Warn("fail to marshal task meta", zap.Error(err))
-		return
-	}
-	gTask.Meta = metaBytes
-}

ddl/ddl.go

@@ -706,6 +706,7 @@ func newDDL(ctx context.Context, options ...Option) *ddl {
 			func(ctx context.Context, taskMgr *storage.TaskManager, serverID string, task *proto.Task) dispatcher.Dispatcher {
 				return newLitBackfillDispatcher(ctx, taskMgr, serverID, task, backFillDsp)
 			})
+		dispatcher.RegisterDispatcherCleanUpFactory(BackfillTaskType, newBackfillCleanUpS3)
 	}
 
 	// Register functions for enable/disable ddl when changing system variable `tidb_enable_ddl`.

disttask/framework/BUILD.bazel

@@ -13,7 +13,7 @@ go_test(
     ],
     flaky = True,
     race = "off",
-    shard_count = 30,
+    shard_count = 31,
     deps = [
         "//disttask/framework/dispatcher",
         "//disttask/framework/handle",

disttask/framework/dispatcher/BUILD.bazel

@@ -37,7 +37,7 @@ go_test(
     embed = [":dispatcher"],
     flaky = True,
     race = "off",
-    shard_count = 13,
+    shard_count = 14,
     deps = [
         "//disttask/framework/mock",
         "//disttask/framework/proto",

disttask/framework/dispatcher/dispatcher.go

@@ -39,8 +39,6 @@ const (
 	MaxSubtaskConcurrency = 256
 	// DefaultLiveNodesCheckInterval is the tick interval of fetching all server infos from etcd.
 	DefaultLiveNodesCheckInterval = 2
-	// defaultHistorySubtaskTableGcInterval is the interval of gc history subtask table.
-	defaultHistorySubtaskTableGcInterval = 24 * time.Hour
 )
 
 var (